1 / 1

Min-SNR Loss Weighting

Diffusion ModelsLoss FunctionsTraining

Medium

Weight a batch of per-sample diffusion losses with the Min-SNR strategy and return the weighted mean loss.

\mathrm{snr}_i = \frac{\bar\alpha_{t_i}}{1 - \bar\alpha_{t_i}}, \qquad w_i = \begin{cases} \min(\mathrm{snr}_i, \gamma) / \mathrm{snr}_i & \text{epsilon} \\ \min(\mathrm{snr}_i, \gamma) & \text{x0} \\ \min(\mathrm{snr}_i, \gamma) / (\mathrm{snr}_i + 1) & \text{v} \end{cases}

L = \frac{1}{B} \sum_{i=1}^{B} w_i\, m_i

Here per_sample_mse holds the unweighted losses $m_i$ of shape (B,), timesteps gives each sample its timestep $t_i$ , and $\bar\alpha_{t_i}$ is alpha_bars at index $t_i$ . The value $\mathrm{snr}_i$ is the signal to noise ratio of sample $i$ , gamma is the cap $\gamma$ , and prediction_type names the network target, which picks the row of the weight formula.

Return a dict with three keys. The keys "snr" and "weights" hold PyTorch tensors of shape (B,) with the dtype of per_sample_mse, and "loss" holds $L$ as a Python float.

Examples

Input: per_sample_mse = [0.5,0.2,0.8], timesteps = [0,1,2], alpha_bars = [0.9,0.5,0.2], gamma = 5.0, prediction_type = "epsilon"

Output: {"snr":[8.999998,1.0,0.25],"weights":[0.555556,1.0,1.0],"loss":0.425926}

Explanation: The first sample has SNR 0.9 / 0.1 = 9, which is clipped to gamma = 5, so its noise prediction weight is 5 / 9, about 0.556. The other two samples have SNR at most 1 and keep weight 1, and the loss is the mean of the weighted errors. The inputs are float32, so the SNR of 9 prints as 8.999998.

Input: per_sample_mse = [1.0,0.4], timesteps = [0,3], alpha_bars = [0.99,0.8,0.5,0.1], gamma = 5.0, prediction_type = "x0"

Output: {"snr":[99.000099,0.111111],"weights":[5.0,0.111111],"loss":2.522222}

Explanation: The inputs are float32, so the first SNR of 0.99 / 0.01 = 99 prints as 99.000099. It is capped at gamma = 5, which is the x0 weight, while the second SNR of 0.1 / 0.9 stays below gamma.

Input: per_sample_mse = [0.3,0.6], timesteps = [1,1], alpha_bars = [0.95,0.75], gamma = 1.0, prediction_type = "v"

Output: {"snr":[3.0,3.0],"weights":[0.25,0.25],"loss":0.1125}

Hint 1

torch.clamp(x, max=gamma) caps a tensor from above.

Hint 2

.mean().item() turns a one-element mean into a Python float.

Requirements

Use PyTorch tensor operations.
Use 0-based indexing into alpha_bars with each sample's own timestep.
Cap the signal to noise ratio from above at gamma, and leave values below gamma unchanged.
Raise a ValueError for any prediction_type other than the exact strings "epsilon", "x0" and "v".
Average the weighted losses over the batch.
Leave all input tensors unchanged.

Constraints

per_sample_mse is a floating-point PyTorch tensor of shape (B,) with B between 1 and 8 and non-negative values.
timesteps is an int64 tensor of shape (B,) with every value between 0 and T - 1.
alpha_bars has shape (T,) with T between 1 and 30, the dtype of per_sample_mse, and every value strictly between 0 and 1.
gamma is a positive Python number between 0.5 and 20.
prediction_type is a Python string.
Supported dtypes are float32 and float64. All values are finite.

Try Similar Problems

Noise Prediction LossEasy Denoising Score Matching LossMedium Conditional Flow Matching LossMedium Posterior Mean and VarianceMedium Prediction ParameterizationMedium Toy Diffusion Training StepHard Gaussian Mixture ScoreMedium Cosine Noise ScheduleMedium Higher-Order Derivatives & Loss CurvatureMedium Loss FunctionsMedium Bradley-Terry Reward Model LossMedium REINFORCE with BaselineMedium

Min-SNR Loss Weighting

Diffusion ModelsLoss FunctionsTraining

Medium

Weight a batch of per-sample diffusion losses with the Min-SNR strategy and return the weighted mean loss.

\mathrm{snr}_i = \frac{\bar\alpha_{t_i}}{1 - \bar\alpha_{t_i}}, \qquad w_i = \begin{cases} \min(\mathrm{snr}_i, \gamma) / \mathrm{snr}_i & \text{epsilon} \\ \min(\mathrm{snr}_i, \gamma) & \text{x0} \\ \min(\mathrm{snr}_i, \gamma) / (\mathrm{snr}_i + 1) & \text{v} \end{cases}

L = \frac{1}{B} \sum_{i=1}^{B} w_i\, m_i

Return a dict with three keys. The keys "snr" and "weights" hold PyTorch tensors of shape (B,) with the dtype of per_sample_mse, and "loss" holds $L$ as a Python float.

Examples

Input: per_sample_mse = [0.5,0.2,0.8], timesteps = [0,1,2], alpha_bars = [0.9,0.5,0.2], gamma = 5.0, prediction_type = "epsilon"

Output: {"snr":[8.999998,1.0,0.25],"weights":[0.555556,1.0,1.0],"loss":0.425926}

Input: per_sample_mse = [1.0,0.4], timesteps = [0,3], alpha_bars = [0.99,0.8,0.5,0.1], gamma = 5.0, prediction_type = "x0"

Output: {"snr":[99.000099,0.111111],"weights":[5.0,0.111111],"loss":2.522222}

Explanation: The inputs are float32, so the first SNR of 0.99 / 0.01 = 99 prints as 99.000099. It is capped at gamma = 5, which is the x0 weight, while the second SNR of 0.1 / 0.9 stays below gamma.

Input: per_sample_mse = [0.3,0.6], timesteps = [1,1], alpha_bars = [0.95,0.75], gamma = 1.0, prediction_type = "v"

Output: {"snr":[3.0,3.0],"weights":[0.25,0.25],"loss":0.1125}

Hint 1

torch.clamp(x, max=gamma) caps a tensor from above.

Hint 2

.mean().item() turns a one-element mean into a Python float.

Requirements

Use PyTorch tensor operations.
Use 0-based indexing into alpha_bars with each sample's own timestep.
Cap the signal to noise ratio from above at gamma, and leave values below gamma unchanged.
Raise a ValueError for any prediction_type other than the exact strings "epsilon", "x0" and "v".
Average the weighted losses over the batch.
Leave all input tensors unchanged.

Constraints

per_sample_mse is a floating-point PyTorch tensor of shape (B,) with B between 1 and 8 and non-negative values.
timesteps is an int64 tensor of shape (B,) with every value between 0 and T - 1.
alpha_bars has shape (T,) with T between 1 and 30, the dtype of per_sample_mse, and every value strictly between 0 and 1.
gamma is a positive Python number between 0.5 and 20.
prediction_type is a Python string.
Supported dtypes are float32 and float64. All values are finite.

Try Similar Problems

per_sample_mse =

Accepts: array

timesteps =

Accepts: array

alpha_bars =

Accepts: array

gamma =

Accepts: number

prediction_type =

Accepts: string