Learning to Discretize Denoising Diffusion ODEs
Authors: Vinh Tong
Dung Trung Hoang
Anji Liu
Guy Van den Broeck
Mathias Niepert
Venue: ICLR 2025 Oral (2025)
Keywords: Diffusion models
Efficient Sampling
Ordinary Differentiable Equations
Primary area: generative models

Abstract: Diffusion Probabilistic Models (DPMs) are generative models showing competitive performance in various domains, including image synthesis and 3D point cloud generation. Sampling from pre-trained DPMs involves multiple neural function evaluations (NFEs) to transform Gaussian noise samples into images, resulting in higher computational costs compared to single-step generative models such as GANs or VAEs. Therefore, reducing the number of NFEs while preserving generation quality is crucial. To address this, we propose LD3, a lightweight framework designed to learn the optimal time discretization for sampling. LD3 can be combined with various samplers and consistently improves generation quality without having to retrain resource-intensive neural networks. We demonstrate analytically and empirically that LD3 improves sampling efficiency with much less computational overhead. We evaluate our method with extensive experiments on 7 pre-trained models, covering unconditional and conditional sampling in both pixel-space and latent-space DPMs. We achieve FIDs of 2.38 (10 NFE), and 2.27 (10 NFE) on unconditional CIFAR10 and AFHQv2 in 5-10 minutes of training. LD3 offers an efficient approach to sampling from pre-trained diffusion models. Code is available at https://github.com/vinhsuhi/LD3.

Published as a conference paper at ICLR 2025
LEARNING TO DISCRETIZE DENOISING DIFFUSION ODES
Vinh Tong' 12 Trung-Dung Hoang", Anji Liu 33, Guy Van den Broeck Mathias Niepert 1.2
University Stuttgart, ?IMPRS-IS, UCLA, *University of Bern
vinh. tong@ki. - -tt.:unninssttttttsstsoe stutgart..de
ABSTRACT
Diffusion Probabilistic Models (DPMs) are generative models showing competitive performance in various domains, including image synthesis and 3D point cloud generation. Sampling from pre-trained DPMs involves multiple neural function evaluations (NFEs) to transform Gaussian noise samples into images, resulting in higher computational costs compared to single-step generative models such as GANs or VAES. Therefore, reducing the number of NFEs while preserving generation quality is crucial. To address this, we propose LD3, a lightweight framework designed to learn the optimal time discretization for sampling. LD3 can be combined with various samplers and consistently improves generation quality without having to retrain resource-intensin neural networks. We demonstrate analytically and empirically that L.D3 improves sampling efficiency with much less computational overhead. We evaluate our method with extensive experiments on 7 pre-trained models, covering unconditional and conditional sampling in both pixel-space and latent-space DPMs. We achieve FIDs of 2.38 (10 NFE), and 2.27 (10 NFE) on unconditional CIFAR1O and AFHQv2 in 5-10 minutes of training. LD3 offers an efficient approach to sampling from pre-traines trained diffusion models. Code is available at https://gtthub.com
vtb.:/gittbb.comm.iiibvhh//... vinhsuh//LD3.
INTRODUCTION
Diffusion Probabilistic Models (DPMs) have emerged as a popular class of generative models, demonstrating competitive performance across various tasks, including image synthesis (Ho et al., 2020; Song et al., 2020b; Dhariwal & Nichol, 2021), text-to-image generation (Nichol et al., 2021; Rombach et al., 2022: Gu et al., 2022), 3D point cloud generation (Luo & Hu, 2021), and molecular generation (Hoogeboom et al., 2022). DPMs learn a multi-step transformation from random (e.g., multivariate Gaussian) noise the data distribution. While they achieve improved sample quality and diversity compared to single-step generative models like GANs (Goodfellow et al., 2014) 4) or VAES (Kingma et al., 2021), their multi-step nature incurs significant computational overhead.
Distillation-based methods are one category of approaches for accelerating DPMs (Meng et al., 2023; Liu et al., 2022b; Salimans & Ho, 2022; Song et al., 2023). These methods refine the denoising network to improve generation quality while reducing the number of sampling steps. Although these approaches can significantly enhance quality, they require costly re-training fine-tuning of the entire network. Moreover, distillation -abelllationbbsedd methods often face challenges such as information loss (Zheng al., 2024) and difficulties with conditional sampling (Meng al., 2023).
The second set of approaches capitalizes on the ability to sample from DPMs by solving corre- sponding diffusion Ordinary Differential Equation (ODE) (Song et al., 2020a: Lu al., 2022a: Zhang & Chen, 2022; Liu et al., 2022a; Song et al., 2020a; Lu et al., 2022b; Zhao et al., 2023; Zheng et al., 2024). These methods use more sophisticated numerical solvers for diffusion ODEs and require no neural network retraining. Solving diffusion ODEs involves mandatory step of selecting discretiza- tion time steps, and the choice of discretization points greatly influences sample quality. However, there does not exist time discretization strategy that works well for all dataset- dataset-model-One model-ODE solver combinations. Here, we propose an effective yet lightweight approach to learn good discretization time steps for any pre-trained DPM.

---

Published as a conference paper at ICLR 2025
Assuming the ODE solving pipeline is differentiable, which holds for most ODE solvers currently in use, our idea is to directly minimize the global truncation error to optimize the time steps, which employ
imput to the solver and hence can be treated as trainable parameters. Specifically, teacher ODE solver that takes small step sizes to approximate the gold standard solution of the ODE. The student with learnable discretized time steps is then tasked to mimic the teacher's s final output given the same input (i.e., the initial condition of the ODE).
One problem with this learning framework is the limited capacity of the student model, which can only optimize few parameters to mimic a stronger teacher, sometimes leading to "underfitting" and suboptimal performance. To address this, we propose a surrogate objective that is easier for the student to optimize. We validate this objective by theoretically proving its "closeness" to the original distillation objective, upper bounding the KI. divergence between the distributions induced by the teacher and the student solvers. The resulting algorithm, termed Learning to Discretize Denoising Diffusion ODEs (LD3), efficiently learns the time discretization by backpropagating through the ODE-solving procedure using the proposed surrogate loss. LD3 can be viewed as an additional step to further improve the sample quality of DPMs after making other design choices such as distilling the denoising model and choosing small set of random samples from a tractable noise distribution. Table summarizes the benefits that LD3 brings compared to related approaches.
ODE solver. Additionally, LD3 is efficient and only requires
Empirically, we test LD3 on both pixel space (Karras et al., 2022) and latent space DPMs (Rombach et al., 2022) across various resolutions, including 32x32 (e.g., CIFAR10 (Krizhevsky & Hinton, 2009)), 256 x56 (e.g., LSUN-Bedroom (Yu et al., 2015), ImageNet (Russakovsky et al., 2015)), and 512 text-to-image generation (Rombach al., 2022: Liu et al., 2023) as well as various types of conditions (e.g., class, text prompt) (Rombach et al., 2022). LD3 performs significantly better than existing time discretization heuristies, especially with small NFE (below 10). Additionally, LD3 can be trained in 5 to 40 minutes on a single GPU.
RELATED WORK
Table I: Comparison of LD3 with existing methods, including Watson et al. (2022), AYS (Sabour al., 2024), GITS (Chen et al., 2024), and DMN (Xue et al., 2024), based on several criteria. The first criterion is training speed: methods that optimize time discretizations in less than one hour are classified as "Fast", while others are considered "Slow". Training stability is true for methods not requiring variance reduction techniques or early stopping to exhibit stable training. A method is, respectively, aaler--deeetalle and network--daptabbl if it considers information about the solver and trained neural networks when optimizing the time steps. Finally, we examine whether model can directly optimize the global truncation error.
Criterion
Watson al. (2022)
AYS
GITS
DMN
LD3 (ours)
Training speed
Slow
Slow
Fast
Fast
Fast
Training stability
X
X
/
Solver adaptability
X
X
Network adaptability
V
V
V
X
Global error optimization
X
X
X
X
well-established approach to accelerating DPMs involves distilling high-quality denoising networks, which typically require many function evaluations. into models that perform the task in fewer steps with minimal performance loss (Meng et al., 2023: Salimans & Ho, 2022; Fan & Lee, 2023). However, this method requires expensive training (Liu et al., 2022b; Luhman & Luhman, 2021; Meng et al., 2023; Song et al., 2023) before the models can be used for sampling. comparison, distillation-basse based methods are significantly slower than approach- by several orders of magnitude- lack the flexibility to be used in plug-and-pll manner.
In addition to fine-tuning the denoising network, many methods focus developing more effective ODE solvers (Lu et al., 2022a; Zhang & Chen, 2022; Liu et al., 2022a; Song et al., 2020a; Lu et al., 2022b; Zhao et al., 2023). The key insight is that truncation errors accumulating during the backward sampling process can significantly degrade the quality of generated images, especially when NFE is reduced. To mitigate this, advanced ODE solvers required. However, since these solvers still rely
2

---

