Multi-Student Diffusion Distillation for Better One-Step Generators

## Abstract

Abstract Diffusion models achieve high-quality sample generation at the cost of a lengthy multistep inference procedure.
To overcome this, diffusion distillation techniques produce student generators capable of matching or surpassing the teacher in a single step.
However, the student model’s inference speed is limited by the size of the teacher architecture, preventing real-time generation for computationally heavy applications.
In this work, we introduce Multi-Student Distillation (MSD), a framework to distill a conditional teacher diffusion model into multiple single-step generators. Each student generator is responsible for a subset of the conditioning data, thereby obtaining higher generation quality for the same capacity.
MSD trains multiple distilled students, allowing smaller sizes and, therefore, faster inference.
Also, MSD offers a lightweight quality boost over single-student distillation with the same architecture.
We demonstrate MSD is effective by training multiple same-sized or smaller students on single-step distillation using distribution matching and adversarial distillation techniques. With smaller students, MSD gets competitive results with faster inference for single-step generation.
Using 4 4 4 4 same-sized students, MSD sets a new state-of-the-art for one-step image generation: FID 1.20 1.20 1.20 1.20 on ImageNet-64 × \times × 64 and 8.20 on zero-shot COCO2014. 1 1 1 Project page: https://research.nvidia.com/labs/toronto-ai/MSD/

## 1 Introduction

Diffusion models are the dominant generative model in image, audio, video, 3D assets, protein design, and more [, , , , ]. They allow different conditioning inputs – such as class labels, text, or images – and achieve high-quality generations. However, their inference process typically requires hundreds of model evaluations – with an often slow and bulky network – for a single sample. This procedure costs millions of dollars per day [, ]. It also prohibits applications requiring rapid synthesis, such as augmented reality. Real-time, low-cost, and high-quality generation will have huge financial and operational impacts while enabling new usage paradigms.

There has been a flurry of work on diffusion-distillation techniques to address the slow sampling of diffusion models [, , ]. Inspired by knowledge distillation [], these methods use the trained diffusion model as a teacher and optimize a student model to match its generated output in as few as a single step. However, most of these methods use the same student and teacher architecture, preventing real-time generation for applications with bulky networks, such as video synthesis [].
While reducing model size can reduce latency, it typically yields worse generation quality, thus presenting a speed-to-quality tradeoff dilemma with existing distillation methods.

Figure: Figure 1: We visualize distilling into multiple students, where each student handles a subset of the input condition. At training, students are trained separately with filtered data. At inference, a single student is retrieved for generation, given the corresponding input condition. We show samples in Fig. 2, with methodological details outlined in Fig. 3, and quantitative results in Tables 2 and 2.
Refer to caption: x1.png

We tackle this dilemma with our method, Multi-Student Distillation (MSD), that introduces multiple single-step student generators distilled from the pretrained teacher. Each student is responsible for a subset of conditioning inputs.
We determine which student to use during inference and perform a single-model evaluation to generate a high-quality sample. This way, MSD enjoys the benefit of a mixture of experts []: it increases the model capacity without incurring more inference cost, thereby effectively pushing the limit of the speed-quality tradeoff.

When distilling into the same-sized students, MSD has the flexibility of conceptually applying to any distillation method for a performance boost. In addition, MSD allows distilling into multiple smaller students for reduced single-step generation time. Using smaller students prevents one from initializing a student from teacher weights, which we solve by adding a relatively lightweight score-matching pretraining stage before distillation (Sec. 4.3).

We validate our approaches by applying MSD to distill the teacher into $4$ same-sized students using distribution matching and adversarial distillation procedures [, ]. The resulting students collectively outperform single-student counterparts, setting new state-of-the-art FID scores of $1.20$ on one-step ImageNet-64$\times$64 generation (Tab. 2) and 8.20 on one-step zero-shot COCO2014 generation (Tab. 2). We also distill the same teacher into $4$ smaller students, which achieve a competitive FID of 2.88 on ImageNet (Tab. 2), with $42\%$ less parameters per student.

Figure: Figure 2: Samples on high guidance-scale text-to-image generations from the SD v1.5 teacher and different sized students, with full training details in App. B. The same-sized student has comparable quality to the teacher. The smaller student, trained on a subset of dog-related data, achieves faster generation while still having decent qualities. The same-sized student is trained with DM stage only, whereas the smaller student is trained with TSM and DM stages (see Fig. 3). See additional samples in Fig. 11.
Refer to caption: x2.png

We summarize our contributions below, which include:

- •
A new framework, MSD, that upgrades existing single-step diffusion distillation methods (Sec. 4.1) by increasing the effective model capacity without changing inference latency.
- •
Demonstrating the effectiveness of MSD by training multiple same sized students using SOTA distillation techniques [, ] in Sec. 4.2, resulting in new record FID scores in ImageNet-64$\times$64 (Sec. 5.2) and zero-shot text-to-image generation (Sec. 5.3).
- •
A successful scheme to distill multiple smaller single-step students from the teacher model, achieving comparable generation quality with reduced inference time.

## 2 Related Work

Diffusion Sampling Acceleration.
While a line of work aims to accelerate diffusion models via fast numerical solvers for the PF-ODE [, , , , ], they usually still require more than $10$ steps. Training-based methods that usually follow the knowledge distillation pipeline can achieve low-step or even one-step generation.
Luhman and Luhman [] first used the diffusion model to generate a noise and image pair dataset that is then used to train a single-step generator. DSNO [] precomputes the denoising trajectory and uses neural operators to estimate the whole PF-ODE path. Progressive distillation [, ] iteratively halves the number of sampling steps required without needing an offline dataset. Rectified Flow [] and follow-up works [, ] straighten the denoising trajectories to allow sampling in fewer steps. Another approach uses self-consistent properties of denoising trajectories to inject additional regularization for distillation [, , , , , , ].

The methods above require the student to follow the teacher’s trajectories.
Instead,
a recent line of works aims to only match the distribution of the student and teacher output via variational score distillation [, , , , , , ]. The adversarial loss [], often combined with the above techniques, has been used to enhance the distillation performance further [, , , , , , , ]. Although MSD is conceptually compatible and offers a performance boost to all of these distillation methods, in this work, we demonstrate two specific techniques: distribution matching [] and adversarial distillation [].

Mixture of experts training and distillation.
Mixture of Experts (MoE), first proposed in [], has found success in training large-scale neural networks [, , , , ]. Distilling a teacher model into multiple students was explored by Hinton et al. [], and after that, further developed for supervised learning [, , ] and language modeling [, , ]. Although several works [, , ] have proposed MoE training for generative adversarial networks, they train the MoE from scratch. This requires carefully tuning the multi-expert adversarial losses. A recent work [] proposes to distill a pretrained diffusion model into an MoE for policy learning, which shares similar motivations with our work. However, to the best of our knowledge, MSD is the first method to distill multi-step teacher diffusion models into multiple one-step students for image generation.

Efficient architectures for diffusion models. In addition to reducing steps, an orthogonal approach aims to accelerate diffusion models with more efficient architectures. A series of works [, , ] introduces vision transformers to diffusion blocks and trains the diffusion model with new architectures from scratch. Another line of work selectively removes or modifies certain components of a pretrained diffusion model and then either finetunes [, , ] or re-trains [] the lightweight diffusion model, from which step-distillation can be further applied [, ]. Our approach is orthogonal to these works in two regards: 1) In our method, each student only handles a subset of data, providing a gain in relative capacity. 2) Instead of obtaining a full diffusion model, our method employs a lightweight pretraining stage to obtain a good initialization for single-step distillation. Combining MSD with more efficient architectures is a promising future direction.

## 3 Preliminary

We introduce the background on diffusion models in Sec. 3.1 and distribution matching distillation (DMD) in Sec. 3.2. We discuss applying adversarial losses to improve distillation in Sec. 3.3.

