Linear Multistep Solver Distillation for Fast Sampling of Diffusion Models

Authors: Yuchen Liang, Xiangzhong Fang, Hanting Chen, Yunhe Wang
Primary area: generative models
Keywords: Diffusion Probabilistic Model, Diffusion Sampler, Solver Schedule

Abstract:
Sampling from diffusion models can be seen as solving the corresponding 
   probability flow ordinary differential equation (ODE). 
   The solving process requires a significant number of function 
   evaluations (NFE), making it time-consuming. 
   Recently, several solver search frameworks have attempted to find 
   better-performing model-specific solvers. However, predicting the impact of 
   intermediate solving strategies on final sample quality remains challenging, 
   rendering the search process inefficient.
   In this paper, we propose a novel method for designing 
   solving strategies. We first introduce a unified prediction formula 
   for linear multistep solvers. Subsequently, we present a solver distillation 
   framework, which enables a student solver to mimic the sampling trajectory 
   generated by a teacher solver with more steps. We utilize the mean Euclidean 
   distance between the student and teacher sampling trajectories as a metric, 
   facilitating rapid adjustment and optimization of intermediate solving strategies.
   The design space of our framework encompasses multiple aspects, 
   including prediction coefficients, time step schedules, and time scaling 
   factors. 
   Our framework has the ability to complete a solver search 
   for Stable-Diffusion in under 12 total GPU hours.
   Compared to previous reinforcement learning-based 
   search frameworks, 
   our approach achieves over a 10$\times$ increase in search efficiency. 
   With just 5 NFE, we achieve FID scores of 3.23 on CIFAR10, 7.16 on ImageNet-64, 
   5.44 on LSUN-Bedroom, and 12.52 on MS-COCO, resulting in a 2$\times$ sampling acceleration ratio 
   compared to handcrafted solvers.

Decision: accept (poster)
Reviewer scores: 8, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730200791065

thread_type:
official_review

<thread_object>
creator: reviewer_dd8u

<content>
summary:
This paper considers diffusion distillation for  a student linear multistep solver from a teach linear multistep solver with more timesteps.  In particular,  the authors propose to train a lightweight neural network to predict the coefficients, time step schedules, and time scaling factors of the student linear multistep solver. The cost function when training the lightweight neural network is taken as the mean squared distance of the difference of diffusion states produced by the student and the teacher linear multipstep solver, respectively. Experiments on FID shows the effectiveness of the new method.

strengths:
The main strength of the paper is that the authors propose to train a lightweight neural network that produce not only the coefficients of the student linear multistep solver but also the timesteps and the scaling factors.  This is based on the assumption that for different ODE trajectories, the optimal coefficients, timesteps and scaling factors are different.

weaknesses:
(0) One weakness is that a small neural network is required to be trained for each particular pre-trained model. Note that not every university or research institute has 8 A00 or H100 GPUs for conducting the training process.  

(1) The literature is not thorough. This work is closely related to a recent paper [1], which is not mentioned at all. The work of [1] considers computing the optimal coefficients of a student linear multistep solver per timestep by solving a quadratic optimization problem. The computational complexity of [1] is negligible as the quadratic optimization problem takes a closed-form solution. The authors should include the performance of [1] in their work. 

(2) One thing that is not clear to me is results for the two experiments of Latent-Diffusion on LSUN-Bedroom and Stable-Diffusion on MS-COCO prompts, where the number of interpolation timesteps M=1. I would think that the teacher ODE solver with two times of the number of times perform betters than the student ODE solver.  Is it the case? If not, explain why. 

(3) I would think that in general, the higher the M value, the better FID score of the student ODE solver. So why in different experimental setups, M were chosen differently? Would higher M value in some cases lead to poor performance? If so, explain why and include the results in the revision. 

(4) Typo:  "can be carry out"

[1] Guoqiang Zhang, Kenta Niwa, W. Bastiaan Kleijn, "On Accelerating Diffusion-Based Sampling Processes via Improved Integration Approximation, ICLR, 2024.

questions:
please check my responses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730571006238

thread_type:
official_review

<thread_object>
creator: reviewer_23f3

