Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed

In plain words

Many image generators start from random noise and clean it up over many small steps. The pictures look good, but every image needs many passes through a large neural network, so generation is slow next to models that need only one pass.

The idea is teacher and student. The slow generator is run on many random noise inputs, and each noise input is paired with the image the slow generator made from it. A second network, the student, starts as a copy of the slow network and is trained to produce that same image from the noise in a single pass. Because the slow generator always gives the same output for the same noise, each target is fixed, so training is a plain regression. There is no discriminator and no adversarial game.

The result is generation as fast as other single-pass generators. On small datasets the samples are reported to be comparable to GAN samples, while larger images come out blurrier. The method needs no change to the network design and applies to any iterative generator that is deterministic.

TL;DR

Distills a 100-step DDIM teacher (or a 50-step one for LSUN) into a single network evaluation §4. The student is initialised from the teacher's network and regresses the teacher's output from the same initial noise, using an L1 or L2 loss and no adversarial training §3.1 app. B. On CIFAR-10 the author-reported FID is 9.36 at one step, against 4.16 for the 100-step teacher tab. 1. Training needs 1.024M teacher samples per dataset, and the repository asks for at least two high-end GPUs or TPUs README.md.

Problem

Score-based and diffusion models produce high-quality samples but need many denoising steps. The paper states they are 2 to 3 orders of magnitude slower than GANs and VAEs §1. DDIM reduces the cost but still needs roughly 20 to 100 network evaluations for good samples §1.

The student has no significant architectural restrictions, so any network can serve as the student §3.2. The authors state that the method applies to any iterative model with a deterministic generative process §6.

Knowledge distillation needs a deterministic teacher. Score-based and energy-based models sample with stochastic MCMC, so they cannot serve as teachers. A GAN is deterministic but already fast, so distilling it gains little §2.1. DDIM is deterministic and is used as the teacher §2.2.

Method

The teacher is the deterministic DDIM sampler, with its noise-prediction network ε_θ taken from a pretrained DDPM (CIFAR-10, LSUN) or DDIM (CelebA) model §2.2 §4. A trained DDPM network can be run as a DDIM by changing only the sampler §2.2.

1. Teacher map. F_teacher(x_T) is the output of the full DDIM sampling chain that starts at noise x_T and ends at x_0 §2.2 §3.1. The chain uses the update in eq. 8.

   x_{t-1} = \sqrt{\tfrac{\alpha_{t-1}}{\alpha_t}}\bigl(x_t - \sqrt{1-\alpha_t}\,\epsilon_\theta(x_t)\bigr) + \sqrt{1-\alpha_{t-1}}\,\epsilon_\theta(x_t) eq. 8

2. Training pairs. Draw x_T from N(0, I) and store the pair (x_T, F_teacher(x_T)). The paper uses 1.024M pairs per dataset and reuses them over several passes §3.1. The code stores them as described in create_dataset.py (see Paper vs code).

3. Student. The student network is initialised with the same architecture and weights as the teacher network ε_θ, so it inherits the teacher's knowledge §3.2. It predicts a noise term, which is subtracted from the input noise to give the sample §3.2. It is conditioned on the highest noise level T §3.2.

4. Loss. Minimise the squared distance between student and teacher outputs, with the constant factor shown in eq. 11 §3.1. The loss is written as a KL divergence between two unit-variance Gaussians, which reduces to a squared error §3.1.

   L_{\text{student}} = \tfrac{1}{2}\,\mathbb{E}_{x_T}\bigl[\lVert F_{\text{student}}(x_T) - F_{\text{teacher}}(x_T)\rVert_2^2\bigr] + C eq. 11

   The paper reports that L1 works better than L2 on CIFAR-10 app. B. The code uses L1 for CIFAR-10 and L2 for the other three datasets training.py.

5. Optimisation. Adam with a linear learning-rate warm-up and no decay, no dropout, no gradient clipping, and an exponential moving average (EMA) of the weights app. B. Hyperparameters per dataset are in the table below tab. 3.

6. Sampling. Draw x_T from N(0, I) and run the student once. Samples are not drawn at reduced temperature §4.1. The code evaluates the EMA weights eval.py.

# Teacher target for one noise input  [§2.2] [code: create_dataset.py#L77-L86]
# noise: (B, H, W, 3) drawn from N(0, I)
# steps: ascending indices into the teacher's 1000-step schedule
def teacher_output(noise, teacher_network, steps):
    # Pair each step with the step before it; the first one gets -1
    previous_steps = [-1] + steps[:-1]
    pairs = list(zip(steps, previous_steps))
    sample = noise
    # Visit the steps from the noisiest to the cleanest  [§2.2]
    for step, previous_step in reversed(pairs):
        # Predict the noise with the teacher network  [§2.2]
        predicted_noise = teacher_network(sample, step)
        # One deterministic DDIM update  [eq. 8]
        sample = ddim_update(sample, predicted_noise, step, previous_step)
    # Clip to [-1, 1] and quantise to 8 bits  [code: create_dataset.py#L92-L95]
    return quantise(sample)