### 3.1 Diffusion Models

Diffusion models learn to generate data by estimating the score functions [] of the corrupted data distribution on different noise levels. Specifically, at different timesteps $t$, the data distribution $p_{\text{real}}$ is corrupted with an independent Gaussian noise: $p_{t,\text{real}}(\bm{x}_{t})=\int p_{\text{real}}(\bm{x})q_{t}(\bm{x}_{t}|\bm
{x})d\bm{x}$ where $q_{t}(\bm{x}_{t}|\bm{x})\sim\mathcal{N}(\alpha_{t}\bm{x},\sigma_{t}^{2}\bm{I})$ with predetermined $\alpha_{t},\sigma_{t}$ following a forward diffusion process [, ].
The neural network learns the score of corrupted data $\bm{s}_{\text{real}}:=\nabla_{\bm{x}_{t}}\log p_{t,\text{real}}(\bm{x}_{t})=-(
\bm{x}_{t}-\alpha_{t}\bm{x})/\sigma_{t}^{2}$ by equivalently predicting the denoised $\bm{x}$: $\bm{\mu}(\bm{x}_{t},t)\approx\bm{x}$. After training with the denoising score matching loss $\mathbb{E}_{\bm{x},t,\bm{x}_{t}}[\lambda_{t}\|\bm{\mu}(\bm{x}_{t},t)-\bm{x}\|_
{2}^{2}]$, where $\lambda_{t}$ is a weighting coefficient,
the model generates the data by an iterative denoising process over a decreasing sequence of time steps.

### 3.2 Distribution Matching Distillation

Inspired by Wang et al. [], a series of works [, , , ] aim to train the single-step distilled student to match the generated distribution of the teacher diffusion model. This is done by minimizing the following reverse KL divergence between teacher and student output distributions, diffused at different noise levels for better support over the ambient space:

$$ $\displaystyle\mathbb{E}_{t}D_{\text{KL}}(p_{t,\text{fake}}\|p_{t,\text{real}}) =\mathbb{E}_{\bm{x}_{t}}\left(\log\left(\frac{p_{t,\text{fake}}(\bm{x}_{t})}{p _{t,\text{real}}(\bm{x}_{t})}\right)\right).$ (1) $$

The training only requires the gradient of Eq. (1), which reads (with a custom weighting $w_{t}$):

$$ $\displaystyle\nabla_{\theta}\mathcal{L}_{\text{KL}}(\theta):=\nabla_{\theta} \mathbb{E}_{t}D_{\text{KL}}\simeq\mathbb{E}_{\bm{z},t,\bm{x}_{t}}[w_{t}\alpha_ {t}(\bm{s}_{\text{fake}}(\bm{x}_{t},t)-\bm{s}_{\text{real}}(\bm{x}_{t},t)) \nabla_{\theta}G_{\theta}(\bm{z})],$ (2) $$

where $\bm{z}\sim\mathcal{N}(0,\bm{I}),t\sim\text{Uniform}[T_{\min},T_{\max}]$, and $\bm{x}_{t}\sim q(\bm{x}_{t}|\bm{x})$, the noise injected version of $\bm{x}=G_{\theta}(\bm{z})$ generated by the one-step student. Here, we assume the teacher denoising model accurately approximates the score of the real data, and a “fake” denoising model approximates the score of generated fake data:

$$ $\bm{s}_{\text{real}}(\bm{x}_{t},t)\approx-\frac{\bm{x}_{t}-\alpha_{t}\bm{\mu}_ {\text{teacher}}(\bm{x}_{t},t)}{\sigma_{t}^{2}},\ \ \bm{s}_{\text{fake}}(\bm{x }_{t},t)\approx-\frac{\bm{x}_{t}-\alpha_{t}\bm{\mu}_{\text{fake}}(\bm{x}_{t},t )}{\sigma_{t}^{2}}.$ (3) $$

The “fake” denoising model is trained with the denoising objective with weighting $\lambda_{t}$:

$$ $\mathcal{L}_{\text{denoise}}(\phi)=\mathbb{E}_{\bm{z},t,\bm{x}_{t}}[\lambda_{t }\|\bm{\mu}_{\text{fake}}^{\phi}(\bm{x}_{t},t)-\bm{x}\|_{2}^{2}].$ (4) $$

The generator and the “fake” denoising model are updated alternatively. To facilitate better convergence of the KL divergence, Distribution Matching Distillation (DMD) and DMD2 [] used two distinct strategies, both significantly improving the generation performance. DMD proposes to complement the KL loss with a regression loss to encourage mode covering:

$$ $\mathcal{L}_{\text{reg}}(\theta)=\mathbb{E}_{(\bm{z},y)\sim\mathcal{D}_{\text{ paired}}}\ell(G_{\theta}(\bm{z}),y),$ (5) $$

where $\mathcal{D}_{\text{paired}}$ is a dataset of latent-image pairs generated by the teacher model offline, and $\ell$ is the Learned Perceptual Image Patch Similarity (LPIPS) [].
DMD2 instead applies a two-timescale update rule (TTUR), where they update the “fake” score model for $N$ steps per generator update, allowing more stable convergence. We use distribution matching (DM) to refer to all relevant techniques introduced in this section.

### 3.3 Enhancing distillation quality with adversarial loss

The adversarial loss, originally proposed by Goodfellow et al. [], has shown a remarkable capability in diffusion distillation to enhance sharpness and realism in generated images, thus improving generation quality. Specifically, DMD2 [] proposes adding a minimal discriminator head to the bottleneck layer of the “fake” denoising model $\mu_{\text{fake}}$, which is naturally compatible with DMD’s alternating training scheme and the TTUR. Moreover, they showed that one should first train the model without GAN to convergence, then add the GAN loss and continue training. This yields better terminal performance than training with the GAN loss from the beginning.
We use adversarial distribution matching (ADM) to refer to distribution matching with added adversarial loss.

## 4 Method

In Sec. 4.1, we introduce the general Multi-Student Distillation (MSD) framework. In Sec. 4.2, we show how MSD is applied to distribution matching and adversarial distillation. In Sec. 4.3, we introduce an additional training stage enabling distilling into smaller students.

Figure: Figure 3: Three-stage training scheme in Eq. 9. Acronym meanings: TSM: teacher score matching (Eq. 8 & Eq. 9); DM: distribution matching (Eq. 9 & Sec. 3.2); ADM: adversarial distribution matching (Eq. 9 and Sec. 3.3). Stage 1 and Stage 2 are techniques from previous works that help with same-sized students; Stage 0 is our contribution, which is required for smaller students who cannot initialize with teacher weights.
Refer to caption: x5.png

### 4.1 Distilling into multiple students

We present Multi-Student Distillation (MSD), a general drop-in framework to be combined with any conditional single-step diffusion distillation method that enables a cheap upgrade of model capacity without impairing the inference speed. We first identify the key components of a single-step diffusion distillation framework and then present the modification of MSD.

In the vanilla one-student distillation, we have a pretrained teacher denoising diffusion model $\bm{\mu}_{\text{teacher}}$, a training dataset $\mathcal{D}$, and a distillation method.
The distillation yields a single-step generator $G(\bm{z};y\in\mathcal{Y})$ via $G=\text{Distill}(\bm{\mu}_{\text{teacher}};\mathcal{D})$. The obtained generator $G$ maps a random latent $\bm{z}$ and an input condition $y$ into an image.
In comparison, in an MSD scheme, we instead distill the teacher into $K$ different one-step generators $\{G_{k}(\bm{z};y\in\mathcal{Y}_{k})\}_{k=1}^{K}$ via

$$ $G_{k}=\text{Distill}(\mu_{\text{teacher}},\mathcal{D}_{k}=F(\mathcal{D}, \mathcal{Y}_{k})),\ \ k=1,...,K$ (6) $$

