Min-SNR-Gamma Loss Weights

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

Problem Statement

The Min-SNR-gamma strategy (Hang et al., 2023) rebalances the multi-task diffusion loss by capping each timestep's weight. Compute the per-timestep weights for the common eps-prediction case.

Background

For eps-prediction, the Min-SNR weight clamps the SNR at gamma and normalizes by the SNR:

$w_t = \frac{\min(\text{SNR}_t, \gamma)}{\text{SNR}_t}$

So high-SNR (low-noise, easy) steps get down-weighted toward gamma/SNR, while low-SNR steps keep weight 1. This stops the near-clean timesteps from dominating training.

Your Task

Implement:

def min_snr_weights(snr, gamma):

Return the list of weights rounded to 4 decimals. A timestep with SNR == 0 gets weight 1.0 (no down-weighting possible).

Input Format

snr: list of non-negative SNRs.
gamma (float): the clamp, typically 5.0.

Output Format

A list of floats rounded to 4 decimals.

Sample

print(min_snr_weights([10.0, 5.0, 1.0], 5.0))

Output:

[0.5, 1.0, 1.0]

Example

Input

print(min_snr_weights([10.0, 5.0, 1.0], 5.0))

Output

[0.5, 1.0, 1.0]

Reasoning

For the first timestep with $\text{SNR} = 10.0$ , the SNR exceeds the clamp $\gamma = 5.0$ , so the weight is calculated as $\min(10.0, 5.0) / 10.0 = 5.0 / 10.0 = 0.5$ .
For the second timestep with $\text{SNR} = 5.0$ , the SNR equals the clamp, resulting in a weight of $\min(5.0, 5.0) / 5.0 = 5.0 / 5.0 = 1.0$ .
For the third timestep with $\text{SNR} = 1.0$ , the SNR is below the clamp, so the weight remains $\min(1.0, 5.0) / 1.0 = 1.0 / 1.0 = 1.0$ .
The computed weights $[0.5, 1.0, 1.0]$ are already within 4 decimal places, so no further rounding changes are needed.
The final output is [0.5, 1.0, 1.0]

Constraints

1 <= len(snr) <= 100000, all snr[i] >= 0, gamma > 0.
w_t = min(SNR_t, gamma) / SNR_t; SNR_t == 0 gives 1.0.
Round to 4 decimals.

Learn this in

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 primerPREVIEW

1. Background Knowledge

In diffusion models, training involves a multi-task objective where the network predicts noise (or signal) at every timestep $t$ . The standard loss is a simple mean over all timesteps, but this is problematic: timesteps with very high Signal-to-Noise Ratio (SNR) correspond to nearly clean images, and the prediction task there is trivially easy yet can dominate the gradient signal. This skews learning toward low-noise regions and degrades generation quality.

The Min-SNR- $\gamma$ strategy (Hang et al., 2023) addresses this by introducing per-timestep weights that cap the effective contribution of high-SNR steps. For the $\epsilon$ -prediction parameterization, the weight at timestep $t$ is defined as:

w_t = \frac{\min(\text{SNR}_t, \gamma)}{\text{SNR}_t}

When $\text{SNR}_t \leq \gamma$ , the weight is $1$ (no change). When $\text{SNR}_t > \gamma$ , the weight becomes $\gamma / \text{SNR}_t$ , which is less than $1$ , thereby down-weighting those easy steps. The parameter $\gamma$ (commonly $5.0$ ) controls the strength of the cap.

A special edge case exists when $\text{SNR}_t = 0$ : division by zero is undefined, and physically a zero-SNR step is pure noise with no signal to predict, so the convention is to assign weight $1.0$ (no down-weighting).

2. Algorithm Approach

This is a straightforward element-wise transformation problem. For each SNR value in the input list, apply a conditional formula:

If the SNR is zero, return $1.0$ .
Otherwise, compute $\min(\text{SNR}, \gamma) / \text{SNR}$ .

The pattern is: iterate, branch on a condition, apply a scalar arithmetic operation, and round the result. No sorting, no accumulation, no data structures beyond the output list are needed.

3. Step-by-Step Strategy

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.