AI Book - Master Artificial Intelligence by Building from Scratch

Learning Objectives

By the end of this section, you will be able to:

Understand how Signal-to-Noise Ratio (SNR) provides a principled framework for designing loss weights
Implement the min-SNR-gamma weighting strategy that balances sample quality and training stability
Apply P2 (Perception Prioritized) weighting to improve perceptual quality
Choose appropriate weighting strategies for different use cases and model architectures

SNR-Based Weighting

The Signal-to-Noise Ratio (SNR) provides a unified framework for understanding and designing loss weighting strategies. At timestep $t$ , the SNR is defined as:

$\text{SNR}(t) = \frac{\bar{\alpha}_t}{1 - \bar{\alpha}_t}$

This ratio captures how much signal (the original data) remains relative to noise at each timestep:

Timestep Region	SNR Value	Interpretation
t near 0	High (>> 1)	Signal dominates, nearly clean data
t in middle	Around 1	Signal and noise are balanced
t near T	Low (<< 1)	Noise dominates, nearly pure noise

Why SNR Matters for Weighting

The VLB-derived weights can be expressed in terms of SNR:

$w_t^{\text{VLB}} \propto \frac{1}{\text{SNR}(t) - \text{SNR}(t-1)}$

This creates extremely large weights where SNR changes slowly (low noise region) and small weights where SNR changes rapidly (high noise region). Alternative weighting strategies modify this relationship to achieve different objectives.

The Core Insight: Different weighting strategies correspond to different assumptions about which timesteps are most important for sample quality. There's no universally "correct" choice - it depends on your evaluation metric and use case.

Common SNR-Based Weights

Uniform: $w_t = 1$ - Equal weight to all timesteps
SNR: $w_t = \text{SNR}(t)$ - Weight proportional to signal strength
1/SNR: $w_t = 1/\text{SNR}(t)$ - Weight inversely proportional to signal (emphasizes noisy timesteps)
Truncated SNR: $w_t = \max(\text{SNR}(t), 1)$ - SNR clamped to minimum of 1

Min-SNR-gamma Strategy

The min-SNR-gamma weighting, introduced by Hang et al. (2023), provides an elegant solution that prevents both extremes of overweighting:

$w_t = \min\left(\text{SNR}(t), \gamma\right)$

where $\gamma$ is a hyperparameter (typically 5). This simple modification has profound effects:

How It Works

At low-noise timesteps (high SNR): Weight is clamped to $\gamma$ , preventing excessive focus on easy predictions
At high-noise timesteps (low SNR): Weight equals the SNR, giving reasonable attention to difficult timesteps
At the crossover (SNR = $\gamma$ ): Smooth transition between regimes

Interactive Visualization

Use the visualization below to explore how different weighting strategies distribute emphasis across timesteps. Pay attention to the min-SNR-gamma curve and how the gamma parameter affects the crossover point.

Loading visualization...

Mathematical Motivation

The min-SNR-gamma weighting can be derived from a modified variational bound. Consider the reweighted objective:

$L = \mathbb{E}_t\left[ \frac{\min(\text{SNR}(t), \gamma)}{\text{SNR}(t)} \cdot L_t^{\text{VLB}} \right]$

This rescales the VLB contribution by a factor that:

Equals 1 for timesteps where $\text{SNR}(t) \leq \gamma$
Equals $\gamma / \text{SNR}(t) < 1$ for timesteps where $\text{SNR}(t) > \gamma$

Choosing Gamma

The original paper recommends

\gamma = 5

as a good default. Lower values (e.g., 1) emphasize high-noise timesteps more, while higher values (e.g., 10-20) behave more like SNR weighting. Experiment based on your dataset and quality metrics.

Perception Prioritized (P2) Weighting

Choi et al. (2022) introduced P2 weighting, which prioritizes timesteps that matter most for perceptual quality rather than likelihood. The key insight is that human perception is most sensitive to certain frequency components that emerge at specific noise levels.

The P2 Weight Formula

P2 weighting is defined as:

$w_t^{\text{P2}} = \frac{1}{(k + \text{SNR}(t))^\gamma}$

where:

$k$ is a constant (typically 1) that prevents division by zero
$\gamma$ controls the decay rate (typically 1)