Specifically, each distilled student $G_{k}$ is specialized in handling a partitioned subset $\mathcal{Y}_{k}$ of the whole input condition set $\mathcal{Y}$. So, it is trained on a subset of the training data $\mathcal{D}_{k}\subset\mathcal{D}$, determined by $\mathcal{Y}_{k}$ via a filtering function $F$. Fig. 1 illustrates this idea.

The partition of $\mathcal{Y}$ into $\{\mathcal{Y}_{k}\}_{k=1}^{K}$ determines the input condition groups for each student. As a starting point, we make the following three simplifications for choosing a partition:

- •
*Disjointness*: This prevents potential redundant training and redundant usage of model capacity.
- •
*Equal size*: Since students have the same architecture, the partitions $\{\mathcal{Y}_{k}\}_{k=1}^{K}$ should be of equal size that require similar model capacity.
- •
*Clustering*: Conditions within each partition should be more semantically similar than those in other partitions, so networks require less capacity to achieve a set quality on their partition.

The first two conditions can be easily satisfied in practice, while the third is not straightforward. For a class-conditional generation, partitioning by semantically similar and equal-sized classes serves a straightforward strategy, though extending it to text-conditional generation is nontrivial. Another promising strategy uses pretrained embedding layers such as the CLIP [] embedding layer or the teacher embedding layer.
One could find embeddings of the input conditions and then perform clustering on those embeddings, which are fixed-length numerical vectors containing implicit semantic information. We ablate partition strategies in Sec. 5.4.

The data filtering function $F$ determines the training subset data $\mathcal{D}_{k}$ from $\mathcal{Y}_{k}$. For example, a vanilla filtering strategy could set $F(\mathcal{D},\mathcal{Y}_{k})=\mathcal{D}_{k}:=\mathcal{D}|\mathcal{Y}_{k}$, where $\mathcal{D}|\mathcal{Y}_{k}$ denotes the subset of the training dataset $\mathcal{D}$ that contains the condition $\mathcal{Y}_{k}$. Empirically, we found that this filtering works in most cases, although sometimes a different approach is justified, as demonstrated in Sec. 4.2.

### 4.2 MSD with distribution matching

As a concrete example, we demonstrate the MSD framework using distribution matching (DM) and adversarial distillation techniques. Inspired by the two-stage framework in [], each of our students is trained with a distribution matching scheme at the first stage and finetuned with an additional adversarial loss at the second stage (adversarial distribution matching, or ADM):

$$ $\displaystyle G_{k}^{(1)}$ $\displaystyle=\text{Distill}_{\text{DM}}\left(\bm{\mu}_{\text{teacher}},F_{ \text{DM}}(\mathcal{D}_{\text{DM}},\mathcal{Y}_{k})\right),\ \ k=1,...,K,$ (7) $\displaystyle G_{k}^{(2)}$ $\displaystyle=\text{Distill}_{\text{ADM}}\left(\bm{\mu}_{\text{teacher}},F_{ \text{ADM}}(\mathcal{D}_{\text{ADM}},\mathcal{Y}_{k});G_{k}^{(1)}\right),\ \ k =1,...,K,$ $$

where we recall that $\bm{\mu}_{\text{teacher}}$ is the teacher diffusion model, $G_{k}^{(i)}$ is the $k$-th student generator at the $i$-th stage, $F$ is the data filtering function, $\mathcal{D}$ is the training data, and $\mathcal{Y}_{k}$ is the set of labels that student $k$ is responsible of. The first stage $\text{Distill}_{\text{DM}}$ uses distribution matching with either a complemented regression loss or the TTUR, with details in Sec. 3.2. These two methods achieve optimal training efficiency among other best-performing single-step distillation methods [, , ] without an adversarial loss, with a detailed comparison in App. A.3. The second stage $\text{Distill}_{\text{ADM}}$ adds an additional adversarial loss (details in Sec. 3.3), by introducing a small prediction head on the fake score model. This introduces minimal additional computational overhead and allows resuming from the first stage checkpoint.

##### Designing the training data

From Sec. 3, the data required for DM and ADM are $\mathcal{D}_{\text{DM}}=(\mathcal{D}_{\text{paired}},\mathcal{C})$ and $\mathcal{D}_{\text{ADM}}=(\mathcal{D}_{\text{real}},\mathcal{C})$, where $\mathcal{D}_{\text{paired}},\mathcal{D}_{\text{real}},\mathcal{C}$ represents generated paired data, real data and separate conditional input, respectively. We now discuss choices for the filtering function.

For the first stage data filtering $F_{\text{DM}}$, we propose $F_{\text{DM}}(\mathcal{D}_{\text{DM}},\mathcal{Y}_{k})=(\mathcal{D}_{\text{
paired}},\mathcal{C}|\mathcal{Y}_{k})$, where $\mathcal{C}|\mathcal{Y}_{k}$ denotes the subset of condition inputs $\mathcal{C}$ that contains $\mathcal{Y}_{k}$. In other words, we sample all input conditions only on the desired partition for the KL loss but use the whole paired dataset for the regression loss. This special filtering is based on the observation that the size of $\mathcal{D}_{\text{paired}}$ critically affects the terminal performance of DMD distillation: using fewer pairs causes mode collapse, whereas using more pairs challenges the model capacity. Naïvely filtering paired datasets by partition reduces the paired dataset size for each student and leads to worse performance, as in our ablation in App. A.2. Instead of generating more paired data to mitigate this imbalance, we simply reuse the original paired dataset for the regression loss. This is remarkably effective, which we hypothesize is because paired data from other input conditions provides effective gradient updates to the shared weights in the network.

For the second stage, we stick to the simple data filtering $F_{\text{ADM}}(\mathcal{D}_{\text{ADM}},\mathcal{Y}_{k})=(\mathcal{D}_{\text{
real}}|\mathcal{Y}_{k},\mathcal{C}|\mathcal{Y}_{k})$, so that both adversarial and KL losses focus on the corresponding partition, given that each student has enough mode coverage from the first stage.

### 4.3 Distilling smaller students from scratch

Via the frameworks presented in the last two sections, MSD enables a performance upgrade over alternatives for one student with the same model architecture. In this section, we investigate training multiple students with smaller architectures – and thus faster inference time – without impairing much performance.
However, this requires distilling into a student with a different architecture, preventing initialization from pretrained teacher weights. Training a single-step student from scratch has previously been difficult [], and we could not obtain competitive results with the simple pipeline in Eq. 7. Therefore, we propose an additional pretraining phase $\text{Distill}_{\text{TSM}}$, with TSM denoting Teacher Score Matching, to find a good initialization for single-step distillation. TSM employs the following score-matching loss:

$$ $\mathcal{L}_{\text{TSM}}=\mathbb{E}_{t}[\lambda_{t}\|\bm{\mu}_{\text{TSM}}^{ \varphi}(\bm{x}_{t},t)-\bm{\mu}_{\text{teacher}}(\bm{x}_{t},t)\|_{2}^{2}],$ (8) $$

where the smaller student with weights $\varphi$ is trained to match the teacher’s score on real images at different noise levels. This step provides useful initialization weights for single-step distillation and is crucial to ensure convergence. With TSM added, the whole pipeline now becomes:

$$ $\displaystyle\bm{\mu}^{(0)}$ $\displaystyle=\text{Distill}_{\text{TSM}}\left(\bm{\mu}_{\text{teacher}}, \mathcal{D}_{\text{real}}\right),$ (9) $\displaystyle G_{k}^{(1)}$ $\displaystyle=\text{Distill}_{\text{DM}}\left(\bm{\mu}_{\text{teacher}},F_{ \text{DM}}(\mathcal{D}_{\text{DM}},\mathcal{Y}_{k});\bm{\mu}^{(0)}\right),\ \ k=1,...,K,$ $\displaystyle G_{k}^{(2)}$ $\displaystyle=\text{Distill}_{\text{ADM}}\left(\bm{\mu}_{\text{teacher}},F_{ \text{ADM}}(\mathcal{D}_{\text{ADM}},\mathcal{Y}_{k});G_{k}^{(1)}\right),\ \ k =1,...,K.$ $$