# One deterministic DDIM update  [eq. 8] [code: models.py#L303-L307]
# alpha, alpha_previous: cumulative products of (1 - beta) at the two steps
def ddim_update(sample, predicted_noise, alpha, alpha_previous):
    # Estimate the clean image implied by the predicted noise  [§2.2]
    clean_estimate = (
        sample - sqrt(1 - alpha) * predicted_noise
    ) / sqrt(alpha)
    # Move the estimate to the previous noise level  [eq. 8]
    return (
        sqrt(alpha_previous) * clean_estimate
        + sqrt(1 - alpha_previous) * predicted_noise
    )

# Student loss for one training batch  [§3.1] [§3.2] [eq. 11]
# noise: (B, H, W, 3) drawn from N(0, I)
# targets: teacher outputs for the same noise, (B, H, W, 3)
def student_loss(noise, targets, use_l2):
    # Predict the noise at the highest step, index 999  [§3.2]
    # [code: models.py#L294-L297]
    predicted_noise = student_network(noise, timestep=999)
    # Subtract the predicted noise from the input  [§3.2]
    # [code: models.py#L300]
    student_output = noise - predicted_noise
    error = targets - student_output
    # Squared error for CelebA and LSUN, absolute error for CIFAR-10
    # [app. B] [code: training.py#L122-L125]
    if use_l2:
        # Summed over pixels, divided by batch size  [code: training.py#L45-L46]
        return sum(square(error)) / batch_size
    # [code: training.py#L48-L49]
    return sum(abs(error)) / batch_size

# Optimisation loop  [app. B] [code: training.py#L189-L222]
# Adam on the raw weights, an EMA of the weights kept alongside
for epoch in range(epochs):
    for shard_index, shard in enumerate(shards):
        for noise, targets in batches(shard):
            # Learning rate rises linearly, then stays at the base rate
            # [app. B] [code: utils.py#L110-L119]
            rate = min(base_rate, step * base_rate / warmup_steps)
            gradient = gradient_of(student_loss, noise, targets)
            weights = adam_step(weights, gradient, rate)
            # [code: training.py#L156-L159]
            ema_weights = ema(ema_weights, weights, ema_decay)
        # Test loss every 80 shards and at the end of each epoch
        # [code: training.py#L200-L220]
        if shard_index % 80 == 79 or shard_index == last_shard:
            test_loss = evaluate(ema_weights, test_set)
            if test_loss < best_test_loss:
                best_test_loss = test_loss
                best_ema_weights = ema_weights

# Sampling with the trained student  [§3.2] [code: eval.py#L23-L27]
# Uses the EMA weights  [code: eval.py#L29-L33]
def sample(student, count, height, width):
    # Initial noise from N(0, I), without temperature reduction  [§4.1]
    noise = normal(0, 1, (count, height, width, 3))
    # One network call gives the whole image  [code: models.py#L294-L301]
    images = student(noise)
    # Clip to [-1, 1] and map to 8-bit pixels  [code: utils.py#L88-L91]
    return to_uint8(clip(images, -1, 1))

Figures are described in words only (no images are reproduced). fig. 1 is a schematic of the teacher-to-student mapping §2.2. fig. 2 to fig. 5 show uncurated samples from the CIFAR-10, CelebA, LSUN Bedroom and LSUN Church models §4.1. fig. 6 and fig. 12 show spherical interpolations between two initial noise vectors §4.2. fig. 11 shows CIFAR-10 samples next to their nearest neighbours in the training set, in Inception feature space app. A. fig. 13 shows 256 × 384 samples from both the teacher and the student, sizes neither model was trained on §4.2 app. A.

Hyperparameters

Values from the paper's table tab. 3 are given as published. Values marked "code" are from the repository. The derived iteration counts are arithmetic on the code's settings, not measured runs.

Not stated in the paper: the numeric value of T; the noise schedule (the code uses β linearly spaced from 1e-4 to 0.02 over 1000 steps create_dataset.py); Adam ε; the number of epochs; the base width; which weights were used for the reported FID.

Results

Author-reported. Setup: CIFAR-10 at 32 × 32, FID and Inception score computed by the authors against the training set tab. 1 eval.py. The Denoising Student uses one network evaluation per sample. The teacher's Inception score is marked with an asterisk, which the paper says means the authors computed it and it is not official tab. 1.

Timing: time to generate 50k CIFAR-10 images, measured on a single NVIDIA V100 at batch size 250 §4.1.

Other reported results:

• CelebA 64 × 64: FID 10.68 for the Denoising Student, given in the text only, with no baseline table §4.1.
• LSUN Bedroom and Church at 256 × 256: no FID or other metric is reported. The paper gives qualitative samples and says the results have blurry textures §4.1 fig. 4 fig. 5.
• Latent-space interpolation is qualitative only §4.2 fig. 6.

Paper vs code

Checked against the official repository at code_commit (MIT). Learning rates, Adam β values, EMA decay, warm-up steps, batch sizes and parameter counts match tab. 3 with utils.py and models.py. The discrepancies are:

