Min-SNR Loss Weights

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

Problem Statement

Compute the Min-SNR- $\gamma$ per-timestep loss weights, for all three of the prediction targets a diffusion model can be trained on.

Background

The DDPM "simple" loss weights every timestep equally in $\varepsilon$ -space. Viewed in $x_0$ -space that is an implicit weight of $\mathrm{SNR}(t)$ , which is enormous at low noise and tiny at high noise. The result is a multi-task optimisation where a handful of easy, low-noise timesteps dominate the gradient and the tasks fight each other.

Efficient Diffusion Training via Min-SNR Weighting (Hang et al., 2023) clamps that weight:

$w_t^{(x_0)} = \min(\mathrm{SNR}(t), \gamma)$

with $\gamma = 5$ working well in practice. Because the three parameterizations are related by fixed scalings of the same error, the clamped weight must be converted into whichever space the loss is actually computed in:

"x0": $\min(\mathrm{SNR}(t), \gamma)$
"epsilon": $\min(\mathrm{SNR}(t), \gamma) / \mathrm{SNR}(t)$
"v": $\min(\mathrm{SNR}(t), \gamma) / (\mathrm{SNR}(t) + 1)$

Note what the epsilon case does: wherever $\mathrm{SNR}(t) \le \gamma$ the weight is exactly 1, i.e. Min-SNR leaves the high-noise steps at the standard simple loss and only down-weights the easy low-noise ones.

Your Task

Implement:

def min_snr_weights(alpha_bars, gamma=5.0, prediction_type="epsilon"):

Return a 1-D NumPy array of weights the same length as alpha_bars. Raise ValueError for an unknown prediction_type.

Input Format

alpha_bars: 1-D NumPy array, entries strictly inside (0, 1).
gamma (float): the clamp, positive.
prediction_type (str): one of "x0", "epsilon", "v".

Output Format

A 1-D NumPy array of weights.

Sample

ab = np.array([0.9, 0.5, 0.1])
print(np.round(min_snr_weights(ab, 5.0, "x0"), 4).tolist())
print(np.round(min_snr_weights(ab, 5.0, "epsilon"), 4).tolist())
print(np.round(min_snr_weights(ab, 5.0, "v"), 4).tolist())

[5.0, 1.0, 0.1111]
[0.5556, 1.0, 1.0]
[0.5, 0.5, 0.1]

The SNRs are [9.0, 1.0, 0.1111]. Clamping at 5 gives [5.0, 1.0, 0.1111] for "x0". Dividing by the SNR gives [0.5556, 1.0, 1.0] for "epsilon" -- only the first, easiest timestep is damped.

Example

Input

ab = np.array([0.9, 0.5, 0.1])
print(np.round(min_snr_weights(ab, 5.0, "x0"), 4).tolist())
print(np.round(min_snr_weights(ab, 5.0, "epsilon"), 4).tolist())
print(np.round(min_snr_weights(ab, 5.0, "v"), 4).tolist())

Output

[5.0, 1.0, 0.1111]
[0.5556, 1.0, 1.0]
[0.5, 0.5, 0.1]

Reasoning

The SNRs are [9.0, 1.0, 0.1111]. Clamping at gamma = 5 gives the x0 weights [5.0, 1.0, 0.1111]. Dividing by the SNR converts to epsilon space: 5/9 = 0.5556, 1/1 = 1.0, 0.1111/0.1111 = 1.0 -- only the low-noise timestep is damped. Dividing by SNR + 1 gives the v-space weights.

Constraints

gamma > 0; entries of alpha_bars lie in (0, 1).
Clamp first, then convert to the requested space.
An unrecognised prediction_type must raise ValueError.
Do not round inside the function.

Learn this in

Forward Diffusion ProcessDiffusion 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

In Diffusion Models, training stability is often compromised by the disparity in difficulty across timesteps. The standard DDPM loss treats all timesteps equally in $\varepsilon$ -space (noise prediction). However, when viewed in $x_0$ -space (image prediction), this corresponds to weighting each timestep by its Signal-to-Noise Ratio ( $\mathrm{SNR}$ ). Since $\mathrm{SNR}$ is very high at early timesteps (low noise) and very low at late timesteps (high noise), the gradient is dominated by the easy, low-noise steps. This causes the model to overfit to clean data while neglecting the harder, noisy steps.

