Accelerating Diffusion Models with Parallel Sampling: Inference at Sub-Linear Time Complexity

Summary

This work improves the sample complexity of diffusion models by using the parallel sampling technique and achieves $\tilde{O}(\text{poly} \log d)$ results for reverse SDE and PFODE settings at the same time. To achieve these results, this work provides a general version of Girsanov’s theorem to deal with the additional dependence introduced by the Picard iteration.

Strengths

1. The general version of Girsanov’s theorem is novel and suitable for the parallel sampling setting, which would raise independent interest.

Weaknesses

1. For the PIADM type algorithm, the framework is similar to Algorithm 2 of [1], which is used to analyze ULMC. It seems that this work replaces the ULMC process with the denoised process. Then, this work relaxes the LSI assumption of data by using a similar technique to [2]. It would be helpful to discuss the technique challenge except for the general Girsanov’s theorem. 2. For the PFODE setting, the PFODE-ULMC type algorithm is introduced and analyzed by [3]. Using this technique, [3] achieve a similar $\sqrt{d}$ improvement. It would be better to discuss the technique novelty compared to [3]. [1] Anari, N., Chewi, S., & Vuong, T. D. (2024). Fast parallel sampling under isoperimetry. arXiv preprint arXiv:2401.09016. [2] Benton, J., Bortoli, V. D., Doucet, A., & Deligiannidis, G. (2024). Nearly d-linear convergence bounds for diffusion models via stochastic localization. [3] Chen, S., Chewi, S., Lee, H., Li, Y., Lu, J., & Salim, A. (2024). The probability flow ode is provably fast. Advances in Neural Information Processing Systems, 36.

Questions

Please see the Weakness part.

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors adequately addressed the limitations and societal impact.

Summary

In this article, the authors propose parallel methods for diffusion models, achieving a poly-logarithmic complexity for the first time under commonly used assumptions. By applying existing Picard Iterations to both the SDE and probability flow ODE of diffusion models, the backward process can be efficiently solved within a period of $O(1)$ length, effectively transferring sequential computations of $O(d)$ into parallelizable iterations with a depth of $O(\log d)$. However, the author's analysis is not trivial. They employ a more intricate mathematical framework for stochastic processes and general forms of Girsanov's theorem to bound the error in each block.

Strengths

- I believe this paper is technically solid and addresses an important open question regarding parallel methods for diffusion models. - This paper is well-written, and the proof is clear and comprehensive. - I think that the sketches and ideas presented in the proof will be useful for future research, although it is similar to the parallel method used in sampling.

Weaknesses

- The authors apply Picard iterations to diffusion models without providing a comparison to other parallel methods, such as randomized midpoint methods. Including such comparisons would enhance the evaluation of their approach. - Please compare the total query complexity of the proposed method to that of sequential methods. Minor: - Line 85: Please add a reference for the Picard iteration. - Line 89: Please include a reference for the claim of exponentially fast convergence. - Last equation on page 4: Specify the distribution that is being assumed for this equation. - Lemma A.6: Ensure that $_F$ is included. - Provide an explanation regarding the different parameter choices for the predictor and corrector in Theorem 3.5.

Questions

See Weaknesses

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Summary

Denoising diffusion models generate samples from a complex target distribution by solving stochastic/ordinary differential initialized at the Gaussian noise. Solving these differential equations is typically done by simulating diffusion solutions. This is inherently sequential since simulating the solutions at time t_i requires the solution at time t_{i-1}. An alternative approach to solve these diffusions is to use Picard Iterations (PI), which involve constructing a contraction mapping where the solution is the fixed point. Notably, a PI iteration updates the entire path, making parallelization amenable. This paper introduces how to sample from a diffusion model incorporating PI into the pipeline, allowing practitioners to utilize their parallel processes efficiently. They show that a PI can speed up the run time in O(poly(log d)) in dimension compared to O(d) runtime of serial implementations of diffusion models.

Strengths

I think the topic is well-motivated, and the authors do a great job of presenting the ideas in the paper. The idea feels quite natural and the algorithms developed have practical utility, which is uncommon for a theory paper. The theory is also strong, I can tell a lot of care has gone into the proofs.

Weaknesses

I understand it is a theory paper, but I was disappointed to see no experiments, or empirical verification of the theory. The main result is the poly(log(d)) run-time, which would have been nice to see how this hold in practice. The algorithms presented in this paper seem fairly easy to implement. Experimenting with even a Gaussian toy model where analytical solutions are known would have made this paper more convincing to practitioners.

Questions

Can you provide some intuition on how the constants scale, and how long does it take for the asymptotic to kick in? Can you provide some intuition on how to tune hyperparameters, including a number of blocks, discretization steps per block, and depth?

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors adequately addressed the limitations.

Summary

