Score-Optimal Diffusion Schedules

Authors: christopher  williams
andrew  campbell
arnaud  doucet
saifuddin  syed
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 18:21:23

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.

Main paper: 
 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 byEQUATIONwith drift f (t)X t , diffusion coefficient g(t) and Brownian noise increment dW 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 p t (x t ) = X0 p 0 (x 0 )p t|0 (x t |x 0 )dx 0 , p t|0 (x t |x 0 ) = N (x t ; s(t)x 0 , σ 2 (t)I).(2)38th Conference on Neural Information Processing Systems (NeurIPS 2024). 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,EQUATIONBy 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 ti and p ti-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 ti and p ti-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) = 1exp -t 0 β(s)ds 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 DDM training. 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. 
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.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 distributionX 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 ,dX t = f (t)X t -1 2 g(t) 2 ∇ log p t (X t ) dt Probability Flow Prediction ODE + -1 2 g(t) 2 ∇ log p t (X t )dt + g(t)d W t . Langevin Correction SDE (4)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 byF 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, dz τ )is the law of Langevin dynamics stationary distribution p t at time τ , initialised at z ∈ X, and running at speed v(t) > 0,Z 0 = z, dZ τ = v(t)∇ log p t (Z τ )dτ + 2v(t)dW τ .(5)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 ′ , Predictor:EQUATIONIn 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 ′ ,EQUATIONPractical algorithms numerically integrate Equation ( 7), e.g. an Euler step with ∆t = t ′ -t,EQUATIONIn such a case, the work done by the corrector depends on the error in the probability flow integrator.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 , dx τ ). 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 predictorEQUATIONNotably, 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,EQUATION)We can approximate Z τ -Zτ using an Euler step, noting that the coupled noise terms cancel,EQUATIONBy substituting Equation (12) in Equation ( 11), we haveEQUATIONwhere 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.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) ,EQUATIONIn 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:EQUATIONThe 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 samplesX 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Ḟt (x) = ∂ ∂t ′ F t,t ′ (x) t ′ =t and Ġ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 ∈EQUATIONTheorem 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 Ḟ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 ,EQUATIONWe 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. For all t, t ′ , if X t ∼ p t and X t ′ ,τ = L t ′ ,τ (F t,t ′ (X t ), dx τ ), then X t ′ ,τ ∼ p t ′ .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. 
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 forEQUATION)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,EQUATION)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 speed1 0 δ(φ(s)) φ(s)ds = 1 0δ(t)dt along a curve φ(s) obtains the "length" Λ, whilst integrating a speed squared, 1 0 δ(φ(s)) φ(s) 2 ds 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 φ,EQUATIONMoreover, E(φ) ≥ Λ 2 , with equality if and only if φ * satisfies,EQUATION)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.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 φ * . We can 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.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 X0,Xt ∥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.Require:Schedule T = {t i } T i=0 , incremental costs {L(t i+1 , t i )} T -1 i=0 1: Λ(t i ) = i-1 j=0L(t j+1 , t j ), i = 0, . . . T ▷ Equation (21); Fritsch and Carlson (1980) 2: Λ = Λ(t T ) ▷ Λ in Equation (19) 3: Λ-1 (•) = Interpolate({( Λ(t 0 ), t 0 ), . . . , ( Λ(t T ), t T )}); ▷ E.g.4: t * i = Λ-1 ( Λ i T ), i = 0, . . . , T ▷ Equation (21) 5: Return: T * = {t * i } T i=0Previous 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 
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 DDM for 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.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. 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 . (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 of schedules during generative sampling with 100 timesteps. "rho=3" and "rho=7" refer to Eq 22 with ρ = 3 and ρ = 7 respectively. LogLinear from Lu et al. (2022) and a convex schedule are also shown. We show our cost optimised schedules for CIFAR10 both using the corrector optimised cost and the predictor optimised cost. 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 scheduleEQUATIONA 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 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. 
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.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.1 Proof of Theorem 2.1Proof. We first note that by the mean value theorem, there exists s(t, t′ ) ∈ [t, t ′ ] L(t, t ′ ) |t ′ -t| 2 = v(t ′ ) 2 E Xt∼pt ∇G s(t,t ′ ) (X t ) t ′ -t 2 (23) = v(t ′ ) 2 E Xt∼pt ∇ Ġs(t,t ′ ) (X t ) 2 (24)Since all t, t ′ we have ∥∇ Ġs (t, t ′ )(x)∥ 2 ≤ V (x), and E Xt∼pt [V (X t )] < ∞, the dominated convergence theorem and the continuity of v(t ′ ) implies,EQUATION)The last equality follows since ∇ Ġt is continuous in t.We will now compute Ġt . Since F t,t (x) = x, we have G t,t (x) = 1 for all x. This implies, we can express Ġt in terms of the derivative of log G t,t ′ ,EQUATIONwherelog G t,t ′ equals, log G t,t ′ (x) = log p t ′ (F t,t ′ (x)) -p t (x) + log det ∇F t,t ′ (x).(28) By combining Equation ( 27) to Equation ( 28), we obtain,EQUATION)For the first term in Equation ( 29), we use chain rule to obtain,EQUATIONwhere '•' denotes a dot product of vectors. By evaluating at t ′ = t, we haveEQUATIONFor the second term in Equation ( 29), we note that as ∆t = t ′ -t → 0F t,t ′ (x) = x + Ḟt (x)∆t + O(∆t 2 ).(33) This implies that the Jacobian determinant admits the following asymptotic expansion,EQUATIONBy substituting in Equations ( 32) and (37) into Equation ( 29) we obtain,EQUATIONA.2 Proof of Theorem 3.1Proof. Let s i = i/T and t i = φ(s i ), Theorem 2.1 impliesEQUATIONEQUATIONwhere ∆ 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 ),EQUATIONIn the limit as T → ∞, this Riemann sum converges to Λ, limEQUATIONWe 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,EQUATIONSubstituting Equation ( 46) into Equation ( 39), we obtain the following estimate for T L(φ, T ),EQUATIONIn the limit as t → ∞, this converges to the integral for E(φ),EQUATIONCombining Jensen's inequality with the fact that φ is increasing implies,EQUATION)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],C = δ(φ * (s)) φ * (s) 2 .(49) By taking square roots and integrating from 0 to s,EQUATIONBy using the substitution in s = 1, along with the constraints Λ(1) = Λ and φ * (1) = 1 we obtain√ C = Λ, √ C = Λ(φ * (1)) = Λ(1) = Λ.(51) Finally, by inverting Equation (50), we conclude our proof,EQUATIONWhen 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 Ġ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), whereEQUATION)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,EQUATIONwhere C(t) > 0 is some constant. Here δ F (t) = Var Xt∼pt ∂ ∂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.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.Require: Initial schedule T = {t i } T i=0 , learning rate γ ∈ (0, 1), score estimate s θ 1: while not converged do 2:for each batch B from data do Fix s θ * and over batch estimate L(t i+1 , t i ), i = 0, . . . , T -1 (Equation ( 13))5:Assign T * ← UpdateSchedule(T , L(t i+1 , t i ))6:Update time locationst i ← γt * i + (1 -γ)t i 7:end for 8: end while 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 Ĵ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],EQUATION)If ∆t is small enough such that,EQUATIONThen, as N → ∞ the following limit exists almost surely,EQUATIONProof. The probability flow ODE update is given throughEQUATIONBy taking the gradient, we have,EQUATION)Since log det(1 + ∆t) = ∆tTrA + O(∆t 2 ) that holds when Tr(A) < 1, our assumption in Equation ( 56) hold, we can apply a Taylor expansionEQUATIONBy taking a gradient, we haveEQUATIONThe 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 Ĵt,N (x).The standard denoising score matching weighting used by Song et al. (2021) EQUATION)Therefore,EQUATIONThe 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,EQUATION)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) 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 tn = (1 -ᾱtn ), where α tn = 1 -β tn and ᾱtn = n i=1 α ti . 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. 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 tn = (1 -ᾱtn ), where α tn = 1 -β tn and ᾱtn = n i=1 α ti . 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 .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. 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 .#  
 
 Appendix:  


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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: limitations: See weaknesses and questions above.
paper_summary: paper_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. 

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