The Min-SNR weighting strategy addresses this by clamping the effective weight. Instead of using the raw $\mathrm{SNR}(t)$ , we use $\min(\mathrm{SNR}(t), \gamma)$ . This ensures that no single timestep contributes disproportionately more to the loss than a threshold $\gamma$ (typically $5.0$ ). High-noise steps (low SNR) retain their original weight, while low-noise steps (high SNR) are down-weighted to $\gamma$ .

Crucially, diffusion models can be trained to predict different targets: the original image ( $x_0$ ), the noise ( $\varepsilon$ ), or a velocity vector ( $v$ ). The relationship between these predictions is linear. Therefore, the loss weight must be adjusted based on the prediction type to maintain the same effective influence on the gradient. The problem provides the specific scaling factors for each type derived from the chain rule of the loss function with respect to the model output.

2. Algorithm Approach

The core algorithm is a vectorized element-wise transformation. Since the input alpha_bars is a NumPy array, we should leverage NumPy's broadcasting and vectorized operations for efficiency. The approach involves three main stages:

SNR Calculation: Convert the provided $\bar{\alpha}_t$ values into Signal-to-Noise Ratios. The relationship is defined by the variance of the signal versus the variance of the noise at timestep $t$ .
Clamping: Apply the $\min(\cdot, \gamma)$ operation to the computed SNR values. This creates the base weight for the $x_0$ prediction target.
Scaling: Depending on the prediction_type, scale the clamped weights using the specific denominators provided in the problem description. This step ensures the loss is correctly normalized for the specific output space of the diffusion model.

This approach avoids explicit Python loops, making it computationally efficient and idiomatic for numerical computing tasks.

3. Step-by-Step Strategy

Validate Input: Check if prediction_type is one of the allowed strings ("x0", "epsilon", "v"). If not, raise a ValueError.
Compute SNR:

Recall that $\bar{\alpha}_t$ represents the cumulative product of noise schedules $\alpha_t$ .
The signal variance is proportional to $\bar{\alpha}_t$ .
The noise variance is proportional to $1 - \bar{\alpha}_t$ .
Calculate $\mathrm{SNR}(t) = \frac{\bar{\alpha}_t}{1 - \bar{\alpha}_t}$ . Use NumPy array operations to compute this for all elements in alpha_bars simultaneously.

Apply Min-SNR Clamping:

Create a new array snr_clamped where each element is $\min(\mathrm{SNR}(t), \gamma)$ .
In NumPy, this can be achieved using np.minimum(snr, gamma).

Apply Prediction Type Scaling:

Use a conditional structure (if-elif-else) to handle the three cases:
If "x0": The weight is simply snr_clamped.
If "epsilon": Divide snr_clamped by the original snr. Note that this results in weights $\le 1$ .
If "v": Divide snr_clamped by $(\mathrm{SNR}(t) + 1)$ .

Return Result: Return the resulting 1-D NumPy array.

4. Common Pitfalls

Incorrect SNR Formula: A common mistake is confusing $\bar{\alpha}_t$ with $\alpha_t$ or misremembering the SNR definition. Ensure you use $\frac{\bar{\alpha}_t}{1 - \bar{\alpha}_t}$ . Using $\frac{1-\bar{\alpha}_t}{\bar{\alpha}_t}$ would invert the logic, weighting high-noise steps more heavily, which is the opposite of the intended effect.
Division by Zero: Although the problem states alpha_bars are strictly inside $(0, 1)$ , numerical precision issues could theoretically lead to $1 - \bar{\alpha}_t$ being extremely close to zero. However, given the constraints, standard floating-point division is usually safe. Be aware that for "epsilon" and "v" types, you are dividing by SNR or SNR+1, which are positive, so no division by zero occurs there.
In-place Modification: Avoid modifying the input alpha_bars array directly. Create new arrays for intermediate calculations to ensure the function is pure and side-effect free.
Wrong Scaling Denominators: Double-check the denominators for "epsilon" and "v".
"epsilon" divides by $\mathrm{SNR}$ .
"v" divides by $\mathrm{SNR} + 1$ .
Confusing these will lead to incorrect gradient magnitudes.
Non-Vectorized Loops: Using a Python for loop to iterate over alpha_bars is significantly slower than using NumPy vectorized operations. Always prefer np.minimum and array arithmetic.

5. Time & Space Complexity

Time Complexity: $O(N)$ , where $N$ is the length of the alpha_bars array. Each operation (division, subtraction, minimum, multiplication) is performed element-wise in constant time per element. NumPy operations are optimized in C, making this very fast in practice.
Space Complexity: $O(N)$ , as we need to store the intermediate SNR array and the final weights array. The space is linear with respect to the number of timesteps.

solution.py

Python 3.13

Test results