Model Parameterization (epsilon-prediction vs. x0-prediction)

A fundamental choice when training a diffusion model is deciding what the underlying neural network should predict. Given the noisy input $x_t$ and the timestep $t$, the model needs to estimate some quantity related to the denoising process. The two most common parameterizations are predicting the noise ($\epsilon$) that was added, or predicting the original clean data ($x_0$). This choice impacts the loss function, training dynamics, and potentially the final sample quality.

Epsilon-Prediction ($\epsilon$-prediction)

This is the standard approach introduced in the original Denoising Diffusion Probabilistic Models (DDPM) paper. The model, denoted as $f_\theta(x_t, t)$, is trained to predict the noise $\epsilon$ that was sampled from a standard Gaussian distribution and added to the original data $x_0$ to create $x_t$ according to the forward process equation:

$$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$$

Here, $\bar{\alpha}_t$ is the cumulative product of the noise schedule variances up to time $t$. The objective function is typically a simplified mean squared error (MSE) loss between the predicted noise $\hat{\epsilon} = f_\theta(x_t, t)$ and the actual noise $\epsilon$ used to generate $x_t$:

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

The expectation is taken over random timesteps $t$, initial data samples $x_0$, and the sampled noise $\epsilon$.

Advantages:

1. Stability: $\epsilon$-prediction is often empirically found to be stable across various datasets and architectures. Predicting a variable ($\epsilon$) with zero mean and unit variance can be a well-behaved regression target for neural networks.
2. Direct Score Connection: Under certain conditions, minimizing $L_{\epsilon}$ is equivalent to score matching, linking it directly to the mathematical foundations of score-based generative models. The score function is the gradient of the log probability density with respect to the data, $\nabla_{x_t} \log p_t(x_t)$, which is closely related to the noise $\epsilon$.
3. Established Practice: It's the most common parameterization used in widely successful models like DDPM, Stable Diffusion (prior to v2.0's v-prediction), and many others, making it a reliable starting point.

During sampling (the reverse process), the predicted noise $\hat{\epsilon}$ is used to estimate the direction towards a less noisy state, often by first estimating the predicted $x_0$ (denoted $\hat{x}_0$) and then using it in the DDPM or DDIM update step.

x0-Prediction ($x_0$-prediction)

An alternative approach is to parameterize the model, let's call it $g_\theta(x_t, t)$, to directly predict the original clean data $x_0$ from the noisy input $x_t$ and timestep $t$. The corresponding MSE loss function aims to minimize the difference between the predicted $\hat{x}_0 = g_\theta(x_t, t)$ and the true $x_0$:

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

Advantages:

1. Intuitive Output: The model's output directly represents the estimated clean data, which can be easier to interpret and potentially debug.
2. Potential for Sharpness: Some studies suggest that directly predicting $x_0$ might lead to sharper image generation in certain contexts, though this is not universally guaranteed.

Connection between Parameterizations:

The two parameterizations are mathematically related. Given the forward process equation $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$, we can express one prediction in terms of the other:

• If the model predicts noise $\hat{\epsilon} = f_\theta(x_t, t)$, the corresponding prediction for $x_0$ is:

  $$\hat{x}_0 = \frac{1}{\sqrt{\bar{\alpha}_t}} (x_t - \sqrt{1 - \bar{\alpha}_t} \hat{\epsilon})$$

• If the model predicts the clean data $\hat{x}_0 = g_\theta(x_t, t)$, the corresponding prediction for $\epsilon$ is:

  $$\hat{\epsilon} = \frac{1}{\sqrt{1 - \bar{\alpha}_t}} (x_t - \sqrt{\bar{\alpha}_t} \hat{x}_0)$$

These relationships show that choosing one parameterization implicitly defines the other. However, training the network to predict one quantity versus the other directly affects the gradients and performance, leading to potentially different training dynamics and final model performance.

Comparison of epsilon-prediction and x0-prediction approaches. The core model architecture is often the same, but the target variable for the loss calculation differs.

Comparison and Approaches