<content>
summary:
This paper proposes DLMS, a flexible solver framework that incorporates the combination of previous model outputs, timestep schedule and timestep scaling factor. The authors further introduce a light weight designer network to dynamically decide the solver strategies for each single trajectory. Experimental results demonstrate DLMS achieves notable improvements over existing solvers and offers faster optimization than search-based methods.

strengths:
1. DLMS offers a flexible framework for diffusion solver, unifying existing methods.
2. DLMS uses dynamic solving strategies for different ODE trajectories, enhacing the potential for diffusion solver.
3. Experimental results demonstrate significant performance improvements compared to existing solvers, and the designer network's training cost is more efficient than that of search-based solvers.

weaknesses:
1. Algorithm 1 implies that each NFE configuration may require independent designer networks, which could limit flexibility in the NFE-performance trade-off. If this is not the case, how does the designer network ensure that $t_N$ becomes a reasonably small when the step count reaches N?
2. The time scaling factor may introduce input distribution misalignment, so further discussion on the motivation and explanation of this would be beneficial.
3. The designer network currently relies on U-Net intermediate feature. As transformers gain popularity in diffusion models, it is uncertain if this approach is adaptable to such architectures.
4. It would be helpful to illustrate differences in solver design choices provided by the designer network across various ODE trajectories to support the claim that a unifed choice for all trajectories is suboptimal.

questions:
1. What is the relationship between the coefficients predicted by the designer network and those derived from Taylor expansion in previous methods? Could you provide a comparison of these coefficients with those from previous methods?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730591133774

thread_type:
official_review

<thread_object>
creator: reviewer_asen

<content>
summary:
The paper proposes a method for learning a trajectory-specific solver for diffusion models. The method suggest to use a small network to predict the best time step size and coefficients of a linear multistep at each step of the solver. The method is tested on a number of image generation tasks.

strengths:
1. The idea to learn a trajectory-specific solver is novel and interesting.
2. The method shows good results on number of benchmark datasets.
3. The method is compared to a number of diffusion dedicated solvers.

weaknesses:
1. The authors make use of what they call "bottleneck feature" without any explanation what are those features, only referencing the relevant paper. The paper should be self contained and the authors should make an effort to give even brief explanation about these features.
2. The method is not compared to any other solver distillation methods such as [1], [2].
3. Discussion and comparison to model distillation is too minimal.
4. The size of the designer network is not provided.

questions:
1. Does a different designer network needs to be learned for each choice of NFE, multistep order?
2. Is the designer network is somehow constrained such that always $t_N=0$?
3. Could the authors provide the size of the designer network?
4. The designer network is dependent only on $h_{t_{n-1}}$ or previous times as well?


[1] Zheng, Kaiwen, et al. "Dpm-solver-v3: Improved diffusion ode solver with empirical model statistics." Advances in Neural Information Processing Systems 36 (2023): 55502-55542.

[2] Shaul, Neta, et al. "Bespoke Non-Stationary Solvers for Fast Sampling of Diffusion and Flow Models." arXiv preprint arXiv:2403.01329 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680398552

thread_type:
official_review

<thread_object>
creator: reviewer_v3hm

<content>
summary:
This paper proposes a Distilled Linear Multistep Solver (DLMS) to learn a faster sampler for diffusion models, requiring fewer function evaluations. The distillation approach is to train a solver that minimizes the Euclidean distance between its trajectory and a teacher solver’s trajectory. DMLS can be trained faster than previous reinforcement learning-based approaches to solver distillation. Experimental results in image generation using unconditional, conditional, latent space, and pixel space diffusion show improved FID scores compared to existing methods, especially in low NFE settings.

strengths:
The strengths of the paper are its well-motivated method and a wide range of experiments.
- The approach can be initialized using existing solvers.
- Unlike model distillation, solver distillation like DLMS can be used for downstream tasks like image restoration.
- The proposed method is simpler than using a reinforcement learning-based approach.
- The method is evaluated in multiple contexts (unconditional, conditional, latent space, and pixel space diffusion) and on multiple datasets with convincing results.
- The paper is overall well-written.

