DDPM posterior from a clean-coordinate prediction

Your task

Function signature

ddpm_posterior(
    x_t: torch.Tensor,
    x0_hat: torch.Tensor,
    epsilon: torch.Tensor,
    beta: torch.Tensor,
    t: torch.Tensor,
    valid_mask: torch.Tensor,
    diffuse_mask: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

A reverse step with nonzero variance

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• x_t, x0_hat and supplied epsilon have shape [B,N,3]. beta has shape [T] and satisfies 0<beta[k]<1; integer t has shape [B] and is ONE-BASED from 1 through T. Define \alpha_k=1-beta[k-1], \bar\alpha_0=1, and \bar\alpha_k=\prod_{j=1}^{k}\alpha_j.
• For each batch item, let \beta_t=beta[t-1]. Compute c_0=\frac{\beta_t\sqrt{\bar\alpha_{t-1}}}{1-\bar\alpha_t}, c_t=\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}, and v=\frac{\beta_t(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}.
• Return posterior_mean = c0x0_hat + ctx_t and x_previous = posterior_mean + sqrt(v)*epsilon on movable valid positions. Return v broadcast to movable residues as posterior_variance [B,N], not its square root.
• At t=1 explicitly use c0=1, ct=0 and v=0, so the movable sample is exactly x0_hat regardless of epsilon or beta. Use a safe positive dummy before any dormant square root so automatic gradients stay finite.
• valid_mask=True marks real positions; diffuse_mask=True permits movement. Fixed valid positions retain x_t in both coordinate returns and have zero variance. Invalid positions return zero coordinates and variance.
• Preserve tiny positive variances with \log\bar\alpha_k=\sum_{j=1}^{k}\operatorname{log1p}(-beta[j-1]) and 1-\bar\alpha_k=-\operatorname{expm1}(\log\bar\alpha_k), or an equally stable equivalent. beta need not be monotone and duplicates are allowed. Invalid beta ranges or time indices raise ValueError.
• B and N are positive; an all-False mask represents an empty example. All values, including padding, are finite and share CPU float32/float64 dtype. There are no additive epsilons, coordinate scales, rotation updates or sampled random values.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def ddpm_posterior(
    x_t: torch.Tensor,  # (B, N, 3)
    x0_hat: torch.Tensor,  # (B, N, 3)
    epsilon: torch.Tensor,  # (B, N, 3)
    beta: torch.Tensor,  # (T,)
    t: torch.Tensor,  # (B,)
    valid_mask: torch.Tensor,  # (B, N)
    diffuse_mask: torch.Tensor,  # (B, N)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    posterior_mean: (B, N, 3).
    posterior_variance: (B, N).
    x_previous: (B, N, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement ddpm_posterior")

Sources and further reading

• RosettaCommons/RFdiffusion · utils.py