Linear Multistep Solver Distillation for Fast Sampling of Diffusion Models

Dear reviewer Dd8U,

Thank you very much for your review. Here are our responses:

W0: A small designer network is required to be trained for each pre-trained models. Not every institute has GPUs for the training process.

If you just want to use DLMS on a public model, you can directly download the pre-trained designer network without the need to retrain it, similar to DEIS, DPM-Solver-v3, and [1], which only require pre-computed values to be downloaded beforehand.

However, if you wish to try training the designer network, the cost of our method is not considered unacceptably high. As shown in Table 2, with only one V100 GPU, similar works like UFS and DPM-Solver-v3 require several days to train on the MS-COCO dataset. In contrast, training the designer network can be completed within half a day, allowing more researchers to delve deeper into our method. Additionally, in situations with more limited resources, the cost of our method, apart from text-to-image generation tasks, is very minimal. While there are works with even lower costs, such as [1], it should also be noted that in our work, these costs bring significant benefits.

W1: This work is closely related to a recent paper [1], which is not mentioned at all.

Thank you for introducing us to this excellent work. [1] proposed a method that utilizes local mean squared error (MSE) to refine existing solvers, which incurs minimal cost but yields limited effectiveness. As we have pointed out in our paper, the choice of time steps can greatly impact the performance of the solvers, with only limited improvements for coefficient optimization. However, the dynamic time steps imply that we need to repeatedly approximate the GT trajectory, incurring some more overheads. Nevertheless, these overheads also bring significant returns. For instance, on CIFAR10, DLMS achieved an FID of 4.52 with 4 NFEs, outperforming the FID of 6.1 with 10 NFEs achieved by IIA-IPNDM [1]. On MS-COCO, DLMS achieved an FID of 12.52 with 5 NFEs of SDv1.5, surpassing the 12.77 achieved by IIA-DPM-Solver with 20 NFEs of SDv2.

Additionally, our method is data-free, while [1] relies on image datasets. Despite this, we believe that the quadratic optimization method used in [1] may further reduce costs for our method. We would be more than happy to reference that article.

W2&3: Would higher $M$ value in some cases lead to poor performance?

In fact, the fitting capacity of student solvers is limited. When the teacher solver is significantly superior to the student, side effects can arise due to fitting difficulties. For models in the latent space, setting $M=1$, meaning the teacher has about twice NFEs of the student, yields the best results. The table below illustrates the effects of increasing $M$ on the LSUN-Bedroom dataset.

Sincerely,

Authors

Dear reviewer 23f3,

Thank you very much for your review. Here are our responses:

W1: Algorithm 1 implies that each NFE configuration may require independent designer networks, which could limit flexibility in the NFE-performance trade-off. How does the designer network ensure that $t_N$ becomes a reasonably small?

The designer network is trained individually for each NFEs. A related question arises: could we attempt training a universal designer network? While this is clearly feasible, it also implies more complex inputs and a multiplied parameter count. Repeatedly invoking a more complex designer network during sampling would directly raise the sampling costs. Therefore, it seems more reasonable to opt for multiple independent designer networks. Methods like USF and AMED also incorporate additional networks, all of which use separate networks.

When the step_index is $N-1$, we constrain the network to directly output $t_N=0$ (which is actually a positive value close to 0). We also enforce that the output $t_n$ must be less than the input $t_{n-1}$.

W2: The time scaling factor may introduce input distribution misalignment, so further discussion on the motivation and explanation of this would be beneficial.

Introducing errors when solving ODEs is inevitable, implying that the estimated value $X_t^S$ contains more noise compared to the groundturth value $X_t$, which is known as exposure bias. In such cases, a time scale slightly greater than 1 is appropriate. Meanwhile, during the solving process from $t_{n-1}$ to $t_n$, what we need to estimate is the integral from $t_{n-1}$ to $t_n$, where a time scale less than 1 can play a role similar to the mean value theorem.

The principle behind the role of the time scale is complex. Taking the experiment with 7 NFEs on Stable Diffusion as an example, the sequence $\{s_n\}$ is $1.013, 1.001, 0.999, 0.935, 0.888, 0.838, 0.880.$ For large $t$, the time scales will be slightly greater than 1, while for $t$ close to 0, the time scales will be less than 1.

W3: 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.

In principle, we aim to extract a particular layer of features from the backbone of the diffusion model to reduce the dimensionality of the features. While the bottleneck layer is a natural choice, it is not mandatory. Additionally, it is worth noting that even without relying on feature information, our method can still significantly improve solving performance through adaptive coefficients, time steps, and time scales. It is also worth mentioning that the U-Net architecture seems to have regained a prominent position in diffusion models [1][2].

W4: 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 unified choice for all trajectories is suboptimal.

The differences in solving trajectories primarily show in the time steps. Taking the case of 10 NFEs on the Stable Diffusion model as an example, our paper notes a basic linear relationship in the learned half-log-snr at each step. However, in practice, there are some variations in the time step across different trajectories. The table below lists the mean and standard deviation of half-log-snr at each step for reference.

Q1: What is the relationship between the coefficients predicted by the designer network and those derived from Taylor expansion in previous methods?

We believe that Taylor expansion is a compromise in the absence of real trajectory values. DLMS, as a data-driven method, does not require excessive emphasis on interpretability. The table below shows a comparison of the inference coefficients for the 5th step at 10 NFEs on the LSUN-Bedroom dataset. It can be observed that DEIS is quite close to the DLMS method, and in many cases, it is also the best-performing handcrafted solver.