• Numerical Stability: Predicting $x_0$ can sometimes be less stable, especially for small $t$ (where $x_t$ is very close to $x_0$) or large $t$ (where $x_t$ is almost pure noise). The scaling factors $\sqrt{\bar{\alpha}_t}$ and $\sqrt{1 - \bar{\alpha}_t}$ in the conversion formulas can become very small or large at the extremes of the timestep range, potentially causing numerical issues if not handled carefully (e.g., via loss weighting or gradient clipping). $\epsilon$-prediction, targeting a variable with fixed variance, often avoids these scaling issues.
• Loss Weighting: When using $L_{x_0}$, it might be beneficial to weight the loss term differently across timesteps compared to the standard $L_{\epsilon}$ loss. The relative importance of predicting $x_0$ accurately might vary depending on the noise level.
• Implementation: Switching between the two involves changing the target variable in your loss computation and ensuring the sampling process uses the corresponding prediction correctly. For example, if your sampler expects an $\epsilon$ prediction, but your model outputs $\hat{x}_0$, you must first convert $\hat{x}_0$ to $\hat{\epsilon}$ using the formula above before proceeding with the sampling step.

import torch
import torch.nn.functional as F

def get_sqrt_alphas_cumprod(alphas):
  """Helper to get cumulative products"""
  return torch.sqrt(torch.cumprod(alphas, dim=0))

def get_sqrt_one_minus_alphas_cumprod(alphas):
  """Helper to get sqrt(1 - alpha_bar)"""
  return torch.sqrt(1.0 - torch.cumprod(alphas, dim=0))

# --- Example Setup ---
T = 1000
betas = torch.linspace(0.0001, 0.02, T) # Example linear schedule
alphas = 1.0 - betas
alphas_cumprod = torch.cumprod(alphas, dim=0)
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)

# --- Assume we have these variables per batch item ---
# x_start: Original clean image [B, C, H, W]
# noise: Sampled Gaussian noise (epsilon) [B, C, H, W]
# t: Sampled timesteps [B]
# model: Your U-Net or Transformer model

# Extract schedule values for batch timesteps t
sqrt_alpha_bar_t = sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
sqrt_one_minus_alpha_bar_t = sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)

# Calculate noisy image x_t
x_t = sqrt_alpha_bar_t * x_start + sqrt_one_minus_alpha_bar_t * noise

# Get model prediction
# model_output shape: [B, C, H, W]
model_output = model(x_t, t)

# --- Loss Calculation ---

# 1. Epsilon-Prediction Loss
target_eps = noise
loss_eps = F.mse_loss(model_output, target_eps)
# Use loss_eps for backpropagation

# 2. x0-Prediction Loss
target_x0 = x_start
loss_x0 = F.mse_loss(model_output, target_x0)
# Use loss_x0 for backpropagation

# --- Sampling (Example DDIM step) ---
# If model predicts epsilon (model_output = predicted_eps):
predicted_x0_from_eps = (x_t - sqrt_one_minus_alpha_bar_t * model_output) / sqrt_alpha_bar_t
# Use predicted_x0_from_eps in the DDIM update formula

# If model predicts x0 (model_output = predicted_x0):
predicted_eps_from_x0 = (x_t - sqrt_alpha_bar_t * model_output) / sqrt_one_minus_alpha_bar_t
# Use predicted_eps_from_x0 in the DDIM update formula (or modify DDIM to use predicted_x0 directly)

Python code snippet illustrating the difference in target calculation for $\epsilon$-prediction and $x_0$-prediction losses, and how to derive one prediction from the other for sampling.

Choosing the Right Parameterization

While $\epsilon$-prediction is the well-established and often more stable default, $x_0$-prediction provides a valid alternative. Experimenting with $x_0$-prediction might be worthwhile if:

• You are encountering stability issues specifically related to the scale of the predicted noise.
• You hypothesize that directly regressing towards the clean image might simplify the learning task for your specific data or architecture.
• You are implementing architectures or techniques specifically designed around $x_0$-prediction.

It's also important to be aware of other parameterizations like $v$-prediction, discussed in the next section, which attempts to combine the benefits of both $\epsilon$ and $x_0$ prediction, particularly improving how the model scales its output across different noise levels. Understanding the implications of predicting $\epsilon$ versus $x_0$ provides a foundation for appreciating these more advanced techniques.