Learning to Discretize Denoising Diffusion ODEs

Thank you for your valuable feedback and for pointing out the areas of improvement in our manuscript.

Weakness 1: Comparision to distillation methods

We thank the reviewer for raising this question. To address this concern, we have expanded Table 9 in the main text into a separate section in Appendix-C (Comparison to Few-Step Generation Methods), where we carefully compare our method to each few-step sampling technique.

In our updated comparison, we ensured that both models operate in a more similar context to enhance the validity of the comparison. While the models are not identical, we have aligned one key factor—the neural network architecture—in the case of Consistency Distillation and Rectified Flow. Although the backbone architecture differs for Progressive Distillation (PD), we believe it is still valuable to report their training cost and performance.

This new setup provides a more equitable basis for comparison, highlighting each model's relative strengths and weaknesses within a consistent framework. We welcome further suggestions for improvements and thank the reviewer again for this helpful feedback.

Comparison to Consistency Distillation

We reference the estimated training time and GPU requirements to train Consistency Distillation from [1]. We train LD3 with NCSN++ architecture [2] to match the same architecture used in CD. The reported results for Consistency Distillation are the best-performing numbers on CIFAR-10, achieved at NFE=2.

Training a consistency model on CIFAR10 requires approximately one day on 8 A100 GPUs. With LD3, we achieve slightly better sample quality than consistency distillation when using 8 NFE. While our inference is 4 times slower, our training is 240 times faster and requires only 1/8 of the GPUs.

Compared to Rectified Flow

We estimate the training time for Rectified Flow using the official code provided by its authors [3], while the FID score is taken from Table 1-a of the original paper.

Rectified Flow seeks to learn a straight flow from the prior distribution to the data distribution. However, the method still requires multiple steps to generate high-quality samples. For instance, achieving a 2.58 FID score necessitates 127 NFE, comparable to LD3 applied to the same diffusion backbone, but comes at a significantly higher computational cost and longer training time.

Compared to Progressive Distillation

Progressive Distillation is a method that trains a student model to reduce the number of sampling steps to half of the teacher model. By repeating this process, Progressive Distillation enables the generation of high-quality images in fewer steps. While we acknowledge that the network architectures used in this comparison are not the same, this comparison can still be a valuable source of information.

All the information in the first row is taken from the original paper [4]. Progressive Distillation proposes some improvements to the architecture of iDDPM [5].

[1] Geng, Z., Pokle, A., Luo, W., Lin, J. and Kolter, J.Z., 2024. Consistency Models Made Easy. arXiv preprint arXiv:2406.14548.

[2] Song, Y., Sohl-Dickstein, J., Kingma, D.P., Kumar, A., Ermon, S. and Poole, B., 2020. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456.

[3] Liu, X., Gong, C. and Liu, Q., 2022. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003.

[4] Salimans, T. and Ho, J., 2022. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512.

[5] Nichol, A.Q. and Dhariwal, P., 2021, July. Improved denoising diffusion probabilistic models. In International conference on machine learning (pp. 8162-8171). PMLR.

Weakness 2: More qualitative comparisons

We thank the reviewer for the question. To address this concern, we have added more qualitative comparisons to our paper (please refer to Figure 3, 4 and Appendix F.5). We include a side-by-side comparison between LD3 and other time optimization methods, such as DMN and GITS. We use the same initial noise, matching NFE, and the same solver across all methods in these comparisons.

Question: LD3 and Distillation methods

We thank the reviewer for this question. The submitted paper does not mention that LD3 and distillation approaches are complementary. If we have overlooked specific sections, please let us know, and we will be happy to provide further clarification.

Since both LD3 and other distillation methods aim to reduce NFE while maintaining sample quality, we have compared LD3 against Consistency Distillation, Rectified Flow, and Progressive Distillation on the CIFAR-10 dataset (see Appendix C in the revised paper). While the comparisons are not fully apples-to-apples, our goal is to address the following question:

“How many sampling steps, with optimized discretization, are needed to match the performance of existing few-step generation methods, and how much training effort can we save?”

In Table 6, we also show that LD3 can further enhance InstaFlow, an already optimized model for few-step generation.

Thank you for your valuable feedback and for pointing out the areas of improvement in our manuscript.

Weakness 1: Evaluation protocol.

