The Training Algorithm

Now that we understand the U-Net architecture, how to integrate timestep information, and the simplified loss function (predicting the noise), let's outline the algorithm used to train this noise prediction network $\epsilon_\theta(x_t, t)$.

The goal is to train the network $\epsilon_\theta$ such that, given a noisy input $x_t$ and its corresponding timestep $t$, it accurately predicts the noise $\epsilon$ that was originally added to the clean input $x_0$ to produce $x_t$. We achieve this using stochastic gradient descent (or variants like Adam) by repeatedly sampling data points and timesteps, computing the loss, and updating the network's parameters $\theta$.

Here is a breakdown of the training loop, typically performed over batches of data:

1. Sample a clean data point: Obtain an example $x_0$ from your training dataset $q(x_0)$. In practice, you'll sample a mini-batch of data points. For simplicity, let's consider a single data point first.
2. Sample a timestep: Choose a timestep $t$ uniformly at random from the range $\{1, 2, ..., T\}$, where $T$ is the total number of diffusion steps defined in the forward process. This ensures the network learns to denoise across all noise levels.
3. Sample noise: Draw a noise sample $\epsilon$ from a standard Gaussian distribution: $\epsilon \sim \mathcal{N}(0, \mathbf{I})$. This is the "ground truth" noise that the network will try to predict.
4. Compute the noisy sample $x_t$: Using the sampled $x_0$, $t$, and $\epsilon$, calculate the corresponding noisy version $x_t$ via the closed-form equation derived from the forward process: $$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$$ Remember, $\bar{\alpha}_t$ is the cumulative product of $(1 - \beta_i)$ up to timestep $t$, based on the predefined noise schedule $\beta_t$.
5. Predict noise using the network: Pass the noisy sample $x_t$ and the timestep $t$ (usually encoded, e.g., using sinusoidal embeddings) through the U-Net model $\epsilon_\theta$ to get the predicted noise: $\epsilon_\theta(x_t, t)$.
6. Calculate the loss: Compute the difference between the actual noise $\epsilon$ (sampled in step 3) and the predicted noise $\epsilon_\theta(x_t, t)$ using the chosen loss function. As discussed previously, this is typically the Mean Squared Error (MSE): $$L = ||\epsilon - \epsilon_\theta(x_t, t)||^2$$ When using mini-batches, this loss is averaged over all samples in the batch.
7. Compute gradients: Calculate the gradients of this loss $L$ with respect to the network parameters $\theta$: $\nabla_\theta L$.
8. Update network parameters: Update the parameters $\theta$ using an optimization algorithm (like Adam) and the computed gradients. For example: $\theta \leftarrow \theta - \eta \nabla_\theta L$, where $\eta$ is the learning rate.

This entire sequence (steps 1-8) constitutes one training step. This process is repeated for many iterations or epochs over the dataset until the model's loss converges and it effectively learns to predict the noise added at various timesteps.

A diagram illustrating the core steps within a single training iteration for the noise prediction network.

By iterating through these steps, the U-Net gradually learns the complex mapping from a noisy image and a timestep to the noise component itself, forming the foundation for the reverse diffusion (denoising) process used during generation.