Skip to main content

Classifier Guidance

ELI5 (Explain Like I'm 5)

  • The Big Idea: To make sure our image generator creates high-quality pictures of a specific class (like a "dog"), we can hire a helper: a separately trained image classifier. At each step of cleaning the noise, the classifier looks at the half-cleaned image and calculates gradients (arrows showing how to change the pixels) to make the image look more like a dog. We use these gradients to steer the diffusion model's steps.
  • Analogy: Imagine drawing a horse with a blindfold on, while a friend who can see stands next to you saying, "Make the ears pointier," "Make the legs longer," or "Nudge your pen upward." The friend is the classifier, and their feedback guides your hand to draw a better horse.
  • Example: If the diffusion model is trying to generate a cat but the face looks a bit blurry, the classifier calculates that adding more contrast around the eyes will increase the "cat-ness" score. It nudges the pixels in that direction, resulting in a much sharper, more recognizable cat face.

Key Insight

Classifier guidance was the first trick that made conditional DDPM samples both sharp and on-target: you train a separate image classifier on noisy images, then during sampling nudge each denoising step in the direction its gradient says will make the chosen class more likely. Intuitively the classifier whispers "a little more cat-ness, this way" at every step, and that push — a score-like signal, the gradient of a log-probability with respect to the image — sharpens the output toward the requested class. The catch is that it needs an extra, specially-trained noisy classifier, exactly the cost that classifier-free guidance (CFG) later removed by folding the same effect into the diffusion model itself. This project builds the original: train the noisy classifier on CIFAR-10 and use its gradients to steer the samples.

What's in this directory

FileRole
noisy_classifier.pyTrains a small time-conditioned CNN on noised digits and measures its accuracy at every noise level
guided_sampling.pyThe guided reverse loop: the DDPM on MNIST project's unconditional DDPM steered by the classifier's gradients, with a guidance-scale sweep

The checked-in demo runs on MNIST so it completes on a CPU alongside project 24's checkpoint; the CIFAR-10 version from the guide is the same two scripts pointed at the DDPM on CIFAR-10 project's model and loader — nothing in the method changes.

Part 1: a classifier that works at every noise level

An off-the-shelf classifier is useless here: during sampling it will be shown images that are 10% signal and 90% static, a distribution it has never seen. So noisy_classifier.py trains exactly the way the diffusion model trains — sample a random t, noise the image with the DDPM on MNIST project's q_sample, and demand the label anyway. The timestep enters through FiLM on each conv block, so one classifier serves all noise levels, just like the U-Net does.

After training, the script measures test accuracy as a function of t:

Classifier accuracy vs noise level

This curve is the guidance budget. Where accuracy is high (low t, nearly clean images) the classifier's gradients are informative; as accuracy decays toward the 10% chance line, its opinions fade into noise. Guidance works because early reverse steps decide the global layout while the classifier still gets a vote through the accumulated denoising — but the strongest, most reliable steering happens in the second half of sampling.

Part 2: steering the unconditional model

The unconditional DDPM proposes a denoising mean; guidance shifts that mean along the classifier's gradient before the noise is re-added (guided_sampling.py):

grad = grad_{x_t} log p(y | x_t, t) # backprop through the classifier
mean <- mean + s * Sigma_t * grad # Sigma_t = the step's posterior variance

Two implementation details that are easy to get wrong:

  • The gradient is taken with respect to the input image, not the weights — x_t is detached, marked requires_grad, and autograd.grad pulls the sensitivity of the selected class's log-probability back through the frozen classifier. Scaling by Sigma_t automatically fades guidance as steps get small and confident.
  • Sampling runs under no_grad for the U-Net but must re-enter enable_grad for the classifier call — see the with torch.enable_grad(): block.

Note what is and is not conditioned: the diffusion model never sees the label at all, at training or sampling time. The label enters only through the classifier's gradient. Class control is bolted onto a finished unconditional model — the exact opposite trade-off from the Class-conditional DDPM project, where conditioning is baked into training.

Run it

python noisy_classifier.py # ~2 min on CPU, incl. the accuracy sweep
python guided_sampling.py # needs the DDPM on MNIST checkpoint, ~3 min

Results

The guidance-scale sweep. Rows top to bottom: s = 0, 1, 5, 20, all asking for the digit 8, all starting from the same noise. At s = 0 you get unconditional samples (whatever the noise wanted to be); as s grows the samples snap to the target class; push far enough and diversity collapses — the classic guidance trade-off of fidelity-to-condition against variety. In the recorded run the classifier scores the rows at 0%, 62%, 88%, and 100% "eights" respectively:

Guidance scale sweep for one target class

guided_sampling.py also scores each row with the classifier itself (fraction of samples it labels as the target at t = 0 — a self-judging metric, but a useful sanity check; see outputs/purity.csv).

All ten classes at a moderate scale (s = 5), one row per requested class, from a model that was never trained on labels:

All ten classes via guidance

Things to try

  • Overdrive it: --target 1 at s = 50. Saturation artifacts and near- duplicate samples show why the scale knob cannot buy unlimited fidelity.
  • Guide with a clean-image classifier (train with --T 1 so it never sees real noise) and watch guidance fail at high t — the reason "trained on noisy images" is in this project's title.
  • Compare against the Class-conditional DDPM project's grid at equal compute. Conditional training wins on purity-per-FLOP; guidance wins on not having to retrain the generator. Holding both trade-offs in your head is the setup for classifier-free guidance (the Classifier-free guidance project), which gets the best of each.