SNR-Weighted Variational Bound Terms

hardFREESignal-to-Noise Ratio and Loss WeightingDiffusion ModelsDiffusion Models 1

Problem Statement

Evaluate the per-sample denoising terms $L_{t-1}$ of the diffusion variational lower bound, using the SNR form that makes the bound's structure obvious.

Background

The VLB of a diffusion model decomposes into one KL per timestep. For the denoising terms the KL is between two Gaussians that differ only in their means, and after substituting the DDPM posterior mean and variance it collapses to a remarkably clean statement (Kingma et al., Variational Diffusion Models):

$L_{t-1} \;=\; \tfrac{1}{2}\left(\mathrm{SNR}(t-1) - \mathrm{SNR}(t)\right)\,\lVert x_0 - \hat{x}_0(x_t, t) \rVert_2^2$

where $\mathrm{SNR}(t) = \bar{\alpha}_t / (1-\bar{\alpha}_t)$ and $\hat{x}_0$ is the model's reconstruction of the clean sample.

Two things to read off this formula:

The weight is the drop in SNR across the step. Since SNR is strictly decreasing in $t$ , the weight is always positive, and it is largest for the low-noise steps where the SNR curve is steepest.
The bound depends on the schedule only through its SNR values at the endpoints of each step. That is the formal reason the noise schedule is a design choice about weighting, not about the model.

The squared norm is a sum over all the data dimensions of one sample, not a mean.

Your Task

Implement:

def vlb_terms(x0, x0_pred, timesteps, alpha_bars):

Return a 1-D NumPy array of length N holding $L_{t-1}$ for each sample in the batch.

Input Format

x0: array of shape (N, D), the clean data.
x0_pred: array of shape (N, D), the model's reconstruction.
timesteps: integer array of shape (N,), one 0-based timestep index per sample.
alpha_bars: 1-D array of cumulative alphas, entries in (0, 1).

Every timestep must satisfy t >= 1 (the $t = 0$ term is the separate reconstruction term $L_0$ , which this function does not handle). If any entry is < 1, raise ValueError.

Output Format

A 1-D NumPy array of N floats.

Sample

ab = np.array([0.9, 0.5, 0.1])
x0 = np.array([[1.0, 0.0]])
x0_pred = np.array([[0.5, 0.0]])
print(np.round(vlb_terms(x0, x0_pred, np.array([1]), ab), 4).tolist())

SNR(0) = 9.0, SNR(1) = 1.0, the squared error is 0.25, so the term is 0.5 * (9.0 - 1.0) * 0.25 = 1.0.

Example

Input

ab = np.array([0.9, 0.5, 0.1])
x0 = np.array([[1.0, 0.0]])
x0_pred = np.array([[0.5, 0.0]])
print(np.round(vlb_terms(x0, x0_pred, np.array([1]), ab), 4).tolist())

Output

[1.0]

Reasoning

Inputs: one sample ( $N = 1$ ) with $D = 2$ dimensions. The clean data is $x_0 = (1.0, 0.0)$ , the reconstruction is $\hat{x}_0 = (0.5, 0.0)$ , the timestep is $t = 1$ , and the schedule is $\bar{\alpha} = (0.9, 0.5, 0.1)$ with 0-based positions.
SNR at both ends of the step: $t-1 = 0$ reads $\bar{\alpha}_0 = 0.9$ , so $\mathrm{SNR}(0) = 0.9 / 0.1 = 9.0$ ; $t = 1$ reads $\bar{\alpha}_1 = 0.5$ , so $\mathrm{SNR}(1) = 0.5 / 0.5 = 1.0$ . The weight is $\tfrac{1}{2}(9.0 - 1.0) = 4.0$ .
Reconstruction error: $\lVert x_0 - \hat{x}_0 \rVert_2^2 = (1.0 - 0.5)^2 + (0.0 - 0.0)^2 = 0.25$ , summed over both dimensions.
Output: $4.0 \times 0.25 = 1.0$ , so the result is [1.0]: one non-negative value per sample in the batch.
Notice: averaging the error over $D$ would give $0.125$ and halve the answer. The same error at $t = 2$ would only get weight $\tfrac{1}{2}(1.0 - 0.1111) \approx 0.4444$ , because the SNR curve is much flatter at the noisy end. A timestep of $0$ is not allowed and must raise ValueError.

Constraints

x0 and x0_pred both have shape (N, D); the output has shape (N,).
Use a sum of squared errors over the D axis, not a mean.
Any timestep < 1 must raise ValueError.
Handle a whole batch of different timesteps at once (fancy-index alpha_bars).
Do not round inside the function.

Learn this in

Introduction to Generative ModelsDiffusion Study Plan Reverse Process & DenoisingDiffusion Study Plan

More from Diffusion Models 1

Classifier-Free Guidance: Uncond/Cond Mixingeasy DDPM Posterior Mean and Variance from x0medium Adaptive CFG Scheduling over the Trajectoryhard Classifier Gradient Guidance Stepmedium Negative Prompt as the Unconditional Branchmedium

Concept primer

1. Background Knowledge

