Faster Diffusion Sampling with Randomized Midpoints: Sequential and Parallel

We thank the reviewer for the review and helpful feedback. We respond to your specific concerns and questions below.

Re: Lipschitz assumption on estimated score: We need this assumption in the proofs so that we get the property ``when the position $\hat x$ is not too far off from the true position $x$, the estimated score is also not too far off". For instance, this property is used in Lemma A.2.

Re: proof of Lemma A.2: You are absolutely right and we apologize for the sloppiness in this proof. Our proof is meant to work with our Assumption 2.4 as stated, and not the stronger point-wise approximation. Specifically, what we meant to say in Line 736 is that 1) at the true position $x$, by Assumption 2.4, the expected difference of the estimated score and true score is bounded; and 2) our estimated $\hat x$ is not too far off from the true $x$, and due to lipschitz continuity assumption on both the estimated score $s$ and true score $\nabla q$, both the estimation score and the true score evaluated at the estimated position $\hat x$ is close to their value evaluated at the true position $x$. However, in this part of the proof, we got confused in the writing, and instead of applying Assumption 2.2, 2.3 and 2.4 together, we applied Assumption 2.2 and the point-wise version of Assumption 2.4. We apologize again for this mistake and have uploaded a corrected version of our submission. We mainly changed Lemma A.2 and Lemma B.2, but also parts of other theorems where these Lemmas are used, the main flow of our argument remain unchanged.

Re: background discussion: We added a definition of the log concave function in section C, and will continue to work on improving the clarity of our discussion for the next version of our paper.

Re: optimality of $d$ dependence for KL and TV: A crucial difference between [1] and our work is that [1] mainly analyze the DDPM setting with SDE, while we analyze the DDIM setting with probability flow ODE. Also, to the best of our knowledge, although the iteration complexity in [1] has a linear dependence on $d$, their optimality discussion does not rule out sublinear (or better than $\sqrt{d}$) iteration complexity. Specifically, in their calculation, the KL error induced by discretizing time is of order $\tilde{O}(d h)$, where $h$ is the step size. They mentioned that in the KL error bound, the linear $d$ dependence in this discretization error bound is optimal, however, the dependence on the step size $h$ is not necessarily optimal and [1] theorized that this could be improved to $O(d h^\beta)$ for $\beta \geq 2$, especially when there is a smoothness assumption for the score function. In this case, one can set the step size to be $h = \Theta(d^{-1/\beta})$ and obtain $d^{1/\beta}$ iteration complexity.

Re: iteration complexity in A.10: We are actually running predictor and corrector step for different amount of time. We are running each predictor step for $1/L$ time (specified in Algorithm 3), but we run each corrector step for $T_{corr} = \Theta(L^{-1/2}d^{-1/18})$ time (as specified in Theorem A.10). Since we take corrector step size $h_{corr} = \Theta(d^{-17/36})$, we need $\Theta(d^{5/12})$ iteration complexity to complete one corrector step, which coincides with the iteration complexity to complete one predictor step.

We thank the reviewer for the review and helpful feedback. We respond to your specific concerns and questions below.

Re: Experiments: We have added an experimental section in our updated pdf (Appendix A), comparing the performance of our proposed randomized midpoint technique with the standard DDIM scheduler. Our experiments generate images using the two schedulers on two diffusion models -- one trained on CIFAR-10 images, and the other on CelebAHQ images. We show that, as predicted by our theory, randomized midpoint consistently outperforms DDIM in terms of FID scores, which measure the quality of generated samples relative to the training dataset.

Re: Space Complexity and Scalability: For the sequential setting, the space complexity is no different than the space complexity of standard schedulers such as DDIM. For the parallel setting, a recent work 1 has shown that an algorithm similar to a parallelized version of DDIM performs well in practice, provided a sliding window method is used to maintain samples $x_{t:t+p}$ for window size $p$, to reduce memory usage. We expect that the same approach would work to implement our randomized-midpoint-based parallel sampling algorithm.

I am assuming the authors have not yet had the time to respond to my review. I will reply once they respond.

We thank the reviewer for the review and helpful feedback. We respond to your specific concerns and questions below.

Re: Experiments: We have added an experimental section in our updated pdf (Appendix A), comparing the performance of our proposed randomized midpoint technique with the standard DDIM scheduler. Our experiments generate images using the two schedulers on two diffusion models -- one trained on CIFAR-10 images, and the other on CelebAHQ images. We show that, as predicted by our theory, randomized midpoint consistently outperforms DDIM in terms of FID scores, which measure the quality of generated samples relative to the training dataset.

We thank the reviewer for the review and helpful feedback. We respond to your specific concerns and questions below.

Re: parallel query complexity: We believe that at least within our proof framework, there is a fundamental tradeoff for parallel algorithms between iteration complexity and total query complexity. Intuitively, randomized midpoint method achieves reduced query complexity by reducing the bias of estimation in each step, and this benefit accumulates as number of iteration grows. For instance, a generalized version of our parallel algorithms can achieve $d^{1/2 - \epsilon}$ (for $\epsilon \leq 1/12$) query complexity with $\log^2 d \cdot d^{\Theta(\epsilon)}$ iteration complexity. In our proof, we set $\epsilon = 0$ to achieve $\mathsf{polylog}(d)$ iteration complexity.

Re: smoothness assumption: Typically, when the smoothness assumption is not present, prior works bound the "effective" smoothness by a function of $\sigma_t$ and $d$. However, the existing bounds ([1], [2]) for this effective smoothness are not strong enough to even achieve a sublinear in $d$ bound on the iteration complexity, and we run into similar barriers here. Within the diffusion community, achieving such a sublinear bound is a well-known question, and is orthogonal to our focus here.

Re: predictor-corrector loop: In both our sequential and parallel algorithm, the predictor-corrector loop runs for multiple ($\tilde{O}(1)$) times. (Note that Algorithm 3 for sequential loops and Algorithm 9 for parallel loops looks almost identical.) However, it is true that in the parallel algorithm, we only take $\tilde{O}(1)$ predictor steps in one predictor-corrector loop; while in the sequential algorithm, we take $\tilde{O}(d^{5/12})$ predictor steps. In the parallel algorithm, we are able to take larger step sizes since we update many randomized midpoints in parallel in one predictor step.

[1] Benton et al. 2024, Nearly-Linear Convergence Bounds for Diffusion Models via Stochastic Localization, ICLR2024