Simplified Training Loss Derivation

The core task of a neural network, often a U-Net, in diffusion models is to learn the reverse diffusion process. This involves estimating the original data distribution by progressively removing noise. While the full probabilistic derivation of diffusion models involves maximizing the data likelihood via the Evidence Lower Bound (ELBO), a common and effective training objective simplifies this significantly.

From ELBO to Simple MSE

The full ELBO for diffusion models contains several terms related to the transitions in the reverse Markov chain. Optimizing it directly can be intricate. However, researchers (notably Ho et al. in the original DDPM paper) found that the training objective can be greatly simplified without sacrificing performance.

Instead of modeling the complex reverse transition probabilities $p_\theta(x_{t-1}|x_t)$ directly, we reparameterize the network $\epsilon_\theta$ to predict the noise $\epsilon$ that was added to obtain $x_t$ from the original image $x_0$ at timestep $t$. Recall the forward process equation that allows sampling $x_t$ directly from $x_0$:

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

Here, $\epsilon$ is a standard Gaussian noise sample $(\epsilon \sim \mathcal{N}(0, I))$, and $\bar{\alpha}_t$ is derived from the noise schedule.

The insight is that if our network can accurately predict this noise $\epsilon$ given the noisy image $x_t$ and the timestep $t$, it effectively learns how to reverse the diffusion step.

The Simplified Loss Function

This leads to a much simpler training objective: we minimize the Mean Squared Error (MSE) between the actual noise $\epsilon$ used to create $x_t$ and the noise predicted by our network $\epsilon_\theta(x_t, t)$.

The loss function $L_{simple}$ is formulated as an expectation over random timesteps $t$, initial data samples $x_0$, and noise samples $\epsilon$:

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

Let's break this down:

1. Sample $t$: Choose a random timestep $t$ uniformly from $\{1, ..., T\}$.
2. Sample $x_0$: Get a data sample (e.g., an image) from your training dataset.
3. Sample $\epsilon$: Draw a noise sample from a standard normal distribution.
4. Compute $x_t$: Calculate the noisy version $x_t$ using the formula $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$.
5. Predict Noise: Feed $x_t$ and the timestep $t$ into the network $\epsilon_\theta$ to get the predicted noise $\epsilon_\theta(x_t, t)$.
6. Calculate MSE: Compute the squared difference between the actual noise $\epsilon$ and the predicted noise $\epsilon_\theta(x_t, t)$.
7. Average: In practice, this expectation is approximated by averaging the MSE over mini-batches during training.

Why This Works

Minimizing this MSE loss encourages the network $\epsilon_\theta$ to become an effective noise predictor for any given noise level $t$. Intuitively, if the model can precisely identify the noise component within $x_t$, it implicitly understands the structure of the underlying data $x_0$ required to reverse the noising process.

While we've skipped the detailed mathematical derivation from the ELBO, it can be shown that this simplified MSE loss corresponds to a specific weighting of the terms in the variational lower bound. This provides theoretical grounding for why this simpler objective is effective for training high-quality diffusion models.

The practical advantage is significant: training boils down to a standard regression problem where the network learns to map a noisy input $x_t$ and timestep $t$ to the noise $\epsilon$ that was added. This is far more straightforward to implement and optimize than dealing directly with the complex distributions of the full ELBO. This simplified loss is the foundation upon which we build the training algorithm described in the next section.