SNR-Weighted Variational Bound Terms

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.