Score-Optimal Diffusion Schedules

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.

Strengths

+ The findings of the paper are interesting. Current schedule designs are heuristic, and this paper provides a theoretical analysis, which is of interest to the community. + The proposed online data-dependent adaptive discretization schedule tuning algorithm is simple and effective. + 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. + The authors provide clear visual comparisons that show the distribution of optimized times in detail.

Weaknesses

+ The experiments in this paper are somewhat weak and lack validation for the LDM model. In fact, the LDM model is widely used, so it is important to validate the proposed method for LDM. + 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

Please see Weakness.

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Please see Weakness.

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.

Strengths

* The problem studied is important for diffusion generative modeling as the manual search for a good sampling schedule can be time-consuming. * The proposed method is simple, despite the rather convolutated presentation. * Empirical evaluation shows that the proposed scheme can have on-par performance with hand-tuned schedules.

Weaknesses

* The presentation of the paper could use a lot of improvement. I found many sections unnecessarily convoluted without conveying much insight. - 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

* 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?

Rating

4

Confidence

3

Soundness

2

Presentation

1

Contribution

2

Limitations

Limitations are discussed.

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

Strengths

The paper is clearly written and derives a principled way to derive an optimal schedule for the popular "predictor-corrector" SDE, of which "annealed Langevin dynamics" and a "reverse diffusion process" are a special case. This is of interest to the sampling and diffusion models communities. The appendix has nice, additional experiments assessing the optimal schedule and 1D and image datasets. The authors are also clear about the limitations of their approach, namely that it assumes perfect estimation of the scores.

Weaknesses

Including in the Appendix something like Figure 10 but for ImageNet, CIFAR, or CelebA would be interesting.

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?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.

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.

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

See weaknesses and questions above.

Answer to authors

Thanks to the authors for their response.

Response to authors

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.