Min-SNR-Gamma Loss Weights
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:
wt=SNRtmin(SNRt,γ)
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:
print(min_snr_weights([10.0, 5.0, 1.0], 5.0))
[0.5, 1.0, 1.0]
- For the first timestep with SNR=10.0, the SNR exceeds the clamp γ=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 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 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, allsnr[i] >= 0,gamma > 0.w_t = min(SNR_t, gamma) / SNR_t;SNR_t == 0gives1.0.- Round to 4 decimals.
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-γ strategy (Hang et al., 2023) addresses this by introducing per-timestep weights that cap the effective contribution of high-SNR steps. For the ϵ-prediction parameterization, the weight at timestep t is defined as:
wt=SNRtmin(SNRt,γ)When SNRt≤γ, the weight is 1 (no change). When SNRt>γ, the weight becomes γ/SNRt, which is less than 1, thereby down-weighting those easy steps. The parameter γ (commonly 5.0) controls the strength of the cap.
A special edge case exists when SNRt=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(SNR,γ)/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
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.