• Training iterations. The paper's table gives 50k, 64k, 320k and 256k iterations tab. 3. Epochs × shards × batches per shard in the code gives 50,000 for CIFAR-10, which matches, and 70,000 for CelebA, 1,600,000 for LSUN Bedroom and 1,280,000 for LSUN Church utils.py training.py create_dataset.py. The LSUN figures are exactly five times the paper's values. This is arithmetic on the settings, not a run. Which is correct is not stated.
• Teacher step schedule. The paper gives only the step counts §4. The code uses a quadratic progression for CIFAR-10 and CelebA that ends at index 800, and a linear one for LSUN that ends at index 980 create_dataset.py create_dataset.py. A code comment says that for CelebA the authors found quadratic gave FID 6.3 and linear gave FID 7.1 in their experiments create_dataset.py.
• Student timestep. The paper conditions the student on "timestep T", the highest noise level §3.2, without a numeric value. The code always uses index 999 models.py. The teacher chain starts at index 800 (CIFAR-10, CelebA) or 980 (LSUN) create_dataset.py, so the student is trained on a different starting index from the one its teacher used.
• Target precision. The paper defines the target as the teacher's real-valued output F_teacher(x_T) §3.1 eq. 11. The code clips the teacher output to [-1, 1] and stores it as 8-bit unsigned integers create_dataset.py create_dataset.py. The loader maps those values back to [-1, 1] training.py. The noise inputs are stored as float16 create_dataset.py.
• Loss normalisation. eq. 11 has a factor of ½ and is a mean over the expectation. The code sums the squared (or absolute) error over all elements and divides by the batch size, with no ½ training.py. Whether this changes the effective learning rate under Adam is not analysed in the paper.
• Early stopping. The paper says it used "test loss evaluated on unseen images for early stopping" app. B. The code does not stop training. It evaluates the EMA model on the test set every 80 shards and at the end of each epoch, and keeps the EMA weights with the lowest test loss training.py.
• Command-line flag. The README documents --use_fewer_lsun_examples False for LSUN README.md. The flag is defined in the parser but not passed to train in the entry point training.py.
• Consistent with the code: the choice of L1 on CIFAR-10 app. B training.py; the teacher for CelebA is DDIM §4 models.py; the student is initialised from the teacher's weights §3.2 models.py.

Pitfalls

• Teacher and student do not share a starting index. As noted above, the teacher's chain starts at index 800 or 980 while the student is always conditioned on 999 create_dataset.py models.py. The paper does not discuss this, and its effect on quality is not tested. An implementation should decide which convention to follow and say so.
• Quantised targets. Stored teacher targets have a resolution of 1/127.5 (about 0.0078) on the [-1, 1] scale, and any teacher output outside [-1, 1] is clipped before storage create_dataset.py training.py. A reimplementation that regresses on raw floats will not reproduce the same targets.
• Expensive data generation. Each training pair needs a full teacher chain of 50 to 100 network evaluations, and the paper uses 1.024M pairs per dataset §3.1 create_dataset.py. The README says the full LSUN set needs about 600 GB of disk README.md and recommends at least 32 GB of RAM and two high-end GPUs or TPUs README.md. The authors' own runs used TPUv2-8 README.md.
• Unstable LSUN gradients. The authors saw gradient norms above 10⁵ on LSUN Bedroom, clipped them by value to ±1, and found that not clipping was more stable app. B. The final recipe has no clipping. The LSUN β1 and learning rate were tuned on Bedroom only, and Church reused them without a sweep app. B.
• Batch size affects reproducibility. The entry point says reproducibility may suffer if the batch size is changed from the default training.py. The paper also says results may differ from the reported ones because training noise is drawn at random README.md.
• Use EMA weights for sampling. The evaluation script loads the EMA checkpoint eval.py. The paper does not say which weights produced its reported FID tab. 1.
• Blurry high-resolution output. The authors attribute the blur on 256 × 256 LSUN to the pixel-wise objective and report it as an observation, not a tested cause §4.1.
• Gap to the teacher. The student's FID is 9.36 against 4.16 for the 100-step teacher on CIFAR-10 tab. 1. The conclusion names closing this gap as future work §6.
• Timing claims. The 100× and 1000× speed-ups are ratios of the times in tab. 2 §4.1. The repository has no timing script, so these numbers cannot be reproduced from the code.
• Metric code is external. The paper says the metrics use the official IS and FID code, and the repository's sampler writes 50k images for that purpose eval.py. The FID for CelebA is compared with a specific preprocessing from another paper eval.py.

Lineage

• Builds on: knowledge distillation with the work of Bucilua et al. and Hinton et al. (arXiv:1503.02531) §2.1; DDIM (arXiv:2010.02502), which supplies the deterministic teacher sampler §2.2; DDPM (arXiv:2006.11239), which supplies the teacher networks §2.2 §4.
• Compared against: the methods listed in tab. 1, including NVAE (arXiv:2007.03898), DDPM, NCSN++, NCSNv2 and StyleGAN2+ADA.
• Related: the probability-flow ODE of the score SDE work (arXiv:2011.13456) is a deterministic multi-step sampler that the authors did not use, because its pretrained models were unavailable when they wrote the paper §2.2 (footnote).
• Superseded by: not stated.