Improved Distribution Matching Distillation for Fast Image Synthesis

Authors: tianwei  yin
michaël  gharbi
taesung  park
richard  zhang
eli  shechtman
frédo  durand
william t freeman
Venue: NeurIPS_cc_2024_Conference
Published: 2024-04-30 20:30:56

Abstract: Recent approaches have shown promises distilling expensive diffusion models into efficient one-step generators.
Amongst them, Distribution Matching Distillation (DMD) produces one-step generators that match their teacher in distribution, i.e., the distillation process does not enforce a one-to-one correspondence with the sampling trajectories of their teachers.
However, to ensure stable training in practice, DMD requires an additional regression loss computed using a large set of noise--image pairs, generated by the teacher with many steps of a deterministic sampler.
This is not only computationally expensive for large-scale text-to-image synthesis, but it also limits the student's quality, tying it too closely to the teacher's original sampling paths.
We introduce DMD2, a set of techniques that lift this limitation and improve DMD training.
First, we eliminate the regression loss and the need for expensive dataset construction.
We show that the resulting instability is due to the "fake" critic not estimating the distribution 
of generated samples with sufficient accuracy and propose a two time-scale update rule as a remedy.
Second, we integrate a GAN loss into the distillation procedure, discriminating between generated samples and real images.
This lets us train the student model on real data, thus mitigating the imperfect "real" score estimation from the teacher model, and thereby enhancing quality.
Third, we introduce a new training procedure that enables multi-step sampling in the student, and
addresses the training--inference input mismatch of previous work, by simulating inference-time generator samples during training. 
Taken together, our improvements set new benchmarks in one-step image generation, with FID scores of 1.28 on ImageNet-64×64 and 8.35 on zero-shot COCO 2014, surpassing the original teacher despite a 500X reduction in inference cost.
Further, we show our approach can generate megapixel images by distilling SDXL, demonstrating exceptional visual quality among few-step methods, and surpassing the teacher. 
We release our code and pretrained models.

Main paper: 
 introduction: 
 Diffusion models have achieved unprecedented quality in visual generation tasks [1] [2] [3] [4] [5] [6] [7] [8] . But their sampling procedure typically requires dozens of iterative denoising steps, each of which is a forward pass through a neural network. This makes high resolution text-to-image synthesis slow and expensive. To address this issue, numerous distillation methods have been developed to convert a teacher diffusion model into an efficient, few-step student generator [9] [10] [11] [12] [13] [14] [15] [16] [17] [18] [19] [20] . However, they often result in degraded quality, as the student model is typically trained with a loss to learn the pairwise noise-to-image mapping of the teacher, but struggles to perfectly mimic its behavior. Nevertheless, it should be noted that loss functions aimed at matching distributions, such as the GAN [21] or the DMD [22] loss, are not burdened with the complexity of precisely learning the specific paths from noise to image because their goal is to align with the teacher model in terms of distribution-by minimizing either a Jensen-Shannon (JS) or an approximate Kullback-Leibler (KL) divergence between the student and teacher output distributions.In particular, DMD [22] has demonstrated state-of-the-art results in distilling Stable Diffusion 1.5, yet it remains less investigated than GAN-based methods [23] [24] [25] [26] [27] [28] [29] . A likely reason is that DMD still requires an additional regression loss to ensure stable training. In turn, this necessitates creating millions of noise-image pairs by running the full sampling steps of the teacher model, which is particularly costly for text-to-image synthesis. The regression loss also negates the key benefit of DMD's unpaired distribution matching objective, because it causes the student's quality to be upper-bounded by the teacher's.In this paper, we show how to do away with DMD's regression loss, without compromising training stability. We then push the limits of distribution matching by integrating the GAN framework into DMD, and enable few-steps sampling with a novel training procedure, which we termed 'backward simulation'. Taken together, our contributions lead to state-of-the-art fast generative models that outperform their teacher, using as few as 4 sampling steps. Our method, which we call DMD2, achieves state-of-the-art results in one-step image generation, setting a new benchmark with FID scores of 1.28 on ImageNet-64×64 and 8.35 on zero-shot COCO 2014. We demonstrate our approach's scalability by distilling from SDXL to produce high-quality megapixel images, establishing new standards among few-step methods.In short, our contributions are as follows:• We propose a new distribution matching distillation technique that does not require a regression loss for stable training, thereby eliminating the need for costly data collection, and allowing for more flexible and scalable training. • We show that training instability in DMD [22] without regression loss stems from an insufficiently trained fake diffusion critic, and implement a two time-scale update rule to address this issue. • We integrate a GAN objective into the DMD framework, where the discriminator is trained to distinguish samples from the student generator vs. real images. This additional supervision operates at the distribution level, which better aligns with DMD's distribution-matching philosophy than the original regression loss. It mitigates approximation errors in the teacher diffusion model and enhances image quality. • While the original DMD only supports one-step students, we introduce a technique to support multi-step generators. Unlike previous multi-step distillation methods, we avoid the domain mismatch between training and inference by simulating inference-time generator inputs during training, thus improving overall performance. 
