Title: Score-Optimal Diffusion Schedules
Venue: NeurIPS 2024
Accepted: yes
Keywords: Diffusion model

adaptive schedule

Annealed Langevin
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

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.

1 Introduction

Denoising Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021) or DDMs are state-of-the-art generative models. They are formulated through considering a forward noising process that interpolates from the target to a 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 p 0 through the forward diffusion process on time interval [0 , 1] , described by

with drift f ( t ) X t , diffusion coefficient g ( t ) and Brownian noise increment d W t . The coefficients f ( t ) and g ( t ) are chosen such that at time t = 1 the distribution of X 1 is very close to a reference Gaussian distribution p 1 in distribution. A sample starting at p 0 and following Equation (1) until time t will be distributed according to p t which is a mollified version of the data distribution

0.6

0.4

0.2

0.0

0.2

0.4

0.6

X

0

1

2

3

4

5

6

Density

Training Data Samples

0.6

0.4

0.2

0.0

0.2

0.4

0.6

X

0

1

2

3

4

5

6

Density

0.6

0.4

0.2

0.0

0.2

0.4

0.6

X

0

1

2

3

4

5

6

Density

The parameters s ( t ) and σ 2 ( t ) define the noising schedule . They can be found in closed form in terms of f ( t ) and g ( t ) (Karras et al., 2022). To obtain samples from p 0 , we can reverse the dynamics of the forward diffusion in Equation (1) to obtain the backward diffusion ,

By simulating Equation (3) backwards in time, we evolve reference samples X 1 ∼ p 1 from t = 1 to t = 0 to obtain samples that are terminally distributed according to the data distribution p 0 . To simulate Equation (3) numerically, we must decide upon a discretisation of time, T = { t i } T i =0 with t T = 1 , t 0 = 0 , which we refer to as the discretisation schedule . For a given noising schedule, it is important to select an appropriate discretisation schedule such that (3) can be simulated accurately, i.e. p t i and p t i -1 should not differ significantly. In this paper, we derive a 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, t i = i/T , with the noising schedule dictating the change between p t i and p t i -1 . Two prominent examples have the form s ( t ) = √ ¯ α ( t ) , σ 2 ( t ) = 1 -¯ α ( t ) with the linear schedule introduced by Ho et al. (2020) having ¯ α ( t ) = 1 -exp ( -∫ t 0 β ( s )d s ) with linear β ( t ) = β min + t ( β max -β min ) . Alternatively, Nichol and Dhariwal (2021) introduce the cosine schedule with ¯ α ( t ) = f ( t ) /f (0) , f ( t ) = cos ( t + ϵ 1+ ϵ π 2 ) 2 for ϵ = 0 . 008 . Karras et al. (2022) 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 a distribution supported on the Cantor set (Figure 1, left), the default linear schedule fails entirely to capture the modes of the data distribution (Figure 1, middle). However, our optimised schedule learned using Algorithm 1 recovers these modes (Figure 1, right). Without such an automatic algorithm, finding performant discretisation schedules often reduces to an expensive and laborious hyperparameter sweep.

We devise a method for selecting a discretisation schedule that yields high-quality samples from p 0 . 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 p 1 to p 0 . Our cost purely depends on the distributional shifts along the diffusion path and assumes perfect score estimation, hence, we refer to our schedules as score-optimal diffusion schedules . The resulting algorithm is cheap to compute, easy to implement and requires no hyperparameter search. Our algorithm can be applied to find discretisation schedules for sampling pre-trained models as well as performed online during DDMtraining. We demonstrate our proposed method on highly complex 1D 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 a notion of cost of traversing from our reference distribution p 1 to the data distribution p 0 through each intermediate distribution p t , 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 (3) from time t to t ′ will always take you from p t to p t ′ , the simulation step will need to do more work to make sure the samples are distributed according to p t ′ if p t and p t ′ are very different distributions rather than if they are close. In the following, we will make this intuition precise.

2.1 Predictor/Corrector Decomposition of the Diffusion Update

To begin, let X = R d and P ( X ) be the space of Lebesgue probability measures on X with smooth densities. For t ∈ [0 , 1] , define the diffusion path p t ∈ P ( X ) as the law of X t satisfying the forward diffusion Equation (1) initialised at the data distribution X 0 ∼ p 0 , or equivalently the law of the X t satisfying the backward diffusion Equation (3) initialised at the reference distribution X 1 ∼ p 1 .

Given a sample X t ∼ p t , a sample X t ′ ∼ p t ′ can be generated by integrating the backward diffusion Equation (3). In Song et al. (2021), it was shown that we can further decompose the backward diffusion Equation (3) at time t into a deterministic flow governed by the probability flow ODE, and the stochastic flow driven by Langevin dynamics targeting p t ,

This decomposition into a deterministic flow and a correction will help us derive our cost in Section 2.2 by analysing the work done by the correction to keep the samples on 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 t ′ and (2) a corrector step that uses Langevin dynamics targeting p t ′ to correct any accrued error. Note we are not advocating for the implementation of such a procedure, only that by imagining simulating with this hypothetical predictor-corrector 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.

̸

