Classifier Guidance: Principles and Implementation

While sophisticated architectures lay the groundwork, the training process itself offers powerful levers to enhance the capabilities of diffusion models. One of the earlier, yet significant, techniques for steering the generation process towards specific desired attributes, such as a particular image class, is classifier guidance. This method modifies the sampling trajectory by incorporating information from an external classifier model.

Imagine you have a diffusion model trained unconditionally on a dataset like ImageNet. During sampling, it generates images representative of the overall dataset. But what if you specifically want an image of a "golden retriever"? Classifier guidance provides a mechanism to inject this conditional information during inference, pushing the sampling process towards images that a classifier recognizes as belonging to the target class.

The Core Principle: Steering with Gradients

Classifier guidance operates entirely during the sampling phase. It doesn't alter the training of the core diffusion model $\epsilon_\theta(x_t, t)$, which is typically trained to predict the noise added at timestep $t$ unconditionally or with simple conditioning. Instead, it uses a separate, pre-trained classifier $p_\phi(y | x_t)$ that has been specifically trained to predict the class $y$ of an image $x$ even when it's noisy (at level $t$).

At each step of the reverse diffusion process (sampling), we start with the noisy image $x_t$. We want to compute the next, slightly less noisy image $x_{t-1}$. The core idea is to adjust the direction of this step based not only on the diffusion model's prediction but also on how we can make $x_t$ more likely to be classified as the target class $y$ according to our classifier $p_\phi$.

This "making more likely" translates mathematically to using the gradient of the classifier's log-probability with respect to the input image $x_t$. That is, we compute $\nabla_{x_t} \log p_\phi(y | x_t)$. This gradient vector points in the direction in input space ($x_t$) that most increases the classifier's confidence in the target class $y$.

Mathematical Formulation

Recall that in a standard DDPM or DDIM sampling step, the diffusion model $\epsilon_\theta(x_t, t)$ predicts the noise $\epsilon$ that was likely added to get $x_t$. Classifier guidance modifies this prediction.

The underlying principle connects to Bayes' theorem and score matching. The score function $\nabla_{x_t} \log p(x_t | y)$ represents the direction to step in to increase the likelihood of $x_t$ given the condition $y$. This can be approximated as:

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

Here, $\nabla_{x_t} \log p(x_t)$ is the score of the unconditional distribution (related to the diffusion model's output), and $\nabla_{x_t} \log p(y | x_t)$ is the score provided by the external classifier $p_\phi$.

In terms of the noise prediction $\epsilon$, the guidance adjusts the diffusion model's output $\epsilon_\theta(x_t, t)$. The modified noise prediction, $\hat{\epsilon}_\theta(x_t, t, y)$, used in the sampling step is calculated as:

$$\hat{\epsilon}_\theta(x_t, t, y) = \epsilon_\theta(x_t, t) - s \cdot \sigma_t \cdot \nabla_{x_t} \log p_\phi(y | x_t)$$

Let's break down the components:

• $\epsilon_\theta(x_t, t)$: The original noise prediction from the unconditionally trained (or base conditional) diffusion model for the current noisy image $x_t$ and timestep $t$.
• $p_\phi(y | x_t)$: The pre-trained classifier's probability estimate that the noisy image $x_t$ belongs to the target class $y$.
• $\nabla_{x_t} \log p_\phi(y | x_t)$: The gradient of the log-probability of the target class $y$ with respect to the input noisy image $x_t$. This is the "guidance signal" from the classifier. It requires computing gradients through the classifier network.
• $s$: The guidance scale (or strength). This is a hyperparameter ($s \ge 0$) that controls how strongly the classifier's gradient influences the noise prediction. A value of $s=0$ recovers the original unconditional sampling. Larger values of $s$ push the generation more strongly towards the target class $y$.
• $\sigma_t$: A scaling factor related to the noise level at timestep $t$. Often, this is related to the standard deviation of the noise, for example, $\sigma_t = \sqrt{1 - \bar{\alpha}_t}$ in DDPM notation, though variations exist. This term helps to balance the magnitude of the gradient relative to the noise prediction at different timesteps.

The resulting $\hat{\epsilon}_\theta$ is then used in the standard DDPM or DDIM update formula to compute $x_{t-1}$.

Implementation Steps

Implementing classifier guidance involves these main components:

1. Pre-trained Diffusion Model: You need a standard diffusion model $\epsilon_\theta(x_t, t)$, trained either unconditionally or potentially with some base conditioning unrelated to the guidance target.
2. Pre-trained Noisy Classifier: This is the critical part. You need a separate classifier network $p_\phi(y | x)$ that is specifically trained to classify images corrupted with noise corresponding to various diffusion timesteps $t$. Training this classifier requires augmenting the training data with noise levels matching the diffusion process schedule. It often shares a similar architecture to the diffusion model's backbone but outputs class probabilities.
3. Sampling Loop Modification:
   • Inside the sampling loop, for each timestep $t$ from $T$ down to 1:
     • Obtain the current noisy sample $x_t$. Ensure $x_t$ requires gradients.
     • Get the unconditional noise prediction $\epsilon_{uncond} = \epsilon_\theta(x_t, t)$.
     • Pass $x_t$ through the noisy classifier $p_\phi$ to get the log-probability $\log p_\phi(y | x_t)$ for the desired target class $y$.
     • Compute the gradient of this log-probability with respect to the input: $g = \nabla_{x_t} \log p_\phi(y | x_t)$. This typically involves a torch.autograd.grad call or equivalent. Remember to detach $x_t$ before feeding it into the diffusion model if you don't want gradients flowing through it.
     • Calculate the guided noise: $\hat{\epsilon}_\theta = \epsilon_{uncond} - s \cdot \sigma_t \cdot g$.
     • Use $\hat{\epsilon}_\theta$ in your chosen sampler (DDPM, DDIM) to compute the denoised estimate for the next step, $x_{t-1}$.
   • Repeat until $x_0$ is obtained.

The following diagram illustrates the data flow during a single guided sampling step:

Data flow for a single step of classifier guidance. The noisy image $x_t$ is processed by both the diffusion model and the noisy classifier. The classifier's output for the target class $y$ is used to compute a gradient, which is then scaled and subtracted from the diffusion model's noise prediction to yield the guided noise $\hat{\epsilon}_\theta$.

Advantages of Classifier Guidance

• Explicit Control: Provides a direct way to control generation towards specific classes or attributes that a classifier can recognize.
• Potential Quality Improvement: Can sometimes enhance sample quality and realism for the target class compared to unconditional generation, especially if the diffusion model struggles with specific modes.

Disadvantages and Challenges

• Requires Separate Classifier: The main drawback is the need to train and maintain an additional classifier model.
• Noisy Classifier Training: This classifier must handle the noise levels encountered during diffusion sampling, making its training challenging and computationally expensive. It needs access to the same noise schedule and data used for the diffusion model.
• Tuning Guidance Scale: Finding the optimal guidance scale $s$ is important. Too low, and the guidance has little effect. Too high, and the process can over-optimize for the classifier, potentially treating the classifier as an adversary and generating samples that look good to the classifier but are unnatural or contain artifacts. This often manifests as overly saturated or strangely textured images.
• Limited Flexibility: While effective for classes, it's less direct to apply to detailed text descriptions compared to methods embedding conditioning directly into the diffusion model architecture (e.g., via cross-attention).

Classifier guidance was an important step in controllable diffusion generation. However, the practical challenges associated with training the noisy classifier led researchers to seek alternative approaches. This paved the way for Classifier-Free Guidance (CFG), a technique discussed in the next section, which cleverly achieves similar guidance effects without needing an external classifier model at all.