related work: 
 Diffusion Distillation. Recent diffusion acceleration techniques have focused on speeding up the generation process through distillation [9, 10, 13-20, 22, 23, 30] . They typically train a generator to approximate the ordinary differential equation (ODE) sampling trajectory of a teacher model, in fewer sampling steps. Notably, Luhman et al. [16] precompute a dataset of noise and images pairs, generated by the teacher using an ODE sampler, and use it to train the student to regress the mapping in a single network evaluation. Follow-up works like Progressive Distillation [10, 13] eliminate the need to precompute this paired dataset offline. They iteratively train a sequence of student models, each halving the number of sampling steps of its predecessor. A complementary technique, Instaflow [11] straightens the ODE trajectories, so they are easier to approximate with a one-step student. Consistency Distillation [9, 12, 19, 26, 31, 32] , and TRACT [33] , train student models so their outputs are self-consistent at any timesteps along the ODE trajectory, and thus consistent with the teacher.GANs. Another line of research employs adversarial training to align the student with the teacher at a broader distribution level. In ADD [23] , the generator, initialized with weights from a diffusion model, is trained using a projected GAN objective with an image-space classifier [34] . Building on this, LADD [24] utilizes a pre-trained diffusion model as the discriminator and operates in latent space, thus improving scalability and enabling higher-resolution synthesis. Inspired by DiffusionGAN [28, 29] , UFOGen [25] introduces noise injection prior to the real vs. fake classification in the discriminator, to smooth out the distributions, which stabilizes the training dynamics. However, purely GANbased methods often struggle to integrate classifier-free guidance directly. For instance, LADD uses diffusion-generated images with classifier-free guidance as real data in its GAN discriminator. Other approaches combine adversarial objectives with a distillation loss to preserve the original guided sampling trajectory. For instance, SDXL-Lightning [27] integrates a DiffusionGAN loss [25] with a progressive distillation objective [10, 13] , while the Consistency Trajectory Model [26] combines a GAN [35] with an improved consistency distillation [9] . In contrast, our approach based on distribution matching [22, 36, 37] inherently integrates classifier-free guidance into the training supervision, significantly simplifying the training process.Score Distillation was initially introduced in the context of text-to-3D synthesis [37] [38] [39] [40] , utilizing a pre-trained text-to-image diffusion model as a distribution matching loss. These methods optimize a 3D object by aligning rendered views with a text-conditioned image distribution, using the scores predicted by a pretrained diffusion model. Recent works have extended score distillation [37, 38, [41] [42] [43] to diffusion distillation [22, 30, 36, [44] [45] [46] . Notably, DMD [22] minimizes an approximate KL divergence, with its gradient represented as the difference between two score functions: one, fixed and pretrained, for the target distribution and another, trained dynamically, for the output distribution of the generator.DMD parameterizes both score functions using diffusion models. This training objective proved more stable than GAN-based methods and has demonstrated superior performance in one-step image synthesis. An important caveat, DMD requires a regression loss for stability, calculated using precomputed noise-image pairs, similar to Luhman et al. [16] . Our work does away with this requirement. We introduce techniques to stabilize the DMD training procedure without the regression regularizer, thus significantly reducing the computational costs incurred by paired data precomputation. Furthermore, we extend DMD to support multi-step generation and integrate the strengths of both GANs and distribution matching approaches [22, 30, 36, 45] , leading to state-of-theart results in text-to-image synthesis. 
background: diffusion and distribution matching distillation: 
 This section gives a brief overview of diffusion models and distribution matching distillation (DMD).Diffusion Models generate images through iterative denoising. In the forward diffusion process, noise is progressively added to corrupt a sample x ∼ p real from the data distribution into pure Gaussian noise over a predetermined number of steps T , so that, at each timestep t, the diffused samples follow the distribution p real,t (x t ) = p real (x)q(x t |x)dx, with q t (x t |x) ∼ N (α t x, σ 2 t I), where α t , σ t > 0 are scalars determined by the noise schedule [47, 48] . The diffusion model learns to iteratively reverse the corruption process by predicting a denoised estimate µ(x t , t), conditioned on the current noisy sample x t and the timestep t, ultimately leading to an image from the data distribution p real . After training, the denoised estimate relates to the gradient of the data likelihood function, or score function [48] of the diffused distribution:EQUATIONSampling an image typically requires dozens to hundreds of denoising steps [49] [50] [51] [52] .Distribution Matching Distillation (DMD) distills a many-step diffusion models into a one-step generator G [22] by minimizing the expectation over t of approximate Kullback-Liebler (KL) divergences between the diffused target distribution p real,t and the diffused generator output distribution p fake,t . Since DMD trains G by gradient descent, it only requires the gradient of this loss, which can be computed as the difference of 2 score functions:EQUATIONwhere z ∼ N (0, I) is a random Gaussian noise input, θ are the generator parameters, F is the forward diffusion process (i.e., noise injection) with noise level corresponding to time step t, and s real and s fake are scores approximated using diffusion models µ real and µ fake trained on their respective distributions (Eq. ( 1)). DMD uses a frozen pre-trained diffusion model as µ real (the teacher), and dynamically updates µ fake while training G, using a denoising score-matching loss on samples from the one-step generator, i.e., fake data [22, 47] .Yin et al. [22] found that an additional regression term [16] was needed to regularize the distribution matching gradient (Eq. ( 2)) and achieve high-quality one-step models. For this, they collect a dataset of noise-image pairs (z, y) where the image y is generated using the teacher diffusion model, and a deterministic sampler [49, 50, 53] , starting from the noise map z. Given the same input noise z, the regression loss compares the generator output with the teacher's prediction:EQUATIONwhere d is a distance function, such as LPIPS [54] in their implementation. While gathering this data incurs negligible cost for small datasets like CIFAR-10, it becomes a significant bottleneck with large-scale text-to-image synthesis tasks, or models with complex conditioning [55] [56] [57] . For instance, generating one noise-image pair for SDXL [58] takes around 5 seconds, amounting to about 700 A100 days to cover the 12 million prompts in the LAION 6.0 dataset [59] , as utilized by Yin et al. [22] . This dataset construction cost alone is already more than 4× our total training compute (as detailed in Appendix I). This regularization objective is also at odds with DMD's goal of matching the student and teacher in distribution, since it encourages adherence to the teacher's sampling paths. 
improved distribution matching distillation: 
 We revisit multiple design choices in the DMD algorithm [22] and identify significant improvements. Our training alternates between 2 steps: 1. optimizing the generator using the gradient of an implicit distribution matching objective (red arrow) and a GAN loss (green), and 2. training a score function (blue) to model the distribution of "fake" samples produced by the generator, as well as a GAN discriminator (green) to discriminate between fake samples and real images. The student generator can be a one-step or a multi-step model, as shown here, with an intermediate step input.The regression loss [16] used in DMD [22] ensures mode coverage and training stability, but as we discussed in Section 3, it makes large-scale distillation cumbersome, and is at odds with the distribution matching idea, thus inherently limiting the performance of the distilled generator to that of the teacher model. Our first improvement is to remove this loss.Naively omitting the regression objective, shown in Eq. (3), from DMD leads to training instabilities and significantly degrades quality (Tab. 3). For example, we observed that the average brightness, along with other statistics, of generated samples fluctuates significantly, without converging to a stable point (See Appendix F). We attribute this instability to approximation errors in the fake diffusion model µ fake , which does not track the fake score accurately, since it is dynamically optimized on the non-stationary output distribution of the generator. This causes approximation errors and biased generator gradients (as also discussed in [30] ).We address this using the two time-scale update rule inspired by Heusel et al. [60] . Specifically, we train µ fake and the generator G at different frequencies to ensure that µ fake accurately tracks the generator's output distribution. We find that using 5 fake score updates per generator update, without the regression loss, provides good stability and matches the quality of the original DMD on ImageNet (Tab. 3) while achieving much faster convergence. Further analysis are included in Appendix F.Our model so far achieves comparable training stability and performance to DMD [22] without the need for costly dataset construction (Tab. 3). However, a performance gap remains between the distilled generator and the teacher diffusion model. We hypothesize this gap could be attributed to approximation errors in the real score function µ real used in DMD, which would propagate to the generator and lead to suboptimal results. Since DMD's distilled model is never trained with real data, it cannot recover from these errors.We address this issue by incorporating an additional GAN objective into our pipeline, where the discriminator is trained to distinguish between real images and images produced by our generator.Trained using real data, the GAN classifier does not suffer from the teacher's limitation, potentially allowing our student generator to surpass it in sample quality. Our integration of a GAN classifier into DMD follows a minimalist design: we add a classification branch on top of the bottleneck of the fake diffusion denoiser (see Fig. 3 ). The classification branch and upstream encoder features in the UNet are trained by maximizing the standard non-saturing GAN objective:EQUATIONwhere D is the discriminator, and F is the forward diffusion process (i.e., noise injection) defined in Section 3, with noise level corresponding to time step t. The generator G minimizes this objective.Our design is inspired by prior works that use diffusion models as discriminators [24, 25, 27] . We note that this GAN objective is more consistent with the distribution matching philosophy since it does not require paired data, and is independent of the teacher's sampling trajectories.With the proposed improvements, we are able to match the performance of teacher diffusion models on ImageNet and COCO (see Tab. 1 and Tab. 5). However, we found that larger scale models like SDXL [58] remain challenging to distill into a one-step generator because of limited model capacity and a complex optimization landscape to learn the direct mapping from noise to highly diverse and detailed images. This motivated us to extend DMD to support multi-step sampling.We fix a predetermined schedule with N timestep {t 1 , t 2 , . . . t N }, identical during training and inference. During inference, at each step, we alternate between denoising and noise injection steps, following the consistency model [9] , to improve sample quality. Specifically, starting from Gaussian noise z 0 ∼ N (0, I), we alternate between denoising updates xti = G θ (x ti , t i ), and forward diffusion steps x ti+1 = α ti+1 xti + σ ti+1 ϵ with ϵ ∼ N (0, I), until we obtain our final image xt N . Our 4-step model uses the following schedule: 999, 749, 499, 249, for a teacher model trained with 1000 steps.Previous multi-step generators are typically trained to denoise noisy real images [23, 24, 27] . However, during inference, except for the first step, which starts from pure noise, the generator's input come from a previous generator sampling step xti . This creates a training-inference mismatch that adversely impacts quality (Fig. 4 ). We address this issue by replacing the noisy real images during training, with noisy synthetic images x ti produced by the current student generator running several steps, similar to our inference pipeline ( § 4.4). This is tractable because, unlike the teacher diffusion model, our generator only runs for a few steps. Our generator then denoises these simulated images and the outputs are supervised with the proposed loss functions. Using noisy synthetic images avoids the mismatch and improves overall performance (See Sec. 5.3). A concurrent work, Imagine Flash [61] , proposed a similar technique. Their backward distillation algorithm shares our motivation of reducing the training and testing gap by using the student-generated images as the input to the subsequent sampling steps at training time. However, they do not entirely resolve the mismatch issue, because the teacher model of the regression loss now suffers from the training-test gap: it is never trained with synthetic images. This error is accumulated along the sampling path. In contrast, our distribution matching loss is not dependent on the input to the student model, alleviating this issue.In summary, our distillation method lifts DMD [22] stringent requirements for precomputed noiseimage pairs. It further integrates the strength of GANs and supports multi-step generators. As shown in Fig. 3 , starting from a pretrained diffusion model, we alternate between optimizing the generator G θ to minimize the original distribution matching objective as well as a GAN objective, and optimizing the fake score estimator µ fake using both a denoising score matching objective on the fake data, and the GAN classification loss. To ensure the fake score estimate is accurate and stable, despite being optimized on-line, we update it with higher frequency than the generator (5 steps vs. 1). 
experiments: 
 We evaluate our approach, DMD2, using several benchmarks, including class-conditional image generation on ImageNet-64×64 [62] , and text-to-image synthesis on COCO 2014 [63] with various teacher models [1, 58] . We use the Fréchet Inception Distance (FID) [60] to measure image quality and diversity, and the CLIP Score [64] to evaluate text-to-image alignment. For SDXL models, we additionally report patch FID [27, 65] , which measures FID on 299x center-cropped patches of each image, to assess high-resolution details. Finally, we conduct human evaluations to compare our approach with other state-of-the-art methods. Comprehensive evaluations confirm that distilled models trained using our approach outperform previous work, and even rival the performance of the teacher models. Detailed training and evaluation procedures are provided in the appendix.Table 1 compares our model with recent baselines on ImageNet-64×64. With a single forward pass, our method significantly outperforms existing distillation techniques and even outperforms the teacher model using ODE sampler [53] . We attribute this remarkable performance to the removal of DMD's regression loss (Sec. 4.1 and 4.2), which eliminates the performance upper bound imposed by the ODE sampler, as well as our additional GAN term (Sec. 4.3), which mitigates the adverse impact of the teacher diffusion model's score approximation error. We evaluate DMD2's text-to-image generation performance on zero-shot COCO 2014 [63] . Our generators are trained by distilling SDXL [58] and SD v1.5 [1] , respectively, using a subset of 3 million prompts from LAION-Aesthetics [59] . Additionally, we collect 500k images from LAION-Aesthetic as training data for the GAN discriminator. Table 2 summarizes distillation results for the SDXL model. Our 4-step generator produces high quality and diverse samples, achieving a FID score of 19.32 and a CLIP score of 0.332, rivaling the teacher diffusion model for both image quality and prompt coherence. To further verify our method's effectiveness, we conduct an extensive user study comparing our model's output with those from the teacher model and existing distillation methods. We use a subset of 128 prompts from PartiPrompts [70] following LADD [24] .For each comparison, we ask a random set of five evaluators to choose the image that is more visually appealing, as well as the one that better represents the text prompt. Details about the human evaluation are included in Appendix K. As shown in Figure 5 , our model achieves much higher user preferences than baseline approaches. Notably, our model outperforms its teacher in image quality for 24% of samples and achieves comparable prompt alignment, while requiring 25× fewer forward passes (4 vs 100). Qualitative comparisons are shown in Figure 6 . Results for SDv1.5 are provided in Table 5 in Appendix D. Similarly, one-step model trained using DMD2 outperforms all previous diffusion acceleration approaches, achieving a FID score of 8.35, representing a significant 3.14-point improvement over the original DMD method [22] . Our results also surpass the teacher models that uses a 50-step PNDM sampler [50] .Figure 5 : User study comparing our distilled model with its teacher and competing distillation baselines [23, 27, 31] . All distilled models use 4 sampling steps, the teacher uses 50. Our model achieves the best performance for both image quality and prompt alignment. In Table 4 , we ablate the influence of the GAN term (Sec. 1 , Figure 2 , and Figure 11 A Limitations While achieving superior image quality and text alignment, our distilled generator experiences a slight degradation in image diversity compared to the teacher models (see Appendix E). Additionally, our generator still requires four steps to match the quality of the largest SDXL model. These limitations, while not unique to our model, highlight areas for further improvement. Like most previous distillation methods, we use a fixed guidance scale during training, limiting user flexibility. Introducing a variable guidance scale [13, 31] could be a promising direction for future research. Furthermore, our methods are optimized for distribution matching; incorporating human feedback or other reward functions could further enhance performance [17, 74] . Lastly, training large-scale generative models is computationally intensive, making it inaccessible for most researchers. We hope our efficient approach and optimized, user-friendly codebase will help democratize future research in this field.We are continually updating the paper and related materials. For access to our code, model, and the latest version of the paper with expanded results and analysis, please visit our project website. Our work on improving the efficiency and quality of diffusion model has several potential societal impacts, both positive and negative. On the positive side, the advancements in fast image synthesis can significantly benefit various creative industries. These models can enhance graphic design, animation, and digital art by providing artists with powerful tools to generate high-quality visuals efficiently. Additionally, improved text-to-image synthesis capabilities can be used in education and entertainment, enabling the creation of personalized learning materials and immersive experiences.However, potential negative societal impacts must be considered. Misuse risks include generating misinformation and creating fake profiles, which could spread false information and manipulate public opinion. Deploying these technologies could result in biases that unfairly impact specific groups, especially if models are trained on biased datasets, potentially perpetuating or amplifying existing societal biases. To mitigate these risks, we are interested in developing monitoring mechanisms to detect and prevent misuse [75, 76] and methods to enhance output diversity and fairness [77] .Table 5 presents detailed comparisons between our one-step generator distilled from SD v1.5 and competing approaches.Qualitative ablation results using SDXL backbone are shown in Figure 7 . Additionally, we compare the image diversity of our 4-step generator with other competing approaches distilled from SDXL [23, 27, 31] . We employ an LPIPS-based diversity score, similar to that used in multi-modal image-toimage translation [85, 86] . Specifically, we generate four images per prompt and calculate the average pairwise LPIPS distance [54] . For this evaluation, we use the LADD [24] subset of PartiPrompts [70] .We also report the FID and CLIP score measured on 10K prompts from COCO 2014 on the side. Table 6 summarizes the results. Table 7 provides further comparisons using image reward [87] and [59] . Our model achieves the best image quality, indicated by the lowest FID and Patch FID scores. We also achieve text alignment comparable to SDXL-Turbo while attaining a better diversity score. While SDXL-Lightning [27] exhibits a higher diversity score than our approach, it suffers from considerably worse text alignment, as reflected by the lower CLIP score and human evaluation (Fig. 5 ). This suggests that the improved diversity is partially due to random outputs lacking prompt coherence. We note that it is possible to increase the diversity of our model by raising the weights for the GAN objective, which aligns with the more diverse unguided distribution. Further investigation into finding the optimal balance between distribution matching and the GAN objective is left for future work.In Section 4.2, we discuss that updating the fake score multiple times (5 updates) per generator update leads to better stability. Here, we provide further analysis. Figure 8 visualizes pixel brightness variations throughout training. The baseline approach, which omits the regression objective from DMD and uses just 1 fake score update, results in significant training instability, as evidenced by periodic fluctuations in pixel brightness. In contrast, our two time-scale update rule with 5 fake score updates per generator update stabilizes the training and leads to better sample quality, as shown in Tab. 3.We further examine the influence of the update frequency for the fake diffusion model µ fake in Figure 9 . An update frequency of 1 fake diffusion update per generator update corresponds to the naive baseline (red line) and suffers from training instability. Although a frequency of 10 updates (magenta line) provides excellent stability, it significantly slows down the training process. We found that a moderate frequency of 5 updates (green line) achieves the best balance between stability and convergence speed on ImageNet. Our approach proves more effective than using asynchronous learning rates [60] (cyan line) and converges significantly faster than the original DMD method that employs a regression loss [22] (dark blue line). For new models and datasets, we recommend adjusting the iteration number to the smallest value that ensures the stability of general image statistics, such as pixel brightness.Additional visual comparisons for the 4-step distilled models are shown in Figure 10 . Sample outputs from our one-step generator are presented in Figure 11 .In Figure 12 , we present qualitative results obtained from our one-step distilled model trained on the ImageNet dataset. This section outlines key aspects of the implementation. All results are reproducible using our open-source training and evaluation code. Generally, we employ the maximum batch size supported by our compute resources. The learning rate is set to the highest stable value, ensuring no divergent loss within the initial 500 iterations. We determine the TTUR iteration count based on the minimum needed for stability, while the guidance scale is selected to optimize performance for the teacher model.Our GAN classifier design is inspired by SDXL-Lightning [27] . Specifically, we attach a prediction head to the middle block output of the fake diffusion model. The prediction head consists of a stack of 4 × 4 convolutions with a stride of 2, group normalization, and SiLU activations. All feature maps are downsampled to 4 × 4 resolution, followed by a single convolutional layer with a kernel size and stride of 4. This layer pools the feature maps into a single vector, which is then passed to a linear projection layer to predict the classification result.Our ImageNet implementation closely follows the DMD paper [22] . Specifically, we distill a one-step generator from the EDM pretrained model [53] , released under the CC BY-NC-SA 4.0 License. For the standard training setup, we use the AdamW optimizer [88] with a learning rate of 2 × 10 -6 , a weight decay of 0.01, and beta parameters (0.9, 0.999). We use a batch size of 280 and train the model on 7 A100 GPUs for 200K iterations, which takes approximately 2 days. The number of fake diffusion model update per generator update is set to 5. The weight for the GAN loss is set to 3 × 10 -3 . For the extended training setup shown in Table 1 , we first pretrain the model without GAN loss for 400K iterations. We then resume from the best checkpoint (as measured by FID), enable the GAN loss with a weight of 3 × 10 -3 , reduce the learning rate to 5 × 10 -7 , and continue training for an additional 150K iterations. The total training time for this run is approximately 5 days.We distill a one-step generator from the SD v1.5 model [1] , released under the CreativeML Open RAIL-M license, using prompts from the LAION-Aesthetic 6.25+ dataset [59] . Additionally, we collect 500K images from LAION-Aesthetic 5.5+ as training data for the GAN discriminator, filtering out images smaller than 1024 × 1024 and those containing unsafe content. Our training process involves two stages. In the first stage, we disable the GAN loss and use the AdamW optimizer with a learning rate of 1 × 10 -5 , a weight decay of 0.01, and beta parameters of (0.9, 0.999). The fake diffusion model is updated 10 times per generator update. We set the guidance scale for the real diffusion model to be 1.75. We use a batch size of 2048 and train the model on 64 A100 GPUs for 40K iterations. In the second stage, we enable the GAN loss with a weight of 10 -3 , reduce the learning rate to 5 × 10 -7 , and continue training for an additional 5K iterations. The total training time is approximately 26 hours.We train both one-step and four-step generators by distilling from the SDXL model [58] , released under the CreativeML Open RAIL++-M License. For the one-step generator, we observed similar block noise artifacts as reported in SDXL-Lightning [27] and Pixart-Sigma [89] . We addressed this by adopting the timestep shift technique from OpenDMD [90] and Pixart-Sigma [89] , setting the conditioning timestep to 399. Additionally, we initialized the one-step generator by pretraining it with a regression loss using a small set of 10K pairs for a short period. These adjustments are not necessary for the multi-step model or other backbones, suggesting this issue might be specific to SDXL. Similar to SD v1.5, we use prompts from the LAION-Aesthetic 6.25+ dataset [59] and collect 500K images from LAION-Aesthetic 5.5+ as training data for the GAN discriminator, filtering out images smaller than 1024 × 1024 and those containing unsafe content. The generator is trained using the AdamW optimizer with a learning rate of 5 × 10 -7 , a weight decay of 0.01, and beta parameters of (0.9, 0.999). The fake diffusion model is updated 5 times per generator update. We set the guidance scale for the real diffusion model to be 8. We use a batch size of 128 and train the model on 64 A100 GPUs for 20K iterations for the 4-step generator and 25K iterations for the 1-step generator, taking approximately 60 hours.For the COCO experiments, we follow the exact evaluation setup as GigaGAN [72] and DMD [22] .For the results presented in Table 5 , we use 30K prompts from the COCO 2014 validation set and generate the corresponding images. The outputs are downsampled to 256×256 and compared with 40,504 real images from the same validation set using clean-FID [91] . For the results presented in Table 2 , we use a random set of 10K prompts from the COCO 2014 validation set and generate the corresponding images. The outputs are downsampled to 512×512 and compared with the corresponding 10K real images from the validation set with the same prompts. We compute the CLIP score using the OpenCLIP-G backbone. For the ImageNet results, we generate 50,000 images and calculate the FID statistics using EDM's evaluation code [53] .To conduct the human preference study, we use the Prolific platform (https://www.prolific.com).We use 128 prompts from the LADD [24] subset of PartiPrompts [70] . All approaches generate corresponding images, which are presented in pairs to human evaluators to measure aesthetic and prompt alignment preference. The specific questions and interface are shown in Figure 13 . Consent is obtained from the voluntary participants. We manually verify that all generated images contain standard visual content that poses no risks to the study participants. L Prompts for Figure 1 , Figure 2 , and Figure 11 We use the following prompts for Figure 1 . From left to right, top to bottom:• a girl examining an ammonite fossil • A photo of an astronaut riding a horse in the forest.• a giant gorilla at the top of the Empire State Building • A close-up photo of a wombat wearing a red backpack and raising both arms in the air. Mount Rushmore is in the background. • An oil painting of two rabbits in the style of American Gothic, wearing the same clothes as in the original. • a portrait of an old man • a watermelon chair • A sloth in a go kart on a race track. The sloth is holding a banana in one hand. There is a banana peel on the track in the background. • a penguin standing on a sidewalk • a teddy bear on a skateboard in times square We use the following prompts for Figure 2 . From left to right, top to bottom:• a chimpanzee sitting on a wooden bench • a cat reading a newspaper • A television made of water that displays an image of a cityscape at night. • a portrait of a statue of the Egyptian god Anubis wearing aviator goggles, white t-shirt and leather jacket. The city of Los Angeles is in the background. • a squirrell driving a toy car • an elephant walking on the Great Wall • a capybara made of voxels sitting in a field • Cinematic photo of a beautiful girl riding a dinosaur in a jungle with mud, sunny day shiny clear sky. 35mm photograph, film, professional, 4k, highly detailed. • A still image of a humanoid cat posing with a hat and jacket in a bar.• A soft beam of light shines down on an armored granite wombat warrior statue holding a broad sword. The statue stands an ornate pedestal in the cella of a temple. wide-angle lens. anime oil painting. • children • A photograph of the inside of a subway train. There are red pandas sitting on the seats. One of them is reading a newspaper. The window shows the jungle in the background. • a goat wearing headphones • motion • A close-up of a woman's face, lit by the soft glow of a neon sign in a dimly lit, retro diner, hinting at a narrative of longing and nostalgia.We use the following prompts for Figure 11 . From left to right, top to bottom:• A close-up of a woman's face, lit by the soft glow of a neon sign in a dimly lit, retro diner, hinting at a narrative of longing and nostalgia. • a cat reading a newspaper • A television made of water that displays an image of a cityscape at night. • a portrait of a statue of the Egyptian god Anubis wearing aviator goggles, white t-shirt and leather jacket. The city of Los Angeles is in the background. • a squirrell driving a toy car • an elephant walking on the Great Wall • a capybara made of voxels sitting in a field • A soft beam of light shines down on an armored granite wombat warrior statue holding a broad sword. The statue stands an ornate pedestal in the cella of a temple. wide-angle lens. anime oil painting. • a goat wearing headphones • An oil painting of two rabbits in the style of American Gothic, wearing the same clothes as in the original. • a girl examining an ammonite fossil • a chimpanzee sitting on a wooden bench • children • A still image of a humanoid cat posing with a hat and jacket in a bar.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We release them at our project website. Guidelines:• The answer NA means that paper does not include experiments requiring code. • The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: The paper does not pose significant risks, and future production systems could implement an output classifier to detect and prevent unsafe content.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: In Appendix I.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [Yes] Justification: We release code and dataset on our website. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: This paper presents DMD2, a few-step distilled generator designed to achieve fast sampling while preserving the high-generation quality of multi-step diffusion models. The paper introduces several key enhancements to the original DMD's training procedure, including: (1) replacing the regression loss with the Two-Time Scale Update Rule to stabilize training, (2) integrating GAN loss to improve generation quality, and (3) employing backward simulation to address potential mismatches between training and inference. Experimental results show that DMD2 achieves state-of-the-art performance, outperforming the original DMD and other competitive models in both image quality and efficiency.

All reviewers unanimously acknowledged the novelty and effectiveness of this work, leading to its acceptance. The authors should also revise their paper according to the reviewers' suggestions in the final version.

Reviews: review: paper_summary: paper_summary: The paper introduces an upgraded version of Distribution Matching Distillation (DMD), i.e., DMD2, which addresses the limitation and inefficiency of previous DMD and improves the performance of efficient and high-quality image synthesis using diffusion models. Specifically, the authors identify the limitations of the original Distribution Matching Distillation (DMD), such as the need for a regression loss and extensive dataset construction. DMD2 eliminates the regression loss, integrates a Generative Adversarial Network (GAN) loss, and introduces a two-time-scale update rule to stabilize training. Additionally, a new training procedure is implemented to simulate multi-step sampling, addressing the training-inference mismatch. Experimental results demonstrate that DMD2 achieves state-of-the-art performance, surpassing the original DMD and other competitive models in image quality and efficiency.
strength_weakness: strength_weakness: I do not find a specific weakness of this paper.
questions: questions: - Training with GAN often entails numerical instability. Does DMD2 have such concerns? If it is true, could the author provide some details in overcoming the instability of DMD2?
- Besides the evaluation metric such as FID and Inception scores, how about some human-related metrics such as ImageReward or aesthetic scores? Does DMD2 show comparable results to teacher SDXL?
limitations: limitations: The paper includes limitations of the proposed method, e.g., requires multiple steps to generate on par with teacher model, e.g., SDXL.
paper_summary: paper_summary: This work addresses identifies reasons for training instability of one of competitive diffusion distillation approaches based on distribution matching, using bi-level optimization and also adopt a GAN based feature space feedback for improved quality. Overall demonstrate very good performance on SDXL, SD checkpoints demonstrating effectiveness on large models too.
strength_weakness: strength_weakness: As the authors discuss on not using real-data within current formulation and setup, there could be a tendency for model to have model collapse? 

As proposed distillation objective is not sampling w.r.t teacher's marginal predictive distribution nor data-distribution. It would be useful to get some diversity metric at per-prompt level using  e.g., LPIPS Diversity or etc comparing to other Distillation approaches and teacher model.

Also it would be useful to understand what stage of training causes this mode collapse, i.e., does small scale training preserve diversity at cost of some quality drop or we observe a consistent drop in diversity?
questions: questions: What is setup of ablation without backward simulation in multi-step setting? Do forward diffusion and query the student generator and based on that query fake score function estimator? If so is the fake score function also trained on equivalent forward diffuse + generator's predictive distribution? As it is currently unclear is it alignment of fake score function alignment to generator or the backward simulation which is resulting in improved performance.

Also given DMD has implicit assumption that the fake score function is capturing predictive distribution of student generator and authors identify its fitting being one of reasons for instability. It might be useful to understand how sensitive is alignment of fake score estimation to student's generator. Also, how good is quality of fake score fun at different stages of training and its implications? 

As at implementation level we are starting with pretrained model and in the limit if distilled student model matches pre-trained model's weights we are asking fake score function to match pre-trained model again.
limitations: limitations: This work build on DMD and improves stability, good engineering practice as primary contribution of this work with limited formulation novelty or insights. So more insights on hyperparameter sensitivity, why and how are different design choices effect final performance, diversity etc as discussed above would make it more useful for community and strong contribution!
paper_summary: paper_summary: This paper introduces DMD2, a few-step distilled generator to achieve fast sampling while maintaining the decent generation quality of the multi-step diffusion models. DMD2 proposes several new improvements to the training procedure of the original DMD, including (1) replacing the regression loss with the Two-Time scale Update Rule (TTUR) to stabilize the training process, (2) incorporating the standard GAN loss to achieve better quality, (3) and utilizing the backward simulation to alleviate the potential mismatch of training and inference. Built upon these modifications, DMD2 achieves excellent results on few-step image generation.
strength_weakness: strength_weakness: 1. It would be better to include a detailed training algorithm to clearly showcase the modifications over the original DMD training process.
2. In practice, the real data used to train the (teacher) diffusion models may not be accessible to the users who hope to distill a small (student) generation model (due to privacy, storage, …). In that case, one limitation of DMD2 is the calculation of GAN loss may become inapplicable. Could you share your opinions about this matter?
questions: questions: 1. In the original DMD paper, the regression loss is capable of mitigating the issue of mode collapse. Would DMD2 also suffer from this issue as the regression loss is removed and an extra GAN loss is introduced?
2. Why the intermediate outputs of DMD2 shown in the right subfigure of Figure 3 are so similar and seem to follow a certain trajectory? As far as I know, distribution matching-based methods do not guarantee the specific paths of the teacher diffusion model and the student generation model are aligned (as mentioned in Lines 35-39). Could authors provide more explanations about this phenomenon (Fig 3)? Note that the samples generated by few-step consistency models may also switch between different paths and generate very different images.
limitations: limitations: Please refer to the weakness and question sections above.
paper_summary: paper_summary: This work proposed an improved training method for distribution matching distillation, named DMD2. Notably, compared to DMD, it does need the regression loss which relies on constructing the synthetic noise-data pairs. Instead, DMD2 introduces three new features: 1) a two time-scale update rule for fake score and generator training, 2) a GAN loss which extracts the bottleneck features of fake score network for a training prediction head as the discriminator, and 3) backward simulation for few-step distillation that produces inference-time generator inputs during training.
strength_weakness: strength_weakness: - The proposed method introduces many hyperparameters and seems to be sensitive to these hyperparameters, such as batch size, guidance scale for teacher score, GAN loss weighting and fake score updating frequency. For instance, in different distillation tasks (EDM on ImageNet, SDv1.5 and SDXL), these hyperparameters are different (as shown in Appendix G). I’m not sure if the hyperparameters need to be specifically tuned for good performance in each distillation task. 
- The authors claimed that “SDXL remains challenging to distill into a one-step generator because of limited model capacity and a complex optimization landscape to the direct mapping from noise to highly diverse and detailed images”. For the point of “limited model capacity”, does it mean that if the student has the same network with SDXL, by any means, we are not able to achieve a one-step generation that matches SDXL’s performance? For the point of “a complex optimization landscape”, it seems that both SDv1.5 and SDXL are trained on LAION dataset, which means they are both trying to learn the same mapping from noise to data. Does it mean the higher-quality generation of SDXL (rather the training data of teacher models) hinders the one-step distillation? On the other hand, I wonder if it is possible to tune the hyperparameters of distilling SDXL for a better one-step generator. For example, in Appendix G, the batch size for distilling SDXL is only 128 while the batch size for distilling SDv1.5 is only 2048. If DMD2 is sensitive to batch size in the large-scale text-to-image case, can we increase the batch size for distilling SDXL to 2048 during training for improved performance? 
- There are some inconsistencies: 1) In Figure 4, the caption says the teacher uses 50 sampling steps while the main text says “while requiring 25x fewer forward passes (4 vs 100)”. 2) EDM (Teacher ODE) originally reported their FID as 2.22, while this work reports 2.32 in Table 1. 3) Both $\mu_{\text{real}}$ and $\mu_{\text{fake}}$ are introduced without definition.
questions: questions: - I’m curious about the extra cost of backward simulation, i.e., producing synthetic images with the current student generator running several steps. Is it possible to compare the training time per iteration with and without backward simulation?
- From the numbers in Table 2, it looks like the 4-step distillation improves Patch FID but gets worse FID and CLIP, compared to 1-step distillation. Any justification?
limitations: limitations: The authors adequately addressed the limitations and potential negative societal impact of their work.
paper_summary: paper_summary: The authors propose an improved way of distilling image-generating diffusion models into fast models capable of generating high-quality images with as few as 1-4 steps. Compared to prior work using distribution-matching (DMD), they do away with the regression loss term that tied the teacher path to the student path. They also introduce a GAN-style loss in the pipeline, and introduce a few other tricks to squeeze out extra performance. They demonstrate SOTA performance among efficient models, and even beat the teacher model (made possible by training using real images and a GAN loss).
strength_weakness: strength_weakness: There's not so much to say here, since I found most aspects of the paper to be excellent. However, I would have liked to see some more details in how their approach relates to competing approaches. Most notably, the line of work by Sauer et al [23, 24] use SDS and a GAN-style loss. Now that a GAN-style loss is introduced in the DMD framework, the gap between these two approaches gets smaller, and it would be nice to see some more explanation about what the core difference is.

