Diffusion Models are Secretly Exchangeable: Parallelizing DDPMs via Auto Speculation

Thank you for your thoughtful review and for taking the time to engage with our work. We hope to address your concerns below:

Comparing ASD with Prior Works

Prior works based on Picard Iterations can only produce approximate samples from the DDPM output distribution. In particular, these approaches need to trade-off sample quality against parallel speedup by tuning the error tolerance hyperparameter of the Picard Iteration. On the contrary, ASD is an error-free parallelization scheme that always produces exact samples from the DDPM output distribution for any choice of hyperparameters (as was also noted by the reviewer). Due to this fundamental difference, we believe the two approaches cannot be compared on equal footing.

Our approach of parallelizing DDPMs via autospeculation represents a novel and fundamentally different "axis of improvement" compared to Picard iterations. In particular, the two approaches are orthogonal to each other and can potentially be combined together for even greater parallel speedups (e.g., by using ASD to sample from each level of the Picard Iteration). This integration presents an exciting direction for future research.

Given these considerations, we believe that the original DDPM sampling method is the most appropriate and fair benchmark for evaluating ASD. We highlight that the primary objective of our experiments was to empirically validate our theoretically guaranteed speedups, while demonstrating a practical application of the hidden exchangeability property of DDPMs (which we believe has potential for broader applications beyond faster DDPM inference).

$\newcommand{\d}{\mathsf{d}}$ $\newcommand{\vx}{\mathbf{x}}$ $\newcommand{\vy}{\mathbf{y}}$ $\newcommand{\vz}{\mathbf{z}}$

Thank you for your thoughtful review and insightful questions which have helped improve our work. We hope to address your concerns below:

Hidden Exchangeability Beyond OU DDPMs

Thanks for the helpful pointer! The hidden exchangeability property is not limited to the OU process and actually holds for a large class of generic DDPM formulations, including both VP-SDE and VE-SDE. This is because both VP-SDE and VE-SDE can be expressed as invertible reparametrizations of the OU process. As a consequence, they are both equivalent to the Stochastic Localization process, and thus, satisfy hidden exchangeability. We shall update our draft to include a proof of this result, which we briefly sketch below:

Consider an arbitrary SDE of the form $\d \vz_t = h(t) \vz_t \d t + \sqrt{g(t)} \d B_t$ where $g, h$ are arbitrary continuously differentiable functions with $g(t) > 0$. Note that $h(t) = -\tfrac{1}{2}g(t)$ recovers VP-SDE and $h(t) = 0$ recovers VE-SDE. Now, consider the OU process $\d \vx_t = - \vx_t \d t + \sqrt{2} \d B_t$ with $\vx_0 = \vz_0$ and let $\vy_t = m(t,\vy_t) \d t + \d B_t$ denote the SL process as defined in Section 3.1 Eqn 4. Now, define the functions $\alpha(t)$ and $r(t)$ as follows: $\alpha(t) = \frac{1}{2} \ln \left(1 + \int_{0}^{t} g(s) \exp({-2 \int_{0}^{s} h(u) \d u}) \ \d s\right)$ $r(t) = \exp(\alpha(t) + \int_{0}^{t} h(s) \d s)$

Since $g(t) > 0$, $\alpha$ is a strictly increasing function with $\alpha(0) = 0$, i.e., $\alpha$ is an invertible function. Furthermore, $r(t) > 0$. One can now Apply Ito's Lemma and the time transformation theorem for SDEs (see [1, Thm 8.5.7]) to prove that $\vz_t = r(t) \vx_{\alpha(t)}$, i.e., $\vz_t$ is reparametrizable to the OU process. As discussed in Section 3.1 of our work, $\vy_t = t e^{s(t)} \vx_{s(t)}$ where $s(t) = \tfrac{1}{2} \ln(\tfrac{t+1}{t})$. It follows that $\vz_t$ also maps to the SL process via the following parametrization: $\vy_t = \frac{t e^{s(t)}}{r(\alpha^{-1}(s(t)))} \vz_{\alpha^{-1}(s(t))}$ Consequently, the hidden exchangeability property also applies to $\vz_t$

Effects of Autospeculation on Sample Quality

As elucidated in Tables 1, 2 and 3, our evaluations demonstrate that ASD consistently achieves the same sample quality as the sequential DDPM implementation (benchmarked via CLIP and FID scores for Image Generation and Task Success Rates for Robomimic Tasks) These findings corroborate our theoretical guarantee that ASD is an error-free parallelization scheme that always produces exact samples from the original DDPM's output distribution (Theorem 3).

References

1. Oksendal : Stochastic Differential Equations

