Min-SNR Loss Weights

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.