Restart Sampling for Improving Generative Processes

Summary

This paper proposes analyzes SDE and ODE-based samplers for diffusion models. Based on the analysis, this paper introduces a new solver, Restart, for sampling from diffusion models. The effectiveness of Restart is demonstrated on various unconditional and conditional generation tasks.

Strengths

- The paper is well-written. - The proposed method shows better performance than EDM at moderate NFE regions.

Weaknesses

I am willing to raise the score by 1 or 2 points if the authors address my concerns satisfactorily. **Weakness 1 : Ambiguity regarding Theorems 1 and 2.** - For Theorem 1, if we set $[t_{\min},t_{\max}] = [0,T]$, the terms for contracted errors $TV(p_T^{ODE_\theta},p_T)$ and $TV(p_T^{SDE_\theta},p_T)$ vanish because $p_T^{ODE_\theta}$, $p_T^{SDE_\theta}$, and $p_T$ are all identically Gaussian distributions. Then, we end up with terms depending on $\delta$, $\epsilon_{approx}$, and $t_{\max} - t_{\min}$ only, so Theorem 1 does not provide any insight into how ODE and SDE have distinct "winning regions", as illustrated in Figure 1 (b). The proof and claim for Theorem 1 should be reformulated such that even with $[t_{\min},t_{\max}] = [0,T]$, Theorem 1 explains how ODE and SDE have winning regions. - Likewise, I think Theorem 2 also should be proven for the entire interval $[0,T]$, so we can directly compare the errors for SDEs, ODEs, and Restart. **Weakness 2 : (Possibly) weak performance on the small NFE regime.** - How does Restart perform in the small NFE regime (NFE $\leq 30$)? In Figure 3, the figure cuts off just before Restart and ODE intersect in the small NFE regime. This seems to contradict the claim that Restart combines the best of both ODE and SDE. Moreover, given the large size of SOTA diffusion models, it is crucial that diffusion samplers work well in the small NFE regime as well. - How does Restart compare to recent fast samplers such as [1], [2], [3] in the small NFE regime? [1] DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps, NeurIPS, 2022. [2] Fast Sampling of Diffusion Models with Exponential Integrator, ICLR, 2023. [3] Denoising MCMC for Accelerating Diffusion-Based Generative Models, ICML, 2023.

Questions