$\newcommand{\bE}{\mathbb{E}}$ Thank you for your review and for taking the time to engage with our work. We hope to address your concerns below:

Analysis of Discrete Steps

We highlight that all our theoretical guarantees for ASD directly consider the discrete-time regime. To this end, our proof of Theorem 2 on the correctness of ASD does not make any use of the continuous time SL process, and relies only on the properties of the Gaussian Rejection Sampler. Our proof of Theorem 3, which analyzes the discrete-time parallel complexity of ASD, uses the SL process only as an analytic tool to interpolate between the ASD increments. Such interpolation arguments are ubiquitous in the sampling literature, particularly in the analysis of diffusion models [1,2] as well as gradient-based sampling algorithms like Langevin Monte Carlo. [3,4]

Impact of Discretization on Exchangeability

While our discussion on hidden exchangeability in Section 3.1 focuses on the continuous case for ease of exposition (to provide intuition behind why the arguments of Anari et. al.'24 can be extended to diffusion models), the impact of discretization on hidden exchangeability is appropriately quantified in the proof of Theorem 3. In particular, the proof of Theorem 3 establishes that, as long as the step-size is not too large, ASD increments aproximately satisfy the hidden exchangeability property. This holds because the discrete trajectory of ASD is statistically indistinguishable from that of the continuous SL process (which exactly satisfies hidden exchangeability as per Theorem 1). We rigorously quantify this in the proof of Theorem 16 in Equations 20 and 21 by proving that the TV distance between ASD and the continuous process is $O(\sqrt{\eta \beta d})$. Hence, the two can be made statistically indistinguishable by choosing $\eta \asymp \tfrac{1}{\beta d}$. We shall update our draft to highlight this point in the main text.

High Dimensional Tasks

Our experiments on Pixel Diffusion demonstrate the effectiveness of ASD on a 196,608 dimensional task. Evaluations on more complex modalities such as video is an interesting avenue of future work which we are unable to pursue at present due to resource constraints.

Unequal Step-Sizes

We respectfully disagree with the claim that the case of unequal step-sizes is not analyzed theoretically. We highlight that Theorem 1 proves a general time invariance property for the distribution of stochastic localization increments, which holds even for unequal step-sizes $\eta_i$. We call this time invariance the hidden exchangeability property as it reduces to exchangeability of increments when the step-sizes are equal. However, the key result of Theorem 1 is agnostic to the choice of step-sizes. Consequently, all our theoretical guarantees for ASD are directly applicable to unequal step-sizes.

Assumption on $\beta$

We first note that $\beta = O(1)$ is a standard assumption in the theory of diffusion models [1,2]. Secondly, for arbitrary $\beta$, the canonical choice of step-size is $\eta \asymp \tfrac{1}{\beta d}$ (e.g. [1] uses $\eta \asymp \tfrac{1}{M_2}$ where $M_2 = \beta d$, see also [2, Appendix C]). Under this setting, the parallel runtime of ASD as per Theorem 4 is $O(K^{2/3})$. Hence, the $K^{1/3}$ parallel speedup guarantee continues to hold.

References

Thank you for the helpful references, which we weren't aware of at the time of writing. We shall update our draft to include them.

Löwner-Heinz Theorem

$\bE[\Sigma_t]$ being a decreasing function in the Löwner order directly follows from the facts that: 1. $\tfrac{d{\bE[\Sigma_t]}}{dt} = -\bE[\Sigma_t^2]$ and, 2. $\bE[\Sigma_t^2]$ is a PSD matrix. To our understanding, the Löwner-Heinz theorem (which deals with the operator monotonicity of $t^p$ for $p \in [-1,1]$) is not relevant here. Please let us know if we have made a mistake.

Comparison to Shih et. al.

Please refer to our response to Reviewer FuEw for a detailed discussion on why our results and that of Shih. et. al. cannot be compared on equal footing (and why our work and Shih. et. al. represent orthogonal axes of improvement for fast DDPM inference)

Quantifying the Acceptance Rates

We highlight that the acceptance rate is quantitatively analyzed in the proof of Theorem 4. In particular, the total probability of rejection is upper bounded by $O(\sqrt{K \theta \eta \beta d})$

Refs

1. Chen et. al. : Sampling is as easy as learning the score
2. Benton et. al. : Convergence Bounds for DDPMs via Stochastic Localization
3. Vempala & Wibisono : Rapid Convergence of the Unadjusted Langevin Algorithm
4. Balasubramanian et. al. : Towards a Theory of Non-Log Concave Sampling
5. Anari et. al. : Parallel Sampling via Counting