Score-Optimal Diffusion Schedules
Authors: Christopher Williams
Andrew Campbell
Arnaud Doucet
Saifuddin Syed
Venue: NeurIPS 2024 poster (2024)
Keywords: Diffusion model
adaptive schedule
Annealed Langevin
Primary area: diffusion_based_models

Abstract: Denoising diffusion models (DDMs) offer a flexible framework for sampling from high dimensional data distributions. DDMs generate a path of probability distributions interpolating between a reference Gaussian distribution and a data distribution by incrementally injecting noise into the data. To numerically simulate the sampling process, a discretisation schedule from the reference back towards clean data must be chosen. An appropriate discretisation schedule is crucial to obtain high quality samples. However, beyond hand crafted heuristics, a general method for choosing this schedule remains elusive. This paper presents a novel algorithm for adaptively selecting an optimal discretisation schedule with respect to a cost that we derive. Our cost measures the work done by the simulation procedure to transport samples from one point in the diffusion path to the next. Our method does not require hyperparameter tuning and adapts to the dynamics and geometry of the diffusion path. Our algorithm only involves the evaluation of the estimated Stein score, making it scalable to existing pre-trained models at inference time and online during training. We find that our learned schedule recovers performant schedules previously only discovered through manual search and obtains competitive FID scores on image datasets.

