Gaussian forward corruption and score targets

Your task

Function signature

gaussian_forward(
    x0: torch.Tensor,
    epsilon: torch.Tensor,
    masks: dict,
    t: torch.Tensor | None = None,
    beta_min: torch.Tensor | None = None,
    beta_max: torch.Tensor | None = None,
    alpha: torch.Tensor | None = None,
    sigma: torch.Tensor | None = None,
    schedule: str = 'linear_vp',
) -> dict[str, torch.Tensor]

Worked example

Explicit three four five

import torch

x0 = torch.tensor([[[5.0, 10.0, 0.0]]], dtype=torch.float64, requires_grad=True)
epsilon = torch.tensor([[[1.0, 0.0, -1.0]]], dtype=torch.float64, requires_grad=True)
t = None
beta_min = None
beta_max = None
alpha = torch.tensor([0.6], dtype=torch.float64, requires_grad=True)
sigma = torch.tensor([0.8], dtype=torch.float64, requires_grad=True)
masks = {
    'valid_mask': torch.tensor([[True]], dtype=torch.bool),
    'diffuse_mask': torch.tensor([[True]], dtype=torch.bool),
}
schedule = 'explicit'

result = gaussian_forward(x0=x0, epsilon=epsilon, t=t, beta_min=beta_min, beta_max=beta_max, alpha=alpha, sigma=sigma, masks=masks, schedule=schedule)

Expected return

{
    'alpha': torch.tensor([0.6], dtype=torch.float64),
    'sigma': torch.tensor([0.8], dtype=torch.float64),
    'noisy': torch.tensor([[[3.8, 6.0, -0.8]]], dtype=torch.float64),
    'score': torch.tensor([[[-1.25, -0.0, 1.25]]], dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Given x0 and supplied epsilon [B,N,3], compute the conditional IID Gaussian corruption x_{t}=\alpha\cdot \mathrm{x0}+\sigma\cdot \epsilon and \mathrm{score}=-(x_{t}-\alpha\cdot \mathrm{x0})/\sigma^{2}=-\epsilon/\sigma at the sampled x_t.
• schedule='linear_vp' (default) takes t [B], beta_min [] and beta_max []; require 0<t\le 1 and 0<\beta_{\mathrm{min}}\le \beta_{\mathrm{max}}.
• With \beta(u)=\beta_{\min}+(\beta_{\max}-\beta_{\min})u, I(t)=\beta_{\mathrm{min}}\cdot t+(\beta_{\mathrm{max}}-\beta_{\mathrm{min}})\cdot t^{2}/2, \alpha=\operatorname{exp}(-I/2), \sigma=\sqrt{1-\operatorname{exp}(-I)}; evaluate the variance as -expm1(-I) to preserve tiny positive times.
• The alternate schedule='explicit' takes alpha [B] and sigma [B] directly, with 0\le \alpha\le 1 and \sigma>0; t/\beta inputs are then absent. Coefficients broadcast over N and xyz and are returned as [B].
• Update only valid_mask & diffuse_mask: fixed valid residues retain x0 and have zero score; invalid residues have zero noisy coordinates and score. No epsilon regularization, coordinate rescaling, centering or COM projection is applied.
• t=0/\mathrm{zero} sigma is excluded because the conditional score has no finite zero-variance limit; t=1 is valid for this forward marginal. schedule is the only key. All masks are bool: valid_mask=True denotes a real residue/atom; diffuse_mask=True denotes a movable one. Their conjunction selects updates.
• B,N (or A) are positive; an all-false mask, not a zero-length tensor, represents an empty example. All supplied numbers, including padding, are finite and share dtype/device; floating tensors are float32 or float64.

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 gaussian_forward(
    x0: torch.Tensor,  # (B, N, 3)
    epsilon: torch.Tensor,  # (B, N, 3)
    masks: dict,
    t: torch.Tensor | None = None,  # (B,)
    beta_min: torch.Tensor | None = None,  # Scalar tensor, shape ().
    beta_max: torch.Tensor | None = None,  # Scalar tensor, shape ().
    alpha: torch.Tensor | None = None,  # (B,)
    sigma: torch.Tensor | None = None,  # (B,)
    schedule: str = 'linear_vp',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    alpha: (B,).
    sigma: (B,).
    noisy: (B, N, 3).
    score: (B, N, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement gaussian_forward")

Related exercises

• Builds onVariance-Preserving Forward Noising
• Builds onGaussian Noising and Score

Sources and further reading

• RosettaCommons/RFdiffusion · diffusion.py
• jasonkyuyim/se3_diffusion · r3_diffuser.py
• microsoft/bioemu · sde_lib.py
• microsoft/Graphormer · main_model.py