Classifier-Free Diffusion Guidance

Jointly training a conditional and unconditional diffusion model and combining their results can attain a trade-off between sample quality and diversity similar to what can be obtained using classifier guidance.

Diffusion Models

  • A diffusion model learns to reverse the noising process
$$ X_T \to X_{T-1} \to ... \to x_0 $$
  • Starting from pure Gaussian noise $X_T$, the model predicts the noise $\epsilon_\theta(x_t, t)$ at each step and removes it
  • For text-to-image models, we want the model to generate $p(x \mid c)$ where $x$ is the image and $c$ is the condition or the text prompt
  • Therefore, we train with $\epsilon_\theta(x_t, t, c)$ and the text-embedding $c$ is fed into the U-Net via cross attention

Classifier Guidance

  • Before classfier-free guidance (CFG), the common approach was classifier guidance in which we train two separate models
    • The diffusion model predicts noise $\epsilon_\theta(x_t, t)$
    • The classifier predicts class probabilities $p_\phi(y \mid x_t)$
  • During sampling, the diffusion trajectory is modified using the classifier gradient
$$ \nabla_x \log p_\phi(y \mid x) $$
  • The intuition is that the noisy image is moved in the direction that makes it more like the caption
  • The guided score becomes
$$ \nabla_x \log p(x \mid y) = \nabla_x \log p(x) + \nabla_x \log p_\phi(y \mid x) $$
  • However, classifier guidance is burdensome because you need both a diffusion model and a separate classifier
  • The classifier must work on noisy images and thus needs to be trained to predict the class label on every noise level $x_t$

Classifier-Free Guidance

  • Instead of training a separate classifier model, we jointly train the unconditional and conditional diffusion model
  • For classifier-free guidance, we drop the condition with some probability during training
  • In the cases where the condition is dropped, we are training the unconditional model and in the cases where the condition is kept, we are training the conditional model
  • During inference time, we run two passes on the model, one without the condition and one with the condition, and take a linear combination of the conditional and unconditional predictions
  • The unconditional prediction is
$$ \epsilon_{uncond} = \epsilon_\theta(x_t, t, \emptyset) $$
  • The conditional prediction is
$$ \epsilon_{cond} = \epsilon_\theta(x_t, t, c) $$
  • The combined prediction is
$$ \epsilon_{guided} = \epsilon_{uncond} + w(\epsilon_{cond} - \epsilon_{uncond}) $$

Conclusion

  • Before CFG, diffusion models were good at generating realistic and diverse images, but were weak at following text
  • CFG gives a knob, the guidance weight $w$, to control how much the model should follow the prompt