Diffusion and flow matching

Variance-Preserving Forward Noising

MediumRaw PyTorchAbout 20 min1 public example plus private tests

FrameDiffDiG proteinBioEmu

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Shrink the clean coordinates by the continuous VP schedule’s alpha and add sigma times the supplied noise at movable valid points. Return noisy coordinates and the conditional score: the gradient of their Gaussian log density with respect to the noisy coordinates, equal to −epsilon/sigma there.

Function signature

vp_forward_noising(
    x0: torch.Tensor,
    epsilon: torch.Tensor,
    valid_mask: torch.Tensor,
    diffuse_mask: torch.Tensor,
    t: torch.Tensor,
    beta_min: torch.Tensor,
    beta_max: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

Integrate a changing schedule before adding noise

Use one batch with x0=[[[2,0,0],[7,0,0],[9,0,0]]], epsilon=[[[1,0,-1],[8,8,8],[8,8,8]]], t=[0.5], beta_min=1, beta_max=3, valid_mask=[[True,True,False]], and diffuse_mask=[[True,False,True]].

I=1\cdot\tfrac12+\tfrac12(3-1)(\tfrac12)^2=\tfrac34,\qquad \alpha=e^{-3/8},\qquad \sigma=\sqrt{1-e^{-3/4}}.

Thus alpha≈[0.687289], sigma≈[0.726384]. The movable point becomes noisy≈[2.100962,0,-0.726384] with score≈[-1.376683,0,1.376683]. The fixed point remains [7,0,0]; the invalid point becomes zero. Both have zero score. Each coordinate output retains shape (1,3,3); no new noise is drawn.

Parameters

Parameter	Shape and meaning
`x0`torch.Tensor	$[B, N, 3]$ Clean positions.
`epsilon`torch.Tensor	$[B, N, 3]$ Supplied IID standard-normal variates.
`valid_mask`torch.Tensor	$[B, N]$ True marks a real point; False returns zero coordinates and score.
`diffuse_mask`torch.Tensor	$[B, N]$ True permits noising at a valid point; False keeps its clean coordinates.
`t`torch.Tensor	$[B]$ Positive continuous diffusion time for each batch item.
`beta_min`torch.Tensor	Scalar tensor, shape $[]$ .Positive scalar rate at time zero.
`beta_max`torch.Tensor	Scalar tensor, shape $[]$ .Scalar end rate, at least beta_min.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Use $I(t)=\beta_{min}t+\tfrac12(\beta_{max}-\beta_{min})t^2$ , $\alpha=e^{-I/2}$ , and $\sigma=\sqrt{-\operatorname{expm1}(-I)}$ . Time t is supplied per batch; beta_min and beta_max are scalar tensors; do not discretize the schedule.
For points where valid_mask and diffuse_mask are both True, compute $x_t=\alpha x_0+\sigma\epsilon$ and score $-\epsilon/\sigma$ . Broadcast coefficients over the point and coordinate axes.
A valid fixed point keeps x0 and has zero score. Invalid points have zero noisy coordinates and zero score; sanitize excluded x0/noise payloads before arithmetic.
Return alpha and sigma as supplied/computed batch vectors. The sampled epsilon is an input; do not sample random noise.
Use CPU float32/float64 tensors with one shared floating dtype/device. Accept noncontiguous input layouts and do not mutate supplied tensors.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.

Why this matters

The score points toward the conditional Gaussian mean, with a magnitude determined by the noise variance.

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 vp_forward_noising(
    x0: torch.Tensor,  # (B, N, 3)
    epsilon: torch.Tensor,  # (B, N, 3)
    valid_mask: torch.Tensor,  # (B, N)
    diffuse_mask: torch.Tensor,  # (B, N)
    t: torch.Tensor,  # (B,)
    beta_min: torch.Tensor,  # Scalar tensor, shape ().
    beta_max: torch.Tensor,  # Scalar tensor, shape ().
) -> 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 vp_forward_noising")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Related exercises

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.