Diffusion models are generative models that learn to reverse a gradual noising process. The training objective is often derived from the Variational Lower Bound (VLB) on the log-likelihood. The VLB decomposes into a sum of terms, one for each timestep $t$ . For $t > 1$ , the term $L_{t-1}$ represents the Kullback-Leibler (KL) divergence between the true posterior $q(x_{t-1} | x_t, x_0)$ and the approximate posterior $p_\theta(x_{t-1} | x_t)$ .

In standard DDPM formulations, this KL divergence simplifies significantly because both distributions are Gaussian with the same variance but different means. The resulting loss term is proportional to the squared difference between the predicted mean and the true mean. However, the proportionality constant depends on the noise schedule parameters $\alpha_t$ .

The problem introduces the Signal-to-Noise Ratio (SNR) formulation, which provides a more intuitive understanding of the loss weighting. The SNR at timestep $t$ is defined as $\mathrm{SNR}(t) = \bar{\alpha}_t / (1 - \bar{\alpha}_t)$ , where $\bar{\alpha}_t$ is the cumulative product of the alphas up to time $t$ . This formulation reveals that the weight of the reconstruction error $\|x_0 - \hat{x}_0\|^2$ is determined by the drop in SNR across the timestep. This highlights that the noise schedule effectively controls how much emphasis is placed on denoising at different levels of noise.

2. Algorithm Approach

The core task is to compute the scalar value $L_{t-1}$ for each sample in a batch using the provided formula:

$L_{t-1} = \frac{1}{2} \left( \mathrm{SNR}(t-1) - \mathrm{SNR}(t) \right) \| x_0 - \hat{x}_0(x_t, t) \|_2^2$

The approach involves three main computational steps:

Validation: Ensure all timesteps are valid ( $t \ge 1$ ).
SNR Calculation: Compute the SNR values for the relevant timesteps ( $t$ and $t-1$ ) using the provided alpha_bars array.
Loss Computation: Calculate the squared Euclidean distance between x0 and x0_pred for each sample, then multiply by the weighted SNR difference.

This is a vectorized operation. Since x0 and x0_pred are of shape (N, D), the squared error should be computed per sample (summing over dimension D). The SNR weights are scalars per sample (indexed by timesteps). The final result is an element-wise multiplication of the weights and the squared errors.

3. Step-by-Step Strategy

Input Validation:

Check if any value in timesteps is less than 1.
If so, raise a ValueError as specified.

Compute SNR Values:

Recall $\mathrm{SNR}(t) = \bar{\alpha}_t / (1 - \bar{\alpha}_t)$ .
Extract the $\bar{\alpha}$ values for the current timesteps t and the previous timesteps t-1 from alpha_bars.
Calculate $\mathrm{SNR}(t)$ and $\mathrm{SNR}(t-1)$ for each sample.
Compute the weight $w_t = \frac{1}{2} (\mathrm{SNR}(t-1) - \mathrm{SNR}(t))$ .

Compute Squared Reconstruction Error:

Calculate the difference: $\text{diff} = x_0 - x_0\_pred$ .
Square the differences element-wise: $\text{diff}^2$ .
Sum along the feature dimension (axis 1) to get the squared L2 norm for each sample: $\| x_0 - \hat{x}_0 \|_2^2 = \sum_{d=1}^D (x_{0,d} - \hat{x}_{0,d})^2$ .

Combine Terms:

Multiply the weight vector $w_t$ by the squared error vector element-wise.
Return the resulting 1-D array.

4. Common Pitfalls

Indexing Errors: The timesteps array is 0-based. Ensure you correctly index alpha_bars for both $t$ and $t-1$ . For a timestep $t$ , you need alpha_bars[t] and alpha_bars[t-1].
Mean vs. Sum: The problem explicitly states that the squared norm is a sum over dimensions, not a mean. Using np.mean instead of np.sum along axis 1 will result in an incorrect scaling factor.
SNR Formula: Ensure the SNR is calculated as $\bar{\alpha} / (1 - \bar{\alpha})$ , not $\bar{\alpha} / (1 - \alpha_t)$ or other variations. The denominator is $1 - \bar{\alpha}_t$ .
Division by Zero: While alpha_bars entries are in $(0, 1)$ , numerical precision issues could theoretically cause $1 - \bar{\alpha}_t$ to be very small. However, given the constraints, standard floating-point arithmetic should suffice.
Broadcasting: Ensure that the weight vector (shape (N,)) and the squared error vector (shape (N,)) are aligned correctly for element-wise multiplication.

5. Time & Space Complexity

Time Complexity: $O(N \cdot D)$ , where $N$ is the batch size and $D$ is the dimensionality of the data. This is because we iterate over each element in the x0 and x0_pred arrays to compute the squared differences and sum them. The SNR calculations are $O(N)$ as they involve simple arithmetic on arrays of size $N$ .
Space Complexity: $O(N)$ for storing the intermediate SNR values and the final result. The squared error computation can be done in-place or with temporary arrays of size $N \cdot D$ , but typically we consider the auxiliary space for the result and weights, which is $O(N)$ . If we count the temporary storage for the squared differences, it is $O(N \cdot D)$ .

solution.py

Python 3.13

Test results