Model Distillation for Diffusion

While advanced samplers like DPM-Solver can reduce the number of steps needed for generation, the underlying diffusion model often remains large and computationally intensive. Model distillation offers a complementary approach to accelerate inference by creating smaller, faster "student" models that approximate the behavior of a larger, pre-trained "teacher" diffusion model. This technique is borrowed from broader deep learning practices where knowledge from a cumbersome model is transferred to a more efficient one.

The fundamental goal of distillation in the context of diffusion models is to train a student network, $\theta_{student}$, to mimic the predictions made by a frozen, pre-trained teacher network, $\theta_{teacher}$. The student model typically has a significantly smaller architecture (e.g., fewer layers, channels, or attention heads), making it faster and less memory-intensive during inference.

Distillation Objectives

Several strategies exist for defining how the student should mimic the teacher:

1. Matching Denoising Predictions: The most straightforward approach is to train the student to predict the same output as the teacher for a given noisy input $x_t$ and timestep $t$. If both models predict the noise $\epsilon$, the distillation loss aims to minimize the difference between their predictions:

   $$L_{distill} = \mathbb{E}_{x, \epsilon, t} [ w(t) || \epsilon_{\theta_{student}}(x_t, t) - \epsilon_{\theta_{teacher}}(x_t, t) ||^2 ]$$

   Here, $x$ is a data sample, $\epsilon$ is sampled noise, $t$ is a timestep, $x_t = \sqrt{\bar{\alpha}_t} x + \sqrt{1 - \bar{\alpha}_t} \epsilon$ is the noisy input, and $w(t)$ is an optional weighting term that can prioritize matching at certain timesteps. A similar objective can be defined if the models predict the denoised sample $x_0$. The teacher model's parameters ($\theta_{teacher}$) remain fixed during this process.

2. Matching Probability Flow ODE Solutions: Instead of just matching the output at discrete steps, the student can be trained to approximate the trajectory defined by the probability flow ODE associated with the teacher model. This often involves matching the predicted score or velocity fields.

3. Feature-Level Distillation: Rather than only matching the final output, the student can be encouraged to replicate the internal feature representations of the teacher model at specific layers. This involves adding auxiliary loss terms that minimize the difference between intermediate activations, potentially providing a richer training signal.

Progressive Distillation

A significant challenge is that a student model performing single-step generation might struggle to replicate the quality of a multi-step teacher. "Progressive Distillation for Fast Sampling of Diffusion Models" (Salimans & Ho, 2022) proposed an effective technique to address this.

The core idea is to iteratively distill the sampling process.

1. Stage 1: Train a student model ($S_1$) to perform in one step what the teacher model ($T$) does in two steps. Specifically, $S_1(x_t, t)$ is trained to predict the output that $T$ would produce after applying its update rule twice, starting from $x_t$. This effectively halves the number of required sampling steps ($N \rightarrow N/2$).
2. Stage 2: Treat the trained student $S_1$ as the new teacher. Train a second student model ($S_2$) to perform in one step what $S_1$ does in two steps. This further halves the sampling steps ($N/2 \rightarrow N/4$).
3. Repeat: Continue this process iteratively. Each stage reduces the number of required function evaluations (NFE) by half. After $k$ stages, the resulting student model requires only $N/2^k$ steps.

This progressive approach allows the student models to learn the complex mapping over multiple intermediate stages, often leading to better sample quality compared to a single-stage distillation attempting a large reduction in steps.

Diagram illustrating the progressive distillation process. Each student model learns to emulate two steps of the previous model (teacher or preceding student), iteratively reducing the required sampling steps.

Relation to Consistency Models

Model distillation, particularly progressive distillation, shares the goal of faster sampling with Consistency Models (discussed in Chapter 5). However, the mechanisms differ:

• Consistency Models: Learn a function that maps any point $(x_t, t)$ directly to the trajectory origin $x_0$ by enforcing a consistency property along the ODE/SDE trajectories. They can be trained via a distillation-like process or independently.
• Progressive Distillation: Focuses on halving the number of discrete sampling steps by training a student to predict the outcome of two teacher steps. It doesn't explicitly enforce the same consistency property across the entire continuous trajectory.

In practice, both approaches can yield significant speedups, often enabling generation in 1 to 8 steps. Progressive distillation might offer more flexibility in retaining aspects of the original multi-step sampling process, while consistency models are specifically designed for extreme few-step or single-step generation based on the consistency property.

Architectural Choices for Student Models

An important advantage of distillation is the flexibility in choosing the student model's architecture. It does not need to match the teacher. Common choices include:

• Using the same architecture type (e.g., U-Net) but with reduced depth, width (channels), or fewer attention heads.
• Employing entirely different, more efficient architectures.

The selection depends heavily on the target application, desired inference speed, acceptable quality trade-off, and available computational resources for inference.

Advantages and Disadvantages

Advantages:

• Speed & Efficiency: Drastically reduces sampling time and computational cost compared to the original teacher model.
• Deployment: Enables deploying diffusion models on devices with limited computational power or strict latency requirements.
• Flexibility: Student architecture can be tailored to specific needs.

Disadvantages:

• Quality Degradation: Distilled models often exhibit some reduction in sample quality or diversity compared to the teacher. This trade-off is usually more pronounced with greater compression (fewer steps).
• Training Cost: The distillation process itself can be computationally expensive, requiring training one or more student models using outputs from the large teacher.
• Hyperparameter Tuning: Finding the optimal distillation objective, student architecture, and training schedule requires careful experimentation.

Model distillation provides a valuable set of techniques for optimizing diffusion models with algorithmic improvements in sampling. By creating smaller, faster student models, distillation makes state-of-the-art generative capabilities more practical for real-world applications, complementing methods like quantization and hardware acceleration discussed later in this chapter.