Okay, we've established that our goal is to train a neural network, typically a U-Net architecture denoted by $\epsilon_\theta$ , to reverse the diffusion process. How do we actually train it? The network needs a clear target, a specific quantity to predict.

Recall the forward process: we start with clean data $x_0$ and progressively add Gaussian noise $\epsilon$ at each timestep $t$ to get $x_t$ . The core idea for training is surprisingly direct: we train the network $\epsilon_\theta$ to predict the exact noise $\epsilon$ that was added to $x_0$ to create $x_t$ , given the noisy input $x_t$ and the timestep $t$ .

The network takes $x_t$ and $t$ as input and outputs a prediction, let's call it $\epsilon_\theta(x_t, t)$ . Our objective is to make this prediction match the actual noise $\epsilon$ that was sampled from a standard Gaussian distribution $\mathcal{N}(0, I)$ during the forward step calculation for $x_t$ .

The Loss Function: Mean Squared Error

To quantify the difference between the predicted noise $\epsilon_\theta(x_t, t)$ and the true noise $\epsilon$ , a common and effective choice is the Mean Squared Error (MSE). Minimizing this error forces the network's output to closely match the target noise.

The training objective, or loss function $L$ , is formulated as the expectation over all possible inputs: the initial data $x_0$ , the randomly chosen timestep $t$ (uniformly sampled between 1 and the maximum timestep $T$ ), and the randomly sampled noise $\epsilon$ . Mathematically, we express this as:

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

Let's break this down:

$\mathbb{E}[...]$ : Denotes the expectation, meaning we average the loss over many examples.
$t \sim \mathcal{U}(1, T)$ : We sample a timestep $t$ uniformly at random from all possible timesteps (1 to $T$ ). This ensures the network learns to denoise at all noise levels.
$x_0 \sim q(x_0)$ : We sample a real data point $x_0$ from our training dataset distribution $q(x_0)$ .
$\epsilon \sim \mathcal{N}(0, I)$ : We sample a noise vector $\epsilon$ from a standard Gaussian distribution.
$x_t$ : This is the noisy version of $x_0$ at timestep $t$ , created using $x_0$ and the sampled noise $\epsilon$ . Recall from the forward process chapter that $x_t$ can be sampled directly using $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$ , where $\bar{\alpha}_t$ comes from the noise schedule.
$\epsilon_\theta(x_t, t)$ : This is the noise predicted by our neural network (parameterized by $\theta$ ) when given the noisy data $x_t$ and the timestep $t$ .
$||\epsilon - \epsilon_\theta(x_t, t)||^2$ : This is the squared L2 norm (Euclidean distance), which calculates the squared difference between the actual noise and the predicted noise across all dimensions (e.g., all pixels in an image).

In practice, during training, we approximate this expectation by averaging the loss over mini-batches of training data. For each item in the batch, we sample a random $t$ , sample a random $\epsilon$ , compute $x_t$ , feed $x_t$ and $t$ to the network $\epsilon_\theta$ , and calculate the MSE between the network's output $\epsilon_\theta(x_t, t)$ and the original $\epsilon$ .

Diagram illustrating the flow for calculating the training loss for a single data point. The network $\epsilon_\theta$ tries to predict the original noise $\epsilon$ based on the noisy data $x_t$ and timestep $t$ . The MSE loss measures the difference between the actual and predicted noise.

Connection to the Probabilistic Objective

While this MSE loss on the noise seems intuitive and works well empirically, it's not just a heuristic. It arises naturally from the more rigorous mathematical objective of diffusion models: maximizing the evidence lower bound (ELBO) on the data log-likelihood.

The full derivation involves variational inference and Bayes' theorem, which we will simplify in the next section. For now, it's sufficient to understand that minimizing $L = ||\epsilon - \epsilon_\theta(x_t, t)||^2$ is directly related to optimizing the ELBO. This connection provides a solid theoretical foundation for why this seemingly simple objective enables the model to learn the complex data distribution required for generation.

By training the network $\epsilon_\theta$ with this objective, we equip it with the ability to estimate the noise present in $x_t$ . This noise estimate is precisely what we need during the reverse process (sampling) to iteratively denoise a random signal back into a sample that resembles our training data.

Was this section helpful?