This 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).
strength_weakness: strength_weakness: Including in the Appendix something like Figure 10 but for ImageNet, CIFAR, or CelebA would be interesting.
questions: 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: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: * The presentation of the paper could use a lot of improvement. I found many sections unnecessarily convoluted without conveying much insight.
  - 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.
  - Section 2.3 and Theorem 2.1 do not seem to be related to the actual algorithm.
* 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.
* 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.
questions: 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)?
* Line 136, what is the point of introducing coupled noises in (9) and (10) if they are going to be canceled anyway?
* Line 218, what is the "distance" referred to here?
* 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)?
* 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?



Minor comments:
* In (3), why is there a tilde on $W_t$?
* Line 75, what are "Lebesgue probability measures"?
* In (4), one could split (3) into a prediction and a correction term differently. Why should this specific way of splitting?
limitations: limitations: Limitations are discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: + 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.
+ 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?

[1] Score-Based Generative Modeling through Stochastic Differential Equations ICLR 2021

[2] Flow Matching for Generative Modeling ICLR 2023

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

[1] The Role of ImageNet Classes in Fréchet Inception Distance ICLR 2023

[2] Diffusion Models Beat GANs on Image Synthesis  NeurIPS 2021

+ 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.
questions: questions: Please see Weakness.
limitations: limitations: Please see Weakness.

