Progressive Distillation for Faster Student Models

1. The Core Idea: Student and Teacher Models

The central problem with few-step sampling is that taking large denoising steps is difficult. If a model trained on small steps is asked to take a giant leap from a very noisy image x_t toward the final clean image x_0, it often produces a blurry, low-quality average of possibilities.

Model distillation offers a clever solution. Instead of training one model to do everything, we use a two-model setup:

• Teacher Model: A large, powerful, pre-trained diffusion model that generates high-quality images using many small, slow steps.
• Student Model: A new model, typically with the same architecture, that learns to produce similar results to the teacher but in far fewer steps.

The key insight of progressive distillation is the student's training objective. The student isn't trained to predict the final clean image. Instead, it's trained to produce, in a single step, the same result that the teacher produces in two steps.

This video provides an excellent summary of the original paper that introduced this idea.

As the video explains, this objective gives the student a sharp, unambiguous target. It doesn't have to guess the single correct x_0 from a noisy x_t; it just has to learn to match the specific, high-quality intermediate output produced by the expert teacher. This makes the learning task much more tractable.

2. The Progressive Distillation Algorithm

The real power of this method comes from its iterative nature. Once a student model is trained to perform N/2 steps with the quality of a teacher that uses N steps, that student can then become the new teacher for the next round of distillation. This process is repeated, halving the steps each time (N -> N/2 -> N/4 -> N/8...), until you reach a very small number of steps (e.g., 4 or 8) while retaining much of the original quality.

Let's break down the algorithm for a single round of distillation:

1. Initialization: The student model's weights (θ_s) are initialized with the trained teacher model's weights (θ_t).
2. Training Loop: For each training batch:
   a. Sample a clean image x_0 and a random timestep t. Create the noisy input x_t.
   b. Teacher Prediction: Use the teacher model and a deterministic sampler like DDIM to perform two denoising steps, starting from x_t. This produces a target output, let's call it x_target.
   c. Student Prediction: Pass the same noisy input x_t to the student model to get its prediction x_pred. This prediction is what the student thinks the denoised image should be.
   d. Calculate Loss: The loss is the mean squared error between the student's prediction (x_pred) and a target derived from the teacher's two-step output (x_target). The video refers to this as "inverting a single step of ddim". In essence, we calculate what the student should have predicted to arrive at x_target in a single DDIM step, and the loss minimizes the difference between that and what it actually predicted.
   e. Update: Update the student model's weights θ_s using gradient descent.
3. Progression: After the student model is fully trained for this stage, its weights replace the teacher's weights, and the process repeats to halve the steps again.

A Note on Parameterization

As you continue your work with diffusion models, you'll find they can be parameterized to predict different things: the noise ε, the clean image x_0, or a value v that combines both. The original progressive distillation paper found that for this task, predicting x_0 or v is more stable than predicting ε, especially at high noise levels. This is a key architectural choice to be aware of when implementing or debugging distillation.

3. Implementation in Practice with PyTorch

Now, let's move from theory to implementation. We will use an unofficial but clear PyTorch implementation of progressive distillation. This repository includes a Colab notebook that lets you run the distillation process without complex local setup.

Hands-on Distillation

Your main task is to work through the distillation process using the provided notebook.

1. Open the Colab: Click the "Open in Colab" badge at the top of the GitHub repository's README.
2. Follow the Notebook: The notebook celeba_u_script.ipynb walks you through the entire pipeline:
   • Training a base model: It first trains a standard DDPM on the CelebA dataset. In a real-world scenario, you would start with an existing high-quality, pre-trained model as your teacher.
   • Running Distillation: The key section is "Distillate". Here, it runs the distillate.py script, loading the base model as the teacher (--base_checkpoint) and training a student model.
   • Iterative Process: It then shows how to run distillation again, this time using the first student as the new teacher to create an even faster second-generation student.

This iterative process is the "progressive" in progressive distillation. The official JAX implementation from the paper makes this loop very explicit.

# A conceptual snippet from the official JAX implementation (resource [LINK](https://colab.research.google.com/github/google-research/google-research/blob/master/diffusion_distillation/diffusion_distillation.ipynb))
# This shows the core 'progressive' loop

end_num_steps = 4
while state.num_sample_steps >= end_num_steps:
    # 1. Train the student model for a set number of iterations
    #    The student's 'num_sample_steps' is half of the teacher's.
    print(f'distilling teacher using {model.teacher_state.num_sample_steps} steps '
          f'into student using {state.num_sample_steps} steps')
    for step in range(steps_per_distill_iter):
        state, metrics = train_step(state, next(train_iter))

    # 2. The trained student becomes the new teacher
    model.teacher_state = get_new_teacher_state(state)

    # 3. Initialize a new student for the next round of distillation
    state = reset_student_for_next_iteration(model.teacher_state)

As you run the PyTorch notebook, compare the generated images from the original model (e.g., 128 steps), the first student (e.g., 64 steps), and the second student (e.g., 32 steps). You should see that the quality remains high despite the drastic reduction in inference time.

4. Advanced Distillation: Adding an Adversarial Loss

Progressive distillation is incredibly effective, but newer methods have pushed the boundaries even further, enabling high-quality generation in as few as 1-4 steps. One prominent technique is Adversarial Diffusion Distillation (ADD), used in models like SDXL-Turbo and Stable Cascade.

ADD enhances progressive distillation by adding a second loss term: an adversarial loss.

The two losses work together:

• The distillation loss ensures the student generates the correct content and follows the teacher's guidance.
• The adversarial loss pushes the student to generate perceptually realistic details and textures, correcting the blurry or "plastic" look that can sometimes appear in few-step generation.

This combination allows for a much more aggressive distillation, collapsing hundreds of steps into just a few while maintaining impressive visual fidelity.

Conclusion

You've now learned and implemented a cornerstone technique for accelerating diffusion model inference. By trading a one-time training cost for a permanent reduction in sampling steps, distillation makes these powerful models far more practical for real-world applications.

Key Takeaways:

• Progressive Distillation is an iterative process that halves the number of required sampling steps in each round.
• The core mechanism is training a student model to replicate, in one step, the output of a teacher model performing two steps.
• This training objective provides a sharp, well-defined target, which is more effective than naively training a model to denoise in large steps.
• The implementation involves a loop where the trained student becomes the teacher for the subsequent round.
• Advanced methods like Adversarial Diffusion Distillation (ADD) build on this by adding a GAN-style loss to improve perceptual quality in very-few-step models.

Next Lesson Preview:

Now that you know how to create a distilled model, how do you know if it's any good? Is a 4-step distilled model better than a 10-step DDIM sample? In the next lesson, we will focus on evaluation. You will learn to compare the performance (speed vs. quality) of distilled models against multi-step samplers, using standard metrics like FID to make quantitative judgments about these optimization trade-offs.