Definition 2.1. A predictor for t, t ′ ∈ [0 , 1] is a smooth bijective mapping F t,t ′ : X ↦→ X , such that det ∇ F t,t ′ = 0 , and the predicted distribution is the pushforward of p t by F t,t ′ denoted F ♯ t,t ′ p t .

Definition 2.2. A corrector for t ∈ [0 , 1] , τ ∈ [0 , ∞ ) is a one-parameter family Markov transition kernel L t,τ : X × P ( X ) ↦→ [0 , 1] such that Z τ ∼ L t,τ ( z, d z τ ) is the law of Langevin dynamics stationary distribution p t at time τ , initialised at z ∈ X , and running at speed v ( t ) > 0 ,

The corrector map L t,τ is specified through an integration time τ and time-dependent speed v ( t ) . We assume τ is fixed and we will describe the appropriate choice of v ( t ) in Section 3.3. The predictorcorrector algorithm, given X t ∼ p t , first applies the predictor Z 0 = F t,t ′ ( X t ) and then uses the corrector to drive the predicted samples towards p t ′ ,

In general, Z 0 will not be a sample from p t ′ exactly because F t,t ′ may not be a perfect transport from p t to p t ′ . In Section 2.2, assessing the work done by L t ′ ,τ to drive the Z 0 towards p t ′ will be key in deriving our cost. Our cost will then depend upon the specific choice for F t,t ′ . Two natural choices for F t,t ′ are apparent. Setting F t,t ′ 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 F t,t ′ 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 is the identity F t,t ′ ( x ) = x . In such a case, the predicted state reduces to the initial state, F t,t ′ ( X t ) = X t ∼ p t . The work done by the corrector step will then be related to the full discrepancy between p t and p t ′ because the predictor provides no help in transporting the sample.

Example 2.2 (Probability Flow ODE) . The predictor step is optimal when F t,t ′ is a transport from p t to p t ′ . In such a case, the predicted state produces a sample from the target distribution, F t,t ′ ( X t ) ∼ p t ′ , and so the corrector step would have to perform no work. An optimal predictor map F t,t ′ can be obtained by integrating the probability flow ODE from time t, t ′ ,

Practical algorithms numerically integrate Equation (7), e.g. an Euler step with ∆ t = t ′ -t ,

In such a case, the work done by the corrector depends on the error in the probability flow integrator.

2.2 The Incremental Cost of Correction

We now focus on deriving a cost related to the work done by the corrector step in the predictorcorrector algorithm. Later, in Section 3, we will find the discretisation schedule that minimises the total cost. To derive the cost, we will analyse the movement of Z 0 under the corrector step's dynamics L t ′ ,τ ( Z 0 , d x τ ) . This requires some care because even if Z 0 is already at stationarity, i.e. perfectly distributed according to p t ′ , applying the Langevin correction step will still result in movement of Z 0 due to the stochasticity of the update. However, the computed work done by the correction step in this case should be 0 . To correctly assign the work done, we will compare two processes. The first is the trajectory of Langevin dynamics, Z τ , defined by the corrector L t ′ ,τ initialised at Z 0 = F t,t ′ ( X t ) targeting p t ′ . The second is a virtual coupled Langevin dynamics ˜ Z τ initialised at F t,t ′ ( X t ) , driven by the same noise and speed but targeting the stationary distribution of the predictor F ♯ t,t ′ p t ,

Notably, Z τ d = X t ′ ,τ and ˜ Z τ d = F t,t ′ ( X t ) share the same law as the corrected sample and predicted sample respectively. Since Z τ and ˜ Z τ are coupled to have the same noise, the difference in their trajectory, Z τ -˜ Z τ , isolates the change in corrector dynamics due to discrepancy between F ♯ t,t ′ p t and p t ′ . If F ♯ t,t ′ p t is very different to p t ′ , then Z τ -˜ Z τ will be large, signifying the corrector is needing to do lots of work to push the distribution of Z towards the target p t ′ . Conversely, if F ♯ t,t ′ p t = p t ′ , then Z τ -˜ Z τ = 0 and no work is done. For small τ , ( Z τ -Z 0 ) /τ is the initial velocity of Z under the p t ′ corrector dynamics, and similarly for ( ˜ Z τ -Z 0 ) /τ . We can then define the incremental cost L ( t, t ′ ) by taking limits as τ → 0 + , measuring the expected L 2 norm ∥ · ∥ of the difference,

We can approximate Z τ -˜ Z τ using an Euler step, noting that the coupled noise terms cancel,

By substituting Equation (12) in Equation (11), we have

where D ( p ∥ q ) = E X ∼ q [ ∥∇ log p ( X ) - ∇ log q ( X ) ∥ 2 ] is a statistical divergence on p, q ∈ P ( X ) , measuring the L 2 distance between the scores of q and p with respect q . D ( p ∥ q ) is referred to as the Stein divergence or the Fisher divergence; see e.g. (Johnson, 2004). For a given choice of v ( t ) and F t,t ′ we now have a cost measuring the change from p t to p t ′ . This cost is intrinsically linked with the effort performed by a DDM sampling algorithm because it is derived through considering the work done by a hypothetical predictor-corrector style update. We note, however, that this general cost can be used to obtain discretisation schedules for use in any style of DDM sampler.

2.3 Corrector and Predictor Optimised Cost

