Score-Optimal Diffusion Schedules

We thank the reviewer for their kind comments on the strengths of our paper. We believe that our schedule tuning algorithm could serve as a strong default for diffusion models and might be applicable more broadly, though a more extensive exploration is difficult to undertake in an initial paper. We hope to use this work as a starting point for further investigation into general schedule tuning in generative models.

In the context of diffusion models, we welcome the reviewer's suggestions for more comprehensive numerical validation. To this end, we have presented preliminary summary statistics for sFID. We are appreciative of the reviewers suggestions for additional metrics, and agree that their addition to a potential camera ready article would increase the strength of the paper.

We believe that our methodology could be applied to flow-matching based models, but this would also require mild adjustment to the analysis we present. We hope our present work could be foundational to a potential future study.

We agree that validation of LDM with our method is a worthwhile and meaningful venture. To this end, we have implemented our method in the HuggingFace LDM library, but due to time constraints do not currently have FID to report. In a potential camera ready article, we would additionally include FID and sFID comparisons on optimised and standard schedules for Latent Diffusion Models.

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.

We thank the reviewer for their suggestion. We would like some assistance from the reviewer, we would like to confirm that sFID refers to the metric computed in:

Nash C, Menick J, Dieleman S, Battaglia PW. Generating images with sparse representations. arXiv preprint arXiv:2103.03841. 2021 Mar 5.

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

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.

We agree with the reviewer that this mathematical content is new to the diffusion literature. In a potential camera-ready version, we would include a brief introduction to metric tensors and charts commonly used in differential geometry, along with their relation to the work we present here. We appreciate the reviewer’s suggestion, which will help us present our work more clearly.

We are grateful for the reviewers for their constructive feedback. We will address each point raised individually below:

• The whole setup of predictor/corrector seems unnecessary to me.

See general rebuttal response.

• Section 2.3 and Theorem 2.1 do not seem to be related to the actual algorithm.

Section 2.3 and Theorem 2.1 are necessary links that take us from the Fisher Divergence at two time points t, t’ to a metric that can be used to derive the optimal schedule. In Section 3.1 when we come to find the costs associated with diffusion paths, we require an infinitesimal notion of cost that can be integrated along the diffusion path. The Fisher Divergence as stated in Section 2.2 is unsuitable for this purpose because it is evaluated between two times t and t’. In Section 2.3 and Theorem 2.1 we derive this corresponding infinitesimal cost. Our algorithm relies entirely on the approximation of this infinitesimal cost along the diffusion path. We appreciate the reviewer highlighting this lack of clarity in the significance of Section 2.3 and Theorem 2.1. In a revised version we will add a clarifying sentence before Theorem 2.1 to emphasise its relevance and importance.

• 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

Our algorithm is grounded in the theoretical insights from Section 3. The relevance of Section 3.1 lies in demonstrating how schedule optimization in diffusion models can be reframed as finding an energy-minimising path with an appropriate metric. The key connection between differential geometry and diffusion models is established in Theorem 2.1, where we derive the Reimannian metric compatible with the objectives of diffusion models, and demonstrate how to approximate the energy using the Fisher divergence. The material from differential geometry included in this section is meant to ensure completeness. To improve the clarity of this section, we will add a linking sentence that explicitly states how the local metric tensor from Theorem 2.1 enables us to use differential geometry to study diffusion schedule design.

• The computation cost could be quite high..

See general rebuttal response.

Questions:

• Line 95

Thank you for pointing out this typo.

• Line 136

The introduction of coupled noises in Section 2.2 is a direct consequence of performing a perturbation analysis on the Langevin dynamics trajectories. Specifically, we are interested in understanding how these trajectories evolve over a small time interval when the score function is slightly perturbed.

By coupling the noise terms in equations (9) and (10), we ensure that any observed differences between these trajectories are solely due to the perturbation and not due to random fluctuations. The noise cancels out naturally because we have isolated the perturbation of the score; this cancellation is a result of the computational approach rather than a manufactured effect. Without coupling the noises, the comparison would be less informative, as it would be unclear whether differences were due to the perturbation or simply random noise.

• Line 218

Here "distance" is intended to mean "length along the path," as referenced in the preceding lines 202-207. We will update the text to more clearly state "length along the path".

• Line 235

We have not claimed that we know $\Delta t$ is sufficiently small, but rather if it is small, then our approximation is valid. One way to assess if a grid is fine enough is by computing the statistical length $\Lambda$, which is independent of the schedule, across a range of grid sizes. When this value stabilises for a particular $\Delta t$, it indicates that the desired integral approximation of Equation 21 has been achieved.

• Simple cubic interpolation in step 3

We also remark that be apprimae $\Lambda^{-1}(t)$ directly by interpolating with knots $\{(\Lambda(t_i), t_i)\}$. This avoids needing to approximate $\Lambda(t)$ and numerically inverting. Since we avoid numerical inversion methods, the choice of interpolant is not overly critical. We use cubic interpolation as an example, but any interpolation method respecting monotonicity (e.g. linear) can also be used. We will add this comment after Algorithm 1.

• In (3)

In the context of the forward/backward SDEs in a diffusion model, the Brownian increments for the forward and backward SDEs are independent of one another. We use a tilde to emphasise that the Brownian increment for the forward diffusion ($\tilde{W}$) is different from the Brownian increment of the backward diffusion ($W$).

• Line 75

