Advanced Loss Function Formulations (v-prediction, L_simple)

While the standard diffusion model training objective, predicting the noise $\epsilon$ added at timestep $t$, works remarkably well, research has explored alternative formulations for the model's prediction target and loss function. These alternatives can offer benefits in terms of sample quality, training stability, and how well the model behaves across different noise levels. Let's examine two significant approaches: the widely used simplified loss $L_{simple}$ (based on $\epsilon$-prediction) and the $v$-prediction objective.

The Standard: $\epsilon$-Prediction and $L_{simple}$

Most diffusion models are trained to predict the noise $\epsilon$ that was added to the original data $x_0$ to produce the noisy sample $x_t$. Recall the forward process:

$$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon \quad \text{where} \quad \epsilon \sim \mathcal{N}(0, \mathbf{I})$$

The model, typically a U-Net or Transformer denoted by $\epsilon_\theta$, takes the noisy input $x_t$ and the timestep $t$ (usually via an embedding) and outputs a prediction of the noise: $\epsilon_\theta(x_t, t)$.

The theoretical loss derived from the variational lower bound (VLB) includes weighting terms that depend on the timestep $t$. However, the DDPM paper found that a simplified, unweighted version of the loss often yields better results in practice:

$$L_{simple} = \mathbb{E}_{t \sim \mathcal{U}(1, T), x_0 \sim q(x_0), \epsilon \sim \mathcal{N}(0, \mathbf{I})} \left[ \| \epsilon - \epsilon_\theta(x_t, t) \|^2 \right]$$

Here, $t$ is sampled uniformly from the total number of timesteps $T$. This $L_{simple}$ objective is computationally straightforward and forms the backbone of many successful diffusion model implementations. It directly optimizes the model to denoise the input by predicting the added noise.

Despite its success, predicting $\epsilon$ can sometimes face challenges, particularly near the boundaries of the diffusion process (very small $t$ or very large $t$) where the signal-to-noise ratio varies drastically.

An Alternative: $v$-Prediction

An alternative formulation proposed to address some limitations of $\epsilon$-prediction is $v$-prediction. Instead of predicting the noise $\epsilon$, the model predicts a different target, $v$, defined as:

$$v = \alpha_t \epsilon - \sigma_t x_0$$

Here, we use the common notation where $\alpha_t = \sqrt{\bar{\alpha}_t}$ (signal scale) and $\sigma_t = \sqrt{1 - \bar{\alpha}_t}$ (noise scale), consistent with the forward process equation $x_t = \alpha_t x_0 + \sigma_t \epsilon$.

Why predict $v$? The intuition is related to signal scaling during the diffusion process. The target $\epsilon$ always has a unit variance. However, the noisy input $x_t$ has variance that changes with $t$. Predicting $v$ can be interpreted as predicting a target that has a more consistent variance across different timesteps, potentially making the learning task easier or more stable for the network. It effectively combines predicting the score function ($\nabla_{x_t} \log p(x_t)$) and the data $x_0$ in a balanced way.

The model architecture remains largely the same (e.g., U-Net), but it's now parameterized to output $v_\theta(x_t, t)$. The corresponding loss function is analogous to $L_{simple}$, but with $v$ as the target:

$$L_v = \mathbb{E}_{t, x_0, \epsilon} \left[ \| v - v_\theta(x_t, t) \|^2 \right]$$

Often, weighting terms similar to those in the full VLB might be reintroduced, or the loss might be formulated based on score matching principles, but the core idea is to predict this $v$ target.

Benefits of $v$-Prediction

1. Improved Sample Quality: Particularly in models like Google's Imagen, $v$-prediction was found to yield higher-quality samples compared to $\epsilon$-prediction, especially when generating high-resolution images or when complex conditioning is involved.
2. Better Scaling Properties: The $v$ target has variance $\alpha_t^2 + \sigma_t^2 = 1$, making the prediction target scale more uniformly across time compared to $\epsilon$, which can simplify network optimization.
3. Handling Boundary Conditions: $v$-prediction can simplify handling the edge cases of $t=0$ (pure data) and $t=T$ (pure noise), potentially leading to more stable training.

Implementation Considerations

When using $v$-prediction, you need to adjust how the model's output is used during sampling. If the model $v_\theta(x_t, t)$ predicts $v$, you can recover the predictions for $\epsilon$ and $x_0$ as needed for the sampling steps (like DDIM or DDPM):

$$\text{Predicted } \epsilon = \alpha_t v_\theta(x_t, t) + \sigma_t x_t$$ $$\text{Predicted } x_0 = \alpha_t x_t - \sigma_t v_\theta(x_t, t)$$

These recovered values can then be plugged into the standard sampling equations.

Comparing $\epsilon$-Prediction and $v$-Prediction

The choice between $\epsilon$-prediction ($L_{simple}$) and $v$-prediction ($L_v$) isn't always clear-cut and can depend on the specific application, dataset, and model architecture.

• $L_{simple}$ ($\epsilon$-prediction):

  • Simpler to implement and understand.
  • Proven effective in a vast number of models (including the original DDPM, Stable Diffusion 1.x/2.x).
  • The prediction target $\epsilon$ has constant unit variance.
  • May require careful noise schedule design and network tuning, especially for extreme SNRs.

• $L_v$ ($v$-prediction):

  • Can lead to improved sample quality, particularly in large-scale models.
  • The prediction target $v$ has properties that might stabilize training across timesteps.
  • Used in models like Imagen.
  • Requires modifying the interpretation of the model output during sampling.

The diagram below illustrates the different targets the network aims to predict in each formulation:

Diagram comparing the prediction targets and loss calculation for $\epsilon$-prediction and $v$-prediction formulations. Both take the noisy data $x_t$ and timestep $t$ as input.

Ultimately, both $L_{simple}$ and $L_v$ are effective objectives for training diffusion models. While $L_{simple}$ remains a baseline, $v$-prediction offers a valuable alternative, particularly when pushing for sample quality or dealing with training dynamics. Experimentation is often necessary to determine the best choice for a given project. Understanding these different formulations provides you with more tools to optimize and control the behavior of your diffusion models during training.