Multistep Distillation of Diffusion Models via Moment Matching

Summary

This paper presents a new method for making diffusion models faster to sample. The method distills many-step diffusion models into few-step models and extends recently proposed one-step methods to the multistep case by moment matching. By using up to 8 sampling steps, the obtained distilled models outperform not only their one-step versions but also their original many-step teacher models, obtaining new SOTA results on the Imagenet dataset.

Strengths

1. This paper presents a new method for making diffusion models faster to sample. The method distills many-step diffusion models into few-step models and extends recently proposed one-step methods to the multistep case by moment matching. 2. The paper is well writing. The proposed moment matching proposed in this paper is a technique worthy of reference and further in-depth study by researchers. 3. Experimental results show the effectiveness of the proposed method and the results of experiments achieve SOTA.

Weaknesses

1. Moment Matching Distillation Part is difficult and too mathematical to understand. The authors could improve the writing to make it easier for the reader to understand. 2. In line 446, "because the first term of $L( \phi)$ can be seen as a standard diffusion loss at $\theta$." I don't understand, the author needs to give more interpretations. 2. The experiment of Class-conditional generation is only on ImageNet. The authors could provide more experiment results on other datasets, such as celeba, lsun and so on in later versions. 3. The experiments lack ablating sampling steps. The authors could provide more experiment results adapting other sampling steps in later versions.

Questions

1. the analysis of Ablating conditional sampling is confusing. It’s not unusual that conditioning leads to much better sample diversity in the distilled model because you use extra label information. Could the authors explain more about the ablating conditional sampling? 2. Why you select 8 sampling steps? How more or less eight steps will affect experimental results, run speeds, etc?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

A limitation of this study is that the theoretical part is difficult and too mathematical to understand.

Summary

This paper proposes an approach to distill a diffusion model into a multi-step generator. Building on previous works that use distribution matching to train a few-step student generator, the paper introduces a novel method of matching the conditional expectation of clean data given noisy data along the sampling trajectory. By focusing on the matching of the first-order moment, the authors derive a gradient for training the generator, which can be approximated using two diffusion models (as seen in Diff-Instruct, DMD, and SwiftBrush) or estimated using batch statistics. The primary contribution is conditioning on noisy data during expectation matching, distinguishing this approach from previous score distillation methods that match unconditioned expectations. Extensive results demonstrate the effectiveness of this method.

Strengths

S1. This paper provides a new perspective on recent distribution matching-based diffusion distillation approaches such as Diff-Instruct, DMD, and SwiftBrush. The moment matching formulation is novel and theoretically sound, leading to practical improvements including conditioning on noisy input to the generator, as well as a different way of training the model that doesn't rely on an auxiliary diffusion model S2. The experimental validation is comprehensive, with the 8-step model achieving exceptionally strong image generation results.

Weaknesses

Minor Typo: There is a typo on line 72: "as described is described." Lack of Detailed Ablation Studies: There are several instances where the authors claim one formulation is better based on their trials but do not provide exact ablation studies. Including these studies would help readers understand the effectiveness of alternative formulations. Specific examples include: - Exclusion of the z_s dependence on the parameters when calculating gradients of the moment matching loss (line 90). - Training the auxiliary diffusion model using a mix of teacher predictions and generated samples (line 110). Comparison Between Optimization Setups: Further comparison between the alternating and instant optimization setups would be beneficial. Currently, alternating optimization appears better, but instant optimization feels more principled. The authors should discuss this discrepancy in more detail and highlight potential challenges to improve performance. Enhancement of Conditional Sampling Ablation: The ablation study on the importance of conditional sampling could be clearer. The paper shows that without conditioning on noisy input, the generator produces less diverse images with drifted tones. However, other methods like Diff-Instruct, DMD, and SwiftBrush, which don't condition on noisy input, work without these diversity/tone issues. This seems contradictory to existing literature. Does this issue arise as the number of sampling steps increases? What is the failure point for other unconditional methods? These experiments and discussions are crucial, as this is the paper's most significant contribution and claim. On L181, the authors mention that "One could imagine fitting x through score matching against the conditional distribution q(z_s | z_t, x_~) but this did not work well when we tried it." Could the authors elaborate on "score matching against the conditional distribution"?

Questions

Please find a set of suggestions in the weakness section above.

Rating

6

Confidence

5

Soundness

2

Presentation

3

Contribution

3

Limitations

Yes.

Summary

The authors proposed a fast sampling method by distilling a diffusion model to model $q(x|z_t)$. This is achieved by matching moments, with two novel approaches proposed to implement it in practice.

Strengths

- The idea of distilling to model the conditional distribution $q(x|z_t)$ by using matching the moments is interesting and novel. - The equivalence between the proposed instantaneous version of the moment matching loss and matching expected teacher gradients is interesting and insightful.