Thank you for your thoughtful comments and for raising this question. The metric we use is the standard FID score (as described in lines 368 and 369) computed against the reference set from each dataset. To improve clarity, we further detailed our evaluation protocol (lines [362-364]) in the revised manuscript.

We appreciate the reviewer's suggestion to add the teacher's performance to Tables 2 and 4. The number of steps of the teacher solver is often less than 20. This is because with ODE solvers, adding more NFE does not necessarily result in a performance gain (see paper [1]-Figure 2 or paper [2]-Figure 7). Appendix B - Experiment Detail already provides more details about the teacher solvers' settings.

[1] Karras, T., Aittala, M., Aila, T. and Laine, S., 2022. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35, pp.26565-26577.

[2] Zhang, Q. and Chen, Y., 2022. Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:2204.13902.

Weakness 2: Number of teacher NFE in Table 6.

Thank you for raising this important question. We selected the 8-step InstaFlow model because it already demonstrates strong performance with an excellent FID score on the MS-COCO dataset. To illustrate this, we evaluated InstaFlow with 10 and 20 NFEs, resulting in FID scores of 13.73 and 13.60, respectively. This indicates minimal improvement despite doubling the number of sampling steps. InstaFlow is specifically designed to generate high-quality samples in just a few steps. In addition to Tables 2 and 3, Table 6 shows that LD3 not only enhances standard diffusion models but also improves few-step generative models that are already optimized for efficient generation.

Weakness 3: The complexity of LD3.

We acknowledge that increasing model complexity, such as designing a more complex neural network with more parameters to predict time steps, might improve performance, as indicated by the results listed in Table 7. However, our primary goal in this work was to develop a model that learns the time discretization as simply and efficiently as possible. Efficiency is our key focus, as it aligns with the practical need for scalable methods in real-world applications. However, making the model that learns the time discretization more complex is an exciting direction for future work. Please let us know if we misunderstood your question or if we need to clarify further.

Weakness 4: Limitation of LD3

Thank you for highlighting this important point. We appreciate your suggestion and have updated the limitations section (lines [534-535]) in the revised manuscript to reflect this consideration, acknowledging it as a valuable area for future research.

Dear reviewer 4e1X,

Since the discussion deadline is approaching, could you kindly confirm if our response has addressed your concerns? If not, we would be happy to answer any follow-up questions and concerns.

Thank you again for your constructive review!

Weakness 3: Beyond image applications.

We thank the reviewer for raising the question about the applicability of LD3 to different tasks.

We aim to demonstrate that using the L2 distance alone already enables LD3 to perform comparably or better than the best heuristic methods. For example, on the FFHQ dataset with DPM_Solver++(3M), LD3 with L2 achieves FID scores of 26.01, 16.33, and 13.47, while the best heuristic baseline achieves 46.14, 22.79, and 14.01, respectively.

Additionally, LPIPS has proven to be an effective distance metric and has been utilized in distillation methods such as [1] and [2]. We acknowledge that LPIPS is specific to the domain of computer vision. Evaluating LD3’s performance on different domain tasks is an excellent approach to demonstrate its generalization capabilities and further strengthen the paper.

To address this, we extended our experiments to other applications, including point cloud generation and molecular docking. Our results indicate that LD3, using the L2 loss, effectively enhances few-step generation performance in these tasks.

We list the following new results in Appendix B of the updated paper. Appendix B also provides additional qualitative examples of these problems.

Application 1: Point Cloud Generation.

Point cloud generation creates 3D representations of objects or scenes, enabling applications in 3D modeling, AR/VR, robotics, and autonomous systems. We evaluate LD3 on a diffusion model [3] trained on the airplane category of the ShapeNet dataset [4].

Using the codebase from [3], we train LD3, generate samples, and assess quality through Chamfer Distance (CD) for reconstruction and metrics like MMD, COV, and JSD for generation. LD3 is trained on 32 noise-sample pairs from a teacher with 100 uniform ODE steps, using MSE loss with student NFEs=${4, 6, 8}$.

Results show LD3 improves generation performance under limited steps and, notably, outperforms the DMPG teacher in the CD-COV metric when NFE=8.

Application 2: Protein-Ligand Docking

