Diffusion and flow matching

Fourier noise features and learned conditioning projection

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

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

Your task

Encode each scalar noise level as a vector of cosine, or cosine-and-sine, values at supplied frequencies and phases. Each Fourier component is one wave evaluated at the transformed noise level; an optional normalized projection turns those waves into learned conditioning features. Return both the waves and the embedding.

Function signature

fourier_noise_embedding(
    noise: torch.Tensor,
    mask: torch.Tensor,
    frequencies: torch.Tensor,
    phases: torch.Tensor,
    options: dict,
    norm_scale: torch.Tensor | None = None,
    norm_bias: torch.Tensor | None = None,
    weight: torch.Tensor | None = None,
    bias: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Apply log scaling before frequencies and phase

With sigma_data=16 and log_divisor=4, the two noise values give $t=0$ and $t=0.5$ . The first row angles are $[0,\pi/2,0]$ , producing [1, 0, 1] up to floating roundoff; the second produces [1, −1, −1]. Projection is disabled, so both returned tensors are equal.

import torch

noise = torch.tensor([16.0, 118.2248975828904], dtype=torch.float64, requires_grad=True)
mask = torch.tensor([True, True], dtype=torch.bool)
frequencies = torch.tensor([0.0, 0.5, 1.0], dtype=torch.float64, requires_grad=True)
phases = torch.tensor([0.0, 0.25, 0.0], dtype=torch.float64, requires_grad=True)
norm_scale = None
norm_bias = None
weight = None
bias = None
options = {
}

result = fourier_noise_embedding(noise=noise, mask=mask, frequencies=frequencies, phases=phases, norm_scale=norm_scale, norm_bias=norm_bias, weight=weight, bias=bias, options=options)

Expected return

(torch.tensor([[1.0, 6.123233995736766e-17, 1.0], [1.0, -1.0, -1.0]], dtype=torch.float64), torch.tensor([[1.0, 6.123233995736766e-17, 1.0], [1.0, -1.0, -1.0]], dtype=torch.float64))

Parameters

Parameter	Shape and meaning
`noise`torch.Tensor	$[\ldots]$ Scalar conditioning values; strictly positive in log_sigma mode and finite in identity mode.
`mask`torch.Tensor	$[\ldots]$ `bool` $[\ldots]$ matching `noise`.
`frequencies`torch.Tensor	$[K]$ Supplied Fourier frequencies; use these values without resampling.
`phases`torch.Tensor	$[K]$ Supplied offsets in cycles, before multiplication by 2π.
`norm_scale`torch.Tensor \| NoneDefault `None`	$[K \mathrm{or2K}]$ Affine LayerNorm multiplier over Fourier channels, only when project=True.
`norm_bias`torch.Tensor \| NoneDefault `None`	$[K \mathrm{or2K}]$ Affine LayerNorm offset over Fourier channels, only when project=True.
`weight`torch.Tensor \| NoneDefault `None`	$[K \mathrm{or2K}, D]$ `float` $[K \lor 2K,D]$ , when `project=True`.
`bias`torch.Tensor \| NoneDefault `None`	$[D]$ `float` $[D]$ , when `project=True`.
`options`dict	Scalar settings. Omitted keys use the defaults below.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

Encode a scalar noise field with supplied Fourier frequencies/phases, optionally followed by learned LN and projection. noise $[\ldots]$ and mask $[\ldots]$ (True means valid) have identical nonempty shape; frequencies $[K]$ ,phases $[K]$ are supplied floating tensors with $K\ge 1$ .
All floats share $\mathrm{float32}/64$ dtype. transform='log_sigma' (default) sets $t=\operatorname{log}(\mathrm{noise}/\sigma_{\mathrm{data}})/\mathrm{log\_divisor}$ , with sigma_data default $16$ and log_divisor default $4$ ; valid noise must be strictly positive.
transform='identity' takes $t=\mathrm{noise}$ and accepts any finite valid value. Compute $\theta[\ldots,k]=2\cdot \pi\cdot (t\cdot \mathrm{frequencies}[k]+\mathrm{phases}[k])$ .
basis='cos' (default) yields $\operatorname{cos}(\theta)$ ; basis='cos_sin' concatenates ALL cos channels then ALL sin channels (a documented extension, not the AF3 default). Return fourier $[\ldots,K \lor 2K]$ . If project=False (default), embedding equals fourier.
If project=True, normalize Fourier channels with population variance and $\varepsilon$ (default $10^{-5}$ ), apply supplied norm_scale and norm_bias, then $\mathrm{embedding}=\mathrm{normalized}\mathbin{@}\mathrm{weight}+\mathrm{bias}$ .
Both outputs are exact zero under the mask; sanitize masked noise to a positive finite value before a log. Frequencies/phases and learned parameters are finite, widths are positive, sigma_data/log_divisor/ $\varepsilon$ are positive.
Required arithmetic intermediates, including phases and projections, must remain finite in the common dtype; divisors, $\varepsilon$ and noise/sigma_data used by the logarithm must remain strictly positive when represented in that dtype. Unknown transform/basis or nonpositive scalar configuration raises ValueError.
Do not sample buffers, train a network, add a residual or perform EDM output preconditioning.

Why this matters

Fixed supplied frequencies turn one scalar conditioning variable into multiple smooth channels that a later model can combine.

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 fourier_noise_embedding(
    noise: torch.Tensor,  # (...)
    mask: torch.Tensor,  # (...)
    frequencies: torch.Tensor,  # (K,)
    phases: torch.Tensor,  # (K,)
    options: dict,
    norm_scale: torch.Tensor | None = None,  # (K or2K)
    norm_bias: torch.Tensor | None = None,  # (K or2K)
    weight: torch.Tensor | None = None,  # (K or2K, D)
    bias: torch.Tensor | None = None,  # (D,)
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    fourier: (...,F).
    embedding: (...,D) when project=True; (...,F) otherwise."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement fourier_noise_embedding")

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.

Sources and further reading

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