Weaknesses

1. Some of the notations should be clarified or made coherent. For instance, do both $g_{\theta}(z_t, t)$ and $g_{\theta}(z_t)$ indicate the parametrized denoiser? What's the definition of $q(z_t)$ in Eq. (6)? In Eq. (4) and Eq. (6), what depends on $\eta$ (please be consistent about dropping this parameter)? 2. Algorithm 2 is somehow heuristics. It may be hard to see if $\mathbb{E}_{g}[\tilde{x}|z_s]$ can be approximated well with the auxiliary network model.

Questions

1. In tables 1 and 2, by increasing NFE to the level of the base model, what the FIDs and IS will be? 2. I understand that the sampling in the base diffusion model may not be optimal. However, which mechanisms in the proposed methods can achieve better FID than the base model with distillation? 3. How is the generation quality and FID by using instantaneous version of the moment matching (algorithm 3) on ImageNet $128\times128$? 4. Can the proposed method be compatible with Consistency Model for further timestep distillation? 5. Also, can the method support training from scratch? 6. In Algorithm 3, storing and computing both teacher gradients and $\eta$-gradients appears costly, particularly as the model's scale (in terms of trainable parameters) increases. I'm curious about the scalability of this distillation method to large latent diffusion models.

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

4

Limitations

The authors provided a discussion on potential limitations on evaluations and 1 or 2 steps generation.

Summary

This paper proposes a diffusion distillation algorithm based on moment matching. The starting point of the paper is to achieve distribution consistency by ensuring that the denoising process $\widetilde{x} = g_\eta(z_t,t)$ with fewer steps conforms to the true distribution, incorporating moment matching. To estimate the first moment of the student models' $x_0$ predictions, the paper suggests using a small surrogate denoising network to learn moment information. Simultaneously, the moment information from the surrogate denoising network is used to optimize the student models. The distillation model proposed in this paper demonstrates good performance with 8-step sampling

Strengths

1. Research based on moment matching requires moment information. This paper proposes a surrogate denoising network for the first moment, based on the principle that one-step sampling in diffusion is unbiased for the first moment. The distillation model needed is trained using this moment network. The approach of using this surrogate is both reasonable and solid 2. The experimental section of the paper demonstrates that, compared to different distillation models, the MM-based distillation model achieves the best results with 8-step sampling

Weaknesses

1. Theoretically, according to the current setting of the paper, the discretization error in the distillation model sampling originates from the gap between the one-step prediction of the distillation model and the first moment of the true teacher models' denoising. It is worth exploring the values of L(ϕ)  at different timesteps after training. 2. The experiments revealed that the performance of the distillation model worsens as the number of steps increases after training (FID of 3.0 for one step and 3.86 for two steps). Could you provide a theoretical explanation for this phenomenon? 3. For the distillation model, the number of steps in the solver of pretrained diffusion models significantly affects the results. This paper lacks an analysis of the solvers' step counts, such as the impact on results when reducing from DDIM 50 steps to DDIM 25 steps. 4. Theoretically, the method proposed in this paper can be applied to any distillation model. If the distillation model uses a pre-trained distillation model for initialization, the accuracy of the x0x0​ predictions will be higher in the early stages of training. Consequently, the optimization effect of the moment estimation surrogate network will be better. Further fine-tuning on this basis could additionally demonstrate the extent to which moment matching improves the distillation model. Could this experiment be supplemented？

Questions

Please see Weakness.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

yes

Summary

The paper presents a novel method for distilling diffusion models to require fewer function evaluations. The method is based on moment-matching, and two practical algorithms are presented and evaluated. Empirical results are state-of-the-art in few-step regimes.

Strengths

The paper presents a well-motivated novel method for distillation fo diffusion models. This method achieves impressive results on standard benchmarks, and represents a solid contribution to the state of diffusion models for image generation. The presentation is clear and complete, and the method is well situated in the literature both through empirical comparison and discussion.

Weaknesses

This is a strong paper overall, without many weaknesses I noticed. I would have liked to see more discussion of why the 1- and 2- step regimes do not perform as well as other approaches. Is this a characteristic of the moment-matching approach?

Questions

1. In Eqn. (6), how is $\mathbb{E}_q[x\mid z_s]$ approximated? Is it through samples from the dataset? If so, are there additional assumptions that need to be made? Naively, I'd expect that, especially as $s \to 1$, the true data distribution would not necessarily be unimodal around the empirical samples in the training set. Am I overthinking this? 1. Why does the instantaneous version of the algorithm have no reported results for 1 and 2 steps?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Although there is no limitations section in the paper, the discussion of the limitations is well addressed throughout the paper.