Although a smaller student may not perfectly match the teacher’s score, it still provides a good initialization for stages 1 and 2. The performance gap is remedied in the latter stages by focusing on a smaller partition for each student. This three-stage training scheme is illustrated in Fig. 3.

## 5 Experiments

To evaluate the effectiveness of our approach, we trained MSD with different design choices and compared against competing methods, including other single-step distillation methods.

In Sec. 5.1, we compare single vs multiple students on a 2D toy problem for direct visual comparison. For these experiments, we used only the DM stage. In Sec. 5.2, we investigate class-conditional image generation on ImageNet-64$\times$64 [] where we have naturally defined classes to partition. Here, we explored training with the DM stage only, with both DM and ADM stages, and with all three stages for smaller students. We then evaluate MSD for a larger model in Sec. 5.3. We explored text-to-image generation on MS-COCO2014 [] with varying training stages.
We use the standard Fréchet Inception Distance (FID) [] score to measure generation quality. Comprehensive comparisons confirm that MSD outperforms single-student counterparts and achieves state-of-the-art performance in single-step diffusion distillation.
Finally, in Sec. 5.4, we summarize our ablation experiments over design choices.

To focus on the performance boost from multi-student distillation, we applied minimal changes to the hyperparameters used by Yin et al. [, ] for their distribution matching distillation implementations. More details on training and evaluation can be found in the App. B and C.

Figure: Figure 4: A 2D toy model. From left to right: teacher (multi-step) generation and student, one-step generation with $1$ and $8$ distilled students, the $\ell_{1}$ distance of generated samples between teacher and students. Takeaway: More students improve distillation quality on this easy-to-visualize setup.
Refer to caption: x6.png

### 5.1 Toy Experiments

In Fig. 4, we show the sample density of MSD with DM stage for a 2D toy experiment, where the real data distribution has $8$ classes, and each class is a mixture of $8$ Gaussians.
We used a simple MLP with EDM schedules to train the teacher and then distill into $1$, $2$, $4$, and $8$ students for comparison.
From the displayed samples and the $\ell_{1}$ distance from teacher generation, we observe that the collective generation quality increases as the number of students increases.

### 5.2 Class-conditional Image Generation

Student architecture the same as the teacher:
We trained $K=4$ students using the MSD framework and the EDM [] teacher on class-conditional ImageNet-64$\times$64 generation. We applied the simplest strategy for splitting classes among students: Each student is responsible for $250$ consecutive classes in numerical order. We compare the performance with previous methods and display the results in Tab. 2. Our DM stage, which uses the complementary regression loss, surpasses the one-student counterpart DMD [], achieving a modest drop of $0.25$ in FID score, making it a strong competitor in single-step distillation without an adversarial loss. We then took the best pretrained checkpoints and trained with the ADM stage. The resulting model achieved the current state-of-the-art FID score of $1.20$. It surpasses even the EDM teacher, StyleGAN-XL [], the multi-step RIN [] due to the adversarial loss. Fig. 5(a) and (b) display a comparison of sample generations, showing that our best students have comparable generation quality as the teacher.

**Table 1: Comparing class-conditional generators on ImageNet-64$\times$64. We measure generation speed with the number of function evaluations (NFE) and quality with FID.**
| Method | NFE ($\downarrow$) | FID ($\downarrow$) |
| --- | --- | --- |
| Multiple Steps |  |  |
| RIN [] | $1000$ | $1.23$ |
| ADM [] | $250$ | $2.07$ |
| DPM Solver [] | $10$ | $7.93$ |
| Multistep CD [] | $2$ | $2.0$ |
| Single Step, w/o GAN |  |  |
| PD [] | $1$ | $15.39$ |
| DSNO [] | $1$ | $7.83$ |
| Diff-Instruct [] | $1$ | $5.57$ |
| iCT-deep [] | $1$ | $3.25$ |
| Moment Matching [] | $1$ | $3.0$ |
| DMD [] | $1$ | $2.62$ |
| MSD (ours): 4 students, DM only | $1$ | $2.37$ |
| EMD [] | $1$ | $2.20$ |
| SiD [] | $1$ | $1.52$ |
| Single Step, w/ GAN |  |  |
| MSD (ours): 4, 42% smaller students, ADM | $1$ | $2.88$ |
| StyleGAN-XL [] | $1$ | $1.52$ |
| CTM [] | $1$ | $1.92$ |
| DMD2 [] | $1$ | $1.28$ |
| MSD (ours): 4 students, ADM | $1$ | $1.20$ |
| teacher |  |  |
| EDM (teacher, ODE) [] | $511$ | $2.32$ |
| EDM (teacher, SDE) [] | $511$ | $1.36$ |

Student architecture smaller than the teacher:
Next, we trained $4$ smaller student models with the prepended teacher score matching (TSM) stage from Sec. 4.3. This achieved a $42\%$ reduction in model size and a $7\%$ reduction in latency, with a slight degradation in FID score, offering a flexible framework to increase generation speed by reducing student size, and increasing generation quality by training more students.
Fig. 5(c) displays sample generation from these smaller students, whereas Fig. 5(d) shows sample generations from an even smaller set of students, with a $71\%$ percent reduction in model size and a $23\%$ percent reduction in latency. We observed slightly degraded but still competitive generation qualities. Using more and larger students will further boost performance, as shown by ablations in Sec. 5.4 and App. A.4. Smaller students without the TSM stage fail to reach competitive performance, indicating its necessity.

Figure: Figure 5: Sample generations on ImageNet-64$\times$64 from the teacher and different sized students, with architecture and latency details in App. B. The same-size students have comparable or better generation quality than the teacher. Smaller students achieve faster generation while still having decent qualities. Same-sized students are trained with DM and ADM stages, whereas smaller students are trained with all three stages (see Fig. 3). See additional samples in Fig. 9 and Fig. 10.
Refer to caption: x13.png

### 5.3 Text-to-Image Generation

Student architecture the same as the teacher: We evaluated the performance of text-to-image generation using the MS-COCO2014 [] evaluation dataset. We distilled $4$ students from Stable Diffusion (SD) v1.5 [] on a 5M-image subset of the COYO dataset []. For splitting prompts among students, we again employed a minimalist design: pass the prompts through the pre-trained SD v1.5 text encoder, pool the embeddings over the temporal dimension, and divide into $4$ subsets along $4$ quadrants. We trained with a classifier-free guidance (CFG) scale of $1.75$ for best FID performance. Tab. 2 compares the evaluation results with previous methods. Our baseline method with only the DM stage again achieved a performance boost with a $0.48$ drop in FID over the single-student counterpart DMD2 without adversarial loss []. Continuing the ADM stage from the best checkpoint yielded a terminal FID of 8.20, again surpassing the single-student counterpart and achieving the current state-of-the-art FID score. In addition, for better visual quality, we also train with a larger CFG scale of $8$, and display corresponding samples in Fig. 2(b) and App. D.2.

Student architecture smaller than the teacher: As a preliminary exploration, with the prepended teacher score matching (TSM) stage, we train a $83\%$ smaller and $5\%$ faster student on a dog-related prompt subset of COYO (containing $\sim$1\,210\,000$$ prompts). We trained with a CFG scale of $8$ and display the samples in Fig. 2. We observed fair generation quality despite a significant drop in model size. Improved training is likely to obtain better sample quality and generalization power.

### 5.4 Ablation Studies

