Classifier Guidance

As mentioned in the chapter introduction, unconditional generation produces samples representative of the overall training data, but we often desire more specific outputs. Imagine wanting to generate images only of cats, or perhaps digits corresponding to the number '8'. How can we steer the diffusion model, which has learned the distribution $p(x)$, to generate samples from a conditional distribution $p(x|y)$, where $y$ represents the desired condition (e.g., $y$= 'cat' or $y$= '8')?

One of the first successful approaches to achieve this is Classifier Guidance. The core idea is to leverage a separate, pre-trained classifier model, let's call it $p_\phi(y|x)$, where $\phi$ represents the classifier's parameters. This classifier is trained to predict the class label $y$ given an input $x$.

However, during the diffusion model's reverse sampling process, we are not dealing with clean data $x_0$ but with noisy intermediate samples $x_t$ at various timesteps $t$. Therefore, for classifier guidance to be effective, the classifier $p_\phi(y|x_t)$ must be trained to recognize the class $y$ even from noisy inputs $x_t$. This means training the classifier not just on the original dataset (like clean images) but also on noisy versions of the data, consistent with the noise levels encountered during the diffusion process.

How does this classifier guide the generation? Recall that the reverse process aims to approximate $p(x_{t-1}|x_t)$. We want to modify this step to sample from $p(x_{t-1}|x_t, y)$. Using Bayes' theorem, we can write:

$$p(x_t | y) = \frac{p(y | x_t) p(x_t)}{p(y)}$$

Taking the logarithm and then the gradient with respect to $x_t$:

$$\nabla_{x_t} \log p(x_t | y) = \nabla_{x_t} \log p(y | x_t) + \nabla_{x_t} \log p(x_t)$$

The term $\nabla_{x_t} \log p(x_t)$ is the score function of the marginal data distribution at time $t$. The diffusion model's noise prediction network $\epsilon_\theta(x_t, t)$ is trained to approximate $-\sigma_t \nabla_{x_t} \log p(x_t)$ (scaled by noise level). The term $\nabla_{x_t} \log p(y | x_t)$ is the gradient of the log-likelihood of the desired class $y$ according to the classifier, evaluated at the current noisy sample $x_t$. This gradient indicates the direction in the input space ($x_t$) that makes the sample look more like class $y$ to the classifier.

Classifier guidance modifies the sampling step by incorporating this classifier gradient. Specifically, when calculating the mean $\mu_\theta(x_t, t)$ for the reverse step $p_\theta(x_{t-1}|x_t)$, we perturb it using the classifier's gradient. The update rule for the mean in the DDPM sampling step can be adjusted as follows:

The original predicted mean is derived from the predicted noise $\epsilon_\theta$:

$$\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right)$$

The guided mean $\hat{\mu}_\theta(x_t, t, y)$ becomes:

$$\hat{\mu}_\theta(x_t, t, y) = \mu_\theta(x_t, t) + s \cdot \Sigma_t \nabla_{x_t} \log p_\phi(y|x_t)$$

Here:

• $\Sigma_t$ is the covariance matrix (variance) of the reverse step $p_\theta(x_{t-1}|x_t)$, often related to $\beta_t$.
• $s$ is the guidance scale, a hyperparameter that controls the strength of the classifier's influence. A higher $s$ pushes the generation more strongly towards the target class $y$.
• $\nabla_{x_t} \log p_\phi(y|x_t)$ is the gradient computed by the classifier $p_\phi$.

The sampling process then proceeds using this adjusted mean $\hat{\mu}_\theta$:

1. Predict the noise $\epsilon_\theta(x_t, t)$ using the U-Net.
2. Calculate the original mean $\mu_\theta(x_t, t)$.
3. Compute the classifier gradient $\nabla_{x_t} \log p_\phi(y|x_t)$ for the desired class $y$.
4. Calculate the guided mean $\hat{\mu}_\theta(x_t, t, y)$ using the guidance scale $s$.
5. Sample $x_{t-1}$ from $\mathcal{N}(x_{t-1}; \hat{\mu}_\theta(x_t, t, y), \Sigma_t)$.

This process is repeated from $t=T$ down to $t=1$.

Diagram illustrating the classifier guidance mechanism within a single reverse step. The U-Net predicts noise, while a separate classifier provides a gradient based on the target class $y$. These are combined, scaled by $s$, to produce a guided mean used for sampling $x_{t-1}$.

Advantages of Classifier Guidance:

• Explicit Control: Provides a direct way to steer generation towards specific attributes that the classifier is trained on.
• Leverages Existing Classifiers: Can potentially use well-trained classifiers off-the-shelf (though they need noise robustness).

Disadvantages of Classifier Guidance:

• Requires a Separate Classifier: Need to train and maintain an additional model $p_\phi(y|x_t)$.
• Classifier Training: The classifier must be robust to the noise levels seen during diffusion, adding training complexity.
• Computational Cost: Requires running inference on both the diffusion model and the classifier at each sampling step.
• Guidance Strength Tuning: Finding the right guidance scale $s$ often requires experimentation. Too low, and the guidance is ineffective; too high, and samples might become unrealistic or suffer from adversarial effects exploited by the classifier gradient.

While effective, the need for a separate, noise-aware classifier led researchers to develop methods that achieve similar guidance without this external dependency. This motivates the next technique we will discuss: Classifier-Free Guidance (CFG).