The progress is mostly empirical in nature. For example, the authors note that the training becomes more stable using the two-scale update rule, but they don't present any theoretical convergence guarantees (which is perfectly fine, the empirical progress is definitely good enough in my opinion).

There are a few minor typos (e.g. line 122, should be "gradient of the data log likelihood"), but few enough not to impair the overall understanding (and I trust the authors to do a final proof reading for the camera ready version).
questions: questions: What is the most significant difference between this work and the line of work by Sauer et al ([23, 24])? Could one sentence about that be added to the "related work" section?

The original DMD paper [22] demonstrated examples of mode collapse when omitting the regression loss. GANs (at least some) are also known to be prone to mode collapse. Could you mention anything about the mode collapse situation in this method? (Since the FID numbers are good, I assume that this is also good, but since this was a main point of analysis in [22]?)
limitations: limitations: Yes, limitation and potential social impact are well-described in section 6 and A.

Comments: title: Thank you for addressing/clarifying most of concerns.
comment: Thank you for the rebuttal. I will maintain my original score.
We are pleased that our responses have addressed your concerns and appreciate your consideration to raise the score!

Regarding your further inquiries, we currently utilize 20K iterations because this represents almost the maximum compute we can afford (64 GPUs over 3 days). However, we have not yet observed the peak performance of our models. For example, in a trial where we extended to 30K iterations, we managed to improve the FID from 19.3 to 18.7. We are eager to explore extending the training duration in our revised version to further confirm the model's stability.

The suggestion to reinitialize the fake diffusion model at a later stage is intriguing, and we look forward to experimenting with this approach. Thank you once again for your invaluable suggestions and insights!
Happy to raise my score.

One thing i want to understand from authors, is how stable is this formulation when number of steps go from 20K to say 60K or more? Do you see performance peak around 20K or further training improves quality? unlike original DMD as authors report only 20K was curious if there were any interesting empirical findings which would provide further insights to community.

This goes back to fake score function alignment, does it help to reinitialize fake score function with teacher model again as student converges well to teacher? As it seems a bit unclear on properties/fit of fake score function and how it effects overall training stability.

Looking forward for more results in later version of paper.
I read the rebuttal and thank the authors for good responses to my questions. I have no further comments, and my "strong accept" recommendation stands. I wish the authors good luck, and I'm looking forward to read the final version!