By inverting Z 0 = F t,t ′ ( X t ) , we can express Equation (13) in terms of an expectation with respect to the reference sample X t ∼ p t , and the score of G t,t ′ : X ↦→ R + , the incremental weight function associated with the transport F t,t ′ from the Sequential Monte Carlo literature (Arbel et al., 2021),

In most cases, it is infeasible to efficiently compute the Jacobian correction in Equation (14). When F t,t ′ ( x ) = x is the identity map corresponding to the corrector optimised update from Example 2.1 Equation (14) reduces a rescaled Stein discrepancy between p t and p t ′ , and G t,t ( x ) = p t ′ ( x ) /p t ( x ) reduces to the likelihood-ratio between p t ′ and p t . We will refer to this case as the corrector-optimised cost denoted L c ( t, t ′ ) , to distinguished it from the predictor-optimised cost L p ( t, t ′ ) derived above, where when relevant, we will use subscripts c and p to distinguish between the two:

The corrector-optimised cost L c ( t, t ′ ) provides meaningful information during the update from reference p t to the target p t ′ . It is worth computing even when the predictor-optimised cost L p ( t, t ′ ) is accessible. L c ( t, t ′ ) measures the change between the reference and target distribution independent of the predictor, whereas L p ( t, t ′ ) measures the residual error between the predictor and target. Notably, L c ( t, t ′ ) encodes information about the incremental geometry of the diffusion path, whereas L p ( 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 t ∼ p t towards p t ′ , we would expect L p ( t, t ′ ) ≤ L c ( t, t ′ ) .

For deriving our optimal discretisation schedule, we require a notion of how L ( t, t ′ ) increases with small increases in t ′ i.e. knowing local changes in incremental cost. In Section 3, we use this local 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 A for a PDE and geometric interpretation.

Theorem 2.1. Suppose p t ( x ) , F t,t ′ ( x ) , v ( t ) and G t,t ′ ( x ) are three-times continuously differentiable in t, t ′ , x and let ˙ F t ( x ) = ∂ ∂t ′ F t,t ′ ( x ) ∣ ∣ t ′ = t and ˙ G t ( x ) = ∂ ∂t ′ G t,t ′ ( x ) ∣ ∣ t ′ = t . Suppose the following hold: (1) for all x ∈ X , t ∈ [0 , 1] , F t,t ( x ) = x and (2) there exists V : X ↦→ R such that for all x ∈ X and t ∈ [0 , 1] , ∥∇ ˙ G t ( x ) ∥ 2 ≤ V ( x ) and sup t ∈ [0 , 1] E X t ∼ p t [ V ( X t )] < ∞ . Then for all t ∈ [0 , 1] , we have L ( t, t ′ ) = δ ( t )∆ t 2 + O (∆ t 3 ) , where

Theorem 2.1 shows that, under regularity assumptions, then the incremental cost is L ( t, t ′ ) ≈ δ ( t )∆ t 2 is locally quadratic and controlled by the local cost δ ( t ) . The δ ( t ) measures the sensitivity of the incremental cost L ( t, t ′ ) to moving samples along the diffusion path to t ′ ≈ t . Notably, δ ( t ) = 0 if and only if the predictor satisfies the continuity equation, ∂ ∂t p t + ∇· ( p t ˙ F t ) = 0 .

3 Score-Optimal Schedules

Given a discretisation schedule T = ( t i ) T i =0 satisfying 0 = t 0 < · · · < t T = 1 , our hypothetical predictor-corrector algorithm recursively uses the predictor and corrector maps to generate a sequence ( X i ) T i =0 starting at X T ∼ p 1 such that the terminal state X 0 approximates samples from p 0 ,

We want to 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 is reasonable if, in our hypothetical corrector steps, τ is set sufficiently large such that the Langevin correction converges to stationarity. We find in our experiments that even if the

schedules derived under Assumption 3.1 are used in sampling algorithms for which Assumption 3.1 does not hold, we still obtain high quality samples. Equipped with Assumption 3.1, we can measure the efficiency of the path update through total accumulated cost L = ∑ T i =1 L ( t i +1 , t i ) , which we will use as our objective to optimise T . In this section, we will identify the optimal schedule T ∗ minimising the cost L by considering an infinitely dense limit. We will then provide a tuning procedure amenable to online schedule optimisation during training. Finally, we will discuss a suitable choice for v ( t ) , the velocity of our hypothetical corrector steps, as well as related work.

3.1 Diffusion Schedule Path Length and Energy

Let φ : [0 , 1] ↦→ [0 , 1] be a strictly increasing, differentiable function such that φ (0) = 0 and φ (1) = 1 . We will say T is generated by φ if t i = φ ( i/T ) for all i = 0 , . . . , T . The schedule generator φ dictates how fast our samples move through their diffusion path. Since every schedule T of size T is generated by some φ , optimising T is equivalent to finding a generator φ minimising L ( φ, T ) , the total cost accumulated by the schedule of size T generated by φ . By Jensen's inequality, we have L ( φ, T ) ≥ Λ( φ, T ) 2 /T , where for t i = φ ( i/T ) ,

As we later prove in Theorem 3.1, in the dense schedule limit as T →∞ , the cost L ( φ, T ) and its lower bound Λ( φ, T ) are controlled by the energy E ( φ ) and length Λ respectively where,

The intuition for why E ( φ ) is an energy, and Λ a length can be gained by first conceptualising the diffusion time t as a spatial variable rescaled by the metric δ ( t ) defined by our cost L . We have φ and ˙ φ are position and velocity, respectively. Integrating the speed ∫ 1 0 √ δ ( φ ( s )) ˙ φ ( s )d s = ∫ 1 0 √ δ ( t )d t along a curve φ ( s ) obtains the 'length' Λ , whilst integrating a speed squared, ∫ 1 0 δ ( φ ( s )) ˙ φ ( s ) 2 d s obtains a 'kinetic energy' E ( φ ) . Note that the length is an invariant of the schedule, whereas the kinetic energy is not. The length Λ measures the intrinsic difficulty of traversing the diffusion path according to the cost independent of φ , whereas E ( φ ) measures the efficiency of how the path was traversed using φ . This geometric intuition hints at the solution to the optimal scheduling problem. The optimal φ should travel on a geodesic path from p 1 to p 0 , at a constant speed with respect to metric δ . For this optimal φ , we then have the kinetic energy being equal to the square of length between p 1 and p 0 . Theorem 3.1 makes the previous discussion precise.

Theorem 3.1. Suppose the assumptions of Theorem 2.1 hold. For all schedule generators φ ,

Moreover, E ( φ ) ≥ Λ 2 , with equality if and only if φ ∗ satisfies,

Notably independent of the choice of φ , as T →∞ , the cost L ( φ, T ) ∼ E ( φ ) /T . This implies that the cost decays to zero at a linear rate, proportional to E ( φ ) and L ( φ, T ) ≳ Λ 2 /T independent of φ . 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 ( φ ∗ ) = Λ 2 . The intuition for the formula φ ∗ ( s ) = Λ -1 (Λ s ) is that this implies Λ( φ ∗ ( s )) = Λ s meaning say 10% of the way through the optimal schedule, we should have traversed 10% of the way along the distance between p 1 and p 0 i.e. 0 . 1 × Λ . This relation holds for constant speed straight lines, meaning φ ∗ is the optimal schedule. For a finite T , Theorem 2.1 implies the optimal schedule T ∗ = { t ∗ i } T i =0 generated by φ ∗ ensures the incremental cost is constant L ( t ∗ i +1 , t ∗ i ) ≈ Λ 2 /T 2 for all i = 0 , . . . , T -1 .

Our geometric intuition in the language of differential geometry is that the diffusion path M = { p t } t ∈ [0 , 1] is Riemannian manifold with metric δ endowed by the incremental cost L ( t, t ′ ) . The schedule generator defines a curve s ↦→ p φ ( s ) ∈ M reparametrising the diffusion path between p 0 and p 1 . Theorem 3.1 shows that φ ∗ is the geodesic of length Λ in M between p 1 and p 0 that traverses the diffusion path at a constant speed √ δ ( φ ∗ ( s )) ˙ φ ∗ ( s ) = Λ with respect to δ and minimises the cost.

3.2 Estimation of Score-Optimal Schedules

Given a schedule T = { t i } T i =0 and estimates of the incremental cost L ( t i +1 , t i ) , Algorithm 1 adapts Algorithm 3 from Syed et al. (2021) to estimate the optimal schedule T ∗ = { t ∗ i } T i =0 generated by φ ∗ . Wecan use Algorithm 1 to refine the schedule for a pre-trained DDM or learn the schedule jointly with the score function. For this joint procedure, we detail in Appendix B.1 how function evaluations can be reused to estimate the cost to minimise computational overhead. For L c ( t, t ′ ) we need only evaluate ∇ log p t ( X t ) and ∇ log p t ′ ( X t ) both available through our model's score estimate. Computing L p ( t, t ′ ) is more challenging since there are Hessian terms that arise in Equation (14). Under the assumption that the step size ∆ t > 0 is sufficiently small, we can approximate ∇ log | det ∇ F t,t ′ ( X t ) | 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 a Hutchinson trace (Hutchinson, 1989) like estimator in Proposition B.1, we compute this quantity memory-efficiently in high dimensions, requiring only standard auto-differentiation back-propagation.

3.3 Choice of Velocity Scaling

Recall that our cost is derived by considering a Langevin dynamics step with velocity v ( t ) . This velocity should be selected so that Langevin dynamics explores the same proportion of our distribution at varying times throughout our diffusion path. Thus, v ( t ) should be on the same scale as the spread of the target, p t . Commonly used noising schedules have s ( t ) ≤ 1 , and our data distribution is normalised so the scale of p t is on the order of σ ( t ) . We therefore set v ( t ) = σ ( t ) . This results in a σ ( t ′ ) -weighted divergence for our incremental cost L ( t, t ′ ) = σ ( t ′ ) 2 D ( p t ′ || p t ) . This can be compared to the weighted denoising score matching loss used to train DDMs (Song et al., 2021), which is also a squared norm of score differences: λ ( t ) E X 0 ,X t [ ∥ s θ ( X t , t ) -∇ log p t | 0 ( X t | X 0 ) ∥ 2 ] for some weighting function λ ( t ) chosen to equalise the magnitude of the cost over the path. In Song et al. (2021), λ ( t ) ∝ 1 / E [ ∥∇ log p t | 0 ( X t | X 0 ) ∥ 2 ] was chosen, which, as we show in Appendix B.3, is λ ( t ) ∝ σ 2 ( t ) . This choice of velocity scaling provides an alternative perspective on this commonly used weighting of squared norms of score differences.

Algorithm 1 UpdateSchedule

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 Kingma et al. (2021), necessitating auxiliary costs and objectives for schedule design. Uniform steps in the signal-to-noise ratio, log ( s ( t i ) /σ ( t i )) , are used by Lu et al. (2022), but this ignores the target distribution's geometry. Watson et al. (2021) 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-stage 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 p t | 0 ( x t | x 0 ) , ignoring the true target distribution. Xue et al. (2024) derive a schedule to control ODE simulation error, but their cost depends only on the ODE solver, and not on the data distribution.

4 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 staircase and its support being the Cantor set, forming a challenging 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 t = 10 -5 . With this mollification, our data density has eight pronounced peaks. We train a one-dimensional DDMfor 150,000 iterations using both a fixed linear schedule and our optimisation algorithm Algorithm 2 initialised at 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 7 we show the complexity of the learned score which displays a self-similar fractal structure.

4.2 Adaptive Schedule Learning for Bimodal Example

We train a DDM on a 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 p 0 ( x ) = 1 2 p left ( x ) + 1 2 p right ( x ) , where p left and p right are normal distributions with means -6 and 6 , respectively, and a common variance σ 2 = 0 . 1 2 .

We learn two diffusion models, one using the linear schedule, and the other using a schedule that is 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. It can be seen in Figure 2 that when the schedule is learned during training, the likelihood evaluation increases and the true score error decreases, in contrast to the linear schedule that remains constant, or worsens, in this regard during training.

0

2000

4000

6000

8000

Training Iterations

0.90

0.95

1.00

1.05

1.10

1.15

1.20

Likelihood Average

Comparison of Linear and Learned Schedules

Linear Schedule

Learned Schedule

0

2000

4000

6000

8000

Training Iterations

300

400

500

600

700

800

900

Score Error

Error Comparison of Linear and Learned Schedules

Linear Schedule Error

Learned Schedule Error

4.3 Scalable Schedule Learning Diffusion

Here we demonstrate that jointly learning the schedule and score using our online training methodology (Algorithm 2) scales to high-dimensional data and converges to a stable solution. We train DDMs on CIFAR-10 and MNIST initialised at the cosine schedule using the codebase from Nichol and Dhariwal (2021). In Figure 3 (left), we show the incremental costs √ L ( t j +1 , t j ) 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.1. Figure 3 (right) shows the learned schedule spends more time at high-frequency details, we visualise a sampling trajectory in Figure 11.

0

20

40

60

80

100

Timestep

10

2

10

1

10 0

10 1

10 2

Noise Level

rho=3

rho=7

LogLinear

CIFAR10 Corrector

CIFAR10 Predictor

0

200

400

600

800

1000

Timestep

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Length Increment

CIFAR Cosine Schedule

0

200

400

600

800

1000

Timestep

0.0

0.2

0.4

0.6

0.8

1.0

(1

t

)

1/2

CIFAR Optimal Scedule

Cosine Initial Scedule

4.4 Sampling Pre-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 pretrained models from Karras et al. (2022), whose DDM is parameterised such that the forward noising distribution is of the form p t | 0 ( x t | x 0 ) = N ( x t ; x 0 , σ 2 t I ) . The scheduling problem then reduces to deciding on a stepping scheme through { σ i } N i =1 , σ N = 0 . Karras et al. (2022) suggest a polynomial based schedule with a parameter ρ that controls the curvature of the schedule

A lower ρ value results in steps near σ min being shortened and steps near σ max being lengthened. Through analysing the truncation error for sampling in Karras et al. (2022), they find that setting ρ = 3 approximately equalises this error, however it is found empirically that ρ = 7 results in better sample quality. We also compare against a schedule that takes uniform steps in log σ space Lu et al. (2022) which we refer to as the LogLinear schedule and a schedule that takes a convex shape in log space. Schedule visualisations are provided in Figure 4 (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 CIFAR10, 40 for FFHQ and AFHQv2, 256 for ImageNet) and solver (Heun second order) as Karras et al. (2022). Our results are shown in Figure 4. Our optimised schedules are able to achieve competitive FID to the best performing ρ = 7 schedule hand-tuned in Karras et al. (2022). This is expected as our schedules take a similar shape to the ρ = 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 trial-and-error.

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 CIFAR10 in Table 2. Notably, the FID decreases with T 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.

We also compare corrector optimised schedules to predictor optimised schedules in Table 1. They provide similar performance so, on image datasets, we encourage the use of the cheaper to compute corrector optimised schedule. Finally, in Figure 4 (left), we report the raw 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 to worse schedules which is 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 is 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.

5 Discussion

We have introduced a 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 p can be computationally expensive due the calculation of second derivatives, however, in Section 4.4 we found L c to provide a cheap and performant alternative. Furthermore, our theory is derived assuming perfect score estimation. Future work can expand on the geometric interpretation of the diffusion path and links to information geometry to further refine the DDM methodology.

References

Arbel, M., Matthews, A., and Doucet, A. (2021). Annealed flow transport Monte Carlo. In International Conference on 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 Abhandlungen 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., Nabiei, 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 hierarchical image database. IEEE Conference on Computer Vision and Pattern Recognition .

Fritsch, F. N. and Carlson, R. E. (1980). Monotone piecewise cubic interpolation. SIAM Journal on Numerical Analysis , 17(2):238-246.

Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems .

Hutchinson, M. (1989). A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in Statistics - Simulation and Computation , 18(3):1059-1076.

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, T. (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, J. (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. In Advances in 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 la physique mathématique. American Journal of Mathematics , pages 211-294.

Sabour, A., Fidler, S., and Kreis, K. (2024). Align your steps: Optimizing sampling schedules in diffusion models. In International Conference on Machine Learning .

Santos, J. E., Fox, Z. R., Lubbers, N., and Lin, Y . 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. In 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 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. In International Conference on Learning Representations .

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.03802 .

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 arXiv:2402.17376 .

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 ′ ) ∈ [ t, t ′ ]

Since all t, t ′ we have ∥∇ ˙ G s ( t, t ′ )( x ) ∥ 2 ≤ V ( x ) , and E X t ∼ p t [ V ( X t )] < ∞ , the dominated convergence theorem and the continuity of v ( t ′ ) implies,

The last equality follows since ∇ ˙ G t is continuous in t .

We will now compute ˙ G t . Since F t,t ( x ) = x , we have G t,t ( x ) = 1 for all x . This implies, we can express ˙ G t in terms of the derivative of log G t,t ′ ,

where log G t,t ′ equals,

By combining Equation (27) to Equation (28), we obtain,

For the first term in Equation (29), we use chain rule to obtain,

where ' · ' denotes a dot product of vectors. By evaluating at t ′ = t , we have

For the second term in Equation (29), we note that as ∆ t = t ′ -t → 0

This implies that the Jacobian determinant admits the following asymptotic expansion,

By substituting in Equations (32) and (37) into Equation (29) we obtain,

Consequentially we have,

A.2 Proof of Theorem 3.1

Proof. Let s i = i/T and t i = φ ( s i ) , Theorem 2.1 implies

where ∆ T = max i | ∆ t i | . By the the mean value theorem, ∆ T ≤ sup s ∈ [0 , 1] ˙ φ ( s ) /T and hence is O ( T -1 ) as T →∞ .

We will first establish the convergence of Λ( φ, T ) . Using Equation (40) we obtain the following estimate for T L ( φ, T ) ,

In the limit as T →∞ , this Riemann sum converges to Λ ,

We will now obtain the limit of T L ( φ, T ) . First denote s i = i/T and ∆ s i = 1 /T . Using the differentiability of φ we have,

Substituting Equation (46) into Equation (39), we obtain the following estimate for T L ( φ, T ) ,

In the limit as t →∞ , this converges to the integral for E ( φ ) ,

Combining Jensen's inequality with the fact that φ is increasing implies,

The last equality follows by substituting t = φ ( s ) . Note a schedule generator φ ∗ obtains the Jensen lower bound if only if there is a C such that for all s ∈ [0 , 1] ,

By taking square roots and integrating from 0 to s ,

By using the substitution in s = 1 , along with the constraints Λ(1) = Λ and φ ∗ (1) = 1 we obtain √ C = Λ ,

Finally, by inverting Equation (50), we conclude our proof,

A.3 Comparison to Fisher Information

When F t,t ′ ( x ) = x , the quantity G t,t ′ ( x ) = p t ′ ( x ) /p t ( x ) reduces to the Radon-Nykodym derivative between p t ′ and p t , and ˙ G t ( x ) reduce to the Fisher score function, ∂ ∂t log p t ( x ) . By Theorem 2.1 the corrector optimised cost satisfies, L c ( t, t ′ ) = δ c ( t )∆ t 2 + o (∆ t ) , where

Suppose p t is sufficiently regular, and the Poincaré inequality (Poincaré, 1890) holds. Since the expectation of the Fisher score with respect to p t is zero, we have,

where C ( t ) > 0 is some constant. Here δ F ( t ) = Var X t ∼ p t [ ∂ ∂t log p t ] is the Fisher information for the diffusion path p t . Suppose we view the diffusion path as a curve in the probability distribution space with metric δ c ( t ) . Equation (54) shows that the topology induced by δ c ( t ) is stronger than the Fisher information, and geometry is more regular.

B Training Algorithms

B.1 Adaptive training

We may incorporate Algorithm 1 into an online algorithm used during training. For a fixed score function along the diffusion path, there is an optimal schedule. Similarly, for a fixed schedule, there is 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 1. 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 γ ∈ (0 , 1) is akin to a learning rate for the schedule optimisation. If γ is set too high, our schedule learning may be overly influenced by the current batch, which could negatively affect the score training performance.

B.2 Estimating predictor optimised cost

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 ∈ N and let ˆ J t,N ( x ) to be the Jacobian of the Hutchinson trace estimator (Hutchinson, 1989) for ∇ ( ∇ log p t ( x ))) at x ∈ X and t ∈ [0 , 1] ,

If ∆ t is small enough such that,

Then, as N →∞ the following limit exists almost surely,

Proof. The probability flow ODE update is given through

By taking the gradient, we have,

Since log det(1 + ∆ t ) = ∆ t Tr A + O (∆ t 2 ) that holds when Tr( A ) < 1 , our assumption in Equation (56) hold, we can apply a Taylor expansion

By taking a gradient, we have

The last line used the fact that the Hessian of log p t is equivalently the Jacobian ∇ ( ∇ log p t ) , the latter we can approximate unbiasedly using the Hutchinson trace estimator to gain ˆ J t,N ( x ) .

B.3 Denoising Score Matching Weighting

The standard denoising score matching weighting used by Song et al. (2021) is λ ( t ) ∝ 1 / E [ ||∇ log p t | 0 ( X t | X 0 ) || 2 ] . For p t | 0 ( x t | x 0 ) = N ( x t ; s ( t ) x 0 , σ 2 ( t ) I ) we have,

Therefore,

The expectation for E [ ||∇ log p t | 0 ( X t | X 0 ) || 2 ] is taken with respect to X t ∼ p t | 0 . We therefore have,

where on the second line we have used the reparameterisation trick. Therefore, we have that the standard weighting for denoising score matching is λ ( t ) ∝ σ ( t ) 2 .

0.0

0.2

0.4

0.6

0.8

1.0

time

0.0

0.2

0.4

0.6

0.8

Schedule (

t

)

Linear Schedule

Predictor Schedule

6

4

2

0

2

4

6

x

0.0

0.5

1.0

1.5

2.0

Density

Predictor Schedule

Linear Schedule

True Density

6

4

2

0

2

4

6

x

500

0

500

Score

Predictor Schedule

Linear Schedule

C Experiment Details

C.1 1D Density Estimation

For all 1D 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 a 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 f ( t ) = β t / 2 and g ( t ) = √ β t using Algorithm 2 with γ = 0 . 1 . We use the weight function v 2 t n = (1 -¯ α t n ) , where α t n = 1 -β t n and ¯ α t n = ∏ n i =1 α t i . We train our model for 5 thousand iterations using both a fixed linear schedule and our optimisation algorithm initialised at the linear schedule. Due to the non-linear dependence during training of the transport-optimised schedule, initialisation in this context plays an important role. We initialise our predictor-optimised schedule with the optimal schedule generated with respect to L c without the Jacobian term as a first approximation. We then add the Jacobian term, optimise the schedule, and train for an additional 5 thousand iterations.

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.

0.4

0.2

0.0

0.2

0.4

x

400

300

200

100

0

100

200

300

400

Score

Corrector Optimised Score

0.0

0.2

0.4

0.6

0.8

1.0

Time

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Schedule (

t

)

Linear Schedule

Predictor Schedule

250

0

250

Score

Predictor Optimised Score

0.4

0.2

0.0

0.2

0.4

x

10

0

10

Score

C.1.2 Mollified Cantor Distribution

We train our model with f ( t ) = β t / 2 and g ( t ) = √ β t using the online Algorithm 2 with γ = 0 . 01 . The weight function is v 2 t n = (1 -¯ α t n ) , where α t n = 1 -β t n and ¯ α t n = ∏ n i =1 α t i . To capture high-frequency details, we train the one-dimensional model for 150,000 iterations using both a fixed linear schedule and our optimisation algorithm initialised at the linear schedule. The difference in sample quality is evident in Figure 1.

C.2 Pretrained Image Models

For sampling pretrained image models, we use the networks and implementation from Karras et al. (2022), https://github.com/NVlabs/edm . 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 et al. (2022). For unconditional CIFAR10 this is a deterministic 2nd order Heun solver with 18 timesteps, therefore a 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) is used with 256 timesteps (NFE=511). The stochasticity settings are left at their default values for this dataset of S churn = 40 , S min = 0 . 05 , S max = 50 , S noise = 1 . 003 .

To compute the corrector optimised schedule for each dataset, we use Algorithm 1. We use 100 data samples for each dataset when computing L c as we find the variation in learned schedule is small between different samples from the dataset. Our initial discretisation schedule used to calculate Λ( t ) is LogLinear with 100 steps. We then fit a monotonic spline to the cumulative estimated Λ( t ) and invert this function to find φ ∗ function from which we can derive our schedules. This takes on the order of 5 minutes to find the Corrector Optimised Schedule for CIFAR10 on a 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 v for each image datapoint when using Hutchinson's trace estimator. It takes 5 GPU hours to compute the Predictor Optimised Schedule for CIFAR10 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 be expected as the schedule is determined through the score which will vary depending on the data distribution. We find that, in general, the higher resolution datasets (FFHQ, AFHQv2 and ImageNet) favour shorter steps nearer the start of the generative process at the expense of larger steps at low noise levels near the end of the generative process.

0.50

0.25

0.00

0.25

0.50

400

200

0

200

400

t = 1

0.50

0.25

0.00

0.25

0.50

300

200

100

0

100

200

300

t = 3

0.50

0.25

0.00

0.25

0.50

200

100

0

100

200

t = 5

0.50

0.25

0.00

0.25

0.50

150

100

50

0

50

100

150

t = 7

0.50

0.25

0.00

0.25

0.50

100

50

0

50

100

t = 9

0.50

0.25

0.00

0.25

0.50

100

50

0

50

100

t = 11

0.50

0.25

0.00

0.25

0.50

75

50

25

0

25

50

75

t = 13

0.50

0.25

0.00

0.25

0.50

50

25

0

25

50

t = 15

0.50

0.25

0.00

0.25

0.50

40

20

0

20

40

t = 17

0.50

0.25

0.00

0.25

0.50

40

20

0

20

40

t = 19

0.50

0.25

0.00

0.25

0.50

20

0

20

t = 21

0.50

0.25

0.00

0.25

0.50

20

10

0

10

20

t = 23

0.50

0.25

0.00

0.25

0.50

20

10

0

10

20

t = 25

0.50

0.25

0.00

0.25

0.50

10

0

10

t = 27

0.50

0.25

0.00

0.25

0.50

15

10

5

0

5

10

15

t = 29

0.50

0.25

0.00

0.25

0.50

10

5

0

5

10

t = 31

0.50

0.25

0.00

0.25

0.50

5

0

5

t = 33

0.50

0.25

0.00

0.25

0.50

7.5

5.0

2.5

0.0

2.5

5.0

7.5

t = 35

0.50

0.25

0.00

0.25

0.50

6

4

2

0

2

4

6

t = 37

0.50

0.25

0.00

0.25

0.50

4

2

0

2

4

t = 39

0

20

40

60

80

100

Timestep

10

2

10

1

10 0

10 1

10 2

Noise Level

CIFAR10 Corrector

CIFAR10 Predictor

FFHQ Corrector

FFHQ Predictor

AFHQv2 Corrector

AFHQv2 Predictor

ImageNet Corrector

10

2

10

1

10 0

10 1

10 2

Noise Level

0

2500

5000

7500

10000

12500

15000

(

)

CIFAR10 Corrector

CIFAR10 Predictor

FFHQ Corrector

FFHQ Predictor

AFHQv2 Corrector

AFHQv2 Predictor

ImageNet Corrector

We analyse the local cost as a function of noise level √ δ ( σ ) for the 4 images datasets within Figure 9. This local cost is used as the metric when determining velocities in the space of schedules from p 1 to p 0 .

C.3 Online Schedule Optimisation of Images

MNIST : For the MNIST experiments, we trained a model with an image size of 32, 32 channels, with a U-Net architecture, with 1 residual block per U-Net resolution, without learning sigma, and

0

10000

20000

30000

40000

50000

60000

Iteration

700

800

900

1000

1100

1200

1300

1400

Length

Cosine Schedule

0

10000

20000

30000

40000

50000

60000

Iteration

1500

2000

2500

3000

3500

4000

4500

5000

Energy

Cosine Schedule

0

100

200

300

400

500

Timestep

0.0

0.2

0.4

0.6

0.8

1.0

(1

t

)

0.5

Cosine Schedule

0000

0

10000

20000

30000

40000

50000

Iteration

3500

4000

4500

5000

5500

6000

6500

Corrector Optimised loss

Schedule lr=0.1

Schedule lr=0.01

Cosine Schedule

0

10000

20000

30000

40000

50000

Iteration

15000

20000

25000

30000

35000

40000

Length

Schedule lr=0.1

Schedule lr=0.01

Cosine Schedule

with a dropout rate of 0.3. The diffusion process was configured with 500 diffusion steps. Training was conducted with a learning rate of 1e-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 was set to 128. We trained two models: one with the schedule optimisation and one without. The schedule training rate was set to γ = 0 . 05 , as stated in Algorithm 2. Training took approximately 12 hours for either model.

CIFAR-10 : For the CIFAR-10 experiments, we trained a model with an image size of 32, 128 channels, and a 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 of 0.3. The diffusion process was configured with 1000 diffusion steps and a cosine noise schedule.

The schedule was initialised with the Cosine schedule and trained for 160,000 iterations on 4 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 an additional 50,000 iterations on a single GPU with a 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

Codebases:

Improved Denoising Diffusion Probabilistic Models Nichol and Dhariwal (2021): MIT License

Elucidating the Design Space of Diffusion-Based Generative Models Karras et al. (2022): Attribution-NonCommercial-ShareAlike 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): Unknown License

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 5.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes] .

Justification: All the proofs are proven in the supplementary material. They are duly cross-referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

Answer: [Yes]

Justification: We provide detailed descriptions of our experimental procedures in Appendix C 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide the code necessary to run our experiments.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide all our experiment details in Appendix C.

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 is FID score for which it is 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 computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide details of the compute resources used in Appendix C.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 of 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 it could improve the quality of generative models and hence been used to generate deepfakes for disinformation.

11. Safeguards

Question: 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, or scraped datasets)?

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 are 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 as details about compensation (if any)?

Answer: [NA] .

Justification: The paper does not involve crowdsourcing 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA] .

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n+ 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.\n\nWEAKNESSES:\n+ 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.\n\nQUESTIONS:\nPlease see Weakness.\n\nLIMITATIONS:\nPlease see Weakness.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nLimitations are discussed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nA 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).\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nIncluding in the Appendix something like Figure 10 but for ImageNet, CIFAR, or CelebA would be interesting.\n\nQUESTIONS:\nIt 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

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

STRENGTHS:
Strengths
* 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
* 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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13