Lebesgue probability measures are those that have a density function with respect to the Lebesgue measure. We've included standard references to make this clearer for readers who may be unfamiliar with probability theory and thank the reviewer for improving the accessibility of our work.

• In (4)

This specific way of splitting is chosen to identify the unique ODE that transports samples from one density $p_{t}$ along the diffusion path to $p_{t'}$. This approach was originally identified and used in the seminal work of Yang Song [reference]. We will add a reference in the text to where this splitting is derived in Yang Song's paper.

We thank the reviewer for their comments. We agree that more visual comparisons of the different schedules would be illustrative. To this end, we will include standard diffusion progression plots, comparing the image diffusion path for CIFAR and CelebA on different schedules. In a potential camera-ready version, we would include one such plot in the main body of the text and all comparisons in the appendix. We greatly appreciate the reviewer’s comment, which has helped improve the presentation of our schedule optimisation.

It is unclear to me if in Eq 57, $J_t$ requires computing the Hessian of $\log p_t$: if so, wouldn't that be an expensive quantity to compute in order to obtain the optimal schedule?

We thank the reviewer for highlighting this important computational point. In our implementation, we have formulated an estimator (Proposition B.1) that completely avoids computing the full Hessian. This proposition is currently in the appendix, but we state it here for completeness:

Proposition B.1: Let $F_{t,t'}$ be the predictor map given by the forward Euler discretisation (8) of the probability flow ODE. For $N \in \mathbb{N}$, let $\hat{J}_{t,N}(x)$ be the Jacobian of the Hutchinson trace estimator (Hutchinson, 1989) for $\nabla(\nabla \log p_t(x))$ at $x \in X$ and $t \in [0, 1]$:

$$\hat{J}_{t,N}(x) = \frac{1}{N} \sum_n \nabla(v_n^T J_t(x) v_n),$$

If $\Delta t$ is small enough such that:

$$\Delta t \cdot \text{Tr} \left( f(t)I - 2g(t)\nabla \log p_t(x) \right) < 1,$$

then, as $N \to \infty$, the following limit exists almost surely:

$$\nabla \log \det \nabla F_{t'}(x) = -\Delta t g(t)^2 \lim_{N \to \infty} \hat{J}_{t,N}(x) + O(\Delta t^2).$$

Here, we only require the trace of the Hessian, so we need to compute only the diagonal entries, ensuring that the memory cost remains linear with respect to the parameters. Additionally, we have utilised a Hutchinson trace estimator for this term, which can be efficiently implemented in standard deep learning libraries.

In our provided code, we have explicitly implemented this in PyTorch and have found that this computation is scalable to image datasets. This approach has enabled us to compute our predictor-optimised schedules, something infeasible if we had needed to compute the full Hessian. Our estimator performs well, producing schedules that achieve good FID scores. This was a subtle computational challenge that was non-trivial to overcome in our work.

In a potential camera-ready version, we will clarify this approximation and emphasise that the computational resources required are significantly less than those needed for computing the full Hessian. We will also make clear in the main text that we have a scalable estimator for this term and do not require computing the Hessian matrix. We thank the reviewer for bringing this point to our attention and helping improve the quality of our paper.

We thank the reviewer for their review and suggested improvements. We agree that providing a simple error plot for cases where the true density is known would be informative. In our diffusion model, we do not have the ability to query the density directly; however, we do predict the score. In Figure 4 of our paper, we make this comparison for our bimodal example of the scores at t=0. It can be noted here that the linear schedule devolves into learning a function dissimilar to the expected shape of the score for a bimodal distribution. To further illustrate this point, we have created a plot that evolves over training, showing the error of the score at time t=0, corresponding to our data distribution. We observe that with our schedule optimisation, we achieve a lower error than with a fixed linear schedule over the course of training. Additionally, we have computed the likelihood of the generated samples in this case where the true density is known. Similarly, we observe that during training, our optimised schedule generates datasets with higher likelihood than the fixed schedule case. We thank the reviewer for their suggestion, which has improved the quality of our paper, and we will include these improvements in a potential camera-ready submission.

Q: It seems the proposed method do not accelerate the backward sampling process since the updated schedule has the same number of timesteps. I wonder if the author have examined the computational efficiency of their Algorithm 2? Since when it comes to SGM, the running time is one of the major concerns.

We thank the reviewer for their question. We do observe in our simple 1D experiment in Figure 1 that an optimised schedule can achieve better results with fewer discretisation points, improving runtime. To better convey this potential use case, we have added follow-up experiments based on the reviewer's previous suggestion to see if we can see a gain in FID with fewer data points for image generation.

Additionally, we have conducted a refinement test to assess the computational benefit of optimising the schedule. For this, we took the rho3 schedule for CIFAR10 and computed the FID score for 18, 20, 30, 50, and 100 discretisation points. In this range, at 18 discretisation points, the rho3 schedule achieves an FID of 5.47, while at 100 discretisation points, it achieves an FID of 2.05. In comparison, the optimised schedule achieves an FID of 1.99 at 18 discretisation points. We further observe the same trend in the sFID, another image performance metric.

In this experiment, we can see that optimising the schedule can lead to improved FID and sFID performance with fewer discretisation points. Furthermore, we compared all schedules with only 10 discretisation points. In this case, the optimised schedule outperforms all others in terms of FID and sFID, demonstrating its superior performance in the few discretisation point regime. We thank the reviewer for their suggestion to perform this interesting computational experiment. In a potential camera-ready submission, we would include this study in the main manuscript.