Summary
The paper proposes a new solver for diffusion SDEs, termed SA-Solver, combining the ideas of a predictor-corrector scheme and the stochastic Adam solvers. The predictor and corrector utilize the Lagrange polynomials for extrapolation to lower the approximation error at future time stamps. Experimentally, SA -Solver improves over previous ODE and SDE across a wide range of image generation benchmarks.
Strengths
- The considered problem of accelerating SDE solvers is of great practical value in the field, as SDE solver often delivers better sample quality than ODE solvers, but are hindered by their slow sampling speed.
- Adapting the idea in Stochastic Adams is interesting, offering higher-order convergence with Lagrange polynomial extrapolation.
- The fusion of the PC sampler and Lagrange polynomial yields superior empirical results on diffusion SDE, tested on a range of dataset resolutions from 32x32 to 256x256. It's nice to see the methods scale to different resolutions and architectures.
Weaknesses
- Throughout the paper, the authors mention several times that empirically the quality of data generated by SDE has a better upper limit. I think some theoretical analysis in the main text to make the paper more self-contained. A concurrent [1] provides some theoretical arguments and comparisons between SDE and ODE, showing that the stochasticity in SDE can reduce overall sample errors. I think it would be beneficial to discuss them to better motivate the idea of the paper.
- The authors didn't provide detailed discussions of the benefit of introducing $\tau(t)$ (at the end the author set $\tau(t)=1$ in experiments). Since the authors observe the benefit of larger $\tau$ on larger NFE, why not use a varying $\tau$ in Fig.2 based on NFE? In addition, the benefit of a larger $\tau$ on larger NFE seems predictable based on the theory in [1].
- The authors did some comparisons of epsilon-prediction and data prediction models. It seems that EDM [2] uses interpolation between these two by pre-conditioning. Does it offer better results?
- Unlike the deterministic path in ODE, the randomness in SDE can change the sample trajectories quite a bit. For example, the trajectories of generating a dog can stray into generating a cat. I wonder if this is an issue for using Lagrange polynomials in SDE, since the preceding predictions could be misleading. If not, could the author provide some intuitions? In addition, the authors mention that "predictor step 2 and corrector step 1 is the most stable setup". such small steps could be caused by the issue I raise.
- The paper separately gives the convergence order of predictor and corrector. Is it possible to give the theoretical order after combing the predictor and corrector?
- I think the author didn't use the "EDM VE" model for CIFAR-10. In Table 2 of EDM [2], the unconditional VE (in config F, EDM VE) can achieve an FID of 1.98 using only 35 NFE, as opposed to the >3 FID in Fig 2 in the current paper. So I guess the author use their VE model in config A (baseline VE). Could the author compare different methods using the VP model (config A), as in [1] and [2], because the VP model is more commonly used in practice and offers better results?
- The stochastic sampler in the concurrent work [1] obtains an FID ~ 1.8 using less than 100 NFE on ImageNet-64 with Pixel DPM, as opposed to the > 2.1 FID in Fig 2 of the current paper. I wonder if one could combine the multi-step idea in this paper with [1] to obtain further improvements.
- Could the author also provide quantitative results on Stable Diffusion experiments, to better showcase the advantage of SA-Solver? Current visualized images do not seem compelling.
### Minors
- line 225, DPM-solver++ and UniPC are *ODE* solvers which would not naturally be special cases of the proposed *SDE* solver.
[1] Restart Sampling for Improving Generative Processes, Xu et al, https://arxiv.org/abs/2306.14878
[2] Elucidating the Design Space of Diffusion-Based Generative Models, Karras et al, https://arxiv.org/abs/2206.00364
# Post rebuttal
Thanks the authors for the rebuttal. After another pass of the paper, it occurred to me that that paper is a direct application of the ideas (variation-of-parameter, or exponential integrator, and Langraian polynomial for approximation) in [1] to SDE samplers. The same formula of the exponential integrated version of SDE (Eq.5) has been proposed in prior works [1] (see their eq.17, arXiv version 1). In addition, [1] also uses a Lagrange polynomial to extrapolate. Due to it's straightforward extension of prior works, I will keep my score unchanged.
[1] Zhang, Qinsheng, and Yongxin Chen. "Fast sampling of diffusion models with exponential integrator." arXiv preprint arXiv:2204.13902 (2022).
Questions
I copy-and-paste some questions from above:
- Since the authors observe the benefit of larger $\tau$ on larger NFE, why not use a varying $\tau$ in Fig.2 based on NFE?
- I wonder if this is an issue for using Lagrange polynomials in SDE, since the preceding predictions could be misleading. If not, could the author provide some intuitions? In addition, the authors mention that "predictor step 2 and corrector step 1 is the most stable setup". such small steps could be caused by the issue I raise.
- Is it possible to give the theoretical order after combing the predictor and corrector?
- Could the author compare different methods using the VP model (config A), as in [1] and [2], because the VP model is more commonly used in practice and offers better results?
- Could the author also provide quantitative results on Stable Diffusion experiments, to better showcase the advantage of SA-Solver?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.