Comments: title: Does reviewer 8ZVm have any further thought?
Response to authors
Answer to authors
comment: Yes，my concerns have been addressed. Thank you for your informative supplementary experiments. I would keep my positive score.
We would like to thank the reviewer once again for their comments, which have led to improvements in our experiments, including a refined analysis that demonstrates a computational gain with our method in the few-point discretisation regime. We would like to confirm if we have addressed all of the reviewer's comments and questions. Thank you!
I thank both the reviewers and the authors for their contributions.
-AC
We thank the reviewer for assessing our work and further thank them for increasing their score. Thank-you!
We thank the reviewer once more for their assessment of our work. Thank-you!
Thank you for your feedback. We would like to specifically address your concern regarding the procedure for inverting $ t \to \Lambda(t) $.

We are very sorry if our previous explanation was unclear. Let us provide some additional clarifying details:

First, $\Lambda$ is an increasing positive function, as the integrand is positive. Furthermore, according to Theorem 2.1, the integrand of $\Lambda$ is a continuous function, making $\Lambda$ a $C^1$ (continuously differentiable) increasing function. Since $\Lambda$ is monotonic increasing, it is injective, and thus an inverse exists. Moreover, this inverse function is also monotonic increasing and $C^1$.

There are two options to obtain $ \Lambda^{-1} $. The first would be to obtain the set of points $ \\{(t_i, \Lambda(t_i))\\}\_{i=1}^N $ and fit a cubic interpolant to this set of points thus obtaining an approximation of $s=\Lambda(t)$. One would then have to perform some numerical strategy to invert this fitted function. The other option (this is the option that we perform in the paper), is to fit a cubic interpolant to the set of points $\\{ (\Lambda(t_i), t_i) \\}\_{i=1}^N$. Notice that we have swapped the order of $ t_i $ and $ \Lambda(t_i) $. When we fit a cubic interpolant to this set of points we directly obtain an approximation of $t= \Lambda^{-1}(s)$ __without ever having to perform any inversion__.

We hope this additional detail clarifies our procedure and addresses your concern about the method of inverting $\Lambda$.

Thank you again for your feedback, and we hope this explanation provides the clarity needed.
Thank you for the detailed response. I remain unconvinced the whole predictor/corrector setup is necessary. For instance, (13) makes sense to me at first glance (as Fisher divergence), and the predictor/corrector setup only adds more confusion. In addition, the authors acknowledged that there is no real benefit from the predictor-optimized schedule as opposed to the simpler corrector-optimized schedule where there is essentially no predictor. The novelty and significance of Section 3.1 remains unclear to me. My other concern about the justification of the ad-hoc way of inverting $\Lambda(t)$ is also not well addressed. As such, I would like to keep my score unchanged.
Thanks to the authors for their response.
I appreciate the author's thorough response to my questions, which are satisfactorily addressed. I find the work to be meaningful and have accordingly increased my score.