Sincerely,

Authors

References

[1] Tian Y, Tu Z, Chen H, et al. U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers. NIPS 2024.

[2] Bao F, Nie S, Xue K, et al. All are worth words: A vit backbone for diffusion models. CVPR 2023.

Dear reviewer AsEn,

Thank you very much for your review. Here are our responses:

Q1: Does a different designer network needs to be learned for each choice of NFE, multistep order?

Yes, the designer network is trained individually for each NFEs. A related question arises: could we attempt training a universal designer network? While this is clearly feasible, it also implies more complex inputs and a multiplied parameter count. Repeatedly invoking a more complex designer network during sampling would directly raise the sampling costs. Therefore, it seems more reasonable to opt for multiple independent designer networks. Methods like USF and AMED also incorporate additional networks, all of which use separate networks.

Q2: Is the designer network is somehow constrained such that always $t_N=0$ ?

Yes, when the step_index is $N-1$, we constrain the network to directly output $t_N=0$ (which is actually a positive value close to 0). We are willing to share further details about the designer network. We also enforce that the output $t_n$ must be less than the input $t_{n-1}$, which is essential in some experiments. Moreover, for positive outputs such as time step $t_n$ and time scale $s_n$, we actually apply exponential transformations.

Q3 Could the authors provide the size of the designer network?

In line 372, we highlighted that the designer network comprises a two-layer MLP with a total parameter count of only 9k. This is negligible compared to the parameter count of diffusion models (e.g., SDv1.5 with a parameter count of 680M). Perhaps we should mention this earlier in our article. Thank you for your suggestion.

Q4: The designer network is dependent only on $h_{t_{n-1}}$ or previous times as well?

Due to the necessity for the designer network to be very lightweight, its input needs to have a low dimensionality. Therefore, we chose to extract features from the bottleneck layer of the U-Net and then compress them into a vector of only 64 dimensions through average pooling. For the Text-to-Image model, the bottleneck features encompass image, text, and time information, ensuring that the input to the designer network is concise yet informative. We only utilize the most relevant $h_{t_{n-1}}$ and discard previous features to simplify the input and save on space costs.

W2: The method is not compared to any other solver distillation methods such as DPM-Solver-v3, Bespoke Solver.

In Tables 1 and 2, our DLMS method demonstrates significant superiority in performance and training costs compared to DPM-Solver-v3. DPM-Solver-v3 performs impressively in scenarios like Latent Diffusion, which are insensitive to time steps. In fact, we could use DPM-Solver-v3 to initialize distillation training, but we prefer our method not to rely on a similar approach.

Bespoke solver is also a type of solver distillation method, but it only uses the final output to adjust hyperparameters, with a deep and complex computation graph that hinders effective hyperparameter tuning. We actively use intermediate values in the trajectory and judiciously stop unnecessary gradient backpropagation, greatly enhancing the effectiveness of solver distillation method. In the publicly available model EDM trained on CIFAR10, our method achieved an FID of 2.53 with 7 NFEs, outperforming their result of 2.75 with 20 NFEs. Nevertheless, we are more than willing to cite this highly relevant article.

W3: Discussion and comparison to model distillation is too minimal.

We agree that a more detailed introduction to model distillation is necessary, and we intended to add relevant content.

Sincerely,

Authors

Dear reviewer V3hM,

Thank you very much for your review. Here are our responses:

Q1: 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.

The primary time consumption is in generating the teacher trajectory, which is roughly proportional to the number of NFEs. In the Stable Diffusion experiments, 5-10 NFEs actually correspond to 0.75 h to 1.5 h. We will revise to a more specific description, such as "about $9 \times NFEs$ mins". We believe that precise statements are crucial, and we sincerely appreciate you pointing this out.

W2: Quantifying the quality of generated images is challenging, so adding multiple metrics like IS or CMMD [1] would increase confidence in the results.

To be frank, while DLMS significantly outperforms handcrafted solvers in terms of FID and visual presentation, CMMD lags behind. It took us several days to understand why. We discovered that, compared to images rich in texture details, CMMD tends to prefer images with simple compositions, clean visuals, and clear semantics.

After comparing CMMD values between DPM-Solver++ at 10 NFE and 20 NFE, we were surprised to find that the CMMD value of 0.578 at 10 NFE is better than the value of 0.583 at 20 NFE. Then everything makes sense now, because our DLMS at 10 NFE emulates the behavior of DPM-Solver++ at 20 NFE as its teacher.

The CMMD scores of more powerful generative models like SDXL and PixArt are significantly worse than the CMMD of SDv1.5. The authors of CMMD attribute this to the COCO dataset. However, we believe that adopting MMD is reasonable, with the issue possibly lying in the CLIP model's excessive emphasis on semantic information.

W3: In the introduction, 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.

Model distillation typically requires thousands of GPU hours. The intuitive difference in magnitude can better help highlight our contribution. Thank you for your suggestion. Additionally, we plan to provide more specific introductions about downstream tasks, including guidance methods and post-processing techniques. This will help understand why solvers can be more flexibly applied to downstream tasks.

W1 & Minor suggestions: 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.

Thank you once again for these helpful writing suggestions. We will make corrections accordingly.

Sincerely,

Authors