Here, we ablate the effect of different components in MSD and offer insight into scaling. Unless otherwise mentioned, all experiments are conducted for class-conditional generation ImageNet-64$\times$64, using only the DM stage for computational efficiency.

**Table 3: Ablation studies on different components of MSD. All experiments are done on ImageNet-64$\times$64, trained with only the DM stage for $20$k iterations, where $B$ is the batch size per student. See App. A.1.**
| Method | FID ($\downarrow$) |
| --- | --- |
| $4$ students, $B=32$ | $2.53$ |
| $1$ students, $B=128$ | $2.60$ |
| $2$ students, $B=128$ | $2.49$ |
| $\bf{4}$ students, $\bf{B=128}$ (baseline) | $2.37$ |
| $8$ students, $B=128$ | $2.32$ |
| $4$ students, $B=128$, $K$-means splitting | $2.39$ |

MSD is still better with the same effective batch size.
To investigate if the performance boost from MSD comes from only a batch size increase over single student distillation, we make a comparison with the same effective batch size. As showcased in Tab. 3, MSD with $4$ students and a batch size of $32$ per student performs slightly better than the single-student counterpart with a batch size of $128$, indicating that MSD likely benefits from a capacity increase than a batch size increase.
As a takeaway, with a fixed training resource measured in processed data points, users are better off distilling into multiple students with partitioned resources each than using all resources to distill into a single student. Although multiple students means multiple model weights to save, storage is often cheap, so in many applications, this cost is outweighed by our improved quality or latency.

Simple splitting works surprisingly well. We used consecutive splitting of classes in Sec.  5.2. In the pretrained EDM model, class labels are fed to the diffusion model via the label embedding layer. Therefore, we investigated another strategy where we performed a $K$-means clustering ($K=4$) on the label embeddings, resulting in $4$ clusters of similar sizes: $(230,283,280,207)$. However, MSD trained with these clustered partitions performs similarly to sequential partition, as shown in Tab. 3. For text-to-image generation, we performed $K$-means clustering on the pooled embeddings of prompts in the training data, resulting in clusters of vastly uneven sizes. Due to computational limitations, we opted for the simpler partition strategies outlined in Sec. 5.

Effect of scaling the number of students. In Tab. 3, we study the effect of increasing $K$, the number of distilled students. We kept the per-student batch size fixed so more students induce a larger effective batch size. We observe better FID scores for more students. We hypothesize that better training strategies, such as per-student tuning, will further improve the quality. Optimal strategies for scaling to ultra-large numbers of students is an interesting area for future work.

## 6 Discussion

### 6.1 Limitations

MSD is the first work to explore diffusion distillation with multiple students, and it admits a few limitations that call for future work.
1) Further explorations could offer more insights into optimal design choices for a target quality and latency on various datasets, such as the number of students, input condition size for each student, and other hyperparameters. This is especially beneficial if the training budget is limited. 2) We apply simple partitioning for both class- and text-conditions and assign them disjointly to different students. Although our empirical study shows that simple alternatives do not offer obvious advantages, more sophisticated routing mechanisms may help. 3) We use simple channel reduction when designing smaller students to demonstrate feasibility. This results in a significantly smaller latency reduction than sample size reduction.
Exploring other designs of smaller students will likely increase their quality and throughput.
4) We train different students separately, but we expect that carefully designed weight-sharing, loss-sharing, or other interaction schemes can further enhance training efficiency. 5) We hypothesize that MSD can be applied to other diffusion distillation methods and other modalities for similar benefits but leave this for future work.

### 6.2 Conclusion

This work presented Multi-Student Distillation, a simple yet efficient method to increase the effective model capacity for single-step diffusion distillation. We applied MSD to the distribution matching and adversarial distillation methods. We demonstrated their superior performance over single-student counterparts in class-conditional and text-to-image generations. Particularly, MSD with DMD2’s the two-stage training achieves state-of-the-art FID scores. Moreover, we successfully distilled smaller students from scratch, demonstrating MSD’s potential in further reducing the generation latency with multiple smaller student distillations. We envision building on MSD to enable generation in real-time, enabling many new use cases.

## Ethics Statements

Our work aims to improve the quality and speed of diffusion models, thus we may inherit ethics concerns from diffusion models and generative models in general. Potential risks include fabricating facts or profiles that could mislead public opinion, displaying biased information that may amplify social biases, and displacing creative jobs from artists and designers.

## Acknowledgements

We thank Amirmojtaba Sabour for their support with the 2D toy model and Marc Law for helpful discussions and feedback.

## Appendix A Additional ablation studies

### A.1 Training curves for Sec. 5.4

Figure: Figure 6: FID comparisons during training for ablations in Table 3.
Refer to caption: x17.png

Fig. 6 displays the training curves for the ablation studies shown in Tab. 3. The relative terminal performances are also reflected in the training process.

### A.2 The effect of paired dataset size on DMD

In Sec. 4.2, we mentioned the special filtering strategy for MSD at DM stage: instead of partitioning the paired dataset for corresponding classes, we choose to keep the same complete dataset for each student. Fig. 7 demonstrates that the alternative strategy discourages mode coverage and leads to a worse terminal performance.

Figure: Figure 7: Comparison of paired dataset filtering strategies for MSD4-DM. Partitioning the paired dataset for each student discourages mode coverage, which results in worse terminal performance. In comparison, keeping the same paired dataset for each student achieves better performance without impairing the convergence speed.
Refer to caption: x20.png

### A.3 Single-step distillation methods comparison

**Table 5: Comparison of various aspects of single-step distillation methods.**
| Method on ImageNet | Terminal FID | Iterated Images | Note |
| --- | --- | --- | --- |
| DMD [] (our reimplementation) | 2.54 | $\sim$ 130M |  |
| DMD2 (w/o GAN) [] | 2.61 | $\sim$ 110M |  |
| EMD [] | 2.20 | $\sim$ 600M |  |
| SiD [] | 1.52 | $\sim$ 1000M | Very unstable |
| CTM [] | $>$5 | $\sim$ 3M |  |

Table 5 justifies our choice of DMD/DMD2 as our first-stage training without adversarial loss. Competitor methods suffer from training instability (SiD, FID fluctuates more than 100%), larger training data size (EMD), or worse quality (CTM and other CM-based methods). DMD/DMD2, on the other hand, strike a good balance. We noticed DMD exhibits more stability for MSD on ImageNet, whereas DMD2 performs better for SD v1.5, which leads to our respective choices. As pointed out in App. B, we used a smaller-sized paired dataset ($10\,000$ images) than the original DMD paper ($25\,000$ images) for ImageNet, which significantly accelerated convergence without impairing the final performance. Moreover, as pointed out in Sec. 4.2, the same paired dataset can be used for all students, eliminating potential additional computation.

### A.4 More results on distilling into smaller students

Figure: Figure 8: The effect of student size, number of classes covered, and training stage on the performance of a single smaller student. XS, XXS, and B refer to model sizes, with precise definitions provided in Tab. 6. Also, see special evaluation details in App. C.
Refer to caption: x21.png

In Sec. 5.2, we trained MSD4-ADM on smaller students to demonstrate the tradeoff between generation quality and speed. Here, we make a more comprehensive ablation study on the interplay between student size, number of classes covered, and training stage, with results displayed in Fig. 8. We observe that generation quality increases with student size and decreases with more classes covered. MSD offers great flexibility for users to make these choices based on computational resources, generation quality, and inference speed requirements.

## Appendix B Implementation Details

### B.1 Toy Experiments