Why Perception Matters

Different noise levels affect different aspects of the image:

SNR Range	Affected Content	Perceptual Impact
Very High	Fine textures, high-frequency details	Subtle, less noticeable
High	Sharp edges, local structure	Important for clarity
Medium	Object shapes, global structure	Critical for recognition
Low	Overall composition, layout	Important but coarse
Very Low	Random structure from noise	Not perceptually meaningful

P2 weighting upweights the medium SNR regime where the most perceptually relevant structure is determined.

Empirical Finding: P2-weighted models achieve better FID and IS scores while maintaining competitive likelihood, suggesting the perceptual prior is well-aligned with sample quality metrics.

Velocity Prediction Weighting

When using v-prediction parameterization (common in latent diffusion models), the natural weighting emerges from the prediction target itself. Recall the v-prediction target:

$\mathbf{v}_t = \sqrt{\bar{\alpha}_t} \boldsymbol{\epsilon} - \sqrt{1-\bar{\alpha}_t} \mathbf{x}_0$

Implicit Weighting

The variance of the v-prediction target depends on the timestep:

$\text{Var}[\mathbf{v}_t] = \bar{\alpha}_t + (1-\bar{\alpha}_t) = 1$

Remarkably, the variance is constant across timesteps when both $\mathbf{x}_0$ and $\boldsymbol{\epsilon}$ have unit variance. This means v-prediction naturally balances the contribution of different timesteps without explicit weighting.

Connection to Other Parameterizations

The implicit weighting of v-prediction can be expressed in terms of epsilon-prediction with explicit weights:

$L^{\mathbf{v}} = \mathbb{E}\left[ \|\mathbf{v}_t - \hat{\mathbf{v}}_\theta\|^2 \right] = \mathbb{E}\left[ w_t^{\mathbf{v}} \cdot \|\boldsymbol{\epsilon} - \hat{\boldsymbol{\epsilon}}_\theta\|^2 \right]$

where the effective weight is:

$w_t^{\mathbf{v}} = \bar{\alpha}_t + (1-\bar{\alpha}_t) \cdot \frac{\text{SNR}(t)}{1 + \text{SNR}(t)} = \frac{\text{SNR}(t) + 1}{1 + \text{SNR}(t)} = 1$

Practical Impact

V-prediction is particularly useful for:

Training at high resolutions where gradient scale matters
Models with very long diffusion schedules
Latent diffusion where the SNR range can be extreme

Comparing Strategies

Different weighting strategies make different trade-offs. Here's a comprehensive comparison:

Strategy	Formula	Best For	Trade-off
Uniform	w = 1	Sample quality	Suboptimal likelihood
VLB	From KL divergence	Likelihood	Poor sample quality
SNR	SNR(t)	Signal preservation	Underweights noise region
1/SNR	1/SNR(t)	Noise region	Underweights clean region
min-SNR-gamma	min(SNR, gamma)	Balance	Requires tuning gamma
P2	1/(k + SNR)^gamma	Perception	Requires tuning k, gamma
v-prediction	Implicit = 1	Stability	Different parameterization

Empirical Results

Based on published benchmarks and practitioner experience:

For FID/sample quality: Uniform or min-SNR-gamma (with $\gamma=5$ ) typically perform best
For likelihood/BPD: VLB weighting gives the best bounds but at the cost of visual quality
For perceptual metrics: P2 weighting can provide improvements especially for high-resolution synthesis
For training stability: V-prediction with its implicit uniform weighting is often most stable

Implementation

Here's a comprehensive implementation of all major weighting strategies:

🐍python