**Question 1** : How does Restart perform on FFHQ-1024 image generation (score function in https://github.com/yang-song/score_sde) compared to Improved SDE or ODE (Heun in EDM)?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

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.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Discussed in Section 6.

Summary

ODE-based samplers plateau in performance while SDE-based samplers deliver higher sample quality. The paper attributes this difference to discretization errors and accumulated errors. Based on these, the authors propose a sampling algorithm called Restart which alternates between the forward diffusion process and backward ODE.

Strengths

1. The authors provide a theoretical explanation of the phenomenon that ODE samplers outperform SDE samplers in the small NFE regime but fall short in the large NFE regime. 2. The experimental results on image generation tasks validate the effectiveness of the method.

Weaknesses

The proposed method relies on several hyperparameters (e.g. $S_{noise}, N_{restart}, i, K_i, t_{min, i}, t_{max, i}$), and the hyperparameters differ in different tasks. It would be hard to effectively tune these parameters in real application.

Questions

Since the derivation of the Restart algorithm is motivated by the theoretical analysis of sampling error, is there a way to choose those hyperparameters based on the error analysis?

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.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The author addressed the limitations of their work.

Summary

The papers propose a new sampling method for diffusion models, termed Restart Sampling. The authors first theoretically analyze the error propagation in diffusion models for stochastic and deterministic samplers under Wasserstein-1 distance and show that ODE samplers have a lower-discretization error but SDE samplers contract the initial distribution error as we run more steps. This agrees with the intuition and the experimental findings that support that ODE samplers are better for low NFEs but their performance flattens for more NFEs. Based on this analysis, the authors propose a method that tries to achieve the best of both worlds: it contracts the initial error with more steps and achieves the same discretization error as the ODE samplers. The implementation of the new method is very straightforward: one runs the ODE sampler and every K steps reverts back to some prior diffusion time using the forward model.

Strengths

The authors study a relevant problem in diffusion models. The proposed solution is simple, effective and novel. The theoretical results motivate the method and show clearly the differences between the SDE and the ODE samplers. The presentation of the paper is excellent. The authors show many experimental results, starting from toy models and going all the way to state-of-the-art text-to-image diffusion models. I think the paper and the method is of interest to the community and the audience of NeurIPS.

Weaknesses

The error propagation of diffusion models has been studied before. The results that I am aware of are from the papers "Sampling is as easy as learning the score", "Restoration-degradation beyond linear diffusions: A non-asymptotic analysis for DDIM-type samplers" and "The probability flow ODE is provably fast". The first studies the error propagation of the SDE sampling method and the latter two the propagation of errors for deterministic samplers. It would be beneficial to compare with these works, highlight potential differences in the approach and the final results, etc. Also, apart from the Stochastic and the ODE samplers, there is a whole family of samplers that satisfy the same Fokker-Planck equations and hence give the same marginals, e.g. see the work "Fast Sampling of Diffusion Models with Exponential Integrator" and also some of the samplers used in the "Elucidating the Design Space of Diffusion-Based Generative Models" (EDM) paper. It would be interesting to compare theoretically and experimentally to these samplers. Another concern I have is that the evaluation is only as thorough as it should have been and it is only done for relative high NFEs. Since evaluating the performance of a trained model is relatively easy, I would expect a more thorough benchmarking. If the performance of Restart sampler breaks for low NFEs, it is useful to know it and acknowledge it in the paper.

Questions

I think it would be useful to know: * how the performance is for lower NFEs, e.g. 20 sampling steps? * the behavior of the sampler across more datasets and NFEs. For example, it would be useful to include more comparisons with the EDM paper (and the references therein) for different datasets. * what are the differences in the theoretical analysis compared to prior results known for error-propagation in diffusion models.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed 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.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The authors adequately addressed the limitations of their work.

Summary

By analyzing of the trade-off between good sample quality and sampling time of both ODE and SDE-based generative models, a restart sampling strategy is proposed by this paper to combine the advantages of ODE and SDE sampling methods. The author proves two theorems that estimate the upper bound on the total error measured by the Wasserstein distance between generated and data distributions of ODE, SDE, and restart sampling methods respectively. It is illustrated that the total error can be decomposed into two parts: additional sampling error generated by discretization error and contracted error generated by the accumulated total error from previous sampling steps. Moreover, it is proved that ODE-based samplers have smaller additional sampling errors and SDE-based samplers have smaller contracted errors. So Comparing the three upper bounds on the total error, it can be proved theoretically that restart sampling yields a smaller total error because its additional sampling error and contracted error are both small. Finally, a range of experiments are done by authors which shows empirically that: 1. The total error of the restart sampler is indeed smaller than that of others. 2. Restart sampler surpasses previous SDE and ODE samplers in both speed and accuracy. 3. Restart sampler better balances text-image alignment/visual quality versus diversity than previous samplers.

Strengths

1. This paper is written with meaningful motivation and a clear structure. The restart sampling method proposed by this paper is innovative, simple, and effective. 2. The reasonableness and effectiveness of the resampling method are proved both theoretically and experimentally. 3. The upper bounds on the total error of the three sampling methods give us an intuitive understanding of the advantages and disadvantages of the three sampling methods.

Weaknesses

1. The effectiveness of restart sampling method on high-resolution image synthesis is not confirmed. For example, a comparison of sampling speed and accuracy on ImageNet 128*128, 512*512 should be added. 2. The paper experiments on the sensitivity analysis of the number of restart iterations K, but there is no experiment on the sensitivity of another hyperparameter: the position and length of the restart interval. 3. As pointed out in the paper, the contracted error further diminishes exponentially with the number of repetitions K though the additional error increases linearly with K. Figure 4 illustrates this trade-off phenomenon, too. So it may be hard to find a suitable K to make the restart algorithm work for different datasets or tasks globally, making it difficult to apply.

Questions

1. Although the sampling step decreases by applying the restart strategy, the repetiness of ODE solver in the restart interval seems time-consuming. Moreover, does NFE computed in the restart interval added to the total NFE in Figure 3? 2. Why does the total error remain the same using the ODE sampler when NFE>20? What are the results of three samplers with NFE<20? Is it consistent with the trend in Figure 1(b)？ 3. Why not compare to DPM-solver? Moreover, why not apply restart to DPM-solver to see if there are consistent benefits?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Summary

The paper proposed a sampler that balances sampling speed and quality by adding noises and restarting the process. They provide theoretical analysis to show a better upper bound of this method compared to original ODE and SDE samplers. Experiments are done to verify their claims.

Strengths

1. Authors identify the main cause of different performances of SDE and ODE in different regimes. And by taking advantage of the contraction by adding noises, they balance the speed and quality of the sampler. 2. The theoretical analysis is clear and well-written. 3. The experiments are thorough with a good explanation of the choices of hyperparameters. 4. The experiments show good results for the proposed method.

Weaknesses

See questions.

Questions

I wonder what the wall clock time looks like between restart, sde, and ode.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

yes

Updated Score

I have raised the score by 1 point. I am not convinced by the result on FFHQ-1024, because I know it is possible to cherry-pick good samples on this dataset. I recommend the authors provide numerical results, e.g., FID, in the revised paper.

Respone to Rebuttal

Thank you for your rebuttal. Please incorporate these discussions in the camera-ready version of your work, if possible. I will increase my score to 7.

Thanks

Thanks for the rebuttal. I like the direct comparison with dpm-solver. I suggest the authors carefully update the paper in the final revision to make clear the points raised in the review. I have increased my score.