Diffusion and flow matching

Gaussian Noising and Score

EasyRaw PyTorchAbout 10 min1 public example plus private tests

RFdiffusionFrameDiffBioEmuChroma

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

Your task

At each movable valid point, combine alpha times the clean coordinate with sigma times the supplied noise. Return this noisy coordinate and its conditional score, the gradient of the Gaussian log density with respect to the noisy coordinate: −epsilon/sigma. Return the supplied batch coefficients unchanged.

Function signature

gaussian_noising_and_score(
    x0: torch.Tensor,
    epsilon: torch.Tensor,
    valid_mask: torch.Tensor,
    diffuse_mask: torch.Tensor,
    alpha: torch.Tensor,
    sigma: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

One point with supplied coefficients

Multiply [5,10,0] by 0.6 and add 0.8*[1,0,-1] to get [3.8,6,-0.8]. The conditional score is -[1,0,-1]/0.8 = [-1.25,0,1.25]. Return alpha=[0.6] and sigma=[0.8] unchanged.

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)
valid_mask = torch.tensor([[True]], dtype=torch.bool)
diffuse_mask = torch.tensor([[True]], dtype=torch.bool)
alpha = torch.tensor([0.6], dtype=torch.float64, requires_grad=True)
sigma = torch.tensor([0.8], dtype=torch.float64, requires_grad=True)

result = gaussian_noising_and_score(x0=x0, epsilon=epsilon, valid_mask=valid_mask, diffuse_mask=diffuse_mask, alpha=alpha, sigma=sigma)

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

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 forces its noisy coordinates and score to zero.
`diffuse_mask`torch.Tensor	$[B, N]$ True allows a valid point to move; False keeps a valid point at x0.
`alpha`torch.Tensor	$[B]$ Supplied mean multiplier for each batch item.
`sigma`torch.Tensor	$[B]$ Supplied positive noise standard deviation for each batch item.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

alpha and strictly positive sigma are supplied as one coefficient per batch; use them directly without constructing a 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 gaussian_noising_and_score(
    x0: torch.Tensor,  # (B, N, 3)
    epsilon: torch.Tensor,  # (B, N, 3)
    valid_mask: torch.Tensor,  # (B, N)
    diffuse_mask: torch.Tensor,  # (B, N)
    alpha: torch.Tensor,  # (B,)
    sigma: torch.Tensor,  # (B,)
) -> 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_noising_and_score")

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.