The real dataset is a mixture of Gaussians. The radius for the outer circle is 0.5, the radius for the 8 inner circles is 0.1, and the standard deviation for each Gaussian (smallest circle) is $0.005$. The teacher is trained with EDM noise schedule, where use $(\sigma_{\min},\sigma_{\max})=(0.002,80)$ and discretized the noise schedule into 1000 steps. We train the teacher for $100\,000$ iterations with AdamW [] optimizer, setting the learning rate at 1e-4, weight decay to 0.01, and beta parameters to (0.9, 0.999). For distillation, we first generated dataset of 1000 pairs, then used DMD [] to train 1, 2, 4, and 8 students, respectively, all for $200\,000$ iterations, with reduced learning rate at 1e-7. We only sample the first $750$ of the $1000$ steps for distillation.

Each subfigure of Fig. 4 is a histogram with 200 bins on $100\,000$ generated samples, using a custom colormap. The loss is the mean absolute difference of binned histogram values.

### B.2 ImageNet

#### B.2.1 Same-sized students

Our ImageNet experiments setup closely follows DMD [] and DMD2 [] papers. We distill our one-step generators using the EDM [] checkpoint “edm-imagenet-64x64-cond-adm”. We use $\sigma_{\min}=0.002$ and $\sigma_{\max}=80$ and discretize the noise schedules into 1000 bins. The weight $w_{t}$ in Eq. (2) is set to $\frac{\sigma_{t}^{2}}{\alpha_{t}}\frac{CS}{\|\bm{\mu}_{\text{teacher}}(\bm{x}_
{t},t)-\bm{x}\|_{1}}$ where $S$ is the number of spatial locations and $C$ is the number of channels, and the weight $\lambda_{t}$ in Eq. (4) is set to $(\sigma_{t}^{2}+0.5^{2})/(\sigma_{t}\cdot 0.5)^{2}$.