1import torch
2import torch.nn as nn
3from typing import Literal
4from enum import Enum
5
6class WeightingStrategy(Enum):
7    UNIFORM = "uniform"
8    VLB = "vlb"
9    SNR = "snr"
10    INVERSE_SNR = "inverse_snr"
11    TRUNCATED_SNR = "truncated_snr"
12    MIN_SNR_GAMMA = "min_snr_gamma"
13    P2 = "p2"
14
15class LossWeighter(nn.Module):
16    """
17    Computes loss weights for different diffusion weighting strategies.
18
19    All strategies are expressed in terms of SNR for unified handling.
20    """
21
22    def __init__(
23        self,
24        strategy: WeightingStrategy,
25        alpha_bar: torch.Tensor,
26        # Strategy-specific parameters
27        gamma: float = 5.0,  # For min-SNR-gamma
28        p2_k: float = 1.0,   # For P2 weighting
29        p2_gamma: float = 1.0,  # For P2 weighting
30        max_weight: float = 100.0,  # Clamp for numerical stability
31    ):
32        super().__init__()
33        self.strategy = strategy
34        self.gamma = gamma
35        self.p2_k = p2_k
36        self.p2_gamma = p2_gamma
37        self.max_weight = max_weight
38
39        # Precompute SNR for all timesteps
40        snr = alpha_bar / (1 - alpha_bar)
41        self.register_buffer("snr", snr)
42        self.register_buffer("alpha_bar", alpha_bar)
43
44        # Precompute weights for efficiency
45        weights = self._compute_weights(snr)
46        self.register_buffer("weights", weights)
47
48    def _compute_weights(self, snr: torch.Tensor) -> torch.Tensor:
49        """Compute weights based on strategy."""
50        if self.strategy == WeightingStrategy.UNIFORM:
51            return torch.ones_like(snr)
52
53        elif self.strategy == WeightingStrategy.VLB:
54            # VLB weight proportional to SNR change rate
55            # Approximate: w_t ~ 1 / (SNR_t - SNR_{t-1})
56            snr_diff = torch.diff(snr, prepend=snr[:1])
57            snr_diff = snr_diff.clamp(min=1e-8)  # Avoid division by zero
58            weights = 1.0 / snr_diff.abs()
59            return weights.clamp(max=self.max_weight)
60
61        elif self.strategy == WeightingStrategy.SNR:
62            return snr.clamp(max=self.max_weight)
63
64        elif self.strategy == WeightingStrategy.INVERSE_SNR:
65            return (1.0 / snr.clamp(min=1e-8)).clamp(max=self.max_weight)
66
67        elif self.strategy == WeightingStrategy.TRUNCATED_SNR:
68            return torch.maximum(snr, torch.ones_like(snr))
69
70        elif self.strategy == WeightingStrategy.MIN_SNR_GAMMA:
71            return torch.minimum(snr, torch.full_like(snr, self.gamma))
72
73        elif self.strategy == WeightingStrategy.P2:
74            return (1.0 / (self.p2_k + snr) ** self.p2_gamma).clamp(max=self.max_weight)
75
76        else:
77            raise ValueError(f"Unknown strategy: {self.strategy}")
78
79    def forward(self, t: torch.Tensor) -> torch.Tensor:
80        """
81        Get weights for given timesteps.
82
83        Args:
84            t: Timestep indices (0-indexed)
85
86        Returns:
87            Weights for each timestep
88        """
89        return self.weights[t]
90
91    def get_effective_weight_ratio(self) -> dict[str, float]:
92        """
93        Compute statistics about weight distribution.
94
95        Returns dict with:
96        - max_min_ratio: Ratio of max to min weight
97        - early_late_ratio: Ratio of early (low noise) to late (high noise) weights
98        """
99        weights = self.weights
100        early_weights = weights[:len(weights)//3].mean()
101        late_weights = weights[-len(weights)//3:].mean()
102
103        return {
104            "max_min_ratio": (weights.max() / weights.min().clamp(min=1e-8)).item(),
105            "early_late_ratio": (early_weights / late_weights.clamp(min=1e-8)).item(),
106        }

Using the Weighter in Training

🐍python

1def weighted_diffusion_loss(
2    model: nn.Module,
3    weighter: LossWeighter,
4    x_0: torch.Tensor,
5    noise_schedule: dict,
6    prediction_type: Literal["epsilon", "x0", "v"] = "epsilon",
7) -> torch.Tensor:
8    """
9    Compute weighted diffusion loss.
10
11    Args:
12        model: Noise prediction network
13        weighter: LossWeighter instance
14        x_0: Clean data batch
15        noise_schedule: Dict with alpha_bar
16        prediction_type: What the model predicts
17
18    Returns:
19        Weighted loss value
20    """
21    batch_size = x_0.shape[0]
22    device = x_0.device
23    T = len(noise_schedule["alpha_bar"])
24
25    # Sample timesteps (0-indexed for weighter)
26    t = torch.randint(0, T, (batch_size,), device=device)
27
28    # Sample noise
29    epsilon = torch.randn_like(x_0)
30
31    # Get alpha_bar for this timestep
32    alpha_bar_t = noise_schedule["alpha_bar"][t].view(-1, 1, 1, 1)
33
34    # Create noisy input
35    x_t = torch.sqrt(alpha_bar_t) * x_0 + torch.sqrt(1 - alpha_bar_t) * epsilon
36
37    # Get model prediction
38    model_output = model(x_t, t + 1)  # Model expects 1-indexed timesteps
39
40    # Compute target based on prediction type
41    if prediction_type == "epsilon":
42        target = epsilon
43    elif prediction_type == "x0":
44        target = x_0
45    elif prediction_type == "v":
46        target = torch.sqrt(alpha_bar_t) * epsilon - torch.sqrt(1 - alpha_bar_t) * x_0
47    else:
48        raise ValueError(f"Unknown prediction type: {prediction_type}")
49
50    # Per-sample MSE
51    per_sample_mse = ((model_output - target) ** 2).mean(dim=(1, 2, 3))
52
53    # Apply weighting
54    weights = weighter(t)
55    weighted_loss = (weights * per_sample_mse).mean()
56
57    return weighted_loss
58
59
60# Example: Comparing different strategies
61def compare_weighting_strategies():
62    """Compare weight distributions across strategies."""
63    T = 1000
64    betas = torch.linspace(1e-4, 0.02, T)
65    alpha_bar = torch.cumprod(1 - betas, dim=0)
66
67    strategies = [
68        (WeightingStrategy.UNIFORM, {}),
69        (WeightingStrategy.SNR, {}),
70        (WeightingStrategy.MIN_SNR_GAMMA, {"gamma": 5.0}),
71        (WeightingStrategy.P2, {"p2_k": 1.0, "p2_gamma": 1.0}),
72    ]
73
74    print("Weight Distribution Analysis")
75    print("=" * 60)
76
77    for strategy, kwargs in strategies:
78        weighter = LossWeighter(strategy, alpha_bar, **kwargs)
79        stats = weighter.get_effective_weight_ratio()
80        print(f"\n{strategy.value}:")
81        print(f"  Max/Min Ratio: {stats['max_min_ratio']:.2f}")
82        print(f"  Early/Late Ratio: {stats['early_late_ratio']:.2f}")
83
84
85if __name__ == "__main__":
86    compare_weighting_strategies()

Practical Recommendations

🐍python

1def get_recommended_weighter(
2    use_case: Literal["quality", "likelihood", "perception", "stability"],
3    alpha_bar: torch.Tensor,
4) -> LossWeighter:
5    """
6    Get recommended weighter for common use cases.
7
8    Args:
9        use_case: What to optimize for
10        alpha_bar: Noise schedule cumulative alphas
11
12    Returns:
13        Configured LossWeighter
14    """
15    recommendations = {
16        "quality": {
17            "strategy": WeightingStrategy.MIN_SNR_GAMMA,
18            "gamma": 5.0,
19        },
20        "likelihood": {
21            "strategy": WeightingStrategy.VLB,
22        },
23        "perception": {
24            "strategy": WeightingStrategy.P2,
25            "p2_k": 1.0,
26            "p2_gamma": 1.0,
27        },
28        "stability": {
29            "strategy": WeightingStrategy.UNIFORM,
30        },
31    }
32
33    config = recommendations[use_case]
34    strategy = config.pop("strategy")
35    return LossWeighter(strategy, alpha_bar, **config)

Key Takeaways

SNR unifies weighting: All strategies can be understood through their relationship to the Signal-to-Noise Ratio
min-SNR-gamma is a strong default: With $\gamma=5$ , it balances sample quality and training stability
P2 targets perception: When perceptual quality is paramount, P2 weighting can provide improvements
V-prediction has implicit weighting: The parameterization itself provides balanced timestep contributions
No universal best: The optimal strategy depends on your evaluation metrics, dataset, and model architecture

Looking Ahead: In the next section, we'll explore the deep connection between diffusion model training and classical denoising autoencoders, providing another lens through which to understand why these losses work.