Score-Optimal Diffusion Schedules
Christopher Williams
Andrew Campbell Department Statistics University of Oxford
Department of Statistics University of Oxford
Arnaud Doucet
Saifuddin Syed Department of Statistics University of Oxford
Department of Statistics University of Oxford
{williams, canmllall,cammelll,, saiieea sww.eddin.iddddnidaa.a.. syyed ac.uk
Abstract
Denoising diffusion models (DDMs) offer a flexible framework for sampling from high dimensional data distributions. DDMs generate path of probability distributions interpolating between a reference Gaussian distribution and data distribution by incrementally injecting noise into the data. To numerically simulate the sampling process, a discretisation schedule from the reference back towards clean data must be chosen. An appropriate discretisation schedule is crucial obtain high quality samples. However, beyond hand crafted heuristics, general method for choosing this schedule remains elusive. This paper presents a novel algorithm for adaptively selecting an optimal discretisation schedule with respect to cost that we derive. Our cost measures the work done by the simulation procedure to transport samples from one point in the diffusion path to the next. Our method does not require hyperparameter tuning and adapts to dynamics and geometry of the diffusion path. Our algorithm only involves the evaluation the estimated Stein score, making it scalable to existing pre-trained models at inference time and online during training. We find that learned schedule recovers performant schedules previously only discovered through manual search and obtains competitive FID scores image datasets.
Introduction
Denoising Diffusion models (Sohl-Dickstein al. 2015 Ho et al. 2020 Song et al. 2021 or DDMs
are offte--o-the-rrr generative models. They are formulated through considering a forward noising process that interpolates from the target reference Gaussian distribution by gradually introducing noise into an empirical data distribution. Simulating the time reversal of this process then produces samples from the data distribution. Specifically, we evolve data distribution po through the forward diffusioo process time interval [0, 1J. described by
d.,- dtt + g(t)dW,,
Xo ~ Po.
(I)
with drift Xe, diffusion coefficient g(t) and Brownian noise increment dW, The coefficients f(t) and g(t) chosen such that at time 1 the distribution of X1 is very close to reference Gaussian distribution p1 in distribution. sample starting at Po and following Equation (I) until time t will be distributed according Pt which is mollified version of the data distribution
(xx) _xx, Po (o(xo)P (ro)
(ejo(xalzo) = Nx (z) (
(2)
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Training Data Samples
Linear Schedule Samples
Optimised Schedule Samples
Density
Density
2-
-04
-14
2.2 0.2 0.6
-2.4 -02 0.0
0.2 0.4 0. -04
-0.
-0.2 0.0
X
X
Figure 1: Density estimates of the mollified Cantor distribution (left) using DDM with schedule
T = [ 100 generated with 100 linearly spaces discretisation times 4 = =/100 (middle), compared to the optimised schedule T {tz { == with 50 discretisation times 1 generated by Algorithm
(right). The eight modes present in our true mollified distribution are shown in grey on each plot.
The parameters s(t) and (1) define the noising schedule. They can found in closed form in terms of (() and g() (Karras et al. 2022). To obtain samples from Po, we can reverse the dynamics of the forward diffusion in Equation to obtain the backward diffusion,
dx, (f(t)x, g(t)²V log (x.)) dt + g(t)dW,
X1 PPI
(3)
By simulating Equation 3 backwards in time, we evolve reference samples X1 ~ p1 from t = 1 to 0 to obtain samples that are terminally distributed according to the data distribution po- To simulate Equation (3) numerically, we must decide upon a discretisation of time, T = {t with br 1, t O, which we refer to as the discretisation schedule. For given noising schedule, is important to select an appropriate discretisation schedule such that (3) can be simulated accurately, i.e Pt. and Pt.., should not differ significantly. In this paper, we derive methodology to compute an optimal discretisation schedule.
Prior work has often joined together the choice of noising schedule and discretisation schedule. A uniform splitting of time would be chosen, 4; = i/T,, with the noising schedule dictating the change between pt, and Pt.-1. Two prominent examples have the form s(t) Va(t), 22 = ā(t) with the linear schedule introduced by Ho et al. (2020) having a(t)
expl fo (s)dss) with linear 3(t) = P.in + t(3max - 3min). Alternatively, Nichol and Dhariwal (2021) introduce the cosine schedule with a()) [(I)/((0)
Nichl
= cos( T+c 5) f 0.008. Karras et al. (0022) refines this approach by splitting the choice of noising schedule from that of the discretisation schedule, however, picking the discretisation schedule is still a matter of hyperparameter tuning.
A good discretisation schedule can drastically impact the efficiency of the training and inference of the generative model, but unfortunately can be difficult to select for complex target distributions. For example, for distribution supported on the Cantor set (Figure I left), the default linear schedule fails entirely to capture the modes of the data distribution (Figure I middle). However, our optimised schedule learned using Algorithm| recovers these modes (Figure I right). Without such an automatic algorithm, finding performant discretisation schedules often reduces to an expensive and laborious hyperparameter sweep.
We devise method for selecting discretisation schedule that yields high-quality samples from Po- Our key contribution is defining an appropriate notion of cost incurred when simulating from one step in the diffusion path to the next. We then choose our discretisation schedule to minimise the total cost incurred when simulating the entire path from P1 to po- Our cost purely depends on the distributional shifts along the diffusion path and assumes perfect score estimation, hence, we refer to our schedules as ecore-optimal diffusion schedules. The resulting algorithm is cheap to compute, easy to implement and requires no hyperparameter search. Our algorithm can applied to find discretisation schedules for sampling tre-trained models as well as performed online during DDM training. We demonstrate our proposed method on highly complex ID distributions and show our method scales to high dimensional image data where it recovers previously known performant discretisation schedules discovered only through manual hyperparameter search. To the best of our knowledge, this is the first online data-dependent adaptive discretisation schedule tuning algorithm.

---

2 The Cost of Traversing the Diffusion Path
To derive our optimal discretisation schedule, we first need to derive notion of cost of traversing from our reference distribution p1 to the data distribution pa through each intermediate distribution Pt, referred to as the diffusion path. We will then later find the discretisation schedule that minimises the total cost of traversing this path. Our notion of cost is based on the idea that while integrating Equation 33 from time to l' will always take you from pt to pe, the simulation step will need to do more work to make sure the samples are distributed according to pr if ph and pr are very different distributions rather than if they are close. the following, we will make this intuition precise.
2.1
Predictor/Corrector Decomposition of the Diffusion Update
To begin, let X = Id and P(X) be the space of Lebesgue probability measures on X with smooth densities. For E [0, 1. define the diffusion path p. E IP(X) as the law of X, satisfying the forward diffusion Equation ( initialised at the data distribution Xo Po. equivalently the law of the X,
satisfying the backward diffusion Equation (3) initialised at the reference distribution X1 ~ pl-
Given sample X₁ ~ Pe. a sample Xe ~ Py can be generated by integrating the backward diffusion Equation (3). In Song al. (2021), it was shown that we can further decompose the backward diffusion Equation (3) at time f into deterministic flow governed by the probability flow ODE, and the stochastic flow driven by Langevin dynamics targeting pt.
dx,- = (f(t)x, - 2 log pa (X.)) dt
2 log (e + g()dW,
(4)
Probubility Prediction
Langevin Correction
This decomposition into a deterministic flow and a correction will help us derive our cost in Section
by analysing the work done by the correction keep the samples the diffusion path. Here, we will first expand upon this decomposition by defining a hypothetical two-step sampling procedure that could be used to sample the DDM. It consists of: (1) a predictor step that generates a deterministic prediction of Xę and (2) corrector step that uses Langevin dynamics targeting Pe to correct any accrued error. Note we are not advocating for the implementation of such a procedure, only that by algorithm, it will be helpful for our theoretical derivation of a cost intrinsically linked to the sampling of DDMs. The two stages of the predictor-corrector algorithm are rigorously defined as follows.
imagining simulating with this hypothetical predictor-coreett
Definition 2.1. A predictor for t, E [0, 1] is a smooth bijective mapping F x X, such that det V Ft. # 0, and the predicted distribution is the pushforward of pt by Fee denoted Ff,
Definition 2.2. A corrector for t E [0, 1]. E [0, oc) is a one-parameter family Markov transition kernel Lt.- X x P(()) + [0, 1] such that Z, L,,(z, dzy) is the law of Langevin dynamics stationary distribution pt at time T, initialised at E X, and running at speed v(t) > O.
Zo 2, az, v(t) log pr (Z,)dr + 2o0()dWW,
(5)
The corrector map L is specified through an integration time T and time-dependenr speed v(t). We assume fixed and we will describe the appropriate choice of e(t) in Section 3.3 The predictor- corrector algorithm, given X, Pt, first applies the predictor Zo = Frr (X.) and then uses the corrector to drive the predicted samples towards pr'
Predictor: Zn = (X,)
Corrector:
Xe.- Le. (Z,,
(6)
In general, Zo will not be sample from Pe exactly because Fir may not be perfect transport from Pt to Pr. In Section 2.2] assessing the work done by Ly to drive the Zo towards pe will be key in deriving our cost. Our cost will then depend upon the specific choice for Ft,rr. Two natural choices for Fi, are apparent. Setting Fe, to the identity means our hypothetical sampling algorithm reduces to the annealed Langevin algorithm for DDMs introduced by [Song and Ermon (2019). The second natural choice is to set Fe,r to the integrator of the probability flow ODE (Song et al. 2021).
Example 2.1 (Annealed Langevin). The predictor step is trivial when the predictor map the identity Fee (x) = In such case, the predicted state reduces to the initial state, (x,) X, Pt- The work done by the corrector step will then be related to the full discrepancy between pt and pr because the predictor provides no help in transporting the sample.
3

---

Example 2.2 (Probability Flow ODE). The predictor step is optimal when Fi, is a transport from Pr to pvr In such case, the predicted state produces a sample from the target distribution, Fee (Xx) ~ pe, and so the corrector step would have to perform no work. An optimal predictor map Fee can be obtained by integrating the probability flow ODE from time 1, 1,1'
dæ dt
= f(t)2; 7 g Vog (x.)).
Practical algorithms numerically integrate Equation 7 e.g. an Euler step with At = t' -t,
(z) = x + (f(t)z - 2 (z))
(8)
In such case, the work done by the corrector depends the error in the probability flow integrator.
2.2 The Incremental Cost of Correction
We now focus on deriving cost related the work done by the corrector step in the predictor- corrector algorithm. Later, in Section we will find the discretisation schedule that minimises the total cost. To derive the cost, we will analyse the movement of Zu under the corrector step's s dynamies ( dr,). ). This requires some care because even if Zu is already at stationarity, i.e. perfectly distributed according to Pe', applying the Langevin correction step will still result in movement of Zo due to the stochasticity of the update. However, the computed work done by the correction step in this case should be O. To correctly assign the work done, we will compare two processes. The first the trajectory of Langevin dynamics, 7-, defined by the corrector Turr initialised at Zo Fr,,((X) targeting pr- The second is virtual coupled Langevin dynamies Z, initialised at Fe,e((X)) driven by the same noise and speed but targeting the stationary distribution of the predictor F Fr,,Pr.
Zo (X,);
dZ, u(t') V Vog Pr (Z, ) + V 2o(t))ww, d² = u(t') Vog Fi )dr 2o(")aaw
(9)
2. (X,)
(10)
Notably, 2, _ Xe and 2, _ Fi, (X.) share the same law as the corrected sample and predicted sample respectively. Since Z, and 2, are coupled to have the same noise, the difference in their trajectory. Z, 2,. isolates the change in corrector dynamics due to discrepancy between Fr,.,Pt and Pr. If F. Pt is very different to pr. then Z, - -2, will be large, signifying the corrector is needing to do lots of work to push the distribution of 2 towards the target pe. Conversely, if F', = pv. and no work is done. For small т, (Z, - Zo)/T is the initial velocity of Z under the Pe corrector dynamics, and similarly for (Z, - Zo)/T. We can then define the incremental cost L(t,t') by taking limits as J 0+, measuring the expected I norm ||- || of the difference,
then Z, 2,
C(t,t)= = lim T FE || (Z+ Zo) - -0z - Zo) 1|] = lim
T 2 E
|| Zr 2,
17]
(11)
T-01
We can approximate Z, 2, using an Euler step, noting that the coupled noise terms cancel,
7, 2, = (t(' (V log pr (Z,,)) - V log (..,)) o(T)-
(12)
By substituting Equation
12 in Equation
11 we have
Pe Pe (2)) ["] = D(p
= 2 log Pe (%))
Fd V log
(13)
where D(pllg) = Ex~alll VI log p(X) - V log q [I is a statistical divergence on p.g g E P(X).
measuring the 2 distance between the scores of q and p with respect q. D(pllg) is referred to as the Stein divergence or the Fisher divergence; see e.g. [Johnson 2004). For given choice of u(t) and Fe. we now have cost measuring the change from pt to pr. This cost intrinsically linked with the effort performed by a DDM sampling algorithm because it is derived through considering the work done by a hypothetical medictor-correctte style update. We note, however, that this general cost can be used to obtain discretisation schedules for in any style of DDM sampler.

---

2.3 Corrector and Predictor Optimised Cost
By inverting Zu Fi. (Xt). we can express Equation 13 in terms of an expectation with respect to the reference sample X; ~ Pt, and the score of Gi X + R.. the incremental weight function associated with the transport Fi, from the Sequential Monte Carlo literature (Arbel et al.) 2021.
Pe (FFe (x(x)
C((,t) - =((')( [IV log Gr.
GLv
det (T).
(14)
In most cases, it is infeasible to efficiently compute the Jacobian correction in Equation (4). When
Fe,e x the identity map corresponding to the corrector optimised update from Example 2.J Equation [4] reduces rescaled Stein discrepancy between and pe, and Gz,(z) = (z) (z) reduces to the likelihood-ratio between pr and Pt. We will refer to this case as the corrector-optiniss cost denoted (t, t'), to distinguished it from the where when relevant, we will use subscripts and p to distinguish between the two:
cost Lp(t, t') derived above,
Co(t,t') = D(p'))(pr)ll))
Cp(t, t') = D(per (......)).
(15)
The corrector-optimised cost Le(4, (1) provides meaningful information during the update from reference Pt to the target pr. It is worth computing even when the predictor-optimiss cost Cp 8') is accessible. Lo(t, t') measures the change between the reference and target distribution independent of the predictor, whereas Cp(t,t') measures the residual error between the predictor and target. Notably, (t, t¹) encodes information about the incremental geometry of the diffusion path, whereas (t,t, t') quantifies information about the incremental efficiency of the predictor. Generally, one does not dominate the other, but if the predictor is well-tuned and the predictor flows samples X₃ Pe towards pe, we would expect (t. < tt.
For deriving our optimal discretisation schedule, we require a notion of how L(, (') increases with small increases in L' i.e. knowing local changes in incremental cost. In Section 3 cost to assign distances to schedules through time, enabling us to find the best schedule. We derive the desired local cost in Theorem 2.1] see Appendix [or PDE and geometrie interpretation.
this local
Theorem 2.1. Suppose Po(x)) Fe,or (x), v(t) and G. (x) are three-times continuously differentiable (
in L,U'x x and let Pr(x)) = 37 Fe (x) n==t and Gz(x)) कe
Gr, (w) Suppose the following
hold: (1) all x X, E [0, 1). (x)) = x and (2) there exists V X + TR such that for all x E X and E [0, 1]. (x)] I (x) and supre(o.1) Exexp. [VX,)] (XXx))
< co. Then for all t E [0, 1), we
have C(1, 1') = (( + O(A where
||
Ø G = at
8(t) = ((t))xxxxx)) -pi
logPt + V Voggp - F, + FrV (16)
Theorem 2.1 shows that, under regularity assumptions, then the incremental cost C(t,t) = 8(t)
is locally quadratic and controlled by the local cost 6(t). The 5(t) measures the sensitivity of the incremental cost C(t, t') to moving samples along the diffusion path to t' Notably, 5(t) = 0 if
and only the predictor satisfies the continuity equation, Pt + V (pr =
Score-Optimal Schedules
Given a discretisation schedule T (t+) satisfying 0 to <... predictor-corrector algorithm recursively uses the predictor and corrector maps to generate sequence (X,) starting at Xy ~ p1 such that the terminal state Xo approximates samples from Po.
< tr = 1, our hypothetical
X ~ Le., (F (X++1), dz;)
(17)
We want identify a discretisation schedule that maximises the efficiency of this iterative procedure. This is not generally possible due to the potential complex interactions that arise from the accrued errors. To simplify our analysis, we make the following assumption.
Assumption 3.1. For all 1', if X, Pt and X L (Fe,e (Xt), dr,)), ), then Xy.r Pe-
Assumption 3. is reasonable if, our hypothetical corrector steps, T is set sufficiently large such that the Langevin correction converges to stationarity. We find in our experiments that even if the
5

---

schedules derived under Assumption 3.1 are used in sampling algorithms for which Assumption 3.]
does not hold, we still obtain high quality samples. Equipped with Assumption 3.I] we can measure the efficiency of the path update through total accumulated cost L Ei=1 ((tr+1, which we will use as objective to optimise T. In this section, we will identify the optimal schedule T minimising the cost & by considering an infinitely dense limit. We will then provide a tuning procedure amenable to online schedule optimisation during training. Finally, we will discuss suitable choice for v((), the velocity of our hypothetical corrector steps, as well as related work.
3.1 Diffusion Schedule Path Length and Energy
Let yp : [0. + [0, 1] be strictly increasing, differentiable function such that (0) = 0 and ((1) 1. We will say T is generated by y if t, ((i for all O..... T. The schedule generator yp dictates how fast our samples move through their diffusion path. Since every schedule T of size T is generated by some yp, optimising T is equivalent to finding a generator &p minimising [(p T), the total cost accumulated lated by the schedule of size T generated by sp By Jensen's inequality, we have ((v,T) 2 A(p, 2 /T, where for 4 = (ii/T),
T
T
T) = Σ ((), A(.TT)=E E VVVVVtttee ((t. i=1
(18)
As we later prove in Theorem 3.1 in the dense schedule limit as T + the cost L(y, T) and its lower bound A(p, T) are controlled by the energy Elpp) and length A respectively where,
[(p)= - S. ((((())))s))dds A A=
A-f(Vo(tjd Vo(t)dt.
(19)
The intuition for why E(p)) is an energy, and A length can be gained by first conceptualising the diffusion time as a spatial variable rescaled by the metric 5(t) defined by our cost L. We have yp and are position and velocity, respectively. Integrating the speed Ju ( = fu (o(r)dt along a curve (ss) obtains the "length" A, whilst integrating a speed squared, Jor (()) obtains a "kinetic energy" E(p). Note that the length is an invariant of the schedule, whereas the kinetic energy is not. The length A measures the intrinsic difficulty of traversing the diffusion path according to the cost independent of (o, whereas E(p) measures the efficiency of how the path was traversed using sp. This geometric intuition hints the solution to the optimal scheduling problem. The optimal should travel on a geodesic path from P to Po, at a constant speed with respect to metric d. For this optimal ys, then have the kinetic energy being equal to the square of length between P1 and pn- Theorem 3.1 makes the previous discussion precise.
Theorem 3.1. Suppose the assumptions of Theorem 2.1 hold. For all schedule generators sp,
lim TL(p,T)- = E(ee), lim A(P,T)) = A.
(20)
T-x
T-o0
Moreover, E(Φ) 2 A with equality if and only if satisfies,
4 (s)) (s) = A- A(I) 1.' Vo(u)du.
(21)
Notably independent of the choice of sp, as T -> oc, the cost C(p, T) E(p)/T. This implies that the cost decays to zero at a linear rate, proportional to E(y)) and L(e, 2 A2/T independent of sp Equation 21) provides an explicit formula for the optimal schedule generator that minimises the dense limit of the total cost and obtains the lower bound E(x*) = A The intuition for the formula (x) A -1Ass) is that this implies Asyp*(s) = As meaning say 10% of the way through the optimal schedule, we should have traversed 10% of the way along the distance between p1 and Po i.e. 0.1 x A. This relation holds for constant speed straight lines, mcaning yo" is the optimal schedule. For a finite T. Theorem 2.1 implics the optimal schedule T = (t, V generated by op ensures the incremental cost is constant ((().)))) R A? /2 for all = . T
Our geometric intuition in the language of differential geometry is that the diffusion path M {pr (o is Riemannian manifold with metric 5 endowed by the incremental cost L(t, t'). The schedule generator defines a curve 8 + Po(s) E -M reparametrising the diffusion path between po and pl. Theorem 3.] shows that s² is the geodesic of length A in M between pi and po that traverses the diffusion path at constant speed (s))
A with respect to 6 and minimises the cost.

---

3.2 Estimation of Score-Optimal Schedules
adapts Algorithm 3 from Syed et al. 2021] to estimate the optimal schedule 7 = { i=o generated by yp" We can use Algorithm| | to refine the schedule for pre-trained DDM or learn the schedule jointly with the score function. For this joint procedure, we detail in Appendix B..ddiix.Boooww.. function evaluations can be reused to estimate the cost to minimise computational overhead. For Le(t, t') we need only evaluate V log Pe (Xt) and V log Pe (X,) both available through our model's s score estimate. Computing Cp(t, t') is more challenging since there are Hessian terms that arise in Equation (14). Under the assumption that the step size åt is sufficiently small, we can approximate V log det F.,, (X,) through Proposition B.1 This approximation only requires us to compute the gradient trace of the Jacobian of our predicted score. With computational cost proportionate to the computational effort for computing the first derivative. Using Hutchinson trace (Hutchinson) 1989) like estimator in Proposition B. we compute this quantity memory-efficiently in high dimensions, requiring only standard outo-differentiation packppropagation.
Given a schedule T= [t, 1- and estimates of the incremental cost L(t tt). Algorithm
3.3 Choice of Velocity Scaling
Recall that our cost is derived by considering a Langevin dynamics step with velocity vt). This velocity should be selected so that Langevin dynamics explores the same proportion of our distribution at varying times throughout our diffusion path. Thus, vtt) should be on the same scale as the spread of the target, Pt- Commonly used noising schedules have s(t S 1, and our data distribution is normalised so the scale of Pt is on the order of a (t). We therefore set u(t) = (t).. This results in (t')) -weightdd divergence for our incremental cost C(t,
(t') D (ppe l|p)). This can be compared to the weighted denoising score matching loss used to train DDMs (Song et al.] 2021)
which is also squared norm of score differences: A())EEx,,x, [llso () V Pogo (X, (x)
for some weighting function A(t) chosen to equalise the magnitude of the cost over the path. In Song
ct al. (2021) X(t) x 1/Ei|| V log Pajo [Xo)) was chosen, which, as we show in Appendix B.3
is A(t) x o (t). This choice of velocity scaling provides an alternative perspective on this commonly used weighting of squared norms of score differences.
Algorithm UpdateSchedule
Require: Schedule T To incremental costs [C(te+1, TT) T-
1: ((t)) = Σj= V((t)-n,t)). O....T
Equation 2
2: Å = -((r))
A in Equation 19
3: (-) = (iteerooaae(((A((t)..)))
((tt),tr){s
b E.g. Fritsch and Carlson
1980
4: t5 = =}=)(At) i= .
> Equation 21
5: Return: T [T
3.4 Related Work
Previous works have devised algorithms and heuristics for designing noising and discretisation schedules. The DDM training objective is invariant to the noising schedule shape, as demonstrated by KKiggma et al. (202T), necessitating auxiliary costs and objectives for schedule design. Uniform steps in the signal- target distribution's 's geometry. Watson et al. (202T) optimise the schedule by differentiating through sampling to maximise quality, but GPU memory constraints necessitate gradient rematerialisation. We avoid this with a simulation-free cost. Closely related to our work is Sabour et al. (2024). who minimise a pathwise KL-divergence between discretised and continuous processes. They require multi- optimisation with early stopping to prevent over-optimisation of their objective which would otherwise result in worse schedules. Amongst the wider literature, various strategies for discretisation schedule tuning have been proposed. Das et al. (2023) derive an equally spaced schedule using the Fisher metric but assume Gaussian data. |Santos et al. (2023) assign time points proportional to the Fisher information of (1oo ao)), ignoring the true target distribution. Xue et al. (2024) derive schedule to control ODE simulation error, but their cost depends only on the ODE solver, and not on the data distribution.
ratio, log (s(1.) (((4)//(( (4:),, are used by Lu et al. (2022). but this ignores the
7

---

Computational Experiments
4.1 Sampling the Mollified Cantor Distribution
The Cantor distribution (Cantor 1884) lacks a Lebesgue density, with its cumulative distribution function represented by the Devil's s starcase and its support being the Cantor set, forming a challeng- ing 1-D test example. When mollified with Gaussian noise, it becomes absolutely continuous and possesses a Stein score. We mollify by running a diffusion with the linear schedule for time = 10-5 With this mollification, our data density has eight pronounced peaks. We train a mieedimeniiona DDM for 150,000 iterations using both fixed linear schedule and our optimisation algorithm Al- gorithm 2 initialised the linear schedule. We find that the non-data-specific default schedule fails to capture these modes, whilst our adaptive method faithfully reproduces the data distribution. In Figure we show the complexity of the learned score which displays self-similar fractal structure.
4.2 Adaptive Schedule Learning for Bimodal Example
We train DDM simple bimodal Gaussian distribution. When the variance of the target bimodal Gaussian is low, it becomes difficult to adequately sample from the target distribution. In our instance, the standard Gaussian reference from the diffusion is given, and the target is the density po(r) = PPee + PP)igh())) where Plan and Pright are normal distributions with means -6 and 6, respectively, and common variance σ² = 0.1
We learn diffusion models, one using the linear schedule, and the other using schedule that learned online during training. We compute the likelihood of the samples generated from either model during training, which is possible in this example because the true probability density is known. that when the schedule learned during training, the likelihood evaluation
can be seen in Figure 2 increases and the true score error decreases, in contrast the linear schedule that remains constant, or worsens, in this regard during training.
Comparison of Linear and Learned Schedules
Error Comparison Linear and Learned Schedules
1.20
Liner Schedule Schodulc
Linea Schedule Frror
Schedule
800
Lke Avenge
700
Score Ener
105
1.00
0.95
400
3.90
300
4000
50:00
4000
6000
Figure 2: Comparison of Linear and Learned Schedules over Training Iterations for the bimodal example. Each point corresponds to 500 training iterations.
4.3 Scalable Schedule Learning Diffusion
Here we demonstrate that jointly learning the schedule and score using our online training method- ology (Algorithm 2) scales to nigh-dimensional data and converges to stable solution. We train DDMs on CIFAR- 10 and MNIST initialised the cosine schedule using the codebase from|Nichol and Dhariwal 2021). ). Figure (left), we show the incremental costs / VI(I++1, 4,) for the cosine schedule and our learned schedule, finding that the increments approximately equalise over the diffusion path as expected by the discussion in Section 3.T Figure 3 (right) shows the learned schedule spends more time at high- frequuryy details, visualise sampling trajectory in Figure []

---

Table Sample quality measured by Fréchet Inception Distance (FID) versus schedule on CIFAR1O (32 x 32), FFHQ, AFHQv2, ImageNet (64 x 64). Pretrained models are used from Karras et al. (2022). All FIDs are calculated using 50000 samples. We highlight the best FID in bold. The ImageNet model lacks second-order differentiation, so no predictor optimised schedule is shown.
Schedule
CIFAR-10
FFHQ
AFHQv2
ImageNet
Eq
p=3
5.47
2.80
2.05
1.46
Eq
p=7
1.96
2.46
2.05
1.42
Logl.inear et al.
2022
2.05
2.42
2.06
1.45
Convex Schedule
22.1
2.43
2.48
1.64
Corrector optimised
1.99
2.46
2.05
1.44
Predictor optimised
1.99
2.48
2.05
19
Schedules from
CO-Cost (x10")
PO-Cost (x10*)
KLUB (x10%)
FID
low to high FID
Noise Level
Eq 22 p
1.96
9.86
2.16
1.75
10
Conee
Lostiterr
CO (ours)
1.99
9.54
2.09
1.39
PO (ours)
1.99
9.61
2.09
1.39
20-1
Logl Jinear
2.05
10.5
2.32
1.05
Eq 22 0
5.47
17.2
4.02
2.82
CFAR10 Corrector
CIARRD Prodictor
Convex Schedule
22.1
45.0
8.17
0.284
Timestep
Figure 4: (Left) Costs associated with different schedule choices for the CIFARIO dataset. Schedules ordered from lowest FID to highest FID. We compare Corrector-optimised (CO) cost and Predictor-optimised (PO) cost versus the Kullback-Leible Upper Bound (KLUB) from Sabour et al. (2024). The minimum value for each cost is highlighted in bold. Note low cost is associated with low FID for our cost and not for the KLUB. (Right) Visualisation schedules during generative sampling with 100 timesteps. "rho=3" and "rho=7" refer Eq with p 3 and p respectively, LogL from Lu et al. (022) and convex schedule are also shown. We show our cost optimised schedules for CIFARTO both using the corrector optimised cost and the predictor optimised cost.
Length Increment
DN4 Contttr Scsohreli- B
12.5
OTA Conactor Scrddali 0.22 Schediis
10.0
CRAR Oitina Stod Capin 112 Screis
00
05
200
400
1000
200
800
3000
T1 mimestop
Timestop
Figure 3: (Left) Incremental costs (1,)+ for the cosine schedule and our online adaptive algorithm. Higher learning rates enforce equalisation of costs more quickly. (Right) Progression of the learned schedule during 40k training iterations, depicted through the standard v1 11- Ot-
4.4 Sampling Fre-Trained Models
In this experiment we demonstrate that our algorithm can recover performant schedules for large image models used in practice and our schedules generate high quality samples. We use the pre- trained models from Karras et al. (2022). whose DDM is parameterised such that the forward noising distribution of the form P:jo (t FO) = [axz; To, of The scheduling problem then reduces to deciding on a stepping scheme through {as (or = 0. Karras et al. (2022) suggest a polynomial based schedule with parameter p that controls the curvature of the schedule
OicN = -Oomax N-I o
(mmax
σ mmnx) )) and ON = 0.
(22)

---

A lower p value results in steps near min being shortened and steps near max being lengthened. Through analysing the truncation error for sampling in Karras al. (2022), they find that setting 3 approximately equalises this error, however it is found empirically ly that p 7 results in better sample quality. We also compare against schedule that takes uniform steps in log space [u et al. (2022) which we refer to as the LogLinear schedule and a schedule that takes convex shape in log space. Schedule visualisations are provided in Figure (right).
We sampled the pre-trained models using these schedules and computed the sample quality using FID. We use the same number of schedule steps (18 for CIFARIO, 40 for FFHQ and AFHQv2. 256 for ImageNet) and solver (Heun second order) as Karras et al. (0022) Our results are shown in Figure 4 Our optimised schedules are able to achieve competitive FID to the best performing p schedule hand-tuned in Karras et al. (2022). This is expected as our schedules take a similar shape to the p = 7 schedule as shown in Figure 4 (right). Therefore, our method provides an entirely automatic and hyperparameter free algorithm to recover this performant schedule that was previously only discovered through -nddrrr
We further analyse how the number of discretisation points, T. used during sampling affects the quality of generated samples for different schedules. We report our results on CIFARIO in Table 2 Notably, the FID decreases with 7 for all schedules and achieves comparable FID once T is large enough. However, when T is small, only the optimised schedules maintain stable performance. This empirically demonstrates an optimised schedule can improve the sampling efficiency by allowing for coarser discretisations and, hence, faster sampling, as predicted by Theorem 3.1
We observe an
identical trend for sFID in Table 3 in the Appendix C.2
points, T
10
20
30
50
100
CO (ours)
2.46
2.02
2.04
2.06
2.07
P 3
50.75
3.92
2.09
2.01
2.05
7
2.70
2.00
2.06
2.05
2.07
P 100
3.09
2.06
2.05
2.06
2.07
Table 2: Comparison of FID across different amounts of discretisation points for different schedules on CIFAR1O. CO stands for our corrector optimised schedule.
We also compare corrector optimised schedules to predictor optimised schedules in Table [] They provide similar performance so, on image datasets, we encourage the use of the cheaper to compute corrector optimised schedule. Finally, in Figure (left), we report the values of our corrector optimised costs and compare these costs to the values of the objective introduced in [Sabour et al.] (2024). Both algorithms aim to find schedules that minimise these costs and therefore it is desirable for low values of cost to be associated with good sample quality (i.e. low FID). We find that low values of our cost correlate much more closely with low FID than the objective introduced by Sabour et al. 2024). Indeed, Sabour et al. (2024) introduce a bespoke multi-stage optimisation for their cost because they found over-optimising their objective can lead worse schedules which explained by the objective not correlating well with FID. We further find that our predictor optimised costs are lower than the corrector optimised costs which to be expected as the predictor reduces the work done by the corrector and thus reduces the incremental cost. The overall shape of schedule, however, between the corrector optimised and predictor optimised costs is similar.
Discussion
We have introduced method for selecting an optimal DDM discretisation schedule by minimising a cost linked to the work done in transporting samples along the diffusion path. Our algorithm is computationally cheap and does not require hyperparameter tuning. Our learned schedule achieves competitive FID scores. Regarding limitations, the computation of L. can be computationally expensive due the calculation of second derivatives, however, in Section [.44] we found L. to provide a cheap and performant alternative. Furthermore, our theory is derived assuming perfect score estimation. Future work can expand the geometric interpretation of the diffusion path and links to information geometry to further refine the DDM methodology.
10

---

References
Arbel, M., Matthews, A., and Doucet, A. (2021). Annealed flow transport Monte Carlo. In International Conference Machine Learning.
Cantor, G. (1884). De la puissance des ensembles parfaits de points: Extrait d' une lettre adressée à l'éditeur. Acta Mathematica, 4:381-392. Reprinted in: E. Zermelo (Ed.), Gesammelte Abhandlun- gen Mathematischen und Philosophischen Inhalts, Springer, New York, 1980.
Choi, Y., Uh, Y., Yoo, J., and Ha, J.-W. (2020). Stargan v2: Diverse image synthesis for multiple domains. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition.
Das, A., Fotiadis, S., Batra, A., Nabici, F., Liao, F., Vakili, S., Shiu, D.-s., and Bernacchia, A. (2023). Image generation with shortest path diffusion. In International Conference on Machine Learning.
Deng, J., Dong. W. Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. (2009). Imagenet: A large-scale hicrarchical image database. IEEE Conference on Computer Vision and Pattern Recognition.
Fritsch, F. and Carlson, E. (1980). Monotone piecewise cubic interpolation. SIAM Journal on Numerical Analysis, (7(2):23--246
Ho, J., Jain, A., and Abbeel, P. (2020). Denoising difffuson probabilistic models. In Advances in Neural Information Processing Systems.
Hutchinson, M. (1989). stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in Statistics -Simulation and Computation, 183))
Johnson, O. (2004). Information Theory and the Central Limit Theorem. World Scientific.
Karras, T., Aittala, M., Aila, T., and Laine, S. (2022). Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems.
Karras, T., Laine, S., and Aila, (2018). A style-based generator architecture for generative adversarial networks. arxiv e-prints. In Conference on Computer Vision and Pattern Recognition (CVPR).
Kingma, D., Salimans, T., Poole, B., and Ho, (2021). Variational diffusion models. In Advances in Neural Information Processing Systems.
Krizhevsky, A., Hinton, G., et al. (2009). Learning multiple layers of features from tiny images.
Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., and Zhu, J. (2022). DPM-solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps. Advances Neural Information Processing Systems.
Nichol, A. Q. and Dhariwal, P. (2021). Improved denoising diffusion probabilistic models. In International Conference on Machine Learning.
Poincaré, H. (1890). Sur les équations aux dérivées partielles de physique mathématique. American Journal of Mathematics, pages 2-1-294.
Sabour, A., Fidler, S., and Kreis, K. (2024). Align your steps:
Optimizing sampling schedules in diffuuson models. International Conference on Machine Learning.
Santos, E., Fox, R., Lubbers, N., and Lin, T. (2023). Blackout diffusion: generative diffusion models in discrete-state spaces. In International Conference on Machine Learning.
Sohl-Dickstein, J., Weiss, E, A., Maheswaranathan, N., and Ganguli, S. (2015). Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning.
Song, Y. and Ermon, S. (2019). Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information P Processing Systems.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021). Score-based generative modeling through stochastic differential equations. International Conference on Learning Representattons.

---

Syed, S., Bouchard-Cóté, A., Deligiannidis, G., and Doucet, A. (2021). Non-reversible parallel tempering: a scalable highly parallel MCMC scheme. Journal of the Royal Statistical Society (Series B), 84:321-350.
Watson, D., Ho, J., Norouzi, M., and Chan, W. (2021). Learning to efficiently sample from diffusion probabilistic models. arXiv preprint arXiv:2106.0380)
Xue, S., Liu, Z., Chen, F., Zhang, S., Hu, T., Xie, E., and Li, Z. (2024). Accelerating diffusion sampling with optimized time steps. arXiv preprint 1rX1:
12

---

A Analysis of incremental cost
A.1 Proof of Theorem 2.1
Proof. We first note that by the mean value theorem, there exists s(t, t')
E [t,
C(t,
VG
27
[t' t/2 = )(t') EXxxpp
(23)
t' -t
22
= =((1))
(24)
Since all t, we have (t)) < V(x), and Exiope Ex [Xx)] < co, the dominated conver- gence theorem and the continuity e(t)) implies.
L(t, )')
lim P- It' -12 = vutt) ((t))'xxxn ExIngi
lim VG Prtt
(25)
= Exxt)2xx.Np
(26)
The last equality follows since VG, is continuous in t.
We will now compute G. Since FI.( = x, we have Gz, = 1 for all x. This implies, we can express G in terms of the derivative of log Gt,
Gx(x)) 0y!
log Geee (I)
(27)
where log G equals,
log Gt. (x) log Per (Fi, (x)) pa(x) + log det (z.
(28)
By combining Equation
27 to Equation 28 we obtain,
(x(x)=
log pe (x)
+ ou log det Fer
(29)
or
===
#=t
For the first term in Equation
29 we use chain rule to obtain.
① log Pe (x))
Opp
(F) V (Fee
OFi.
(0))
(30)
ay
Pe (Feer (x))
or
0y
where s denotes a dot product of vectors. By evaluating
t, we have
i
1 dp:
1
log Pe (x))
+ (t(
Vpe(z) F((x))
(31)
11' (((x) Ot
log pe (x) V (o) F(()).
(32)
Ot
For the second term in Equation 29h. we note that At = t
(z) + (z) + (
(33)
This implies that the Jacobian determinant admits the following asymptotic expansion,
log det VP.. = log det + V + o(tt)
(34)
= =log(1 + TrV
+ o(At))
(35)
+ = 1
(36)
Consequentially we have,
log det (x)
(37)
ov
/=t
By substituting in Equations
32 and 37 into Equation
29 we obtain,
Gx(x) Øt
log + Vlogpt()) F(())) + Fr))
(38)
13

---

A.2
Proof of Theorem 3.1
Proof. Let s; = =i/T and t, (s;), Theorem 2.1 implies
(11441,4) 6((44)) + oA)7),
(39)
Vi(t41, (((t++1)At, + o(Ar),
(40)
where AT = max, At,] By the the mean value theorem, År S .elo ((s)/T and hence is O(T T co.
We will first establish the convergence of A(p, T). Using Equation 40 we obtain the following estimate for T
A(p, T) Σ ((li11,44
(41)
In the limit as T + 0o, this Riemann converges to A,
lim A T) = S Vo(()dl.
(42)
T-x
We will now obtain the limit of TL T). First denote 8; i/T and As; 1/T. Using the differentiability of &o we have,
At ((s++)) - (ss) = ((((+1)
o (T
(43)
T
Substituting Equation 466 into Equation 339
we obtain the following estimate for TL(s, T),
TL( T) = T
4(1+1,4)
(44)
1=
T-1
T
Σ
8(t++1) + o(T
(45)
T-1
= Σ (((((4+1))((((4+))) 2 + o1).
(46)
i=0
In the limit as -> oo, this converges to the integral for E((),
TC(,,T L. ((s(ss) E(y).
(47)
Combining Jensen's 's inequality with the fact that is increasing implies,
E(p)- = L' ((s)ss (f, (X(p(s))((sdds)
-s. Vo(t)at) = -
(48)
The last equality follows by substituting L= p(s). Note a schedule generator yo obtains the Jensen lower bound if only if there is C such that for all [0,1]
C= (()) (s)){²((s)²
(49)
By taking square roots and integrating from to
VCs /-" (s(l) (s')ds'
* (*))
Vo(t)dt = (s)).
(50)
By using the substitution = 1, along with the constraints A(1) = A and (1) = we obtain VC=A.
VC = (1)) A(1) = A.
(51)
Finally, by inverting Equation (50). we conclude our proof,
(s)= A (As).
(52)
14

---

A.3
Comparison to Fisher Information
When Fi, (z) = z, the quantity (z) = pe (x) reduces to the Radon- Nykkody derivative between Py and pt, and (x(z) reduce to the Fisher score function, कt log (xx). By Theorem 2.1 the corrector optimised cost satisfies, (t t') = 6t(t) + o(At), where
(t) = u() Ex Xpmpi V | = v(t) Xampp
V
log Pt
(53)
Øt
Suppose Pt is sufficiently regular, and the Poincaré inequality (Poincaré expectation the Fisher score with respect to Pt is zero, we have,
1890 holds. Since the
[ logp:
6c(1) 2 C(()o(()EX)M))
C(t)o(t)ā=((),
(54)
where C(t) > is some constant. Here 6p (t) = Varx, at log Pt. is the Fisher information for the diffusion path pa. Suppose view the diffusion path curve in the probability distribution space with metric (t) Equation (54) shows that the topology induced by 6c(t) is stronger than the Fisher information, and geometry more regular.
B Training Algorithms
B.1 Adaptive training
We may incorporate Algorithm| [ into online algorithm used during training. For fixed score function along the diffusion path, there is an optimal schedule. Similarly, for a fixed schedule, there an optimal score function that can be learnt from the data. To incorporate these two steps, we propose a two-step algorithm for online training.
First, for a fixed schedule, we optimise the score function. Then, using this estimated score function, we compute the optimal schedule through Algorithm T To add regularity throughout training, as our score predictions are over batches rather than the entire dataset, we do not replace the current schedule with our computed optimal one. Instead, we take a weighted combination of the current schedule with the computed optimal one.
The weighting factor 7 (0, 1) akin learning rate for the schedule optimisation. If is set high, our schedule learning may be overly influenced by the current batch, which could negatively affect the score training performance.
Algorithm 2
Require: Initial schedule T (4.) learning rate 7 E (0, 1), score estimate 88
1: while not converged do
2: for each batch from data do
3: Fix T and assign 0 a (0. B, (,B,T) T)
4: Fix Sø* and over batch estimate [(t;+ 12), i 1= 0,..., T -
(Equation 13
5: Assign T +- UpdateSchedulee
L(t441,44)
6:
Update time locations t; + -t; + 101-11 7);
7: end for
8: end while
B.2 Estimating predictor optimised cost
Proposition B.1. Let Fi,e be the predictor map given by the forward Euler discretisation S of the probability flow ODE, For N E I and let (x) to be the Jacobian of the Hutchinson trace estimator (Hutchinson 1989 for V(V (x))) E X and E [0, 1].
1 N N
(x =
Σ
(z))))) en ~(0
(55)
15

---

If At is small enough such that.
1 ATT (1W)- ((0)I-591 -001- logp((a) (l
< 1.
(56)
Then, as N the following limit exists almost surely.
At g()) lim JNN + 0(AI))
v log det (Fr
(57)
2
N
Proof. The probability flow ODE update is given through
1 (x) AI (Mxe) = log
(58)
By taking the gradient, we have,
1 - logp((t)
V (LN(( 1+ Al
(59)
Since (1 A1) AITrA + 0 that holds when Tr(A) < 1, our assumption in Equation (56) hold, we can apply a Taylor expansion
log det (x)e = AITI (1O)) (1())- 1
(o0gggl) + O(A())
(60)
By taking gradient, we have
At V log del Fey v(()2VTr (V2TTT(2110))) logp+ + 0(A12)
(61)
=-E(ve)) At
g TTr log (X)) + O(AP))
(62)
At (( VTr (v (o 0(Ar).
(63)
The last line used the fact that the Hessian of log P: equivalently the Jacobian V(V log the latter we can approximate unbiasedly using the Hutchinson trace estimator gain Jr (z).
B.3 Denoising Score Matching Weighting
The standard denoising score matching weighting used by Song et al. (2021) log Pogll///ooooooooo/Xooo For (x (mro = (20e; s())oo, o² (7)7) we have,
X(I) x
log P: (x) ||2, s(t)xoll? d
log (t)².
(64)
o(1)² 2
Therefore,
s())z0 xt x(t)
V nx log (x0
(65)
The expectation for E[[] V log Proo Xx)x)) is taken with respect to X, ~ Pejo- We therefore have,
s(t) - X,
E|| oo = = (x
(66)
ott)
s(t) - s(t) (t)e
EN(e:D.)
(67)
a(1)
-6
EN(eoo))
(68)
(t)
Ev(Co.) [||e|]]]
(69)
(1)
(70)
(t)
where on the second line we have used the reparameterisation trick. Therefore, we have that the standard weighting for denoising score matching is X(t) x att)?
16

---

Schedule (Be)
0.8
Linear Schedule
Corrector Schedule
Predictor Schedule
0.4
0.2
0.0
0.0
time
2.0
Corrector Schedule
Predictor Schedule
Density
Linear Schedule
1.0
True Density
0.0
-2
X
Corrector Schedule
Predictor Schedule
Score
Linear Schedule
-500
-2
X
Figure 5: Schedules and density estimates for: linear (blue); Stein score optimised (green); and predictor optimised (red) schedules. The predictor optimised schedule identifies bump along the diffusion path where the reference Gaussian density splits into two modes. In the regions where the score is evaluated (around +-6), our trained score accurate compared to the linear schedule score, which fails to match the slope of the true score, resulting in a wider variance density estimate.
C Experiment Details
C.1 ID Density Estimation
For all ID experiments, we train a one spatial dimension model with continuous time encoding via Gaussian Fourier features to embed time values into a higher-dimensional space. The model architecture includes hidden dimension of 128, five layers, an embedding dimension of 12, and one residual time step. It features a combination of residual blocks, both incorporating linear layers with GELU activation and LayerNorm, tailored to integrate time embeddings.
C.1.1 Bimodal Example
We train our model in the form of ftt) = 3:/2 and g(t) = BB, using Algorithm 2 with 7 = 0.1. We use the weight function of = ). where at., train our model for thousand iterations using both a fixed linear schedule and our optimisation algorithm initialised the linear schedule. Due to the non-linear dependence during training of the vransport-optimisee schedule, initialisation in this context plays an important role. We initialise our predictor-optimise schedule with the optimal schedule generated with respect to L. without the Jacobian term as a first approximation. We then add the Jacobian term, optimise the schedule, and train for additional 5 thousand iterations.
Si, and ā,, Th 0t2, We
In Figure 5 we see that the linear schedule forms an estimate with greater variance than the true data and the predicted densities of the optimised schedules.
17

---

Corrector Optimised Score
Predictor Optimised Score
Score
Corrector Sctwddlle Prodicter Scheduli
Xt
Scheoule
Score
Linear Scheoule Score
Scare
02
Time
Figure 6: Comparison of sampling from a mollified Cantor distribution using DDMs with two different schedules: linear (blue) and optimised (green). The optimised schedule enables the DDM to capture eight distinet data modes centered on the mollified Cantor distribution (grey), whereas the linear schedule does not have clear mode separation. The optimised schedule diffusion accurately predicts the score for the mollified Cantor distribution, being near vertical lines interweaving the Cantor set, where the linear schedule fails to adequately approximate the score.
C.1.2 Mollified Cantor Distribution
We train our model with [(t) 9 and g(t) v3. using the online Algorithm 2 with 0.01. The weight function is vf. (1 āt,)). where Ctt,, = Bt., and ãt,, -II Ote, To capture high-frequency details, we train the one-dimensional model for 150,000 iterations using both fixed linear schedule and our optimisation algorithm initialised at the linear schedule. The difference in sample quality is evident Figure
C.2 Pretrained Image Models
For sampling pretrained image models, we use the networks and implementation from Karras et al. (2022), hhttps://githtii
Sonn/NVlabs/eda We compute the FID using the provided FID script within the codebase with the standard 50, 000 samples using the same fixed seeds 0 49999 for all schedules. For each image dataset, we use the default sampling strategy included in the codebase by Karras al. 2022). For unconditional CIFARIO this deterministic 2nd order Heun solver with 18 timesteps, therefore total of 35 NFE (number of function evaluations) for the underlying denoising network. For both unconditional FFHQ and unconditional AFHQv2 the same deterministic 2nd order Heun solver is used with 40 timesteps (NFE 79). For class conditional ImageNet, the bespoke stochastic solver from Karras et al. (2022) used with 256 timesteps (NFB=511). The stochasticity settings are left at their default values for this dataset of Sehum 40, Smin 0.05, Smax 50, Smoise 1.003.
To compute the corrector optimised schedule for each dataset, we use Algorithm I We use 100 data samples for each dataset when computing L. as we find the variation in learned schedule small between different samples from the dataset. Our initial discretisation schedule used to calculate A(t) is LogLinear with 100 steps. We then fit monotonic spline to the cumulative estimated (t) and invert this function find (p function from which we can derive our schedules. This takes on the order of 5 minutes to find the Corrector Optimised Schedule for CIFAR 10 single RTX 2080Ti GPU. For predictor optimised schedules we repeat the same procedure however also include estimation of the Hessian term. We use 5 samples of for each image datapoint when using Hutchinson's 's trace estimator. takes GPU hours to compute the Predictor Opti- mised Schedule for CIFARIO due to the extra computational cost of computing the second derivatives.
We compare the corrector optimised and predictor optimised schedules for the 4 image datasets in Figure 8 We find that all of the schedules have the same general shape with increasing step sizes in log space as the generative process approaches clean data. However, the curvature of the schedule varies between datasets which is to expected as the schedule determined through the score which will vary depending
the data distribution. We find that, in general, the higher resolution datasets (FFHQ, AFHQv2 and ImageNet) favour shorter steps nearer the start the generative process the expense of larger steps at low noise levels near the end of the generative process.
18

---

1
5
100
300
200
130
100
52
200
-100
-200
400
-10
325
350
50-
25
-50
50
75
-0.25
0.00
0.25
0.0
-2.50
-2.25
0.00
0.25
0.30
-0.30
-0.25
0.00
0.22
0.50
-030
-022
3.50
3.3
20-
W
20
20
10-
20
40
6.50
0.25
0.50
8.00
0.5
.0.50
.025
0.00
0.25
0.50
3.5
15
10-
10
-10
-20
-10
-0.50 -0.25 0.00
025
030
-3.50
-375
0.00
025
0.50
-0.50
-0.25
0.00
0.25
0.50
-0.5
-025
3.0 325
350
t
t
t
t 39
7.
5.-
0.0
5.0
-4
0.50 0.25
0.50
0.25
0.50
2.50
3325
0.00
0.25
0.50
0.50 -025 0000
0.25 0.50
-050 -025
300
a2s
a50
Figure 7: Evolution of the estimated score for the mollified Cantor distribution Section 4.1 with Corrector Optimised Schedule. In this case the linear schedule fails to evenly progress the progression of the score, see Figure showing the terminal score estimate this case. The estimated score exhibits self-similar nature of interweaving roots around the centers of mass of the mollified Cantor distribution. Identification of these roots amounts to estimated modes in our density estimate, see Figure
19

---

10a
10
MNIST: For the MNIST experiments, we trained model with an image size of 32, 32 channels, with U-Net architecture, with residual block per U-Net resolution, without learning sigma, and Noise Level o
10
CIFAR10 Corrector
CIFAR10 Predictor
10-1
FFHQ Corrector
FFHQ Predictor
AFHQv2 Corrector
10
AFHQv2 Predictor
ImageNet Corrector
20
60
80
100
Timestep
Figure 8: Corrector optimised and predictor optimised schedules for the image datasets, CIFARIO. FFHQ, AFHQv2 and ImageNet.
CIFARIO Corrector
CIFARLO Predictor
15000
Corrector
FFHO Predictor
12500
APHQV2 Corrector
ima geNeet Corrector
10000
7500
5000
2500
10
10-1
100
10a
102
Noise Level
Figure 9: Local cost / /o()) versus for image datasets and using the corrector optimised versus predictor optimised costs.
We analyse the local cost as function of noise level v(a) for the 4 images datasets within Figure This local cost used the metric when determining velocities in the space of schedules from Pr to Po-
# points, T
10
20
30
50
100
CO (ours)
3.94
3.78
3.80
3.81
3.81
P 3
24.08
4.90
3.80
3.77
3.80
P =
4.02
3.76
3.78
3.80
3.81
P 100
4.31
3.81
3.81
3.81
3.81
Table 3: Comparison of sFID across different amounts of discretisation points for different schedules on CIFARIO. CO stands for our corrector optimised schedule.
C.3 Online Schedule Optimisation of Images
20

---

1403
5002
Cosine Schedule
-007
Casine Schedule
64-
1223
Length
1140
Energy
2002
201
003
52
Schedule Ir- 0.05
Cosine Schedule
1502
2 10000 2000 28000 52000
50392
S5000
22000
20000
33000
40000
50020
50020
280
202
100
402
500
Iteration
Iteration
Timestop
Figure 10: Progression of the length and energy Equation (18) over training of MNIST. Both models are trained from initialisation, one with adaptive schedule learning (red) and one without (blue). We can see that the energy and length quantities increase during training. Recall that for fixed path of scores + V log Pr that the length A is constant. As are learning the score, this value is not constant during training. Interestingly, by optimising the schedule during training we observe larger length value, possibly indicating that the diffusion path learned with the optimised schedule differs greatly from the path learned without.
Sample progression of MNIST digits for the standard cosine schedule with 0.008 (top) against our optimised schedule (bottom). As we can see, the cosine schedule spends more time near the Gaussian reference distribution whereas the optimised schedule quickly determines large scale features and spends more time toward the data distribution.
Figure 11:
2520
Schedule 1r=0.1
40000
Schedule Ir=0.1
Corrector Optimited LooS
5000-
Schedule Ir=0.01
Schedule Irw0.01
Cosine Schedule
Cosine Schedule
3000
3520
3550
15000 1,000
33838
20030
53003
Iteration
Figure 12: Progression of the length A and cost through online training for CIFAR-10. For the larger learning rate, Algorithm seems to garner a larger A value at a faster rate that the lower schedule learning rate. For the fixed cosine schedule run, A is stable, perhaps because the score estimate has stabilised for the fixed cosine schedule already during the model training burn-in phase.
21

---

with dropout rate of 0.3. The diffusion process was configured with 500 diffusion steps. Training was conducted with learning rate of le-4.
The schedule was also initialised with the Cosine schedule and trained for 60,000 iterations on an NVIDIA 1080 Ti GPU with 12 GB RAM. The batch size for MNIST
128. We trained two models: one with the schedule optimisation and one without. The schedule training rate was set to 7 0.05, as stated in Algorithm 2
Training took approximately 12 hours for either model.
CIFAR-10: For the CIFAR-10 10 experiments, we trained a model with an image size of 32, 128 channels. and U-Net architecture with 3 residual blocks per multiplier resolution (described in codebase Nichol and Dhariwal (2021)). without learning sigma, and with a dropout rate 0.3. The diffusion process was configured with 1000 diffusion steps and cosine noise schedule.
The schedule was initialised with the Cosine schedule and trained for 160,000 iterations on NVIDIA A40 GPUs, each with 48 GB RAM. The batch size was set to 1,536. After training for 160,000 iterations, we trained two models online: one with the corrector schedule optimisation and one without, for additional 50,000 iterations on single GPU with batch size of 384. The burn in phase for training took approximately 50 hours, with the individual schedule optimisations after this taking approximately 24 hours.
C.4 Licenses
Codehases:
Improved Denoising Diffusion Probabilistic Models Nichol and Dhariwal (2021): MIT License
Elucidating the Design Space of Diffusion-Basse Generative Models Karras et (0222)
Craibution-ooCoommmen -haraal----------e-e 4.0 International
Datasets:
CIFAR-10 Krizhevsky et al.
2009); MIT license
FFHQ Karras et al.
(2018): Creative Commons BY-NC-SA 4.0 license
AFHQv2 Choi et al.
(2020): Creative Commons BY-NC 4.0 license
ImageNet Deng et al. 2009): y: Unknown License
D
Acknowledgements of Funding
CW acknowledges support from DST Australia. AC acknowledges support from the EPSRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/1). acknowledges support from the NSERC Postdoctoral Fellow Program.
22

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes].
Justification: Our main theoretical and experimental contributions are clearly stated in the abstract and demonstrated in the paper. They reflect the paper's contributions and scope.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitations of our work in Section S
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
Answer: [Yes].
Justification: All the proofs are proven the supplementary material. They are duly cross-referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide detailed descriptions of our experimental procedures in Appendix and provide the code to run our experiments.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We provide the code necessary
experiments.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide all our experiment details in Appendix
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Our main metric FID score for which standard practice to report it calculated on the first 50,000 images generated from the model. Standard deviations could be bootstrapped from this set but this is not standard practice. Furthermore, the computational cost to sample large image models 50,000 times multiple times is prohibitive for our academic compute cluster.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed reproduce the experiments?
23

---

Answer: [Yes]
Justification: provide details of the compute resources used in Appendix C
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes].
Justification: After careful review of the NeurIPS Code of Ethics, it is clear that the research presented in this paper conforms with the Code of Ethics in every respect.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper is mostly theoretical and methodological. We do not see immediate societal impact this work. However, we acknowledge that large scale implementation of our algorithm might suffer from the same societal biases as any other generative models. Indeed could improve the quality of generative models and hence been used to generate deepfakes for disinformation.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We provide the references and licenses for the codebases and datasets used in our work in Appendix C.4
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not introduce any new assets.
Guidelines:
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA].
Justification: The paper does not involve crowdsoureing nor research with human subjects. 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human
Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA].
Justification: The paper does not involve crowdsourcing nor research with human subjects.


Reviews:
[{"review_id": "vCOC6USLC8", "reviewer": "Reviewer_KRNd", "summary": "The paper proposes an improved discretization schedules based on a novel cost measure that they proposes. The proposed method can update a given schedule to achieve greater sample quality as is demonstrated with solid experiments.", "questions": "* It seems the proposed method do not accelerate the backward sampling process since the updated schedule has the same number of timesteps. I wonder if the author have examined the computational efficiency of their Algorithm 2? Since when it comes to SGM, the running time is one of the major concerns.", "limitations": "See weaknesses and questions above.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strengths\n* The proposed method is novel and can adapt to most given SGM model to update its discretization schedule to achieve better sample quality. The computation of the proposed method is also tractable.", "weaknesses": "Weaknesses\n* While the paper provides heuristic theoretical justification to the design of the cost function, a better way to justify that the updated schedule provides better sample may be to directly measure the difference between the learned density (using a discretization schedule) and the true density $p_0$ and point out that updated schedule can reduce this difference.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their review and suggested improvements. We agree that providing a simple error plot for cases where the true density is known would be informative. In our diffusion model, we do not have the ability to query the density directly; however, we do predict the score. In Figure 4 of our paper, we make this comparison for our bimodal example of the scores at t=0. It can be noted here that the linear schedule devolves into learning a function dissimilar to the expected shape of the score for a bimodal distribution. To further illustrate this point, we have created a plot that evolves over training, showing the error of the score at time t=0, corresponding to our data distribution. We observe that with our schedule optimisation, we achieve a lower error than with a fixed linear schedule over the course of training. Additionally, we have computed the likelihood of the generated samples in this case where the true density is known. Similarly, we observe that during training, our optimised schedule generates datasets with higher likelihood than the fixed schedule case. We thank the reviewer for their suggestion, which has improved the quality of our paper, and we will include these improvements in a potential camera-ready submission.\n\n\n\n_Q: It seems the proposed method do not accelerate the backward sampling process since the updated schedule has the same number of timesteps. I wonder if the author have examined the computational efficiency of their Algorithm 2? Since when it comes to SGM, the running time is one of the major concerns._\n\nWe thank the reviewer for their question. We do observe in our simple 1D experiment in Figure 1 that an optimised schedule can achieve better results with fewer discretisation points, improving runtime. To better convey this potential use case, we have added follow-up experiments based on the reviewer's previous suggestion to see if we can see a gain in FID with fewer data points for image generation. \n\nAdditionally, we have conducted a refinement test to assess the computational benefit of optimising the schedule. For this, we took the rho3 schedule for CIFAR10 and computed the FID score for 18, 20, 30, 50, and 100 discretisation points. In this range, at 18 discretisation points, the rho3 schedule achieves an FID of 5.47, while at 100 discretisation points, it achieves an FID of 2.05. In comparison, the optimised schedule achieves an FID of 1.99 at 18 discretisation points. We further observe the same trend in the sFID, another image performance metric.\n\nIn this experiment, we can see that optimising the schedule can lead to improved FID and sFID performance with fewer discretisation points. Furthermore, we compared all schedules with only 10 discretisation points. In this case, the optimised schedule outperforms all others in terms of FID and sFID, demonstrating its superior performance in the few discretisation point regime. We thank the reviewer for their suggestion to perform this interesting computational experiment. In a potential camera-ready submission, we would include this study in the main manuscript.", "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": "cC3Wa5b7EQ", "reviewer": "Reviewer_hc4h", "summary": "A popular way to sample a target distribution is to run a \"predictor-corrector\" SDE, which additively combines two processes whose stationary distribution is the target: a Langevin (or \"corrector\") process and a reverse-diffusion (or \"predictor\") process. \n\nHow to discretize this SDE without deterioriating the approximation quality of its terminal distribution is the topic of this paper. The authors derive a discretization cost (Eq 13), which is the Fisher divergence between two distributions $\\textemdash$ one obtained running the \"predictor-corrector\" process in continuous-time and the other obtained by running the \"predictor\" process only in discrete-time $\\textemdash$ multiplied by the velocity of the \"corrector\" process.\n\nThis cost can be approximates using the estimated scores (section 3.2) and the authors optimize this cost with respect the schedule. They illustrate numerically the optimized schedules on image datasets and verify that they yield better approximations of the target distribution (measured in FID).", "questions": "It is unclear to me if in Eq 57, $J_t$ requires computing the Hessian of $\\log p_t(\\cdot)$: if so, wouldn't that be an expensive quantity to compute in order to obtain the optimal schedule?", "limitations": "Yes.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clearly written and derives a principled way to derive an optimal schedule for the popular \"predictor-corrector\" SDE, of which \"annealed Langevin dynamics\" and a \"reverse diffusion process\" are a special case. This is of interest to the sampling and diffusion models communities. \n\nThe appendix has nice, additional experiments assessing the optimal schedule and 1D and image datasets.\n\nThe authors are also clear about the limitations of their approach, namely that it assumes perfect estimation of the scores.", "weaknesses": "Including in the Appendix something like Figure 10 but for ImageNet, CIFAR, or CelebA would be interesting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments. We agree that more visual comparisons of the different schedules would be illustrative. To this end, we will include standard diffusion progression plots, comparing the image diffusion path for CIFAR and CelebA on different schedules. In a potential camera-ready version, we would include one such plot in the main body of the text and all comparisons in the appendix. We greatly appreciate the reviewer’s comment, which has helped improve the presentation of our schedule optimisation.\n\n_It is unclear to me if in Eq 57, $J_t$ requires computing the Hessian of $\\log p_t $: if so, wouldn't that be an expensive quantity to compute in order to obtain the optimal schedule?_\n\nWe thank the reviewer for highlighting this important computational point. In our implementation, we have formulated an estimator (Proposition B.1) that completely avoids computing the full Hessian. This proposition is currently in the appendix, but we state it here for completeness:\n\n**Proposition B.1:** Let $ F_{t,t'} $ be the predictor map given by the forward Euler discretisation (8) of the probability flow ODE. For $ N \\in \\mathbb{N} $, let $ \\hat{J}_{t,N}(x) $ be the Jacobian of the Hutchinson trace estimator (Hutchinson, 1989) for $ \\nabla(\\nabla \\log p_t(x)) $ at $ x \\in X $ and $ t \\in [0, 1] $:\n \n$$\n\\hat{J}_{t,N}(x) = \\frac{1}{N} \\sum_n \\nabla(v_n^T J_t(x) v_n),\n$$\n\nIf $ \\Delta t $ is small enough such that:\n\n$$\n\\Delta t \\cdot \\text{Tr} \\left( f(t)I - 2g(t)\\nabla \\log p_t(x) \\right) < 1,\n$$\n\nthen, as $ N \\to \\infty $, the following limit exists almost surely:\n\n$$ \n\\nabla \\log \\det \\nabla F_{t'}(x) = -\\Delta t g(t)^2 \\lim_{N \\to \\infty} \\hat{J}_{t,N}(x) + O(\\Delta t^2).\n$$\n\nHere, we only require the trace of the Hessian, so we need to compute only the diagonal entries, ensuring that the memory cost remains linear with respect to the parameters. Additionally, we have utilised a Hutchinson trace estimator for this term, which can be efficiently implemented in standard deep learning libraries.\n\nIn our provided code, we have explicitly implemented this in PyTorch and have found that this computation is scalable to image datasets. This approach has enabled us to compute our predictor-optimised schedules, something infeasible if we had needed to compute the full Hessian. Our estimator performs well, producing schedules that achieve good FID scores. This was a subtle computational challenge that was non-trivial to overcome in our work.\n\nIn a potential camera-ready version, we will clarify this approximation and emphasise that the computational resources required are significantly less than those needed for computing the full Hessian. We will also make clear in the main text that we have a scalable estimator for this term and do not require computing the Hessian matrix. We thank the reviewer for bringing this point to our attention and helping improve the quality of our paper.", "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": "H7idRTE0vf", "reviewer": "Reviewer_8ZVm", "summary": "This paper proposes a method to optimize for the discretization schedule of the diffusion sampling process. The main idea is to minimize a surrogate for the total length of the diffusion sampling path between two consecutive time steps, where the length is defined in terms of the Fisher divergence between two steps. Experiments show that the proposed scheme can find schedules competitive with the ones discovered with manual search.", "questions": "* Line 95, why is the Markov kernel taking in a probability measure as the second argument, as opposed to the usual definition where it takes in a Borel set (so that conditioned on the first argument you get a probability measure)?\n* Line 136, what is the point of introducing coupled noises in (9) and (10) if they are going to be canceled anyway?\n* Line 218, what is the \"distance\" referred to here?\n* Line 235, can the authors comment on why we know $\\Delta t$ is sufficiently small a priori (since we are optimizing the time intervals themselves)?\n* Can the authors justify the usage of simple cubic interpolation in step 3 of Algorithm 1 in order to invert $\\Lambda$? Why would interpolation work well?\n\n\n\nMinor comments:\n* In (3), why is there a tilde on $W_t$?\n* Line 75, what are \"Lebesgue probability measures\"?\n* In (4), one could split (3) into a prediction and a correction term differently. Why should this specific way of splitting?", "limitations": "Limitations are discussed.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "* The problem studied is important for diffusion generative modeling as the manual search for a good sampling schedule can be time-consuming.\n* The proposed method is simple, despite the rather convolutated presentation.\n* Empirical evaluation shows that the proposed scheme can have on-par performance with hand-tuned schedules.", "weaknesses": "* The presentation of the paper could use a lot of improvement. I found many sections unnecessarily convoluted without conveying much insight.\n  - The whole setup of predictor/corrector seems unnecessary to me. If I understand correctly, Section 2.1 and 2.2 are used to introduce the cost which turns out to be just the Fisher divergence either between $p_{t_{i+1}}$ and $p_{t_{i}}$ (the so-called corrector optimized cost) or $F_\\sharp p_{t_i+1}$ and $p_{t_i}$ for $F$ equal to the one-step of ODE integration.\n  - Section 2.3 and Theorem 2.1 do not seem to be related to the actual algorithm.\n* The theoretical insights in Section 3 are unclear to me. Section 3.1 seems to be regurgitating the fact in differential geometry that the energy-minimizing path connecting two points is a constant-speed geodesic. In particular, I'm unable to identify interesting theoretical contributions related to the setting of diffusion generative modeling.\n* The computation cost could be quite high, especially using the predictor-optimized schedule, which does not really offer a noticeable advantage over the cheaper corrector-optimized schedule. This seems to suggest one should always choose the simpler corrector-optimized schedule. Yet with the corrector optimized schedule, the picture of predictor/corrector seems irrelevant since the predictor is just the identity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for the reviewers for their constructive feedback. We will address each point raised individually below:\n\n - _The whole setup of predictor/corrector seems unnecessary to me._\n\nSee general rebuttal response. \n\n- _Section 2.3 and Theorem 2.1 do not seem to be related to the actual algorithm._\n\nSection 2.3 and Theorem 2.1 are necessary links that take us from the Fisher Divergence at two time points t, t’ to a metric that can be used to derive the optimal schedule. In Section 3.1 when we come to find the costs associated with diffusion paths, we require an infinitesimal notion of cost that can be integrated along the diffusion path. The Fisher Divergence as stated in Section 2.2 is unsuitable for this purpose because it is evaluated between two times t and t’. In Section 2.3 and Theorem 2.1 we derive this corresponding infinitesimal cost. Our algorithm relies entirely on the approximation of this infinitesimal cost along the diffusion path.\nWe appreciate the reviewer highlighting this lack of clarity in the significance of Section 2.3 and Theorem 2.1. In a revised version we will add a clarifying sentence before Theorem 2.1 to emphasise its relevance and importance.\n\n- _Section 3.1 seems to be regurgitating the fact in differential geometry that the energy-minimizing path connecting two points is a constant-speed geodesic_\n\nOur algorithm is grounded in the theoretical insights from Section 3. The relevance of Section 3.1 lies in demonstrating how schedule optimization in diffusion models can be reframed as finding an energy-minimising path with an appropriate metric.  The key connection between differential geometry and diffusion models is established in Theorem 2.1, where we derive the Reimannian metric compatible with the objectives of diffusion models, and demonstrate how to approximate the energy using the Fisher divergence. The material from differential geometry included in this section is meant to ensure completeness. To improve the clarity of this section, we will add a linking sentence that explicitly states how the local metric tensor from Theorem 2.1 enables us to use differential geometry to study diffusion schedule design.\n\n- _The computation cost could be quite high.._\n\nSee general rebuttal response.\n\nQuestions:\n\n- _Line 95_\n\nThank you for pointing out this typo.\n\n- _Line 136_\n\nThe introduction of coupled noises in Section 2.2 is a direct consequence of performing a perturbation analysis on the Langevin dynamics trajectories. Specifically, we are interested in understanding how these trajectories evolve over a small time interval when the score function is slightly perturbed.\n\nBy coupling the noise terms in equations (9) and (10), we ensure that any observed differences between these trajectories are solely due to the perturbation and not due to random fluctuations. The noise cancels out naturally because we have isolated the perturbation of the score; this cancellation is a result of the computational approach rather than a manufactured effect. Without coupling the noises, the comparison would be less informative, as it would be unclear whether differences were due to the perturbation or simply random noise.\n\n- _Line 218_\n\nHere \"distance\" is intended to mean \"length along the path,\" as referenced in the preceding lines 202-207. We will update the text to more clearly state \"length along the path\".\n\n- _Line 235_\n\nWe have not claimed that we know $\\Delta t$ is sufficiently small, but rather if it is small, then our approximation is valid. One way to assess if a grid is fine enough is by computing the statistical length $\\Lambda$, which is independent of the schedule, across a range of grid sizes. When this value stabilises for a particular $\\Delta t$, it indicates that the desired integral approximation of Equation 21 has been achieved.\n\n- _Simple cubic interpolation in step 3_\n\nWe also remark that be apprimae $\\Lambda^{-1}(t)$ directly by interpolating with knots $\\{(\\Lambda(t_i), t_i)\\}$. This avoids needing to approximate $\\Lambda(t)$ and numerically inverting. Since we avoid numerical inversion methods, the choice of interpolant is not overly critical. We use cubic interpolation as an example, but any interpolation method respecting monotonicity (e.g. linear) can also be used. We will add this comment after Algorithm 1.\n\n- _In (3)_\n\nIn the context of the forward/backward SDEs in a diffusion model, the Brownian increments for the forward and backward SDEs are independent of one another. We use a tilde to emphasise that the Brownian increment for the forward diffusion ($\\tilde{W}$) is different from the Brownian increment of the backward diffusion ($W$). \n\n- _Line 75_\n\nLebesgue probability measures are those that have a density function with respect to the Lebesgue measure. We've included standard references to make this clearer for readers who may be unfamiliar with probability theory and thank the reviewer for improving the accessibility of our work.\n\n- _In (4)_ \n\nThis specific way of splitting is chosen to identify the unique ODE that transports samples from one density $p_{t}$ along the diffusion path to $p_{t'}$. This approach was originally identified and used in the seminal work of Yang Song [reference]. We will add a reference in the text to where this splitting is derived in Yang Song's paper.", "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": "5yBqXiWwlA", "reviewer": "Reviewer_g4tL", "summary": "This paper proposes a novel algorithm for adaptively selecting an optimal discretisation schedule in the training of diffusion model. The proposed method does not require hyper-parameter tuning and adapts to the dynamics and geometry of the diffusion path. This method also achieves competitive FID scores on image datasets.", "questions": "Please see Weakness.", "limitations": "Please see Weakness.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "+ The findings of the paper are interesting. Current schedule designs are heuristic, and this paper provides a theoretical analysis, which is of interest to the community.\n+ The proposed online data-dependent adaptive discretization schedule tuning algorithm is simple and effective.\n+ The mathematical analysis presented in this paper is exceptionally detailed. Meanwhile, adequate related work was discussed in detail. The authors also provide corresponding code that makes the work in this paper very solid.\n+ The authors provide clear visual comparisons that show the distribution of optimized times in detail.", "weaknesses": "+ The experiments in this paper are somewhat weak and lack validation for the LDM model. In fact, the LDM model is widely used, so it is important to validate the proposed method for LDM.\n+ Some other diffusion frameworks also use schedule discretization in the same way, e.g., VP-SDE [1] and Flow Matching [2], is the method proposed in this paper robust for them as well?\n\n[1] Score-Based Generative Modeling through Stochastic Differential Equations ICLR 2021\n\n[2] Flow Matching for Generative Modeling ICLR 2023\n\n+ The authors should provide additional quantitative and visual comparisons, as the FID metric alone is insufficient for assessing generation quality [1]. Low FID scores do not necessarily correlate with superior generative performance. Therefore, incorporating other metrics such as sFID, Recall, and Precision would offer a more comprehensive evaluation. This enhanced analysis will help further substantiate the validity of the proposed method.\n\n[1] The Role of ImageNet Classes in Fréchet Inception Distance ICLR 2023\n\n[2] Diffusion Models Beat GANs on Image Synthesis  NeurIPS 2021\n\n+ The mathematical parts of this paper are somewhat difficult to understand, especially the differential geometry part of it, and I would suggest that the authors add some background and notation notes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their kind comments on the strengths of our paper. We believe that our schedule tuning algorithm could serve as a strong default for diffusion models and might be applicable more broadly, though a more extensive exploration is difficult to undertake in an initial paper. We hope to use this work as a starting point for further investigation into general schedule tuning in generative models.\n\nIn the context of diffusion models, we welcome the reviewer's suggestions for more comprehensive numerical validation. To this end, we have presented preliminary summary statistics for sFID. We are appreciative of the reviewers suggestions for additional metrics, and agree that their addition to a potential camera ready article would increase the strength of the paper. \n\nWe believe that our methodology could be applied to flow-matching based models, but this would also require mild adjustment to the analysis we present. We hope our present work could be foundational to a potential future study. \n\nWe agree that validation of LDM with our method is a worthwhile and meaningful venture. To this end, we have implemented our method in the HuggingFace LDM library, but due to time constraints do not currently have FID to report. In a potential camera ready article, we would additionally include FID and sFID comparisons on optimised and standard schedules for Latent Diffusion Models.\n\n_The authors should provide additional quantitative and visual comparisons, as the FID metric alone is insufficient for assessing generation quality [1]. Low FID scores do not necessarily correlate with superior generative performance. Therefore, incorporating other metrics such as sFID, Recall, and Precision would offer a more comprehensive evaluation. This enhanced analysis will help further substantiate the validity of the proposed method._\n\nWe thank the reviewer for their suggestion. We would like some assistance from the reviewer, we would like to confirm that sFID refers to the metric computed in:\n\nNash C, Menick J, Dieleman S, Battaglia PW. Generating images with sparse representations. arXiv preprint arXiv:2103.03841. 2021 Mar 5.\n\nWe have implemented this metric and computed preliminary results for sFID on the CIFAR dataset which are included in our rebuttal document. We have found that the results for sFID correlate strongly with the results for FID in this case. In a potential camera-ready version, we would incorporate this metric across all the datasets presented. We appreciate the reviewer’s input, which has helped substantiate the validity of our method and improve the results presented in our paper.\n\n_The mathematical parts of this paper are somewhat difficult to understand, especially the differential geometry part of it, and I would suggest that the authors add some background and notation notes._\n\n We agree with the reviewer that this mathematical content is new to the diffusion literature. In a potential camera-ready version, we would include a brief introduction to metric tensors and charts commonly used in differential geometry, along with their relation to the work we present here. We appreciate the reviewer’s suggestion, which will help us present our work more clearly.", "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:
We thank all the reviewers for taking the time to read our paper and for their kind comments on the strengths of the presentation and methodology. We also thank the reviewers for providing constructive feedback to improve the paper for the potential camera-ready version. We will address each point made by the reviewers individually. Below, we would like to summarize the additional experiments we conducted based on the reviewer feedback comparing (1) performance metrics and (2) schedule refinements. Please see the rebuttal document for the results. 

(1) Performance metrics: 
We have implemented the sFID metric in addition to FID and KLUB on the CIFAR dataset, which is included in our rebuttal document. We have found that the results for sFID correlate strongly with the results for FID in this case. In a potential camera-ready version, we would incorporate this metric across all the datasets presented. We appreciate the reviewer’s input, which has helped substantiate the validity of our method and improve the results presented in our paper.

(2) Schedule refinements:
Additionally, we have conducted a refinement test to assess the computational benefit of optimising the schedule. For this, we took the rho3 schedule for CIFAR10 and computed the FID score for 18, 20, 30, 50, and 100 discretisation points. In this range, at 18 discretisation points, the rho3 schedule achieves an FID of 5.47, while at 100 discretisation points, it achieves an FID of 2.05. In comparison, the optimised schedule achieves an FID of 1.99 at 18 discretisation points. We observe the same trend in the sFID, another image performance metric. Notably, the optimised schedule maintains stable performance even with 10 discretisation points compared to a suboptimal schedule.

We would further like to briefly comment on importance of the predictor/corrector framework used in our work and the reasoning for presenting our work within it. We further would like to address why we include analysis for both predictor and corrector optimised schedules in our work.

Our predictor/corrector framework is based on the seminal work of Yang Song et al. [1,2] (also see [2,3]). It is crucial to justify the final cost that we use to find optimal schedules. We agree that the final cost coincides with the standard Fisher Divergence; however, without the predictor/corrector formalism, it would be unclear whether this cost is justified and what relation controlling the Fisher Divergence would have in controlling error along the trajectory of a diffusion model. Our work provides this link by showing that the Fisher Divergence cost can be derived by minimising the work done by the diffusion model update steps under the predictor-corrector framework, thereby providing the relationship between the Fisher Divergence and diffusion model sampling. Furthermore, without understanding this link between predictor/corrector and the Fisher Divergence, it would be unclear what weighting should be applied to the Fisher Divergence when computing the cost for different time points along the trajectory. In Section 3.3, we find that the variance of the applied noise is the appropriate scaling, which is derived by noting that Langevin corrector steps should have steps on the order of the scale of the target distribution, which directly implies a meaningful scaling for the Fisher Divergence. Without the predictor/corrector framework, it would be unclear which scaling to use.

Corrector-optimised schedules are indeed more computationally efficient, requiring only one additional function evaluation per training step without gradient tracking compared to using no schedule optimisation. Despite this minimal overhead, they provide performance similar to the predictor-optimised schedule, which requires higher-order gradients. We explicitly mention this in our experimental section (Lines 311-312), where we recommend the use of the more cost-effective corrector-optimised schedule for image datasets. However, we believe it is important to include the analysis of the predictor-optimised schedule for two main reasons: firstly, it is a natural extension of the corrector schedule, and understanding its computational trade-offs could be valuable for the research community; secondly, while the corrector-optimised schedule performs similarly on image datasets, the predictor-optimised schedule may offer advantages in other application domains. We agree that the corrector-optimised schedule's benefits should be more prominently highlighted for its simplicity, lower cost, and competitive performance. We propose making this recommendation earlier in the paper, particularly in Section 2.

- [1] Generative Modeling by Estimating Gradients of the Data Distribution, Yang Song, Stefano Ermon
- [2] Score-Based Generative Modeling through Stochastic Differential Equations, Yang Song et al
- [3] Elucidating the Design Space of Diffusion-Based Generative Models, Term Karras et al
- [4] The probability flow ODE is provably fast, Sitan Chen et al

Meta-review: This work proposes an approach to select a good discretization schedule for simulating the backward generation process of diffusion model. The selection was guided by an theoretically grounded approximation of discretization cost. Although there were some concerns about implementation cost and idealized assumptions such as perfect score, overall the merits overweight the limitations, and I'm delighted to recommend an acceptance. The authors are encouraged to account for the discussions in a revision.
Decision: Accept (poster)