Molecular docking predicts how small molecules bind to proteins, a key step in drug discovery. We evaluate LD3 with DiffDock [5], a diffusion model that outperforms traditional docking methods. DiffDock uses 18 steps (NFE) of reverse SDE for sampling; we adapt this by converting the SDE to an ODE and training LD3 on 10 noise-sample pairs with 18 ODE steps. Our student model uses 6 NFE, and we compare its performance to the default uniform discretization.

Using the evaluation protocol from [5], we sample 40 candidates per protein-ligand pair, rank them with a pre-trained confidence model, and measure RMSD against ground truth coordinates from the PDB test set. Predictions with RMSD < 2Å are considered good. The results show that LD3 significantly reduces the performance gap between a 6-step student and an 18-step teacher.

[1] Song, Y., Dhariwal, P., Chen, M. and Sutskever, I., 2023. Consistency models. arXiv preprint arXiv:2303.01469.

[2] Kang, M., Zhang, R., Barnes, C., Paris, S., Kwak, S., Park, J., Shechtman, E., Zhu, J.Y. and Park, T., 2025. Distilling diffusion models into conditional gans. In European Conference on Computer Vision (pp. 428-447). Springer, Cham.

[3] Luo, S. and Hu, W., 2021. Diffusion probabilistic models for 3d point cloud generation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 2837-2845).

[4] Chang, A.X., Funkhouser, T., Guibas, L., Hanrahan, P., Huang, Q., Li, Z., Savarese, S., Savva, M., Song, S., Su, H. and Xiao, J., 2015. Shapenet: An information-rich 3d model repository. arXiv preprint arXiv:1512.03012.

[5] Corso, G., Stärk, H., Jing, B., Barzilay, R. and Jaakkola, T., 2022. Diffdock: Diffusion steps, twists, and turns for molecular docking. arXiv preprint arXiv:2210.01776.

Thank you for your valuable feedback and for pointing out the areas of improvement in our manuscript.

Weakness 1: Typos in the paper.

Thank you for pointing out several typos in our paper. We have corrected these errors and updated the submission accordingly. We appreciate your careful review and attention to detail.

Weakness 2: Invertibility assumption of the theorem.

Thank you for your insightful observation about the invertibility assumption. The teacher solver $\Psi_*$, which employs small step sizes, can be assumed to be invertible. The challenge lies with the student solver, especially when it operates with a small number of NFEs.

The theorem bounds the KL divergence between the student distribution and teacher distribution when the loss in Equation (6) is zero. In this scenario, the student solver is already surjective since the image of $\Psi_\mathbf{\xi}$ and the image of $\Psi_*$ is identical (i.e., Im($\Psi_\mathbf{\xi}$) = Im($\Psi_*$)).

To further evaluate its invertibility (i.e., ensuring it is both injective and surjective), we conducted an experiment to test the injectivity of the student solver with a small NFE (e.g., NFE=2). Specifically, we generated 1M samples from the CIFAR-10 diffusion model using the student solver with NFE=2 and checked whether any two samples were identical. We observed no such cases. While we acknowledge that this experiment is insufficient to conclude that the student solver is always injective, it provides empirical evidence of how closely our solver aligns with the invertibility assumption required by the theorem.

The theorem provides a theoretical insight into the relationship between the relaxation of the requirement that the start points of the solvers are identical and the KL divergence of the student and teacher distributions. We believe it is an important addition to the paper, even if the invertibility assumption is not always met in practice.

Dear reviewer 4pzZ,

Since the discussion deadline is approaching, could you kindly confirm if our response has addressed your concerns? If not, we would be happy to answer any follow-up questions and concerns.

Thank you again for your constructive review!

I appreciate the authors providing additional explanation. However, I believe the point I raised about the invertibility assumption has not been fully addressed. Could you clarify the experiment described further? Specifically:

1. How were the initial noises for the 1M samples generated? Were they IID Gaussian samples, or were they chosen to be in close proximity to each other?
2. Was a class-conditional model used for this experiment? And if so, were all samples generated from the same class?
3. What method was used to determine if two samples were identical?

The additional experiments on point-clouds and protein-docking are very appreciated. However, I would like to request an expansion of the results to provide a clearer understanding of the benefits. Specifically, only a single heuristic schedule (linear schedule) is used as the baseline, yet the linear schedule is often far from the optimal choice. It would be helpful to include additional options, such as EDM's polynomial schedule, quadratic schedule, and cosine schedules as well.

Since the method shows promise for use in other domains, which was one of my main concerns, I’ll be raising my score.