weaknesses:
The weaknesses of the paper include its figure and table captions and reliance on FID as the sole quantitative metric.
- In general, the figures and table are not self-contained. Figure 4, for example, could be made more interpretable by describing the significance of dashed lines.
- FID scores are the only quantitative metric used to evaluate the method. Quantifying the quality of generated images is challenging, so adding multiple metrics like IS or CMMD [1] would increase confidence in the results. 
- In the introduction (line 047) it is argued that distillation is expensive, requiring multiple GPU days of training. The reported training times for DLMS are still more than ten hours, so these methods could still be compared, if not to better understand their respective strengths and weaknesses.
- The time comparisons in Table 2 are hard to draw conclusions from considering the reported times are compared to those from previous papers that were run on different systems with different GPUs and software environments. It is stated that this is due to limited code availability. If they are possible to obtain, FLOP counts (or an estimate of them) would be more comparable.

**References**

[1] Jayasumana, S., Ramalingam, S., Veit, A., Glasner, D., Chakrabarti, A. and Kumar, S. “Rethinking FID: Towards a Better Evaluation Metric for Image Generation”, CVPR 2024

questions:
- On line 398 it is reported that the distillation time of DMLS is approximately 1.5 * 8 = 12 hours for stable diffusion, while the abstract claims that the framework “has the ability to complete a solver search for Stable-Diffusion in less than 10 total GPU hours”. Is this difference due to a rounding error or do these claims refer to different times? Clarify this discrepancy and ensure consistency between the abstract and results.

**Minor suggestions that do not individually affect the score**
- Line 129: “can be carry out” -> “can be carried out”.
- Line 189: Remove “precious”.
- Line 263: Introduce the strop gradient operation.
- Line 284: Reformulate.
- Line 292: “PLMS(iPNDM)” -> “PLMS (iPNDM)”.
- Line 340: “AMED-Plugin(Zhou et al., 2024)” -> “AMED-Plugin (Zhou et al., 2024)”.
- Line 363: Specify “various aspects” and “as well as the ablation…” -> “as well as ablation…”.
- Line 394: “MS-COCO(2014)” -> “MS-COCO (2014)”.
- Line 485: “Handcrafted(best)” -> “Handcrafted (best)”.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732107934952

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revision Summary

comment:
Dear AC and reviewers,

We have revised the paper according to the reviews. The modifications are as follows:

 1. We find that adding constraint $\sum_{k=1}^p a_{k}=1$ ensures the unbiasedness of the estimation. This significantly improves the performance of the DLMS.  We will update new results gradually.  
The table below shows the new results on Stable Diffusion v1.5.
|NFE | 5 | 6 | 7 | 8 | 9 | 10 | 
|--|--|--|--|--|--|--|
| DEIS | 15.16 | 13.95 | 13.61 | 13.58  |13.46 | 13.38 |
| DLMS w/o constraint | 15.69 | 13.54 | 13.20  | 13.07 | 12.69 | 11.72 |
| DLMS w constraint | **12.52**  | **11.87** | **11.65** | **11.63** | **11.52** | **11.45** |
 
 2. A more detailed discussion on the limitations of model distillation in image processing problems has been added as Sec. 2.2. 
  
 3. Added two related recent works [1] [2] in introduction and related works. And added the comparison with Bespoke solver in Sec 3.4.

 4. A sensitivity analysis of number of interpolation time steps $M$ is added in Sec. 4.3.
 
 5. Add more description to the tables and figures.

 6. The typos mentioned have been corrected.

Sincerely,

Authors

References

[1] Shaul, Neta, et al. "Bespoke Non-Stationary Solvers for Fast Sampling of Diffusion and Flow Models.", ICML, 2024.

[2] Guoqiang Zhang, Kenta Niwa, W. Bastiaan Kleijn, "On Accelerating Diffusion-Based Sampling Processes via Improved Integration Approximation, ICLR, 2024.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces DLMS, a novel method that trains a lightweight neural network to generate optimal coefficients, timesteps, and scaling factors for a linear multistep ODE solver, tailored to specific trajectories. This approach offers a flexible framework unifying existing diffusion solvers and demonstrates significant performance improvements over them.  Key strengths highlighted by reviewers include:

Reviewers agreed that the proposed approach for learning trajectory-specific solvers demonstrates strong empirical results on benchmark datasets across diverse diffusion contexts and is effective in lowering training cost compared to search-based solvers.