For the DM stage, we prepare a distillation dataset by generating $10\,000$ noise-image pairs using the deterministic Heun sampler (with $S_{\text{churn}}=0$ over 256 steps. We use the AdamW optimizer [] with learning rate 2e-6, weight decay 0.01, and beta parameters (0.9,0.999). We compute the LPIPS loss using a VGG backbone from the LPIPS library [], and we upscale the image to $224\times 224$ using bilinear upsampling. The regression loss weight is set to 0.25. We use mixed-precision training and a gradient clipping with an $\ell_{2}$ norm of 10. We partition the 1000 total classes into consecutive blocks of 250 classes and trained 4 specialized students using $\text{Distill}_{\text{DM}}{}$ and $F_{\text{DM}}{}$ defined in Sec. 7. Each student is trained on 4 A100 GPUs, with a total batch size of 128, for $200\,000$ iterations. This yields the MSD4-DM checkpoint in Tab. 2.

For the ADM stage, we attach a prediction head to the middle block of the fakescore model. The prediction head consists of a stack of $4\times 4$ convolutions with a stride of 2, group normalization, and SiLU activations. All feature maps are downsampled to $4\times 4$ resolution, followed by a single convolutional layer with a kernel size and stride of 4. The final output linear layer maps the given vector to a scalar predicted probability. We load the best generator checkpoint from DM stage, but re-initialize the fakescore and GAN classifier model from teacher weights, as we observed this leads to slightly better performance. We set the GAN generator loss weight to 3e-3 and the GAN discriminator loss weight to 1e-2, and reduce the learning rate to 5e-7. Each student is trained on 4 A100 GPUs, with a total batch size of 192, for $150\,000$ iterations. This yields the MSD4-ADM checkpoint in Tab. 2.

#### B.2.2 Smaller students

Following a minimalist design, we pick our smaller student’s architecture by changing hyperparameter values of the “edm-imagenet-64x64-cond-adm” checkpoint architecture. See details in Tab. 6.

**Table 6: Hyperparameter details for different sized student models of the ADM architecture. Unspecified hyperparameters remain the same as the teacher. Latency is measured on a single NVIDIA RTX 4090 GPU.**
| Model identifier | # channels | Channel multipliers | # Residual blocks | # parameters | latency |
| --- | --- | --- | --- | --- | --- |
| B (teacher) | 192 | [1,2,3,4] | 3 | 296M | 0.0271s |
| S | 160 | [1,2,2,4] | 3 | 173M | 0.0253s |
| XS | 128 | [1,2,2,4] | 2 | 86M | 0.0209s |
| XXS | 96 | [1,2,2,2] | 2 | 26M | 0.0192s |

For the MSD4-ADM-S checkpoint in Tab. 2, we train the TSM stage using the model architecture S with the continuous EDM noise schedule with $(P_{\text{mean}},P_{\text{std}})=(-1.2,1.2)$ and the weighting $\lambda_{t}=(\sigma_{t}^{2}+0.5^{2})/(\sigma_{t}\cdot 0.5)^{2}$. We use a learning rate of 1e-4. Each student is trained on 4 A100 GPUs, with a total batch size of 576, for $400\,000$ iterations. Then DM and ADM stages were trained using a total batch size of 160, following otherwise the same setup as Sec. B.2.1.

For the ablation study in A.4, we instead train a common TSM stage for all students for computational efficiency. We train this common stage using 16 A100 GPUs, with a total batch size of $3584$ and $4864$ for architecture XS and XXS, respectively. The DM and ADM stages are followed by specialized students with filtered data and 4 A100 GPUs each, with a total batch size of 224 and 256 for architecture XS and XXS, respectively, and using the same learning rate of 2e-6.

### B.3 SD v1.5

#### B.3.1 Same-sized students, CFG=1.75

Our SD v1.5 experiments setup closely follows DMD2 [] paper. We distill our one-step generators from the SD v1.5 [] model, using a classifier-free guidance scale of $1.75$ for the teacher model to obtain the best FID score. We use the first 5M prompts from the COYO dataset [], and the corresponding 5M images for the GAN discriminator. We apply the DDIM noise schedule with 1000 steps for sampling $t$. The weight $w_{t}$ in Eq. 2 is set to $\frac{\sigma_{t}^{2}}{\alpha_{t}}\frac{CS}{\|\bm{\mu}_{\text{teacher}}(\bm{x}_
{t},t)-\bm{x}\|_{1}}$ where $S$ is the number of spatial locations and $C$ is the number of channels, and the weight $\lambda_{t}$ in Eq. 4 is set to $\alpha_{t}^{2}/\sigma_{t}^{2}$.

For the DM stage, we use the AdamW optimizer [] with learning rate 1e-5, weight decay 0.01, and beta parameters (0.9,0.999). We use gradient checkpointing, mixed-precision training, and a gradient clipping with an $\ell_{2}$ norm of 10. We partition the prompts and corresponding images by the 4 quadrants formed by the first 2 entries of the embeddings, where the embeddings are pooled from the outputs of the SD v1.5 text embedding layers. We choose not to include a regression loss but instead use a TTUR, which updates the fakescore model 10 times per generator update. Each of the 4 students is trained on 32 A100 GPUs, with a total batch size of 1536, for $40\,000$ iterations. This yields the MSD4-DM checkpoint in Tab. 2.

For the ADM stage, we attach a prediction head that has the same architecture (though with a different input size) as Sec. B.2. We load both the best generator checkpoint and the corresponding fakescore checkpoint from DM stage. We set the GAN generator loss weight to 1e-3 and the GAN discriminator loss weight to 1e-2, and reduce the learning rate to 5e-7. Each student is trained on 32 A100 GPUs, with a total batch size of 1024, for $5000$ iterations. This yields the MSD4-ADM checkpoint in Tab. 2.

#### B.3.2 Same-sized students, CFG=8

The above CFG=1.75 setting yields sub-optimal image quality. Similar to previous works [, , ], we choose CFG=8 for enhanced image quality. Due to time and computational resource limitations, we only train with the ADM stage. Each of the $4$ students is trained on 32 A100 GPUs, with a learning rate of 1e-5 and a batch size of 1024 for both the fake and the real images, for $6000$ iterations. This yields the checkpoint that is used to generate Fig. 2 (b) and Fig .11. Longer training with the added DM stage can likely further improve the generation quality.

#### B.3.3 Smaller student, CFG=8

We again pick our smaller student’s architecture by changing the hyperparameter values of the SD v1.5 architecture. See details in Tab. 7.

**Table 7: Hyperparameter details for different sized student models of the SD v1.5 architecture. Only the diffusion model part is measured since the text encoder and the VAE remain frozen. Unspecified hyperparameters remain the same as the teacher. Latency is measured on a single NVIDIA RTX 4090 GPU.**
| Model identifier | # block_out_channels | # parameters | latency |
| --- | --- | --- | --- |
| B (teacher) | [320,640,1280,1280] | 860M | 0.041s |
| S | [160,320,320,640] | 142M | 0.039s |

To create a subset of dog-related data, we first selected $\sim$1\,210\,000$$ prompts in the COYO [] dataset whose embeddings are closest to “a dog.” We then created an equal number of noise-image pairs from the SD v1.5 teacher using these prompts with CFG=8. We train the TSM stage using the model architecture S with the 1000-step DDIM noise schedule and the weighting $\lambda_{t}=\alpha_{t}^{2}/\sigma_{t}^{2}$. We use a learning rate of 1e-4. We then continue to the DM stage with the paired regression loss, using a learning rate of 1e-5, and finally continue to the ADM stage using generated paired images as ”real” images with a learning rate of 5e-7. We use 16 A100 GPUs. The TSM stage is trained with a total batch size of $1536$ for $240\,000$ iterations. The DM stage is trained with a total batch size of $512$ for both the paired and the fake images for $20\,000$ iterations, and the ADM stage is trained with the same batch size for $6000$ iterations. This yields the checkpoint used to generate Fig. 2(c). Longer training and better tuning are again likely to improve the generation quality further.

## Appendix C Evaluation Details

For zero-shot COCO evaluation, we use the exact setup as GigaGAN [] and DMD2 []. Specifically, we generate $30\,000$ images using the prompts provided by DMD2 code. We downsample the generated images using PIL to $256\times 256$ Lanczos resizer. We then use the clean-fid [] to compute the FID score between generated images and $40\,504$ real images from the COCO 2014 validation dataset. For ImageNet, we generate $50\,000$ images and calculate FID using EDM [] evaluation code. When selecting the best checkpoints for partitioned students, the same procedure is applied only for prompts/classes within respective partitions. For the ablation study in Sec. A.4, $10\,000$ images are generated for only the first $10$ classes for an apple-to-apple comparison.

## Appendix D Additional qualitative results

### D.1 Additional Imagenet-64 × \times × 64 results

Figure: Figure 9: Collective one-step samples from $4$ same-sized students trained with MSD-ADM on ImageNet (FID=1.20).
Refer to caption: x22.png

Figure: Figure 10: Collective one-step samples from $4$ smaller students trained with MSD-ADM on ImageNet (FID=2.88).
Refer to caption: x23.png

In Fig. 9, we present more ImageNet-64$\times$64 qualitative results collectively generated by our $4$ same-sized students trained with MSD-ADM. In Fig. 10, we display corresponding generations from $4$ smaller students with architecture S (see Tab. 6).

### D.2 Additional Text-to-image Synthesis Results

Figure: Figure 11: Collective one-step samples from $4$ SD v1.5 students trained with MSD-ADM on COYO with CFG=8.
Refer to caption: x24.png

In Fig. 11, we present more text-to-image qualitative results collectively generated by our $4$ students trained on the COYO dataset with MSD-ADM. These students are trained with a teacher classifier-free guidance (CFG) scale of $1.75$, which yields sub-optimal visual qualities despite having a good FID score. Generation with better qualities can be obtained with a higher CFG scale.

## Appendix E Prompt details

### E.1 Prompts for Fig. 11

We use the following prompts for Fig. 11, from left to right, top to bottom:

- •
wise old man with a white beard in the enchanted and magical forest
- •
a high-resolution photo of an orange Porsche under sunshine
- •
Astronaut on a camel on mars
- •
a hot air balloon in shape of a heart. Grand Canyon
- •
transparent vacation pod at dramatic scottish lochside, concept prototype, ultra clear plastic material, editorial style photograph
- •
penguin standing on a sidewalk
- •
border collie surfing a small wave, with a mountain on background
- •
an underwater photo portrait of a beautiful fluffy white cat, hair floating. In a dynamic swimming pose. The sun rays filters through the water. High-angle shot. Shot on Fujifilm X
- •
3D animation cinematic style young caveman kid, in its natural environment
- •
robot with human body form, robot pieces, knolling, top of view, ultra realistic
- •
3D render baby parrot, Chibi, adorable big eyes. In a garden with butterflies, greenery, lush, whimsical and soft, magical, octane render, fairy dust
- •
a chimpanzee sitting on a wooden bench
- •
a capybara made of voxels sitting in a field
- •
a cat reading a newspaper
- •
a squirrell driving a toy car
- •
close-up photo of a unicorn in a forest, in a style of movie still

### E.2 Prompts for Fig. 2

We use the following prompts (same for all three models), from left to right, top to bottom:

- •
dog on a bed
- •
Your Puppy Your Dog
- •
Trained Happy Dog
- •
Very handsome dog.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Why is atleast CLIP score not reported on either COCO-2017 or 2014 which could be informative as FID has its deficiencies, could consider HPSv2 or other metrics too for completeness.

What is total training compute required for proposed method? How does it compare to previous methods which do not specialize to sub-sets of data?

To better justify and understand motivation of this work, it might be useful to consider pruning or smaller architecture of already distilled one-step model as a baseline or initialization in their work? How much of training compute can be exploited with better initialization compared to distilling from scratch, such analysis would better benefit community as it currently lacks novel insights to adopt broadly for practical applications too. 

Authors cite EM Distillation as justification to emphasize difficulty of training one-step model from scratch? While it is known from consistency distillation, rectified flow and other works too that training a one-step models is hard not sure why cite distillation method to justify training from scratch as this is also not focus of this work.
soundness: 3
strengths: Paper is easy to read and understand the setting, focusing on domain specific student (partition of dataset).
Different objectives and better initialization to perform distillation makes sense and resultant effectiveness is demonstrated empirically. 

Demonstrates finetuning with adversarial training further improves quality of distilled model, which makes sense.
summary: In this work authors propose a way to distill a pre-trained diffusion model into multiple student where each student is specialized for sub-domain or specific partition of data. To perform distillation authors propose different objectives and also consider smaller architectures for student guided by target from pre-trained diffusion model.
weaknesses: Currently this work lacks strong motivation or useful analysis. 
There are previous works like eDiff which specialize different diffusion models per timestep and also works exploring MoE for efficient inferen w.r.t efficiency as motivation more effective pruning, efficient architectures, caching across timesteps etc. have been proposed to achieve smaller models and/or lower latency. 

This work explores splitting student into multiple models w.r.t dataset, while that is practical this work does not provide any novel insights nor significant performance boost.In case of text to image with SD1.5 FID boost only marginal by 0.15 combining all 3 objectives and 4 sets of parameters instead of one, which asks for more memory, more complicated orchestration etc. 

While FID is evaluated, it is unclear how well MSD recovers marginal data distribution i.e., diversity of generation and resultant sampled/recovered distributions (posterior) w.r.t conditional i.e., something like LPIPS_Diversity and aggregated distribution Precision-Recall or other metrics. This helps understand if there is any feature collapse, mode collapse etc?

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: See Weaknesses.
soundness: 4
strengths: 1. The paper is well written and presented. I enjoyed reading the paper. Though MoE is not a new idea, using it for Distillation is new, further using it to accelerate inference is commendable.
2. The idea of using Multiple-Students for distillation for inference time-quality tradeoff is quite intuitive. Moreover, assigning a student to a subset of conditions is a smart choice to increase the capacity of overall model.
3. Authors solve the obvious problem with above choice - initialization from scratch - by introducing an additional TSM stage which gives a good initialization, allowing for further distillation stage.
4. The empirical results are quite strong and encouraging. The proposed method achieves SoTA FID on Imagenet 64x64. Further, it shows encouraging results on distilling StableDiffusion performing better than several one-step generation methods.
summary: This paper introduces a 'Multi-Student Diffusion Distillation' framework. The core idea behind the proposed method stems from Mixture-of-Experts. Particularly, the paper proposes to distill a pre-trained Diffusion model (Teacher model) into multiple  Student model, where each Student is responsible for learning of a subset of conditions. This effectively increases the model capacity by amortizing the set of conditions into smaller subsets where a smaller Student model is responsible for corresponding subset. There are few keypoints pertaining to the proposed method: (a) partitioning/filtering function to partition the set of conditions into subsets, some of the desired features of such function are described in Section 4.1; (b) distillation into multiple Student models, where each Student is responsible for a subset of conditioning variables; (c) support for smaller-sized Student model unlike previous methods which employ same-sized Student model; (d) a teacher score matching phase for smaller-sized Student networks for initialization and better training. The paper primarily deals with Distribution Matching Distillation (DMD) and its extension Adversarial Distribution Matching (ADM). The proposed method SoTA FID on ImageNet 64x64 for one-step generation. The paper is well written and presented. The idea is very intuitive, however, it is interesting to see it working in practice on models like StableDiffusion.
weaknesses: 1. The paper focuses exclusively on DMD (Distribution Matching Distillation) and its extension ADA, which limits the demonstration of the method's generality. While the authors acknowledge this limitation, can the authors demonstrate preliminary results with other distillation approaches, particularly Consistency Distillation [1-3], on simple datasets like Mixture-of-Gaussian. Such experiments would better establish MSD's generality beyond DMD/ADA.
2. There is insufficient clarity regarding the text condition partitioning process in the latent space of the text-encoder during inference.  As I understand, the authors partition the text conditions in latent space of text-encoder. In that case during inference, how is the appropriate Student model selected during inference? Specifically, given that text conditions are not naturally disjoint (unlike ImageNet-style datasets), could the authors provide details on how they determine which Student to use during inference for text-to-image generation? Do they use same text-encoder partitioning technique as in training, or is there a different mechanism?
3. The authors outline several desired properties for the partitioning function in Section 4.1, yet the implemented solution simply uses consecutive classes as partitions (validated in Section 5.4). Could you compare a random partitioning strategy with your current approach? This would be valuable to determine whether the specific partitioning method offers advantages over any balanced data division.
4. The central contribution of the paper is that 'it offers a flexible framework to increase generation speed by reducing student size, and increasing generation quality by training more students. This is seen in Table 3 as well. In fact, in Table 1, the authors show that the Students outperform the Teacher. Does this observation also hold for text-to-image SD models? 
5. Minor:
	1. The partition function notation $F(\cdot) = (\cdot, \cdot | \cdot)$ needs proper definition as it resembles conditional probability notation.
	2. The MSD results appear to use a Student of equal size to the Teacher. Please include results for smaller-sized Students (as used in Fig. 5c) or explain their omission.
	3. Just for clarity: In Table 1, a single Student is used for generation (the Student responsible for a particular prompt), that is why the NFE is 1, right? 


Overall, the paper has merit. Albeit a simple idea, using MoE for distillation and accelerating inference is commendable. However, I am worried about the scope of the paper (see Weaknesses). I would like the authors to address the questions/doubts listed above. I am resorting to score of 6, I am open to increase it once these comments are addressed.

[1] Song, Yang, et al. "Consistency models." arXiv preprint arXiv:2303.01469 (2023).

[2] Zheng, Jianbin, et al. "Trajectory consistency distillation." arXiv preprint arXiv:2402.19159 (2024).

[3] Luo, Simian, et al. "Latent consistency models: Synthesizing high-resolution images with few-step inference." arXiv preprint arXiv:2310.04378 (2023).


----------------------------------------

**Post Rebuttal**

I am satisfied with the author’s response and rebuttal, as they address my concerns. While the method may initially appear straightforward, the paper tackles subtle yet significant challenges, which leads me to support its acceptance.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Please refer to the Weaknesses and Questions.
soundness: 2
strengths: Overall, this paper is well-written and easy to follow, with relatively new comparison methods.
summary: This paper addresses the high computational cost associated with multistep inference in diffusion models by focusing on the speed-quality tradeoff in distillation. The authors propose a Multi-Student Distillation (MSD) framework to enhance both generation speed and output quality. In this framework, a teacher model is distilled into several single-step student models, each specialized for generating data under specific input conditions.
weaknesses: 1. The authors state in Line 257 that 'Conditions within each partition should be more semantically similar than those in other partitions, so networks require less capacity to achieve a set quality on their partition.' However, there are no experiments presented to support this claim. I believe that implementing this idea is challenging and will demand additional computational resources. I recommend including relevant experiments and source code to facilitate a comprehensive review.

2. The statement in Line 15 that 'the student model’s inference speed is limited by the size of the teacher architecture' is misleading, as the inference speed of the student model is independent of the teacher model; the student only depends on the teacher during the distillation training phase. I recommend proofreading the entire paper to ensure clarity and professionalism.

3. The proposed method introduces multiple student models; therefore, comparisons and analyses of the model parameters should be a focal point of the paper.

4. The proposed method leverages adversarial distillation with the expectation of enhancing the distillation effect. However, there is no comparison with standard distillation methods or other variants to validate the adversarial distillation’s anticipated advantages.

5. In the ablation studies section, only a quantitative analysis of the generation effect is presented. I believe that a qualitative analysis should also be included, as the paper aims to enhance generation quality.

6. I can not conduct a comprehensive review of the technological accuracy of this paper, as it is empirical rather than theoretical, and the implementation code is not provided.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: Does this method require more model storage for practical deployment? Do the authors have any solutions to address this issue?
soundness: 2
strengths: + The organization and writing of the paper are clear, making it easy to understand and follow, and the review of related work is thorough.
+ The discussion on data partitioning is valuable and aligns well with the positioning of the proposed method. I also suggest the authors conduct more comprehensive experimental validation on this aspect.
summary: The paper introduces Multi-Student Distillation (MSD), an approach for diffusion model distillation, which improves existing single-step methods by increasing effective model capacity without added inference latency. MSD uses multiple student models, each optimized for a subset of conditioning inputs, to generate samples in a single step. This framework enhances flexibility by supporting multiple smaller student models to reduce generation time and enables initialization without requiring teacher weights. The authors validate MSD through experiments, achieving improved FID scores on various benchmarks with reduced parameters and comparable generation quality.
weaknesses: - The technical novelty is limited, as the approach simply extends the distillation of a teacher diffusion model to multiple student diffusion models. The distillation methods used, DM and ASD, are existing techniques, so the improvement offered by this approach is marginal.
- The performance improvement is also marginal. As seen in Tables 1 and 2, the improvement of MSD over DMD2 is small, yet it requires more training and model resources. Although the paper explores smaller student models, Table 1 shows a significant drop in performance, which reduces the practical contribution of this work in terms of lowering inference costs.

Meta-review: This paper presents a method for single step diffusion model distillation by turning the student to a MoE style conditioned on the data partition. It has received mixed reviews -- reviewer igyA is overall positive about the contributions, however other three reviewers challenged its novelty, generality and effectiveness. I like the the simplicity aspect of the work, and I don't particularly agree that it should be heavily criticized for lack of novelty due to the simplicity, however I do agree that certain critiques are valid. Especially, the question regarding the generality of the method wrt distillation baselines and introduced complexities and parameter overhead are valid concerns. Based on these considerations, I think this work is not ready but I encourage the authors to keep improving it.

Decision: Reject