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 = jnp.tile(jnp.arange(10), 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).clip(-1, 1)
    agree = (classifier(guided[gamma]).argmax(1) == y).mean()
    flat = guided[gamma].reshape(80, -1)
    dist = jnp.mean(jnp.array([
        jnp.linalg.norm(flat[y == c][:, None] - flat[y == c][None], axis=-1).sum()
        / 56 for c in range(10)]))  # 8 samples: 56 ordered pairs per class
    print(f'{gamma:>5.1f} {agree:>16.3f} {dist:>22.2f}')
gamma  class agreement  within-class distance
  1.0            0.650                  16.46
  2.0            0.850                  15.24
  4.0            0.875                  14.47

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.