In this paper, the authors propose to analyse the error in parallel sampling for diffusion models. This represents the first theoretical analysis of parallel sampling for diffusion models. The initial methodology was proposed in ParaDiGMS [1]. In this paper, the authors propose some incremental improvements on the samplers used in ParaDiGMS. The main contribution resides in the control of the error in parallel sampling. Notably, instead of deriving a global time complexity which depends on $O(d)$ they derive rates that are of order $O(\mathrm{poly log}(d))$. This improvement comes from the parallel sampling nature of the scheme they analyse. The authors not only deal with the SDE (Stochastic Differential Equation) sampler but also with ODE (Ordinary Differential Equation) samplers. Parallel sampling for ULA (Unadjusted Langevin Algorithm) was already theoretically analysed in [2] but to the best of my knowledge this is the first rigorous work on the complexity of parallel sampling for diffusion models. [1] Shih et al. (2024) -- Parallel Sampling of diffusion models [2] Anari et al. (2024) -- Fast parallel sampling under isoperimetry

Strengths

* This paper is very well-written. I found the flow of the paper to be quite easy to follow. * I have read the proofs in the case of the SDE samplers and they are correct. I also commend the authors on adding a proof sketch in the main paper. This makes the whole paper much easier to read and the proofs become more intuitive. * The rates obtained by the authors are quite compelling. In particular the $O(\mathrm{poly log}(d))$. time complexity echoes the rates obtained in [1]. * The authors not only analyse the SDE sampler but also the ODE sampler which is usually more complicated than the SDE one. I am pleasantly surprised by the gains they obtain in terms of memory complexity. [1] Anari et al. (2024) -- Fast parallel sampling under isoperimetry

Weaknesses

* My main complain is that the results of Theorem 3.3 and Theorem 3.5 are given in terms of a given level of optimisation $\delta^$. This gives a specific choice of constant for the discretisation stepsize $\epsilon$, the time $T$, the block size $h$, the number of blocks $N$, the number of parallel iterations $K$ and $M$ the number of steps in each block. However, it would also be interesting to give a bound that depends on those quantities. What I am thinking about is that then one could optimize the bound under time complexity and space complexity constraints (for example with a requirement that $d M \leq M_0$ where $M_0$ is the available memory of the system). Such a bound would be more interesting to gain insights on what we really gain by moving to parallel sampling. * While I'm interested in the ODE framework I am a bit confused by the explanation on how the authors have managed to reduce the $O(d^2)$ memory complexity to $O(d^{3/2})$. The explanation of l.281 "The reduction of space complexity by the probability flow ODE implementation is intuitively owing to the fact that the probability flow ODE process is a deterministic process in time rather than a stochastic process as in the SDE implementation." is not very intuitive to me. It would be nice to point out where the improvement comes from in the proof sketch (3.2.4). * Regarding the ODE sampler, I am not sure that the bounds with SDE are directly comparable since the proof techniques use different metrics (even though the KL can be related to the TV using Pinsker's inequality). To get comparable rates it would have been better to provide a unified framework in which both the SDE and the ODE samplers are analysed under the same techniques (and assumptions). * Assumption 3.4 is extremely strong and not satisfied in most applications. It seems that compared to [1] which analyses ODE samplers, Assumption 3.4 is not required there. Similarly for Assumption 3.3. These assumptions prevent the model to explode near the data which is something that is observed in practice (and would happen theoretically). I would have liked a more detailed comparisons between the differences of assumptions required for the analysis of sequential and parallel samplers of diffusion models. [1] Chen et al. (2023) -- "Restoration-Degradation Beyond Linear Diffusions: A Non-Asymptotic Analysis For DDIM-Type Samplers"

Questions

* "We propose new parallel inference algorithms for diffusion models using parallel sampling" (in the TLDR). It seems that one of the claim of the authors is that they also propose a new sampler. Compared to ParaDiGMS [1], the methodological innovation seems pretty incremental (exponential integrator, shrinking step size, early stopping, see Section 1.1 "Contributions"). Is there something more? If there is a key methodological contribution, this should be assessed rigorously against ParaDiGMS in experimental settings. I realize that this is a theoretical paper so I would suggest to tame down the claims of methodological novelty and instead focus on the theoretical analysis. * I would be interested to see if the improvements obtained by the authors can be translated to the manifold hypothesis setting of [2, 3]. In particular, can the $O(\mathrm{poly log}(d))$ rate be kept for this setting and/or improved to replace the dimension of the space by the dimension of the implicit manifold? * See "Weaknesses" for other questions and concerns. [1] Shih et al. (2024) -- Parallel Sampling of diffusion models [2] De Bortoli (2022) -- Convergence of denoising diffusion models under the manifold hypothesis [3] Yang et al. (2023) -- The Convergence of Variance Exploding Diffusion Models under the Manifold Hypothesis

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

In the paper the limitations are discussed in Section 4, i.e. "Although we anticipate implementing diffusion models in parallel may introduce engineering challenges, e.g. scalability, hardware compatibility, memory bandwidth, etc., we believe that our theoretical contributions lay a solid foundation that not only supports but also motivates the empirical development of parallel inference algorithms for diffusion models since advancements continue in GPU power and memory efficiency." I think there are other limitations that need to be discussed regarding the current work and the theoretical assumptions that the authors needed to make in order to derive their results. While I don't think these limitations hinder the results of the paper they should be clearly laid out.