Diffusion Models for Images

Dive into Deep Learning · §17.8

Diffusion models for images
a U-Net noise predictor · class conditioning with a null label · classifier-free guidance

The Same Model, an Image-Shaped Network

• Fashion-MNIST in [-1, 1]; T = 1000, \beta_t linear from 10^{-4} to 0.02;
• the noise-conditional U-Net of the earlier sections predicts the noise, with the step index embedded;
• a class embedding is added to the time embedding, plus a null label that carries no class;
• label dropout with probability 0.1: one network, two predictors.

Sampling uses an exponential moving average of the weights.

Unconditional Samples after Ten Minutes of Training

Many recognizable garments from about 340 thousand parameters; detail limited by the budget.

Classifier-Free Guidance

\tilde{\boldsymbol{\epsilon}} = \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_t, t, \varnothing) + \gamma\, \big( \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_t, t, c) - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_t, t, \varnothing) \big)

Bayes’ rule twice: the class gradient is proportional to the difference of the two predictions the network already provides.

Classifier Guidance on the Same Requests

\hat{\boldsymbol{\epsilon}} = \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_t, t, \varnothing) - s\, \sqrt{1 - \bar{\alpha}_t}\; \nabla_{\mathbf{x}_t} \log p_{\boldsymbol{\phi}}(c \mid \mathbf{x}_t, t)

• the class direction comes from a classifier trained on noisy images;
• agreement rises and within-class distance falls with s, with more malformed samples at large s;
• the cost: a second network, trained on noisy data, with its own confusions.

Agreement Rises, Within-Class Distance Falls

y = torch.arange(10, device=device).repeat(8)  # eight samples per class
guided = {}
print(f'{"gamma":>5} {"class agreement":>16} {"within-class distance":>22}')
for gamma in (1.0, 2.0, 4.0):
    guided[gamma] = sample(ema, y, gamma, seed=2).clamp(-1, 1)
    with torch.no_grad():
        agree = (classifier(guided[gamma]).argmax(1) == y).float().mean()
    flat = guided[gamma].flatten(1)
    dist = torch.stack([torch.pdist(flat[y == c]).mean() for c in range(10)]).mean()
    print(f'{gamma:>5.1f} {agree:>16.3f} {dist:>22.2f}')

gamma  class agreement  within-class distance
  1.0            0.625                  16.63
  2.0            0.875                  15.33
  4.0            0.950                  13.44

Measured with a classifier trained on clean images: agreement with the requested class increases with \gamma, within-class distance decreases.

The Same Seed at Two Guidance Scales

first grid: guidance scale 1; second grid: guidance scale 4; columns are the ten classes

Columns are classes. At \gamma = 4 the instances are cleaner, more prototypical, and more alike; some saturate to white.

Recap

• Images change the network, not the model: a U-Net predicts the noise.
• A null label makes one network conditional and unconditional at once.
• Classifier-free guidance extrapolates between the two; \gamma trades diversity for class fidelity. Classifier guidance does the same with a noisy-image classifier’s gradient.
• Beyond: fast samplers, distillation, latent diffusion, transformer backbones, and the continuous-time view of the appendix.