Published as a conference paper at ICLR 2025
on multi-step sampling, selecting an appropriate strategy is crucial. Current approaches often rely on handcrafted schedules, which may not be optimal.
Recent work has focused on optimizing time schedules. Xue et al. (2024) formulate an optimization problem aimed at identifying the optimal time discretization. They derive an upper bound for the global truncation error under the assumption that the score prediction error of the pretrained model is uniformly bounded by a small constant. However, this assumption is quite strong, as it leads to an optimization problem that depends solely on the noise schedule parameters, ignoring the influence of both the solver and the neural network. While this allows for a fast solution, typically found in a matter of seconds, it overlooks critical information about the pretrained model (trained dataset) and solver design. Furthermore, minimizing the upper bound does not necessarily equate to minimizing the actual global error.
Recent work by Sabour et al. (2024) empirically observe this problem when they derive bound to the divergence between the analytical ODE solution distribution and the numerical solution distribution. Their objective is challenging to optimize that they need to simulate many sampling trajectories and use a large batch size when performing optimization to reduce the variance and early stopping to prevent divergence. Consequently, their proposed approach is slow and hard to use. Instead of optimizing the global truncation error, Chen et al. (2024) optimizes the local truncation errors. However, their method ignores the information about the solver being used to solve the ODE and is not guarantee to optimize the glohal truncation error. Watson et al. (2022; 2021) propose the Differentiable Diffusion Sampler Search (DDSS) method, which aims improve the Kernel Inception Score by optimizing time discretization. By leveraging Kernel Inception Score (KID) to guide the optimization process, DDSS aims to enhance the quality of generated samples. However, their method requires a large amount of training samples and needs over 50k iterations with batch size 512 converge. We summarize some key differences between LD3 and similar approaches in Table 1.
3 BACKGROUND
Diffusion Probabilistic Models (DPMs) (Ho et al., 2020; Song et al., 2020b) involve a forward diffusion process that gradually converts samples following data distribution into samples from pre-specified noise distribution. Specifically, given a sample xo from the data distribution g, the forward process gradually perturbs it by adding Gaussian noise, which is chosen such that the distribution at time step T is a Gaussian distribution: XT (0, I)). For any 1 E [0, TJ. the Gaussian transition kernel is defined as
q(x:|xx) (aaxx,, II
where a₁ and de are two noise schedule hyperparameters designed so that the signal-to-noins ratio (SNR) ag /of is strictly decreasing when increasing t. This ensures that more information about Xo is discarded as increases.
To learn the data distribution g(xn), DPMs are tasked to recover information discarded by the forward process. This results in a so-called backward process that starts from the noise distribution Xy ~ q(xT) and moves backward through time to reconstruct Xo- Specifically, a neural network eo (x f) is trained to predict the noise added by the forward pass given x; and è by minimizing
Ex,,e,
- <l
(2)
where x, XgX0 OLE with (0 I) is the noisy sample at time t. t t~ T] is the time step. and (t) E + is a time-dependent weight for time step t. (x t) is a learnable deterministic function parameterized by 0 that predicts the noise added X,-
One way to sample from a trained DPM is to first draw Xy randomly and then use as the initial condition to solve the following diffusion ODE, whose solution is proven to match the data distribution ((Xo) if e (xe, 1) is optimal (Song et al., 2020b; Lu et al., 2022a; Zhang & Chen, 2022; Liu et al., 2022a; Song al., 2020a; Lu al., 2022b; Zhao al., 2023; Zheng al., 2024):
dx.
= f(1)x, + g2(
alog Oi: (1)=
do
20, ca <x 1), where f(1)
25(1))²
(3)
0t
Øt
Øt
3

---

Published as a conference paper at ICLR 2025
x M2
0-
6g v,(0)
ODE solver steps
-2-
-2
1 T2(1)
-4
0.0
0.25
0.5
0.75
1.0
fg 7y(2)
0.25
0.5
0.75
1.0
Time
Time
Xy
1/ (3)
(g(x5))
9.(x2)
vx(xx)
pd
(g(xx)
a (9,((x)) (x(x))
Integral curves
Teacher trajectories
Student injooctorics
Figure 1: Motivation and elaboration of LD3. (a) Directly optimizing the global truncation error loss Chard by minimizing the teacher and student outputs improves sample quality. (b) The surrogate objective Laoft that allows discrepancies in the initial condition (i.e., xT) between the teacher solver and the student solver is easier to optimize. (c) By optimizing the surrogate objective, LD3 learns better discretization strategies.
LEARNING TO DISCRETIZE DENOISING DIFFUSION ODES
Our main goal is reduce the computational burden of DPMs while maintaining their generation quality. In particular, we focus the small number of neural function evaluation (NFE) regimes with at most 10 evaluations of the denoising network generate a sample.
The ODE view of DPMs allows us to treat the problem of DPM sampling solving class of ODEs, where choosing discretization points is critical. Although one may hope that a universal discretization strategy exists for various DPMs, we find that both the structure of the DPM and the training data influence the optimal time discretization. with no single strategy working well for all cases. Instead, propose general yet efficient algorithm to learn discretization strategies for DPMs.
In the following, we first approach this problem under the framework of global truncation error optimization (Sec. 4.1). We then elaborate on potential underfitting problem of this optimization problem and propose a surrogate objective to improve sampling performance (Sec. 4.2). Finally, we provide details of the algorithm and the training process (Sec. 4.3).
4.1 LEARNING TO DISCRETIZE BY OPTIMIZING THE GLOBAL TRUNCATION ERROR
Given pre-trained denoising network ce(-.-) and an initial state X ~NO, 1f time step T, we solve the diffusion ODE stated in Equation (3), typically by applying numerical method such as Euler or Heun's method. This process is carried out over sequence of decreasing time steps 1 where T = to 1 dynamics, guided by the denoising network. The final solution, computed at time iN = 0 denoted as V (xx,, {t. ) ee). Recall that goal is to learn the set of time steps {t, j To ensure these time steps remain monotonic (i.c., -on-increasing). we encode the time steps by a monotonic function Te parameterized by such that for each i E N, 1000 notation, we now express V (xx,, {Te (i)) Ea) as Ve (xr). Additionally, we denote the distribution induced by the transformed state W (XT) as pe (Xo), which is obtained by first sampling Xy and then passing it through the network (e(:))
IN 0. At each time step, the state is updated according to the ODE
N, we have T (i) t,.' For simplicity in
Define I, (XT, Ee) as a teacher ODE solver that accurately solves the diffusion ODE and the distribution induced by V, as q(xo). We aim to minimize the KI. divergence between the teacher distribution g(xp) and the student distribution (x())):
((xo) Pe (xx) )
min (xx(((Xo)) |l pe(x))) = min Ex, xx) log
To minimize the objective mentioned above, we focus reducing the global truncation error. which entails training the student ODE solver Ve to closely mimic the behavior of the teacher ODE solver V_. We refer to this as hard teacher forcing. Formally, this gives rise to the following optimization
Mooe details about Fy will he discussed in Section 4.3.
4

---

Published as a conference paper at ICLR 2025
problem we propose:
min Lhard (E) := min Wa v. (xT))]]
(5)
$
where d(..-) .) is differentiable function that satisfies y d(x.y) > 0 and d(x, y) 0 if and only if x = y. Some examples are squared 12 distance yxx,y) = |x y and Learned Perceptual Image
Patch Similarity (LPIPS) (Zhang et al., 2018).
For any valid loss in the form of Equation (5), its global optimal solutions are also global optimal solutions to Equation (4). Our proposed objective improves upon previous works by directly opti- mizing the global truncation error and, consequently, the Kullback-Leible (KL) divergence between the student and teacher distributions. Unlike existing approaches that focus minimizing local truncation error (Chen et al., 2024) or optimizing an derived upper bound (Sabour et al., 2024; Xue et al., 2024), our method addresses the fundamental error more effectively. Additionally, our approach leverages information from both the student solver and the trained neural network, which are often overlooked in similar studies (Chen et al., 2024; Xue et al., 2024). This holistic consideration enhances the accuracy and performance of our solution.
4.2 OPTIMIZING DISCRETIZATION POINTS BY SOFT TEACHER FORCING
Despite having the same global optimum as the KI. divergence between the teacher-induced and the student-induced distribution, directly optimizing hardd (E) could lead to underfitting the objective, we need to ensure Vr (xr) = Ve (xxx for any Xy, which is hard as we are only allowed to optimize &. which typically contains no more than 20 parameters for student ODE solvers with low NFE. This issue is illustrated by the ID ODE shown in Figure I(a), where the green curves are the ground truth integral curves, and the red trajectories are from the teacher. The teacher matches the ground truth closely, it can take fine-grained steps solve the ODE. However, given the restriction that the student solver can only evaluate the ODE at three time steps before generating the output, an inevitable truncation error exists between the teacher and the student.
minimize
One way around this problem is to optimize the parameters of the student denoising network Ea in addition to &. However, this will significantly increase the sample complexity and the training time, limiting the method's efficiency and portability. Instead, we propose to relax the "hard" teacher forcing criterion. Specifically, for any XT and the corresponding output of the teacher Xu == V (xT). we only require the existence of an input Xyy that "close" to Xy, such that the student's 's output given X (i.e., Ve (xx))) matches Xg. We define B Bxx, roT) = {x' ||| x' S ror} the L.2 ball of radius rar around x. Take the ID ODE in Figure ((a) as an example. Although it may be impossible force the student solvers map the same point xy to the teacher's output, we can instead find nearby input X that does map to the teacher's output V (xT) (i.e., () The goal is that, by finding such Xy points close to X7,, as illustrated in Figure I(c), the student solver's distribution at t == O will still closely match the teacher's 's distribution. Note that xy. is only used during training to find the optimal discretization points. Formally, we relax the objective defined in Equation (5) into the following for 0:
min Cson(() min E ((0.00)
min Ex,,((x.rrr) ror)
a d(Vs( (xT))
(6)
The effectiveness of this relaxed objective depends on two ques- tions: (i) compared to (E), how much easier is it to optimize Looft ( given the fact that we only have a handful of learnable parameters (i.e., {); (ii) whether minimizing Csot (E) leads to stu- dent distributions that have small KL divergence with the teacher distribution (i.e., the ground-truth objective in Eq. (4)).
(an(E)
0.16
0.14
0.12
0.10
0.08
We start by showing positive evidence the first question em- pirically, small r suffices to ensure Looit (E) being much smaller than Lhara(E) after training. We experiment on a pre-trained DPM (Karras et al., 2022) AFHQv2 (Choi et al., 2020). We optimize the discretization points (i.e., І) with respect to L (E)
r=0.0
0.5
Figure 2: Cuon (E) drops signifi- cantly as we increase r.
using different and plot the training loss. We use the LPIPS distance (Zhang al., 2018) for d(.,) ). As shown in Figure 2, compared to r = 0.0, where Ccon. and Lhard are identical, the loss is significantly reduced relax the optimization problem by increasing
5

---

Published as a conference paper at ICLR 2025
We now move to the second question: if we can effectively minimize (E), can establish some form of guarantee of the student in terms of its KL divergence with the teacher distribution? We confirm this with the following theoretical result.
Theorem Ler V. and e be a teacher and student ODE solver each with noise distribution NO. I) E IRA,, and with, respectively, distributions q and pç. Assume both Vp and V are invertible. Let r > 0, if the objective from Equation (6) has an optimal solution E for with objective value 0, we have
(g(x) |I Pe- (x)) S 2 +rvd++++ Ex
2
where C(Ve (x)) = log det Juc- g (xx)).
The proof is provided in Appendix A.1. Intuitively, the theorem states that, if we can find an optimal solution 5* that minimizes Coofs then the KL divergence between the teacher (i.e., q(x)) and the student (i.e., Pe- (x)) can be upper bounded. The first two terms depend mainly on and the square root of the dimensionality: +d+1. I. Since is chosen to be quite small (e.g., 0.19 for CIFAR 10, NFE), the first two terms are effectively tight in practice. While it hard to establish an analytic bound for the third term, conduct numerical evaluations estimate its magnitude in practice and observe that it reduces with r. See Appendix A.2 for more details.
4.3 PRACTICAL IMPLEMENTATION
Alg Algorithm LD3
Require:
Student solver Vgg, teacher solver V. and
1: D t- [xxr, Xr, (x(x)) |xr ~ ~V(0 of x, = xr}
Generate data D
2: while not converged do
(Xy, Xx, V. (xx)) ~
4: CIE,E 1.55,,. xT) = L.IPP(( ee (x))) V. (xx))) subject xyy E B(xr, ror)
5: Update &, E" and xyy using the corresponding gradients V V.(. 5", x,)))
6: X/. Xy |||x;- xrlla r] r
x -x
Projected SGD
- xrll2
7: Update D with new x/
8:
end while
Now that we have justified the effectiveness of Loott we are left with the question of how to optimize it in practice. This is achieved by treating Looft (E) as jointly optimizing & and xy with the constraint that x/ within the r-ball of xy:
C(E, x²) = 1xx=~N10. (LPIPS(VE V. (XT))] subject to X/ E B(xT,ToT).
(7)
Note that we choose d LPIPS (Zhang et al., 2018) as the distance metric in setting. LPIPS is also a common choice in many distillation- models (Song et al., 2023; Salimans & Ho. 2022). As illustrated in Figure ((b), given input-output pair (Xr, V (x7)) from the teacher, forward through the student ODE solver steps with learnable parameters & and xy-- Backpropagation is then performed get the gradients w.r.t. LPIPS( W (x,)), (xr)). Finally,
projected SGD to enforce the constraint on x/ and use SGD to update &. Specifically, let x be the center of a r-sphere and x point. Then the projection of x/, which is the intersection of the line between x7 and x/r
with the sphere's surface, can be computed as Xp = XT -x,--
x/ XT
Parameterization. Starting from trainable vector & E N+1 we model T as a strictly mono- tonically decreasing function using a cumulative softmax function, followed by renormalization to the range [tmin: T):
(i( _ Emin (T - (min) + 4min where
N
(i)
T<(1)
Σ m=1
coftmax($))
{ Emax
Emin
Here, fmin is often utilized in training and sampling as a substitute for 0 in the diffusion model to mitigate numerical instability problems (Karras al., 2022; Song et al., 2023).

---

Published as a conference paper at ICLR 2025
During training. DPMs denoise images from real data, while during inference, they use predictions from previous steps. leading to discrepancies and errors (Ning et al., 2023; Li et al., 2023). As suggested by (Li et al., 2023), learn decoupled time steps as imput to the denoising model, while i determines the solver's step size. We parameterize if as = = (e + c. For simplicity, we still use Ve to refer to the student solver with decoupled time step variables.
Training. The final LD3 algorithm is shown in Algorithm 1. In line 1, we first generate training samples by first sample Xr ~ (O, IrI) and then compute the corresponding teacher outputs (x(x).. Initially, the same starting samples are used for both the student and the teacher (i.e., X7 Xr): resulting in dataset D = {(T. Xr, V, (xT))). In every iteration, we first compute the objective shown in Equation (7) (line 4), and then apply gradient-based updates to & and x/p (line 5). An additional projected SGD) step is applied to X/. to bound the distance between X/. and Xy (line 6). The algorithm terminates after convergence.
Since computing V (x() involves evoking the denoising network Eo multiple times, naively storing all intermediate outputs for efficient backpropagation would lead to memory overhead that scales linearly w.r.t. NFE. To reduce the memory overhead, we use the rematerialization technique proposed in (Watson al., 2022) only cache the intermediate xg. This leads to an almost constant memory overhead w.r.t. NFE. Please refer to Appendix D.3 for details.
5 EXPERIMENTS
Experiment setup. We evaluate pre-trained diffusion models across different domains. For pixel space models, we include CIFARIO (32x32) (Krizhevsky & Hinton, 2009), FFHQ (64x64) (Karras et al., 2019), and AFHQv2 (64x64) (Choi et al., 2020). For latent space models, we assess LSUN- Bedroom (256x256) (Yu et al., 2015) and class-conditional ImageNet (256x256) (Russakovsky et al., 2015) with a guidance scale of 2.0. Additionally, we consider text-to-image generation models, including Stable Diffusion v1.5 (Rombach et al., 2022) at 512x512 pixels with guidance scale of 7.5, and InstaFlow Liu et al. (2023).
We primarily assess LD3 using advanced diffusion ODE solvers, including DPM__olvers (Lu et al., 2022b), Uni PC (Zhao et al.. 2023), and iPNDM (Zhang & Chen, 2022). Additionally, we test LD3 with Euler, standard black-box ODE solver. To evaluate the performance of our learned discretization method, we compare it against 8 existing discretization methods. For commonly used discretization choices, we include time uniform (Lu al., 2022b; Ho et al., 2020), time quadratic (Song al., 2020a), time EDM (Karras et al., 2022), and time logSNR (uniform in X) (Zhang & Chen, 2022), with details provided in Appendix D.2. For more recent advanced discretization methods, we compare LD3 to DMN (Xue al., 2024), GITS (Chen et al., 2024), AYS (Sabour et al., 2024), and Watson et al. (2022). Since AYS (Sabour et al., 2024) and Watson et al. (2022) do not have published code, adhere to their settings and compare results with their reported metrics in Appendix E.
For CIFARIO, FFHQ, and AFHQv2, we use 100 samples for both training and validation and train LD3 for 7 epochs with a batch size of 2. We set proportional the dimensionality d and inversely proportional to the squared NFE: r = y x NFE2 where y = 0.001 in all experiments. For Latent Diffusion (Rombach al., 2022) on ImageNet and LSUN-Bedroom, we use 100 samples for both training and validation, with the training conducted over 5 epochs. Unless stated otherwise, we draw 50k samples for the evaluation using the FID score (Heusel et al., 2017) against reference data set, where lower scores indicate better quality. For text-to-image generation, we train Stable Diffusion (Rombach et al., 2022) and InstaFlow Liu et al. (2023) by randomly selecting 5 prompts from the MSCOCO dataset (Lin et al., 2015) and generate 10 training pairs for each prompt. We train them for 5 epochs. Details of experiment settings can be found in Appendix
LD3
DMN
GITS
LD3
DMN
GITS
"A medieval knight standing in custle courtyard." Figure 3: Side-by-side comparison of selected images generated with Stable Diffusion [iPNDM).
"A red apple on white table."
Left: NFE=6, Right: NFE=5.
7

---

Published as a conference paper at ICLR 2025
Main results. LD3 consistently improves generation quality across all solvers (Table 2), particularly
in low NFE settings. For instance, on the AFHQv2 dataset with NFE=4, iPNDM [LD3] achieves an FID score of 9.96, outperforming other approaches, which achieve a best score of 12.89. At NFE=10. LD3 delivers the best performance across all datasets, with FID scores of 2.38 on CIFAR 10, 2.27 on AFHQv2, and 3.25 on FFHQ. Additional results, including for the pixel space L.SUN-Bedroon (256x256) dataset and various NFE settings, are available in Appendix F.1 and Appendix F.4. Table compares LD3's optimized time steps with common time discretization methods, demonstrating clear advantage of the optimized steps over standard choices.
Table 2: FID comparison on CIFARIO, AFHQv2,
Table FID score at small NFE regimes on latent diffusion models. We investigate and ImageNet datasets. FID scores are computed based on 50k samples using the reference data set. Additional results can be found in Appendix F.4.
and FFHQ and the two solvers Uni PC and iP- NDM. We compare LD3 and two different time discretization optimization methods, DMN and GITS. FID scores are computed with 50k samples using the reference data set.
Method
NFE=4
NFE-5
NFE-6
NFE-7
Method
NFE=4
NFE=6
NFE-8
NFE=10
LSUN Bedroom-22555 256 (latent space)
CIFARIO
Uni [ni_PCCC (3M)
39.78
13.88
6.57
4,56
PC (3M)
43.92
4.41
3.16
UC [GITS]
70.93
47.37
22.33
17.27
Uni_CC [GITS]
25.32
11.19
5.67
3.70
Uni Uni_PC [DMN]
29.22
8.21
4.40
Uni PC [DMN]
26.35
8.09
5.90
2.45
Uni [L.33]
20.15
4,98
Uni PC [LDD3]
13.72
5.92
3,42
2.87
iPNDM (3M)
11.93
6.38
5.08
4.39
iNNDM (3M)
35.04
5.67
3.69
iPNDM [GITS]
76.86
59.17
28.09
19.54
iPNDM [GITS]
15.63
6.82
4.29
2.78
iPNDM [DMN]
11.82
6.15
4.71
5.16
iPNDM [DMN]
28.09
9.24
7.68
3.31
iPNDM [LD3]
8.48
4.52
4.31
iPNDM [LD3]
9.31
2.81
2.38
Teacher
3.06
Teacher
2.08
Imagenet-255 256 (latent space)
AFHQQ2
Uni PC (3M)
20.01
8,51
5.92
5,20
Uni_CC PC (3M)
33.78
3.81
Uni Uni_PC [GITS]
54.88
34.91
14.62
Uni PC [G[TS]
12.20
7.26
3.86
2.88
[Ci_PC [DMN]
16.72
7.96
7.54
7,81
Uni PC [DMN]
30.32
14.46
6.85
2.94
Uni PC [LD3]
9.89
5.03
4.46
4.32
PC [LD3]
12.99
2.90
2.84
iPNDM (3M)
13.86
7,80
6.03
5.35
iNDDM
23.20
iPNDM [GITS]
56.00
43.56
19.33
10.33
iPNDM [GITS]
12.89
6.10
4.03
3,26
iPNDM [DMN]
10.15
7.25
iPNDM [DMN]
33.15
16.01
10.12
3.22
iPNDM (LD3)
9.19
6.03
5.09
4.68
iPDDM [LD3]
9.96
3.63
2.63
2.27
Teacher
4.17
Teacher
2.11
Table 5: FID scores on Stable Diffusion v1.5.
FFH
We follow the standard FID evaluation with 30k
Uni_ PC (3M)
53.25
5.59
3.90
captions from MS-COCO (Lin al., 2015).
Uni_PC PC [GITS]
21.38
12.21
7.84
4.46
Uni. PC [DMN]
25.82
9.47
6.85
3.54
Uni PC [LD33]
21.00
5.97
3.50
3,27
Method
NFE=44
NFE=5
NFD=
NFE=7
iNDDM (3M)
36.54
16.44
8.11
5.39
iPNDM (2M)
17.76
14.41
13.86
13.76
iPNDM [GITS]
18.05
9.38
5.72
3,96
iPNDM [GITS]
18.05
14.11
12.10
11.80
iPNDM [DMN]
31.30
12.12
11:00
5.24
iPNDM [DMN]
21.70
17.30
13.68
11.88
iPNDM [D33]
17.96
6.47
3.97
3.25
iPNDM [L.33]
17.32
13.07
12.40
11.83
Teacher
2.54
Table 6: LD3 improves InstaFlow, few-step text-to-image generation model. We compare the FID score with 10k captions sampled from MS- COCO (Lin et al., 2015).
Table Comparison of FID scores on CIFARIO using iPNDM solver. We compare LD3's dis- cretization with commonly selected heuristies.
Discretization type
NFE=4
NFE=6
NFE=8
NFE=10
Method
NFE=2
NFE=4
NFE=6
Time LogSNR
35.04
11.80
3.69
InstaFlow
22.56
16.04
14.78
Time Uniform
266.20
229.39
205.24
185.28
InstaFlow [LD3]
14.12
Time Quadratic
139.72
68.82
37.82
23.40
Time EDM
29.78
9.95
5.41
3.80
Teacher (NFR 8, Uniform)
LD3
9.31
3.35
2.81
2.38
For latent space diffusion models, compare method to the best discretization approaches using Latent Diffusion models (Rombach et al., 2022) trained on LSUN-Bedrood and ImageNet, as shown in Table 4. Our model consistently matches outperforms baseline methods, particularly
8

---

Published as a conference paper at ICLR 2025
in low NFE scenarios. For example, at NFE=4, Uni PC [LD3] achieves an FID score of 20.15 on -UN--edroom, around 9 points better than the next best model, Uni PC [DMN].
We also test the performance of LD3 on text-to-image generation. Table shows performance of LD3 on Stable Diffusion. Generally, LD3 outperforms the default Time Uniform discretization while GITS and DMN only improve the performance given enough number of steps. InstaFlow (Liu al., 2023) a few-step eext-to-image generation, it can generate high-quality images in just a few steps using simple Euler solver. We further show that LD3 significantly boost the generation quality by optimizing the time step (Table 6). For example, LD3 improves FID score from 22.56 to 15.49 with only 2 NFE. One might question how different training prompts affect the final results. Interestingly, when we train our model using different sets of prompts, we achieve similar FID scores with small variance. Please refer to Appendix F.3 for more details. For qualitative comparisons, please refer to Figures 3 and 4 and Appendix F.5.
LD3
DMN
GITS
Time
Uniform
AFHQv2 DPM__olver-- NFE=6
ImageNet- Uni_PC, NFE=6
LD3
DMN
GITS
Time
Uniform
FFHQ DPM_Solver++, NFE=6
d) LSUN-Bedroom: iPNDM, NFE=6
Figure 4: Side-by-side comparison of random images generated by different pre-trained models across four datasets: AFHQv2, ImageNet, FFHQ, and L.UM-BBdroom. We compare LD3 with DMN (Xue et al., 2024), GITS (Chen et al., 2024), and Time Uniform discretization. For each dataset, samples from each column are created using the same initial noise, solver, and the number of NFE. We provide more side-by side-by-side bid comparisons in Appendix F.5
Ablation study. We investigate the importance of different components in our model as shown Table 7. Initially, compare variants with and without training E (cf. the decoupling technique described in Section 4.3) and then examine the effects of other factors on performance. This is particularly important because the variant without & is easily integrated into any ODE samplers, making detailed study of its behavior beneficial to the community. We observe that e" significantly contributes to the performance of LD3, which is expected as it effectively doubles the number of trainable parameters. For both versions, optimizing Lhard results in worse performance compared to optimizing Lsoft. For instance, on CIFAR10, with NFF=6, using Loott improves the FID score by 44% (from 13.46 to 7.51). The effect of Csoft is less pronounced without ; with the same NFE, Looft boosts the FID by 15% (from 14.20 to 12.03).
Table 8 highlights the importance of learning time discretization tailored to the ODE solver type. In particular, we train LD3 using DPM Solvertt, utilizing the optimized time steps for testing on both irrelevant, the optimized time steps would
DPM Solver++ and Euler. If the choice of solver perform well on both solvers. However, this is not the case. Optimized time steps for DPM Solver++ perform poorly for Euler and vice versa. For instance, the optimized time steps for

---

Published as a conference paper at ICLR 2025
yield an FID score of 42.44 on Euler, whereas those optimized for Euler achieve of 25.28 on the same solver.
Table 7: Ablation study on CIFARIO and FFHQ We conduct the experiment with (PM_Soover++<3MM
Table 8: We examine the importance of tailor- ing the time discretization to ODE solvers. We use Euler and DPM_ SMM_ooever++ with LD3 on the AFHQv2 dataset for this experiment. Each row shows the FID scores when LD3 is trained with a specific solver, while each column presents
Setting
NFEa4
NFE
NFE
Full setring
19.39
9.08
7.51
w/ Ş
Chard
19.62
10.17
13.46
33.25
15.13
11.78
CIFARIO
the FID scores when sampled with specific
Full setting
33.34
19.39
12.03
w/o C
w/ Chard L2
33.75
19.50
14.20
solver.
56.49
21.65
15.36
Full setting
27.99
13.32
7.53
w/ 5"
w/ Crard
30.52
14.20
9.12
Solver name
DPM_ Solver++
Fuler Solver
26.01
16.33
13:47
Solver++
FIIQ
NFE=
Euler Solver
Euler_Solver
231.00
25.28
Full setting
30.94
16.76
11.05
w/o E
w/ Ciard
31.04
17.64
12.58
NFE=6
DPM Solver++
14.16
38.73
21.97
16.57
Fuler_Solver
46.63
11.08
Efficiency and performance. LD3 demonstrates significantly faster training times compared to AYS (Sabour et al., 2024). As shown in Figure (a), LD3 can be optimized in less than an hour on a single GPU, while AYS typically requires several hours on multiple GPUs. For instance, at 10 NFE, our model needs approximately 36 minutes on single NVIDIA A100 GPU, whereas AYS requires to hours on 8 NVIDIA RTX6000s.
We generally only need 100 teacher noise-sample pairs for training and validation. Figure 5 (b) demonstrates that increasing the training size reduces the FID score. Since the performance plateaus at around 100 samples, we use (1 samples in practice to balance performance and efficiency.
mins
200
LO2 Irage Nen-35
A3CO)
A30:
FID-
40-
ARHOUO
CIARSU
0.16
DPM [DMN]
DPM Solvare (3R)
150
LB2 Stabre ci
AYS CIARLO 32
(1 A120)
30-
0.12
DPM_Solvere+ (LO3]
(1 R7X00000)
0.10
100
AYS Inagehet 250
(8 (U300000))
hour
0.08
50-
0.06
0
0
NFE=4
10
30 50
0.02
a) Number of Function Evaluations
b) Data Size
NFE-4
NFE-6
NFÉ-8
NFE-10
Figure 5: (a) Training time for LD3 and AYS (Sabour et al., 2024) on various NFE; (b) The effect of data size model performance (FID is evaluated with 5K samples, DPM_Solver++<3MM, NFF=4).
Figure 6: Comparison of global truncation error across time dis- cretization methods. RMSD is com- puted between student and goal samples using 100-step DDIM.
6 CONCLUSION AND LIMITATIONS
We introduced LD3, lightweight framework designed to reduce computational costs when sampling from pre-trained DPMs. LD3 learns time discretization for ODE sampling, significantly lowering the NFE needed to generate high-quality images with minimal training overhead. Our experiments across various datasets demonstrate that LD3 consistently improves sampling quality. For instance, LD3 reduced the FID score on CIFAR 10 (4 NFE) from 35.04 to 9.31 with 5 minutes of training on GPU. These results indicate that LD3 provides a more efficient approach to sampling from pre-trained diffusion models, with promising applications in image synthesis and beyond.
Limitations and broader impact. Despite significant performance improvements for a few NFE, LD3 still falls short of distillation-based methods regarding sample quality. A limitation of LD3 is that it needs to be trained separately for each given number of NFEs. Additionally, our model necessitates differentiable solver. which may not always be feasible. Using LD3 with advanced diffusion ODE solvers to generate fake content could also exacerbate the potential risks of LD3 being used for malicious purposes.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This work was funded in part by Deutsche Forschungsgemeinnchatt (DFG, German Research Foun- dation) under Germany's 's Excellence Strategy EXC 2075 390740016, the DARPA ANSR program under award FA8750022-2-0004, the DARPA CODORD program under award HR001 12590089, the NSF grant #IIS 1943641, and gifts from Adobe Research, Cisco Research, and Amazon. We acknowledge the support of the Stuttgart Center for Simulation Science (SimTech). VT and MN thank IMPRS-IS (International Max Planck Research School for Intelligent Systems) for the support. This work also received partial support from the Diabetes Center Berne.

---

Published as a conference paper at ICLR 2025
REFERENCES
Angel X Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, al. Shapenet: An information-rich 3d model repository. arXiv preprint arXiv:15/2.03012 2015.
Defang Chen, Zhenyu Zhou, Can Wang, Chunhua Shen, and Siwei Lyu. On the trajectory regularity of ode-based diffusion sampling. arXiv preprint arXiv:2405. 1326, 2024.
Yunjey Choi, Youngjung Uh Jaejun Yoo, and Jung-Woo Woo Ha. Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 8188-8197, 2020.
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. arXiv preprint arXiv:2210001776 2022.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794 2021.
Jiyu Fan, Ailing Fu, and Zhang. Progress molecular docking. Quantitative Biology, 7:83-89, 2019.
Ying Fan and Kangwook Lee. Optimizing ddpm sampling with shorteut fine-tuning. arXiv preprint 133.:22301..3362, 2023.
Zhengyang Geng, Ashwini Pokle. William Luo, Justin Lin, and J Zico Kolter. Consistency models made easy. arXiv preprint wxXiv:2466.15548 2024.
Ian Goodfellow, Jean Pouget-Abadie. Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. In Proceedings the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 10696-10706, 2022.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 23:6840-6851, 2020.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International conference machine learning, pp. 8867-8887. PMLR, 2022.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEF/CVE conference on computer vision and pattern recognition, pp. 4401-4410, 2019.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in Neural Information Processing Systems, 35:26565-26577 2022.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696-21707. 2021.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, University of Toronto. 2009.
Ravi Kumar, Manish Purohit, Zoya Svitkina, Erik Vee, and Joshua Wang. Efficient rematerialization for deep networks. Advances in Neural Information Processing Systems, 32, 2019.
12

---

Published as a conference paper at ICLR 2025
Mingxiao Li, Tingyu Qu, Wei Sun, and Marie- Franne-r Moens. Alleviating exposure bias in diffusion models through sampling with shifted time steps. arXiv preprint arXiv:23305.1558,, 2023.
Tsung- Yi Lin, Michael Maire, Serge Belongie, Lubomir Bourdev, Ross Girshick, James Hays, Pictro Perona, Deva Ramanan, C. Lawrence Zitnick, and Piotr Dollár. Microsoft coco: Common objects in context, 2015.
Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. arXiv preprint arXiv:2202.09778 2022a.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 03003, 2022b.
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, et al. Instaflow: One step enough for high-quality diffusion-based text-to-image generation. In The Twelfth International Conference Learning Representations, 2023.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu, Dpm-solver: fast ode solver for diffusion probabilistie model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint 2rX0v:22 2022b.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02338 2021.
Shitong Luo and Wei Hu. Diffusion probabilistic models for 3d point cloud generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2837-2845, 2021.
Chenlin Meng. Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint arXiv:2 12. 2021.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171, PMLR, 2021.
Mang Ning, Enver Sangineto, Angelo Porrello, Simone Calderara, and Rita Cucchiara. Input perturbation reduces exposure bias diffusion models. arXiv preprint arXiv:2301.11700 2023.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10685,10695, 2022.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115:211-252. 2015.
Amirmojtaba Sabour, Sanja Fidler, and Karsten Kreis. Align your steps: Optimizing sampling schedules in diffusion models. arXiv preprinl 1.Xi:2244411507, 2024.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005// 2022.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502. 2020a.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 1. 13446, 2020b.
13

---

Published as a conference paper at ICLR 2025
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469, 2023.
Daniel Watson, Jonathan Ho, Mohammad Norouzi, and William Chan. Learning to efficiently sample from diffusion probabilistic models. arXiv preprint arXiv:2106.0380 038000 2021.
Daniel Watson, William Chan, Jonathan Ho, and Mohammad Norouzi. Learning fast samplers for diffusion models by differentiating through sample quality. arXiv preprint arXiv:2202.05830,
2022.
Shuchen Xue, Zhaoqiang Liu, Fei Chen, Shifeng Zhang, Tianyang Hu, Enze Xie, and Zhenguo Li. Accelerating diffusion sampling with optimized time steps. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8292-8301, 2024.
Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of large-scale image dataset using deep learning with humans in the loop. arXiv 2015.
preprint
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:2204. /3902, 2022.
Richard Zhang. Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features perceptual metric. In CVPR, 2018.
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu. Unipe: A unified predictor- corrector framework for fast sampling of diffusion models. arXiv preprint arXiv:2302.04866 2023.
Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Dpm-solver-v3: Improved diffusion ode solver with empirical model statistics. Advances in Neural Information Processing Systems, 36, 2024.
14

---

Published as a conference paper at ICLR 2025
LEARNING TO DISCRETIZE DIFFUSION ODES
ADDITIONAL MATERIAL
A Analytical proofs
16
A.I Proof of theorem
16
A.2 Empirical measurement the bound.
17
Applying L.33 to different domains
18
B.1 Point cloud generation
18
B.2 Molecular Docking
20
C Comparison with few-step generation methods
21
C.I Comparison with Consistency Distillation
C.2 Comparison with Rectified Flows
21
Comparison with Progressive Distillation
D Experiment Details
22
D.1 Practical Implement Details
................................
22
D.2 Baseline discretization heuristics
25
................................
D.3 The Rematerialization Trick
26
E
More comparison with time discretization optimization methods
26
E.1 Compare to AYS on CIFAR10 and FFHQ
26
E.2 Compare AYS and GGDM-PRED+TIME on ImageNet-64
27
F
Additional Results
27
F.1
LD3 on pixel space SUUN-Bedroom-256 trained with iDDPM
27
F2 FID Progression during training
28
F3
Stable diffusion trained with different prompt
28
F4 Full result tables
29
F.5
Additional Samples
30
15

---

Published as a conference paper at ICLR 2025
A ANALYTICAL PROOFS
A.1 PROOF OF THEOREM
Theorem Let V and Ve be a teacher and student ODE solver each with noise distribution NO II., 0 E IRs, and with, respectively, distributions q and pç. Assume both v and Ve are invertible. Letr > O, if the objective from Equation (6) has an optimal solution &x for with objective value 0, we have
x2 2
DKL. (g(x) Pe- (x:) S +rvd+++1 Ex [[C(.(x)) C( (x))].
where C(Ve (x)) = log det Juc- (v (x))..
Proof. By assuming the invertibility of the solvers and the loss of Equation (6) having an optimal (zero loss and satisfying all røy-ball constraints) solution g", we have for every x ~ q(x) exactly one with V (x) = b and exactly one corresponding a with V (x) = a. Moreover, by definition of the loss objective and the fact that a an optimal and, therefore, also feasible solution, we have B(b, rar) and thus ||a - bll2 S ror.
Using the density function of the normal distribution, we can now write
gxx)
[xmg((x)
log
(x)
((b) det
(b) db
-xx log
(aa det dV)
[xx log(N(()) log (scee dVV((b)
log(W(a)) - log dett dVve- (a)
db
da
log (lact (b)
db 11) log (lace (V
da }
Ex [lo((N(b)) lo)()(a))] + Exmg(x))
I Vr
exp
(-2 1207
1 6a
da Excwy(x) [log (laoa d've (a) ) log (laee dVF((b)
)
-xxg(x)
log
| 0r 2m
(-1
db
exp
(8)
with (x) = a and 1/ (x) = b.
Now, simplifying the left expression from above, we obtain
I orV or exp
(-2 &²
)
Exxg(x) log
2
Exng(x)
Σ
(a
12)
(-2 a²
20z jm1
exp
-2
Rewriting 0; by 6 for 6 IR we have
d
d
Ex-g(x)
207 i=1
(2c,b; + 6))
=
of
1
B Bxxg(x)
Σ cgby
+
207
E
Ex-g(x)
Σ
Since l/a b1l2 E rar, have that (a= b,) < r20 and again with 0g Eg have
that Edh of < rz
Hence, have that
,
E
Ex-o(x)
Σ
2
Ex
[r202] =
20
20
2
i=
16

---

Published as a conference paper at ICLR 2025
where the last equality follows from the independence of the random variables in the multivariate distribution. Moreover, applying the Cauchy-Socn Schwarz inequality, we have:
d
d
1/2
1/2
[xx
Ex
gxx
Σ
c
2
o7
T
2
d
1/27
Ex xxx((x)
roT Σ
(9)
2
T
1/2
d
T Eb-N(0,021)) OT
where the second inequality follows from the definition of r. Since bj itd _N o2), the sum of
squares follows a Chi-squared distribution scaled by of (i.e., Σ by x Thus:
1/27
r
E
Σ
E
1o7x2
TTE
OT
T
To this end, applying Gautschi's inequality, we have:
ا=ههه r
d+1
S
(
which gives:
I' rv
d+1
S 2
rv
r (
Thus, the first term of Equation (8) is upper bounded by:
x² 2 +
2
DKL (g(x) pe- (x)) S
+rvd+1.
Combining with the second term of Equation (8), we get the final bound:
2²
DKL. (g(x) || pe- (x)) S +rvd+I+ Exxg(x) [[C(V.()) - C(Ve--()))l
2
where C(VE (x)) log det Jwg- (v=
A.2 EMPIRICAL MEASUREMENT THE BOUND.
50K
40K
(x)) C(V
30K
20K
IC(V, [C(V,(x))
10K
FFHQ AHFQV2
0
r=0.01
2.0
4.0
6.0
8.0 10.0
Figure 7: Estimation of CC(V. (x)) C(T (x)) across various values.
17

---

divergence (JSD). We use the same reference set Published as a conference paper at ICLR 2025
We compute the first two terms of the KD bound in Theorem since we know r and d. Although it would be hard to establish a tight analytic bound for the third term, we can still empirically estimate its magnitude in practical situations.
We test the empirical effect of r on the AFHQv2 (Choi et al., 2020) and FFHQ (Karras et al., 2019) datasets. We change from 10.0 down to 0.0 and approximate (C((V(())) C(V (x))x by first sampling b (0, I).... We then approximate
dV.(() db
(x), ~ log det
using teacher solver V (.), here the Uni_ solver with 20 NFE. Next we randomly sample an a E (bb, ror), to approximate
C(VE. (x)) = log ((cet d (a)
da
with the student solver Vg (-), here the Uni_PC solver with 7 NFE. Finally, we take the absolute difference
H(x) = (C(. ((Ve-) log (laet d'v
( log
dV (b)
db )
det
da
Since the computation of the Jacobian is very slow, repeat this for 100 samples and take the average to approximate Ex--((x)
[xxx)())[((xx]] (x)], which is the third term of the bound in Theorem Figure shows how [xx [x [x)] evolves as changes. We observe that as decreases so does Ex>g((x) ( (x)] Hence, we empirically verified that the overall bound in Theorem tightens with smaller r as desired.
However, we also observe that Ex~q(x<) [H (x)] does not converge to 0 for r + O. We hypothesize
that this related to the fact that we only compute an approximation of [xx) (H()) for randomly
sampled and b with maximum I.2 distance of and not for a and b for which we have that
Ve- (a) (b) as assumed by the theorem.
B APPLYING LD3 TO DIFFERENT DOMAINS
LD3 is applicable not only image synthesis but also various generation tasks. We demonstrate that LD3 can be seamlessly integrated into diverse applications, including Point Cloud Generation (Luo Hu, 2021) and Molecular Docking (Corso al., 2022). is important to note that our goal is not to compete with state-of sate-of-the-arr models for each task. Instead, we select one representative method per task to showcase how LD3 significantly enhances performance while using limited number of sampling steps.
B. POINT CLOUD GENERATION
Point cloud generation the process creating 3D representations of objects or scenes using discrete points in space. This task is crucial for wide range of applications, including 3D modeling, virtual and augmented reality, robotics, and autonomous systems. We test LD3 on a diffusion model named DMPG (Luo & Hu, 2021) trained on the airplain category of the ShapeNet dataset (Chang al., 2015).
We codebase from (Luo & Hu, 2021) to train LD3, generate samples, and evaluate the quality of these samples which allow us to evaluate the reconstruction quality of the point clouds using Chamfer Distance (CD). The generation quality measured with minimum matching distance (MMD), the coverage score (COV), and the following the original paper.
We generate 32 noise-sample pairs to train LD3 using 100 uniform ODE steps (teacher), We train three LD3 students with NFE=(4, 6, 8) using MSE loss function.
Table 9 shows that LD3 consistently improves point cloud generation performance when the number of steps is limited. Interestingly, when NFE=8. LD3 even surpasses DMPG teacher in CD-COV metric.
Code for DMPG available https: //gi wttp:h..c.mo.tt..//o/// ffusi on-pointtttttooo nt-cloud.
18

---

Published as a conference paper at ICLR 2025
DMPG
DMPG [LD3]
a) NFE=4
DMPG
DMPG (LD3)
b) NFE-6
DMPG
DMPG [LD3]
c) NFE-8
Figure 8: Selected point cloud generation examples. From top to bottom, we visualize student models with NFE values of 4, 6, and 8. The left column shows DMPG with the default time discretization, while the right column displays DMPG with LD3-optimized steps.
19

---

Published as a conference paper at ICLR 2025
Table Comparison of point cloud generation performance. We use the public code" provided by
(Luo & Hu, 2021) for sampling and evaluation. CD is multiplied by 10 and JSD is multiplied by 103.
Model DMPG
CD-MMD() CD-COV (%, 1)
JSD (1)
w/ Uniform
3.82
34,93
4.08
w/ EDM
3.50
42.83
4.38
NFE=4
w/ Quadratic
3.42
46,62
2,68
w/ DMN
3.44
45,30
3.04
w/ LD3
3.42
47.28
2.66
w/ Uniform
3.49
43.66
2.31
w/ EDM
3.37
47.28
2.51
NFE=6
w/ Quadratic
3.34
45,80
1.82
DMN
3,51
43.16
2.42
w/ LD3
3.34
46.62
1.79
w/ Uniform
3.39
46.63
1.71
w/ EDM
3.35
45.96
2.02
NFE-8
w/ Quadratic
3.34
46.95
DMN
46.62
1.61
w/ LD3
3.34
47.12
1.61
NFE=100
Teacher
3.27
47.12
1.03
B.2 MOLECULAR DOCKING
Molecular docking predicts how small molecules (ligands) bind to proteins, a critical process in drug discovery for identifying potential therapeutics by understanding atomic-level interactions (Fan et al., 2019). We evaluate LD3 using DifiDock (Corso et al., 2022). an advanced diffusion model that surpasses traditional docking methods. DifiDock's sampling process involves steps (NFEs) of a reverse SDE. We adapt this process by converting the SDE to an ODE and by generating 10 noise-sample pairs with 18 ODE steps (NFE) to train LD3. Our student model uses 6 NFE, and we compare the performance of LD3 with that of the default uniform discretization method.
Green Ground truth
Red DiffDock[[D3]
Blue DiflDocc(Unitorm)
Figure 9: Selected molecular docking example. Comparison of predicted structures from Diff- Dock[LD3] (red) and DiffDock[Uniform] (blue) with the ground truth crystal structure (green). The structures are generated using NFE. Left: Both models correctly predict the binding site. Right: DiftDock[LD3] (red) shows better alignment with the ground truth conformer (green) than DiffDock(Uniform) (blue). The RMSD values for DiffDock[LD3) and DiffDock(Uniform are 1.70 and 3.53, respectively.
We follow the evaluation protocol of prior work Corso al. (2022). For each protein-ligand pair, sample 40 candidates and rank them using a pre-trained confidence model. The primary evaluation metric the RMSD between the predicted ligand coordinates and the ground truth ligand from the PDB set. A prediction is considered accurate if its RMSD is less than 2 (Corso et al., 2022).
As shown in Table 10, LD3 significantly narrows the gap between 6-step student and an 18-step teacher.
20

---

Published as a conference paper at ICLR 2025
Table 10:
Comparison of molecular docking performance. We train LD3 and evaluate sample quality using the public code provided by (Corso et al., 2022).
DiffDock
Top-1 RMSD: %<2λ (1)
Top-5 RMSD: %<2Å (1)
w/ DMN, NFE=6
33.62
40.40
w/
EDM, NFE=6
33.43
39.66
w/
Quadratic, NFE=6
33.43
39.94
w/
Uniform, NFE=6
32.77
39.55
w/ LD3, NFE=6
35.31
40.68
Teacher, NFE=18
38.46
46.15
C COMPARISON WITH FEW-STEP GENERATION METHODS
While LD3 significantly improves generation performance in few-step regimes, it still falls short of 1-step distillation-bast methods regarding generation efficiency. However, for users with access to pre-trained diffusion model and limited computational resources, LD3 provides practical solution for efficiently optimizing sampling performance, without having to retrain the diffusion model. LD3 is especially beneficial when a slightly slower sampling process (compared to distillation) is acceptable, it requires only single GPU, drastically reduces training time by hundreds of times, and maintains high efficiency.
We compare LD3 against Consistency Distillation (Song et al., 2023), Rectified Flow (Liu al., 2022b), and Progressive Distillation (Salimans & Ho, 2022) on the CIFAR-I0 dataset. Our experiment aims answer 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?"
C.1 COMPARISON WITH CONSISTENCY DISTILLATION
Table 11: Comparison with Consistency Distillation: We take the estimated training time and the
number of GPUs required train the model from (Geng al., 2024). Both diffusion models use the NCSN++ architecture (Song et al., 2020b). The reported results for Consistency Distillation are the best-performing numbers CIFAR- 10, achieved NFE=2. For LD3, we use the iPNDM solver.
Architecture
Training time
#GPUs
FID/NFE
CD
NCSN+
day
A100s
2.93 / NFE=2
LD3 NCSN++]
NCSN+
<3 3 minutes
1 A100
15.61 / NFE=4
LD3 NCSN++ NCSN+
<3 minutes
1 A100
7.29/ NFE=5
LD3 [CCN+++ NCSN+
R 3 minutes
I A100
5.60 / NFE=6
LD3 [NCS++ NCSN+
< 6 minutes
1 A100
3.19 / NFE=7
LD3 [NCSN++] NCSN+
6 minutes
1 A100
2.90 / NFE=8
LD3 [NCSN++]
NCSN+
< minutes
1 A100
3.59 / NFF=9
LD3
[NCSN++]
NCSN+
6 minutes
1 A100
2.87 NFE=10
Consistency distillation (Song et al., 2023) diffusion-based distillation method designed to produce high-quality samples in just one or two steps. Training a consistency model on CIFARIO requires approximately one day 8 A100 GPUs. For fair comparison, train LD3 under the same conditions, using the same NCSN++ diffusion model backbone (Song et al., 2020b). With LD3, we achieve slightly better sample quality than consistency distillation when using NFE. While our inference is 4 times slower, training is 240 times faster and requires only 1/8 of the GPUs.
C.2 COMPARISON WITH RECTIFIED FL.WSS
Rectified Flow (Liu et al., 2022b) seeks to learn a straight flow from the prior distribution to the data distribution. However, the method still requires multiple steps generate high-quality samples. For
21

---

Published as a conference paper at ICLR 2025
Table 12: Comparison with Rectified Flow: We estimate the training time for Rectified Flow using the official code provided by its authors (Liu et al., 2022b), while the FID score is taken from Table 1-a of the original paper. For LD3, we employ the iPNDM solver.
Architecture
Training time
#GPUs
FID/NFE
RF
+DPM+
day
8 A100s
2.58 NFE=117
LD3 [DDPM++]
+DPM++
< 3 minutes
A100
9.31 N
LD3 [DDPM++]
+
< minutes
A100
8.17 NFE=5
LD3 [DDPM++]
+DPM++
R minutes
A100
3.35 N
LD3 [DDPM++]
+DPM+
< minutes
A100
3.97 NFE=7
LD3 [DDPM++]
+DP
<6 minutes
A100
2.81 NFE=8
LD3 [DDPM++]
DDPM++
<6 minutes
A100
2.51 NFF=9
LD3 [DDPM++]
DDPM+
23 6 minutes
A100
2.38 NFE-10
instance, achieving 2.58 FID necessitates 127 NFE, which is comparable LD3 applied to the same diffusion backbone but comes at a significantly higher computational cost and longer training time.
C.3 COMPARISON WITH PROGRESSIVE DISTILLATION
Progressive Distillation (Salimans & Ho, 2022) is a distillation method that trains 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 differ, include this experiment to strengthen our argument that training LD3 is significantly more efficient than training a distillation- based method.
Table 13: Comparison to Progressive Distillation. The data in the first row is sourced from the original paper by Salimans Ho (2022). Progressive Distillation introduces architectural enhancements to iDDPM (Nichol & Dhariwal, 2021).
Architecture
Training time
#GPUs
FID/NFE
PD
iDDPM*
1 day
8 TPUs
9.12 / NFE=1
PD
iDDPM"
1 day
8 TPUSs
4.51 / NFE=2
PD
iDDPM*
1 day
8 TPUs
3.00 /НE=4
PD
iDDPM*
day
8 TPUs
2.57 /NFE=8
[D3[DDPM+*
DDPM+
3 minutes
AA100
9.31 NFFE=
LD3 [DDPM++]
DDPM+
< minutes
A100
8.17 / NFE=5
LD3 [DDPM++]
DDPM+
23 3 minutes
A.00
3.35 NFE=6
LD3 [DDPM++]
+
6 minutes
1 A1100
3.97 / NFE=7
LD3 [DDPM++]
DDPM+
6 minutes
A1000
2.81 NFE=8
LD3 [DDPM++]
DDPM+
6 minutes
A
2.51 NFE=9
LD3 [DDPM++]
DDPM+
minutes
AA10
2.38 / NFE=10
D EXPERIMENT DETAILS
In this section, we provide more experiment details for each setting, including the codebases and the configurations for evaluation.
D.1 PRACTICAL IMPLEMENT DETAILS
We denote W. H. and C as the width, height, and number of channels of image, respectively. Thus we have C' x W x H. Similarly, W H', and C' represent the corresponding dimensions in the latent space for the Latent Diffusion model (Rombach et al., 2022), d = C' x W x H'.
22

---

Published as a conference paper at ICLR 2025
Teacher Solver and Data Generation. Generally, we select the rest-performing solver 20 NFE as our teacher solver, except for ImageNet and -oxttt generation tasks. For ImageNet, we select 10 NFE as the teacher because it performs better than higher NFE, as noted in (Sabour et al., 2024). For text-to-image generation, we also use 10 NFE, since increasing NFE does not significantly enhance our model's performance. We generate 50 to 100 samples for training. Detailed choices of solvers and NFE for the teacher solver are provided in Appendices D.1.1 to D.1.3.
Optimizer and Trainable Parameters. We update three primary parameter sets during training: the step-size parameters &, the coupling time step imput to the denoising model & and the starting point x-. Their optimizers are RMSprop for & and SGD for both e and X7. The learning rates are denoted as Ig. lgs, and lx, We set 0.005 for pixel space datasets and lg = 0.001 for latent space datasets, while le- and lx, are NFE-dependent. Further details are available in Appendices D.1.1 to D.1.3.
Initialization. We evaluate several baseline heuristics for & by computing their validation loss (a process taking only few seconds) and choose the one that minimizes this loss for parameter initialization. We initialize ç" as zero vector.
Training. Initially, freeze ee and optimize only & and x during the first two epochs (we call it phase 1). After this phase, we jointly update all parameters (we call it phase 2). At the end of each epoch, we update x for all samples in the validation set, without updating & and x/y, with respect to Leoft- We save the checkpoint that minimizes Loon. for the validation set after each training iteration.
We evaluate our model using the FID score with 50,000 randomly generated samples. For ImageNet, we generate an equal number of samples for each class to ensure a balanced FID evaluation.
Evaluation.
We provide more detail about our training setting in Appendices D.1.1 to D.1.3.
D.1.1 PIXEL SPACE DIFFUSION ON CIFARIO, FFHQ, AND AFHQv2
Pre-trained model:
EDM (Karras et al., 2022)
Image resolution:
W H 32, = 3 for CIFAR10.
- W = - 64, = 3 for FFHQ and AFHQv2.
Training/Validatee data:
Size: 50/50.
-The data are generated with the best baseline at NFE of 20, which is Uni_PC PC solver with Time LogSNR discretization for all three datasets.
-
Phase and Phase 2:
Phase 1: 2 epochs.
Phase 2: 5 epochs.
Starting point x/, and r:
- SGD optimizer, Lx, 32.0 NFE' momentum 0, weight decay = 0.
- 0.001x CXC NFE
Time discretization parameters
R.Mpprpp optimizer, 1g = 0.005, momentum = 0.9, weight decay = 0.
Gradients are clipped by the norm of 1.0.
If the validation loss does not improve for consecutive iterations (patience 5).
reduce the learning rate by a factor of 0.8. Stop decay if the learning rate reaches minimum of 5 x 10-5
23

---

Published as a conference paper at ICLR 2025
Time coupling parameters E-:
SGD optimizer, 1ę- NFE'
0.1
Gradients are clipped by the norm of 1.0.
If the validation loss does not improve for 5 consecutive iterations (patience = 5),
reduce the learning rate by factor of 0.8. Stop decay if the learning rate reaches minimum of 1 X 10-6
D.1.2 LATENT SPACE DIFFUSION ON BUUN-BEDROOM AND IMAGENET
Pre-trained model:
- Latent Diffusion (Rombach et al., 2022)
Image resolution:
W H 256, C C=3 3.
W" H' 64, C' 3.
Guidance scale: 2.0 (for ImageNet).
Training/Validation data:
Size: 50/50.
The data are generated with the best baseline at NFE of 20 for and 10 for ImageNet, using the Uni PC solver with Time uniform discretization for LSUN- Bedroom and Time quadratic for ImageNet.
Phase and Phase 2:
Phase 1: epochs.
Phase 2: 3 epochs.
Starting point x and r:
SGGD optimizer, 1x, NFE momentum 0, weight decay = 0.
12.0
r- 0.001 xW H''''' NFF
Time discretization parameters
RMSprop optimizer, Ig 0.001, momentum 0.9, weight decay 0.
Gradients are clipped by the norm 1.0.
If the validation loss does not improve for consecutive iterations (patience = 5).
reduce the learning rate by factor of 0.8. Stop decay if the learning rate reaches minimum of 5 x 10-5
Time coupling parameters --:
SGD optimizer, 1ę= 0.001
NFE-
Gradients are clipped by the norm of 1.0.
If the validation loss does improve for consecutive iterations (patience = 5),
reduce the learning rate by factor of 0.8. Stop decay if the learning rate reaches minimum of x 10
D.1.3 TEXT-TO-IMAS GENERATION WITH STABLE DIFFUSION
We randomly sample 5 captions from the MS COCO dataset (Lin al., 2015), which
1. "Two individuals learning to ski along with an instructor."
2. "A man sitting on a chair that is on a deck over the water."
3. "A dog sitting at a table in front of a plate."
4. "Four people sit around eating food outside together." "
5. "A cut dips its paws into cup nightstand."
We utilize Stable Diffusion v1.5 (Rombach et al., 2022) with the iPNDM solver and the NFE equal to the number of NFE of the student plus one. We used optimized GITS timestep teacher.
24

---

Published as a conference paper at ICLR 2025
Pre-trained model:
-table Diffusion (Rombach et al., 2022)
Image resolution:
W 512, C 3.
W" H' 64. C' 4.
Guidance scale: 7.5.
Training/Validation data:
- -Size: 25/25.
The data are generated with more than one NFE compared to student. The discretization of teacher is optimized GITS.
Phase and Phase 2:
Phase 1: 2 epochs.
Phase 2: 3 epochs.
Starting point X² and r:
12.0 SGD optimizer, bx NFE momentum = 0, weight decay = O.
0.001 "''
NFE
Time discretization parameters E
RMSprop optimizer, 1g = 0.001, momentum = 0.9, weight decay = 0.
- Gradients are clipped by the norm 1.0.
If the validation loss does not improve for 5 consecutive iterations (patience = 5).
reduce the learning rate by factor of 0.8. Stop decay if the learning rate reaches minimum of 5 x 10-5
Time coupling parameters 5":
SGD optimizer, 1ę- NFE
0.001
Gradients are clipped by the norm 1.0.
the validation loss does not improve for consecutive iterations (patience = 5).
reduce the learning rate by factor of 0.8. Stop decay if the learning rate reaches minimum of 1 x 10-5,
D.2 BASELINE DISCRETIZATION HEURISTICS
We compare our learned time discretization with the following discretizations heuristics:
Polynomial discretization (Time quadratic, Time uniform): function. Specifically:
This discretization is polynomial
(tmux Emin)
fmin
tmax
T, i 0,
(10)
here p is often set to or 2 (Song et al., 2020b; Ho et al., 2020; Lu et al., 2022b; Song et al., 2020a) which corresponds time quadratic and time uniform discretization.
Time EDM discretization: First introduced by (Karras et al., 2022), Time EDM discretization has been shown be effective with Heun's solver on EDM pre-trained model (Karras et al., 2022):
(t4))
max
N
(
mmax)
(11)
Time LogSNR: This schedule uniformly separate the logSNR. Specifically: -
a.t.))
A(4)
N (Amax -maa) + Amin where A(4,)
(12)
o(4))
this schedule is often used with solvers from (Lu et al., 2022a;b; Zheng al., 2024).
25

---

Published as a conference paper at ICLR 2025
D.3 THE REMATERIALIZATION TRICK
Although the student model has few trainable parameters, the memory cost of maintaining the forward pass state scales linearly with the number of inference steps when taking gradients concerning model samples. This scaling can quickly become unfeasible given the large size of DPM architectures. To address this, we adopt the gradient rematerialization technique proposed in (Watson et al., 2022), following (Kumar al., 2019). Rather than storing specific forward pass outputs for backward pass computations, we recompute them as needed, trading (NN) memory cost for ( computation time. This involves rematerializing calls to the pre-trained DDPM while keeping all progressively denoised images from the sampling chain in memory. Our model exhibits rapid convergence, even with a small batch size of or 2.
E MORE COMPARISON WITH TIME DISCRETIZATION OPTIMIZATION METHODS
In this section, we further compare our method with some time discretization methods such as AYS (Sabour al., 2024) and GGDM PRED (Watson et al., 2022).
AYS (Sabour al., 2024) focuses aligning the discretized sampling process with the continuous- timm process in diffusion models. The authors introduce the KL-divergence Upper Bound (KLUB) objective to formalize the optimization of the sampling schedule. This objective measures the discrepancy between the discrete and continuous processes, guiding the optimization toward schedules that minimize this discrepancy. However, optimizing this objective function is challenging due to its high variance, often requiring hours of running time on multiple GPUs.
GGDM PRED TIME (Watson et al., 2022) a Differentiable Diffusion Sampler Search method designed to improve the Kernel Inception Score by optimizing time discretization. However, their method requires 50,000 training iterations with a batch size of 512 to achieve convergence.
We compare our method with AYS in Appendix E.1 on CIFARIO and FFHQ Since the official implementation of AYS has not been published, we have used the numbers reported in their paper for the settings we have in common. We also compare our method with AYS and GGDM PRED + TIME on an unconditional ImageNet-64 dataset Appendix E,2 trained with a diffusion model from (Nichol & Dhariwal, 2021).
E.1 COMPARE TO AYS ON CIFARIO AND FFHQ
Table 14: Comparison of LD3 to AYS (Sabour al., 2024) (results sourced from (Sabour et al., 2024)) and the best baseline (Best) among time uniform, time quadratic, time EDM, and time LogSNR. Unlike the experiments Table 2, which use DPM_Solver++ (3M) (order = 3), we use DPM_Solver++ (2M)
(order = 2) following (Sabour al., 2024).
Pre-trained DPM (Karras et al., 2022)
CIFAR10
FFHQ
NFE
10
20
10
20
DPM_ Solver++ (2M) (Best)
5.07
2.37
7.07
3.41
(2M) (AYS)
2.98
2.10
5.43
3.29
(2M) LD3
3.38 = 0.64
2.36 = 0.02
3,98 + 0.10
2,89 + 0.03
We compare LD3 with AYS (Sabour et al., 2024) CIFARIO and FFHQ Since the implementation of AYS has not been published, we use their reported result and run LD3 on CIFAR 10 and FFHQ with pre-trained EDM model (Karras et al., 2019) using DPM Solver++ (2M). It can be seen from Table 14 that LD3 achieves competitive results with AYS Specifically, YAS outperforms LD3 on CIFAR 10 while our model achieves significantly better results on the FFHQ dataset.
26

---

Published as a conference paper at ICLR 2025
E.2 COMPARE TO AYS AND GGDM-PRED+TIME ON IMAGENET-64
Table 15: Comparison of LD3 AYS (Sabour et al., 2024) and GGDM-PRED+TIME (Watson et al.,
2022) on unconditional ImageNet-64 dataset. Baseline performance is sourced from (Sabour et al., 2024).
Model
Solver
Discretization method
NFE=5
NFE=10
NFE=15
NFE=20
NFE-25
DDIM
Time Uniform
135.4
40.70
28.54
24.23
22.13
3M steps
DDIM
Time Quadratic
409.1
148.6
67.65
45.60
36.11
GGDM PRED + TIME
Learned Schedule
55.14
37.32
24.69
20.69
18.40
DDIM
Uniform
145.01
42.51
30.32
26.60
24.77
1.5M steps
DDIM
50.38
29.23
24.21
22.26
21.42
DDIM
1.D3
49.79
28.61
22.93
21.16
19.81
Following Sabour al. (2024), compare our model with GGDM + PRED TIME (Watson et al., 2022) on the ImageNet-64 dataset trained using iDDPM (Nichol & Dhariwal, 2021). The sampler being used is DDIM. Since the code for GGDM + PRED TIME has been published, we rely on the reported numbers in their paper. We ran our model using the published checkpoint from Nichol Dhariwal (2021), which was trained for 1.5 million steps. Table 15 shows that, despite testing on less trained model, LD3 significantly outperforms GGDM PRED + TIME in the small steps regime (NFE < 20). Additionally, we compared LD3 with AYS on this dataset, and our method shows better results in all settings.
F ADDITIONAL RESULTS
F.1 LD3 ON PIXEL SPACE
-256 TRAINED WITH IDDPM
To prove the effectiveness of our method on high-resolution pixel space diffusion models, we evaluate dataset trained with iDDPM (Nichol & Dhariwal, 2021). We use their implemented solver, DDIM sampler, and compare the FID score generated using different discretization strategies. Table compares LD3 with commonly used discretization strategies. LD3 consistently demonstrates superior performance across most NFEs (Number of Function Evaluations), achieving significantly lower FID scores than the other methods. For instance, at NFE 4, LD3 records an FID of 92.43, outperforming Time Uniform (135.57), Time Quadratic (129.88) and Time LogSNR (178.69). This trend continues as NFE increases, with LD3 maintaining the lowest FID scores up NFE 9. However, at NFE 10, Time Uniform slightly surpasses LD3, achieving an FID of 22.62 compared to LD3's 23.03, marking the only instance where LD3 is not the top performer.
LD3 on
Table 16: FID comparison between LD3 and commonly used discretization strategies LSUN- Bedroom-256, pixel space pre-trained model iDDPM (Nichol Dhariwal, 2021), ) and DDIM sampler.
NFE
Solver
Discretization
6
7
8
10
Time LogSNR
178.69
142.71
106.64
85.24
67.72
56.67
47.80
Time Uniform
135.57
88.04
61.06
44.30
33.66
26.81
22.62
DDIM
Time Quadratic
129.88
101.76
81.18
65.94
52.77
45.86
39.47
Time EDM
294.54
282.52
273.05
265.34
258.88
252.15
245.67
LD3
92.43
65.95
49.30
37.85
31.73
25.82
23.03
27

---

Published as a conference paper at ICLR 2025
F.2 FID PROGRESSION DURING TRAINING
FID
CIFARIO DPM,
50
Uni_PC
40
30
20
0
20
40
60
80
Training iteration
Figure 10: FID progression during training. We report the FID values during training for CIFAR10 and FFHQ datasets, using 4 NFE with the DPM Solver++ (3M) for CIFARIO and the Uni PC (3M) solver for FFHQ. The FID, computed based on 5000 generated samples, is evaluated each time the validation loss decreases.
In our method, we save the optimal parameters based on the validation loss criterion. Figure 10 shows a decreasing trend in FID values throughout the training process. This trend suggests correlation between FID and validation loss, indicating that FID generally decreases as the validation loss decreases. This observation further validates the effectiveness of our method.
F.3 STABLE DIFFUSION TRAINED WITH DIFFERENT PROMPT SETS
To demonstrate the robustness and independence of the method with respect to the training prompts, different subsets of 5 prompts for training. Insterestingly, despite varying the prompts, the performance remains consistent, as shown in Table 17.
Table 17: FID score averaged over 5 runs on Stable Diffusion v1.5 using the iPNDM [LD3] solver.
NFE=4
NFE=5
NFE=6
NFE=7
FID 17.34
0.02
13.22 + 0.12
12.48 + 0.08 12.04 + 01.2
28

---

Published as a conference paper at ICLR 2025
F.4 FULL RESULT TABLES
Table 18: Full FID comparison on CIFARIO 32x32 (Krizhevsky & Hinton, 2009).
NFE-4
NFF-8
NFF=
Time LogSNR
46.59
24.99
12.16
6.88
4,62
3.54
3.08
Time Uniform
282.94
263.80
249.75
238.26
227.79
217.86
208.34
Time Quadratic
170.95
124.73
91.56
69.93
54.97
44.59
37.21
(3M)
Time EDM
50,91
32.13
18.64
11.71
8.44
6.51
5.16
DPM
GITS
31,58
16.28
11,07
6.04
5.61
3.96
3.51
DMN
39.31
22.88
12.37
8.21
7.23
4.35
2.67
LD3
19.39
9.08
7.51
5,90
3.42
2.90
3.09
Time LogSNR
35.04
20.52
11.80
7.69
5.67
4.37
3.69
Time Unifoem
266.26
242.94
229.39 39
217.13
205.24
194.72
185.28
iPNDM
(3M)
Time FDM
Time Quadratic
139.72
29.78
96.67
17.35
68.82
9.95
52.32
7.61
37.82
5.41
27.51
5.00
23.40
3.80
GITS
15.63
10.12
6.82
4.48
4,29
3.10
2.78
DMN
28.09
16,76
9.24
6.85
7.68
4,75
3.31
LD3
9.31
8.17
3.35
3,97
2.81
2.51
2.38
Time LogSNR
43.92
24.01
13.12
6.63
4.41
3.16
Time Uniform
282.77
263.61
249.42
237.81
227.21
217.09
207. 40
Time Quadratic
164.78
117.28
85.14
65.25
51.91
42.48
35.52
(3M)
Time EDM
50.65
34.24
19.56
12.68
9.65
7.83
6.12
Uni
GITS
25.32
13.79
11.19
6,80
5.67
4,40
3,70
DMN
26.35
12.93
8.09
5.42
5.90
3.62
2.45
LDD3
13.72
6.39
5.92
2.98
3.42
2.81
2.87
Table 19: Full FID comparison on AFHQv2 64 (Choi et al., 2020).
Solver
Discretization
NFE=
NFE=5
NFE=6
NFE=7
NFE=8
NFE=9
NFE=10
Time LogSNR
27.82
17.88
10.72
6.15
4,28
3.64
3.19
Time Uniform
153.20
143.45
134.46
125.39
116.13
107.01
98.25
Time Quadratic
68.78
42.48
30,92
24.70
20.94
18.48
16.73
DPM (3M)
Time EDM
19.65
13.34
9.64
7,82
6.54
4.11
GITS
16.79
12.62
7.75
5.55
4.37
3.98
DMN
37.54
30.09
18.41
14.66
9,10
5.94
3.24
LD3
13.86
9.48
5.47
4.27
3.63
2.73
2.68
Time LogSNR
23.20
15.63
9.55
5.98
4.49
3.75
3.19
Time Unifoem
159.22
139.47
125.17
113.68
102.99
92.57
82.89
iPNDM
(3M)
Time Quadratic
Time FDM
53.67
15.35
32.27
9.01
23,80
6.26
20.32
4.73
18.18
3.83
16.52
3.55
15.10
3.02
GITS
12.89
8.72
6.10
5.48
4.03
3.47
3.26
DMN
33.15
26.11
16.01
13.03
10.12
6.68
3.22
LD3
9.96
6.09
3.63
2.97
2.63
2.25
2.27
Time LogSNR
33.78
13.01
8.27
5.07
4.60
4.46
3.81
Time Uniform
153.24
143.46
134.27
124.98
115.52
106.21
97.30
Time Quadratic
66.38
41.56
30.62
24.26
20.47
18.04
16.33
(3M)
Time EDM
23.74
15.84
10.24
8.37
7.75
6.26
Uni
GITS
12.20
9.45
7.26
4,85
3.86
3.26
2.88
DMN
30.32
23.04
14.46
11.53
6.85
4.36
2.94
I.D3
12.99
6.35
3.81
3.18
2.90
2.87
2.84
29

---

Published as a conference paper at ICLR 2025
Table 20: Full FID comparison on FFHQ 64x64 (Karras al., 2019).
Solver
Discretization
NFE=4
NFE=5
NFE=6
NFE=7
NFE=8
NFE=9
NFB=10
Time LogSNR
46.14
22.79
14.01
8.63
6.18
4,86
4.18
Time Uniform
179.07
165.50
153.56
142.80
133.07
124.26
116.30
Time Quadratic
94.11
70.17
54.93
44.64
37.41
32.12
28.30
(3M)
Time EDM
39.59
23.81
15.29
11.17
9.67
8.45
7.01
DPM
GITS
29.09
17.54
12.73
9.72
7.60
6.30
4.99
DMN
40.23
26.31
14.73
11.21
10.30
7.37
4.29
LD3
27.99
13.32
7.53
4,86
4.42
3.37
3.55
LogSNR
Uniform
39.49
17.52
Quadratic
71.51
51.66
39.17
31.32
26.50
21.22
iPNDM
(3M)
Time EDM
29.35
17.52
11.44
8.76
6.86
4.93
GITS
18.05
12,91
9.38
7.56
5.72
4.75
3.96
DMN
31.30
20.93
12.12
10.17
11.00
8,36
5.24
LD3
17.96
11.98
6.47
4.88
3.97
3.57
3.25
Time LogSNR
53.25
20.20
11.24
7.09
5.59
3,90
Time Uniform
178.68
165.03
152.97
142.09
132.22
123.28
115.20
PC
(3M)
Time Quadratic
89.75
65.94
41.47
34.67
26.45
EDM
47.74
21.38
25.82
LD3
Table 21: Full FID comparison on LSUN- Bedroom 256x256 (Yu et al., 2015).
Table 22: Full FID comparison ImageNet
256x256 (Russakovsky et al., 2015). We cor- rected a typo in the table: The results for the first baselines of DPM Solver++ and iPNDM solvers in this table were previously swapped.
Sulver
Discrotization
NFT=4
NFT==5
NFT=66
NFE=7
Solver
Discretizaioon
NFE:4
NEL-5
NFE=6
NFE=?
Time LogSNR
$4.61
23.24
11.52
7.26
Time LogSNR
80.44
35.81
16.95
11.38
Time Uniform
26.07
7.51
Time Quadiatic
Time Quadratic
47,61
21.29
12:42
9.68
Time FDM
244.49
233.18
221.56
210.02
(3M)
Time FDM
324.1
294.61
268.96
243.91
71.77
23.28
DPM
DMN
21.48
9.36
7.77
7.61
DMN
35.93
11.13
4.97
4.21
LD3
16.92
6.74
4.74
4.44
LIX
28.83
12.17
5:33
4.16
Time LogSNR
$1.35
24.93
13.94
9.11
Time LogSNR
55.77
32.51
20.26
14.52
Uniforas
Time Quadratic
27.44
11.93
18.77
14.39
11.71
PANM
Time FDM
Time Quadratic
237 68
28.54
223.29
15.98
210.10
9.50
195.99
6.76
(3M)
Time FDM
312.44
284.15
252.37
221.56
76.96
59.17
DMN
10.15
7.33
7.25
7.40
DMN
11.82
6.15
4.71
5.16
LD3
9.19
6,03
5.09
4.68
LDD3
8.48
5.93
4.52
4.31
Time LogSNR
50.26
19.22
9.08
5.87
Time LugSNR
73.87
34,06
17.18
12.05
Uniform
39.78
13.88
Ust PC (M)
Time Quadratic
30.66
13.71
7.15
5.17
Quadratic
35.97
15.93
9.43
Time EDM
235.51
218.15
203.26
186.88
(AM)
Time EDM
297.83
259.90
232.79
204.89
GITS
70.93
47.37
22.33
17:27
DMN
16.72
7.96
7.54
7,81
DMN
29.22
8.21
LD3
LDX
20.15
9.09
4.98
4.18
F.5 ADDITIONAL SAMPLES
We offer additional visual examples in Figures to 23 to illustrate the qualitative effectiveness of LD3. The visual quality of LD3 surpasses discretization baselines. Our method can generate images with more visual details and less severe contrast.
30

---

Published as a conference paper at ICLR 2025
(a) Time Uniform
(b) DMN
(c) GITS
(d) LD3
Figure 11: Text prompt: "An expressive painting of a basketball player dunking, depicted as an
explosion of nebula" Stable Diffusion v1.5 (Rombach et al., 2022), iPNDM (2M) solver, NFE = 5.
(a) Time Uniform
(b) DMN
(c) GITS
(d) LD3
Figure 12: Text prompt: "A plush toy robot sitting against a yellow wall" Stable Diffusion v1.5 (Rom- bach et al., 2022), iPNDM (2M) solver, NFE- = 5.
(a) Time Uniform
(b) DMN
(c) GITS
(d) LD3
Figure 13: Text prompt: "A hand drawn sketch of a Porsche 911" Stable Diffusion v1.5 (Rombach et al., 2022), iPNDM (2M) solver, NFE 55.
(a) Time Uniform
(b) DMN
(c) GITS
(d) LD3
Figure 14: Text prompt: photo of a Samoyed dog with tongue out hugging a white Siamese cat"
Stable Diffusion v1.5 (Rombach et al., 2022), iPNDM (2M) solver, NFE=
31

---

Published as a conference paper at ICLR 2025
(a) Time Uniform
(b) DMN
(c) GITS
(d) LD3
Figure 15: Text prompt: "An expressive oil painting of basketball player dunking, depicted as an
explosion of a nebula" Stable Diffusion v1.5 (Rombach et al., 2022), iPNDM (2M) solver, NFE 10.
(a) Time Uniform
DMN
(c) GITS
(d) LD3
Figure 16: Text prompt: "A van Gogh style painting of an American football player" Stable Diffusion v1.5 (Rombach et al., 2022), iPNDM (2M) solver, NFE- 10.
(a) Time Uniform
(b) DMN
(c) GITS
(d) LD3
Figure 17: Text prompt: "A portrait of a dog" Stable Diffusion v1.5 (Rombach et al., 2022), iPNDM
(2M) solver, NFE 10.
32

---

Published as a conference paper at ICLR 2025
Time Uniform
DMN
GITS
LD3
(a) NFE 4
Time Uniform
DMN
GITS
LD3
(b) N
Time Uniform
DMN
GITS
LD3
(c) NFE =6
Time Uniform
DMN
GITS
LD3
(d) NPE =
Figure 18: Side-by-side comparison. Random samples on
Qualitative comparison
of methods across NFES using the iPNDM solver, with the same randomly sampled initial noises.
33

---

Published as a conference paper at ICLR 2025
Time Uniform
DMN
GITS
LD3
(a) NFE 4
Time Uniform
DMN
GITS
LD3
(b) NFE 5
Time Uniform
DMN
GITS
LD3
(c) NFE
Time Uniform
DMN
GITS
LD3
(d) NFE =
Figure 19: Side-by-side comparison. Random samples on class-conditional ImageNet. Qualitative comparison of methods across NFEs using the Uni solver, with the same randomly sampled initial noises.
34

---

Published as a conference paper at ICLR 2025
Time EDM
DMN
GITS
LD3
(a) NFE 4
Time EDM
DMN
GITS
LD3
(b) NFE =6 6
Time EDM
DMN
GITS
LD3
(c) NFE
Time EDM
DMN
GITS
LD3
(d) NFE 10
Figure 20: Side-by-side comparison. Random samples on AFHQv2. Qualitative comparison of
methods across NFEs using the DPM_Solver++ solver, with the same randomly sampled initial noises.
35

---

Published as a conference paper at ICLR 2025
Time EDM
DMN
GITS
LD3
(a) NFE
Time EDM
DMN
GITS
LD3
(b) NFE
Time EDM
DMN
GITS
LD3
(c) NFE =
Time EDM
DMN
GITS
LD3
(d) NFE = 10
Figure 21: Side-by-side comparison. Random samples FFHQ Qualitative comparison of methods
across NFEs using the DPM UPMM_Solver+ solver, with the same randomly sampled initial noises.
36

---

Published as a conference paper at ICLR 2025
Time EDM
DMN
GITS
LD3
(a) NFE
Time EDM
DMN
GITS
LD3
(b) NFE 6
Time EDM
DMN
GITS
LD3
(c) NFE
Time EDM
DMN
GITS
LD3
NEE 10
Figure 22: Side-by-side comparison. Random samples on CIFAR10. Qualitative comparison of
methods across NFEs using the PC solver, with the same randomly sampled initial noises.
UC
37

---

Published as a conference paper at ICLR 2025
(a) Baseline, FID 22.56
(b) LD3, FID=15.49.
Figure 23: Side-by-side comparison. Text prompts, from left to right, top to bottom: "A large grizzly bear with grass in the background.", "A man standing in front of microwave, next to pots and pans.", "A bedroom scene with bookcase, blue comforter, and window.", "A herd of elephants walking through a water-fille filled lake." Instaflow (Liu et al., 2023), NFE 2.
38

---

Published as a conference paper at ICLR 2025
39


Reviews:
[{"review_id": "LTdONNr4Qa", "reviewer": "Reviewer_4e1X", "summary": "The paper introduces a sampling method for diffusion models in order to reduce the sampling time required to generate an image. The paper proposes to learn the sampling steps from a teacher model that accurately solves the ODE by taking small step sizes. Extensive experiments show the effectiveness of the method while only requiring small amounts of training time.", "questions": "I would like to hear the opinion of the authors on the concerns I raised in the weaknesses section.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to follow.  \n- It presents an easy solution to the sampling problem of diffusion models that only requires limited training time while obtaining. \n- The soft teacher loss is effective and simple to implement. \n- The evaluation is thorough and includes multiple models, multiple datasets, and multiple sampling strategies.\n\nIn general, I liked the paper and I lean toward acceptance. However, since this is not my area of expertise, I would wait for the discussion with the authors and other reviewers to increase the score to Accept and recommend borderline Accept for now.", "weaknesses": "Although I liked the paper, there are some concerns that, if addressed, would improve the paper. In the following paragraphs, I describe my concerns in detail:\n\n- In the table with the main results, sometimes it is not clear what the metrics are computed against. I suppose the metrics in table 2, 3, 4, and 5 are computed against random samples of the model using the accurate estimation of the ODE. However, if this metric is computed against the true distribution, the performance of the teacher with the accurate computation of the ODE should be shown (1000 steps). I think the evaluation protocol needs to be more clearly defined.\n\n- In a similar direction, Table 6 shows the performance of a teacher model using 8 steps. Why only 8 steps are used here? Would not the teacher use a higher number of samples?\n\n- The model used is quite simple being only composed of a single vector (or two in the decoupled version). From the results in Table 7, increasing the number of parameters leads to better results. Would increasing the complexity of the model lead to better results?\n\n- In the limitations section I found missing that the proposed method needs to be retrained for the target number of sampling steps. One model trained to generate images with 2 samples, would not be useful for 3 and a new model would need to be trained. This might be a problem since different images might necessitate a different number of steps to achieve good quality.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iPG9QpC4M1", "reviewer": "Reviewer_4pzZ", "summary": "The paper tackles the challenge of optimizing timestep schedules for sampling in diffusion and flow-based generative models. The authors propose optimizing the schedule in order to minimize the global truncation error of multi-step sampling by reducing the discrepancy between a high-NFE teacher model and a low-NFE student model. They demonstrate that a relaxed version of this optimization problem yields even better results. Impressively, their proposed algorithm runs very quickly, achieving convergence in under 1 hour on a single GPU. A comprehensive evaluation is done using multiple pretrained diffusion models and ODE solvers. The proposed algorithm is compared against several hand-designed and learnt sampling schedules, and is shown to considerably improve the image quality in the low-NFE regime.", "questions": "- Can the invertibility assumption in Theorem 1 be relaxed and still achieve an upper bound on the KL divergence?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The LD3 algorithm is extremely lightweight, requiring only 100 samples and less than 1 hour on a single GPU to learn optimized sampling schedules.\n- The method is evaluated on a comprehensive set of pretrained models and compared against several baseline, showing improved quality in the majority of cases\n- A proper ablation study is done on the various choices/hyperparameters.", "weaknesses": "- There are several typos in the paper. See some examples below:\n    - Algorithm 1 Line 6: $x'_T ← x'_T + ...$  must be $x'_T ← x_T + ...$\n    - Line 251: $x'T \\rightarrow x'_T$\n    - Line 251: $\\Psi*(x_T) \\rightarrow \\Psi_*(x_T)$\n- Theorem 1 requires more explanation on its invertibility assumption. Specifically, if the NFE is small, functions $\\Psi_*, \\Psi_\\xi$ invertibility is a non-trivial fact which requires some justification on its assumption. \n- The method relies on a learned perceptual distance (LPIPS) to achieve optimal results, as shown by the significant quality drop in Table 7 when switching to a standard Euclidean loss. This raises questions about how well the method might generalize to other data types beyond images.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "onj2EEF9HY", "reviewer": "Reviewer_YeX6", "summary": "This paper proposes LD3, a method to discover timesteps that will yield good samples for diffusion models when the number of forward passes is very small at inference time, through optimization. The authors propose a distillation-like objective where a global notion of distance between the original model and a student model with learned timesteps is minimized. Two variants are presented, one \"soft\" objective achieving best results that only requires samples to be within a ball of hyperparameter radius $r$ to the teacher's samples. The soft objective additionally is linked theoretically to also upper bounding the KL divergence between the two models. Various experiments are conducted to demonstrate that LD3 achieves better FID scores than prior methods on few-step sampling, and an ablation study is included to demonstrate the importance of different components of LD3, the most important seemingly being the decoupling of timestep choice and step size, and the use of a perceptual distance (LPIPS) as opposed to pixel-based L2.", "questions": "Near the introduction, the paper suggests the approach should be seen as complementary to distillation methods as opposed to going head-to-head, but the goal of both said approaches and LD3 is to achieve good sample quality with the least number of steps possible. Why do the authors argue this is the case? It is not clear to me (and it is not demonstrated in the paper) that distillation methods are compatible with optimizing the timesteps used, e.g. progressive distillation is trained to match two steps of the teacher with the student, so it's not clear a priori whether changing the timesteps later will completely break a model distilled this way.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper includes proofs of soundness of their proposed minimization objectives, going beyond purely empirical contribution.\n- The number of experiments is substantial across both datasets, baselines from prior work, and choice of pretrained models.\n- The experiments conducted include notoriously difficult datasets in the literature of diffusion step reduction like ImageNet, and shows improvement in more complex settings such as a text to image model.\n- The objective is cheap to train compared to prior work; the key being that a very low batch size of 2 is permissible to use.\n- Ablation studies demonstrate the importance of the proposed changes separately.\n- The samples presented qualitatively look very reasonable and show clear improvement over the usual hand-crafted timestep schedules, and they fix the random seed so the same samples can be comapred.", "weaknesses": "There are two main areas around which the paper could be much stronger. The first is in comparisons to distillation methods, which are among the strongest in the literature. The paper includes a comparison to progressive distillation and consistency distillation in Table 9, but it is really difficult to compare these methods apples-to-apples. There are details missing (please correct me if I missed these e.g. in the supplementary material) such as what models were compared and where are the baseline scores taken from; ideally the same model should be post-trained with the different techniques. The number of forward passes across methods also doesn't match, making it difficult to draw any conclusions. One conclusion that can be drawn, however, is that progressive distillation remains better than LD3 in FID score at NFE=8, albeit requiring much more compute to distill.\n\nThe other major weakness is the lack of careful qualitative comparisons to other step reduction methods. The vast majority of the qualitative samples are compared to hand-crafted schedules, which are the weakest baselines. This is really important, especially because prior work has shown that very low FID scores can be achieved somewhat adversarially, resulting in strange samples (e.g., consider the CIFAR10 samples in the GGDM paper), so quantitative results are insufficient to truly demonstrate that LD3 improves over all prior work. Careful side-by-side comparisons of different step reduction methods, derived from the same pre-trained model and using the same initial noise and matching NFE would be significantly more convincing.\n\nOverall, the work is strong, and the quantitative results already put this paper as a valuable contribution to the literature that should be accepted at the conference. I am opting for a weak accept, because the comparisons to distillation methods seem improper and incomplete, and the qualitative comparisons require more care. But even if so, due to the very low cost of the proposed technique and the achieved scores, the work already has intrinsic value. I strongly encourage the authors to address the concerns outlined above as it would make the work excellent.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)