Accelerated Diffusion Models via Speculative Sampling

We would like to thank the reviewer for their positive assessment of our work.

This paper shares a similar idea to speculative decoding for LLMs [1], but to my knowledge, it is the first to apply this idea to diffusion model sampling. T-stitch [2] leverages the idea of using a smaller draft model for fast sampling, but from a totally different perspective. It didn't share the same theoretical foundation as speculative decoding.

We thank the reviewer for pointing us to [1]. While [1] also leverages two models, their approach differs from ours. Specifically, they employ a cheap model for early denoising stages and an expensive model for later stages. This strategy does not guarantee the same quality as the original superior model. In contrast, our method samples exactly from the superior model, utilizing the cheap model only for parallel verification, thus maintaining the desired quality. We will however include a discussion regarding this very related work in our updated version of the manuscript.

[1] Pan et al. (2024) – T-Stitch: Accelerating Sampling in Pre-Trained Diffusion Models with Trajectory Stitching

Some experimental settings are not clearly described.

Below we answer the concerns of the reviewer.

What kind of sampler is used for 30-step sampling in Table 1? And what is the original total number of sampling steps of the target model?

In Table 1 (and all of our experiments), we used the stochastic sampler of flow matching models, see [1] for instance. An explicit form of the sampler is given in Equation (37), line 1843. In the revised version of the paper, we will highlight Equation (37) in the main paper. We emphasize that our method is compatible with other stochastic methods such as DDPM [2].

[1] Albergo et al. (2023) – Stochastic Interpolants: A Unifying Framework for Flows and Diffusions

[2] Ho et al. (2020) – Denoising Diffusion Probabilistic Models

The idea of applying speculative decoding to diffusion model sampling is interesting. However, since there are already many fast samplers like DPM-Solver [1] for diffusion models, what are the main advantages of speculative sampling over the other samplers? In this paper, there is no experiment comparing them.

We thank the reviewer for this insightful remark. We emphasize that our method is largely orthogonal to other acceleration techniques. In particular we now combine our approach with [1,2], see our detailed answer to Reviewer zMQr, which yields comparable NFE improvements (including against DPM-Solver++ [3]). Furthermore, our method's applicability extends beyond accelerating denoising diffusion models; it can also be used to accelerate the simulation of any diffusion equation with a computationally expensive drift, such as Langevin diffusion to sample from unnormalized target distributions, see our detailed answer to Reviewer E771. We have added experiments as well as an Appendix in the revised version demonstrating the efficiency of the methodology in this context.

[1] Tong et al. (2024) – Learning to Discretize Denoising Diffusion ODEs

[2] Shih et al. (2023) – Parallel Sampling of Diffusion Models

[3] Lu et al. (2022) – DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models

In Table 1, is the draft model independent or frozen? If independent, how is NFE computed? Is it only dependent on the calls to the target model?

To clarify, we employed the frozen strategy for all Image Space and Robotics experiments reported in Table 1. The independent strategy was used exclusively in the low-dimensional setting to demonstrate the superior performance of the frozen draft approach. Consequently, in that specific case, only the target model was evaluated, and the NFE reported corresponds solely to that model. We will explicitly address this distinction in the revised manuscript.

Versatility of the method

We have now included additional experiments regarding comparison with DPM++ solvers, see answer to Reviewer eBDt, as well as an experiment showcasing the efficiency of the method in a sampling context (i.e. a non generative modeling setting), see answer to Reviewer E771.

We would like to thank the reviewer for their insightful comments.

how to think about $g_t$ [...] for classical schedules.

Let $A(t) = g_t^2 ((1/\sigma_t - \sigma_t)^2 + \alpha_t^2)$ and consider a few schedules:

• Rectified flow [1]: $\alpha_t = 1-t$, $\sigma_t = t$. We have that $A(t) = 2 t (1-t) (1 + (1 + 1/t)^2)$.
• Cosine [2]: $\alpha_t = \cos(\pi t /2)$, $\sigma_t = \sin(\pi t /2)$. We have that $A(t) = \frac{\pi}{2}\sin(\pi t) (1 + \mathrm{cotan}(\tfrac{\pi}{2} t))$.

Note that $A(t) \to 0$ as $t \to 1$ and $A(t) \to +\infty$ as $t \to 0$. Hence the acceptance is close to 1 around $t \approx 1$ and worsens when $t \to 0$. This is consistent with our empirical results. We will include a plot of $A(t)$ as a function of $t$ for those schedules.

[1] Liu et al. (2022) – Flow Straight and Fast

[2] Nichol and Dhariwal (2021) – Improved Denoising Diffusion Probabilistic Models

difficult to parse how useful Theorem 4.3 is.

We believe that Theorem 4.3 is useful because, although we do not expect our bound to be tight, it captures rigorously the impact of the various hyperparameters of the algorithm on the expected acceptance rate. In particular:

• that decreasing the stepsize increases the acceptance rate at the cost of increasing the number of total steps.
• The influence of the schedule (see above).
• The existence of an optimal value for $\varepsilon$ is also observed in practice, see our experimental section.

needs more support is the sentence "Our approach complements [other approaches for acceleration] and can be combined with...better integrators." [...] I would be more convinced about the proposed approach if, combined with one of these existing approaches, it gets a tangible improvement.

We have compared our approach with LD3 [1] and AYS[2]. In addition, to further substantiate our claim we have combined our acceleration method with parallel sampling [3] and observed benefits of the method, see the answer to Reviewer zMQr for more details. We compare the results on CIFAR-10 as reported in LD3 [1]. Our best speculative sampling method outperformed both LD3 and AYS. We also included our best results obtained with a uniform timesteps spacing and EDM timestep spacing [4]. These results are based on the same model as “Best speculative”. We sweep over $\rho = [1.0, \dots, 8.0]$ in the case of EDM timestep spacing. This improves the quality of the samples but they remain inferior in quality to the ones obtained with our best speculative model. We re-implemented LD3 [1] in our setting and used it to learn a timestep spacing. Our setting is similar to the one of [1]. Finally, we compare our approach with a distilled generator trained on top of our best model. We focus on Multistep Moment Matching Distillation (MMD) [5]. We sweep over several hyperparameters for MMD and report the best result.

We will include those results in the revised version of our paper. We will also comment on LD3 [1] and AYS [2] in the related work section as these works are indeed complementary to ours.

Finally we note that our method is not restricted to diffusion models. It can be applied to accelerate simulation to sample from unnormalized distributions. We have added an experiment on the $\phi_4$ potential demonstrating the efficiency of the method, see the answer to Reviewer E771.

[1] Tong et al. (2024) – Learning to Discretize Denoising Diffusion ODEs

[2] Sabour et al. (2024) – Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

[3] Shih et al. (2023) – Parallel Sampling of Diffusion Models

[4] Karras et al. (2022) – Elucidating the Design Space of Diffusion-Based Generative Models

[5] Salimans et al. (2024) – Multistep Distillation of Diffusion Models via Moment Matching

[is the] method actually orthogonal to other accelerations. This is tricky as many few-NFE methods are based on ODEs. [...] setup you could try where you could achieve comparable FID with existing acceleration baselines using [...]?

It is indeed correct that our approach does not apply to deterministic samplers, but we adapt common acceleration methods to the stochastic case. For example in the case of LD3 [1] we train LD3 in our setting by considering the noise added during the sampling to be frozen. Similarly, in [2] (see end of Section 3.1), the authors freeze the noise. Following those principles we are able to combine several acceleration techniques with our approach.

[1] Sabour et al. (2024) – Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

[2] Shih et al. (2023) – Parallel Sampling of Diffusion Models

We would like to thank the reviewer for their very positive assessment of our manuscript.

Overall, their claims are well-supported by theory and experiments. Thus I think the contributions of the paper are solid and beneficial to the community of both speculative decoding and diffusion models.

We appreciate that the reviewer recognizes the value of our work and its benefits for both the LLM and diffusion models communities.

How much does the proposed method accelerate inference of diffusion models when combined with other acceleration methods?

Below, we combine our speculative sampling method with parallel sampling [1].

We report FID score and NFE for CIFAR-10 with a number of steps of $30$. We vary the temperature parameter $\tau$, the churn parameter $\varepsilon$ as well as the number of parallel iterations, see [1,2]. For each combination of hyperparameters we also consider window sizes $5$, $10$ and $20$ and report the best run (in terms of FID).

The original speculative sampling procedure corresponds to $p=0$. The best FID number that can be achieved with this configuration is $2.23$ with a NFE of $15.69$. However, by combining our speculative sampling procedure with parallel sampling then we can reach a FID of $2.07$ with a NFE of $15.42$. This shows the benefits of combining our speculative sampling procedure with other acceleration methods. We will report those results in the updated version of our manuscript.

[1] Shih et al. (2023) – Parallel Sampling of Diffusion Models

[2] Tang et al. (2024) – Accelerating Parallel Sampling of Diffusion Models

Captions for the tables can be improved.

We are willing to improve the captions of our table if the reviewer has more feedback for us.

Versatility of the method

We have now included additional experiments regarding comparison with DPM++ solvers, see answer to Reviewer eBDt, as well as an experiment showcasing the efficiency of the method in a sampling context (i.e. a non generative modeling setting), see answer to Reviewer E771.

We would like to thank the reviewer for their positive assessment of our paper.

The approach assumes parallel evaluation allowing lower latency, but the method increases compute and memory requirements, which are also critical components that are not handled by the approach.

Our approach indeed trades memory for latency. We believe this trade-off is advantageous for certain applications. For example, in applications requiring extremely low latency, increased memory consumption is often acceptable. Note that this fact is also emphasized in [1] “Currently, for large models like SD, ParaTAA requires the use of multiple GPUs to achieve considerable speedup in wall-clock time. Nonetheless, as advancements in GPU technology and parallel computing infrastructures evolve, we anticipate that the cost will be significantly lower”

[1] Tang et al. (2024) – Accelerating Parallel Sampling of Diffusion Models

The evaluation is conducted using CIFAR10 and LSUN datasets. More thorough evaluation is needed using ImageNet-1k classes.

We have now evaluated our method on ImageNet (64x64x3) and show similar improvements in terms of NFE. With a baseline at 250 steps (corresponding to a FID/IS of 4.18/49.54), our method achieves a NFE of 95.38 corresponding to a halving of the NFE. The FID/IS score of the speculative sampling approach is 4.35/48.28 (which is of the same order as 4.18/49.54). We acknowledge that our baseline FID/IS is not state-of-the-art, but this was the only model we had the time to retrain by the rebuttal deadline. We are currently training another model in order to improve the base FID/IS score.

The parameters of the base sampler are given by $\varepsilon = 0.25$ and the guidance value is set to $0.25$. In our speculative sampling experiment we consider a temperature of $\tau=1.0$ and a window size of $10$.

Need to compare with other approaches that allow efficient sampling of diffusion models. o “Fast Sampling of Diffusion Models via Operator Learning” o “Parallel Sampling of Diffusion Models” o “Accelerating Parallel Sampling of Diffusion Models”

We do not compare with [1] as this approach is very different from ours and does not yield state-of-the-art results. However, we have now compared with [2] and refer to the answer to Reviewer zMQr for more details. In short, our method outperforms [2] and, combining our speculative approach with [2], leads to improved results which also outperforms strong baselines such as [4,5]. However, we emphasize that our approach is not directly comparable to [2,3], since 1) these methods do not sample exactly from the original denoising diffusion models and 2) only converge in the infinite number of steps limit since they rely on the fixed point property of the SDE integrators.

[1] Zheng et al. (2022) – Fast Sampling of Diffusion Models via Operator Learning

[2] Shih et al. (2023) – Parallel Sampling of Diffusion Models

[3] Tang et al. (2024) – Accelerating Parallel Sampling of Diffusion Models

[4] Tong et al. (2024) – Learning to Discretize Denoising Diffusion ODEs

[5] Sabour et al. (2024) – Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

Supplementary material Section C – Algorithm 7. The return statement on Line 980 is probably incorrect.

Thanks, this is indeed a typo.

Versatility of the method

To showcase the applicability of our method beyond generative modeling, we demonstrate its efficiency in the context of Monte Carlo sampling. We consider the $\phi_4$ model [1,2] which defines $\pi(x) \propto \exp[-U_\beta(x)]$ for

$U_\beta(x) = (\beta/2) \sum_{|i-j| = 1} (x_i - x_j)^2 + \sum_{i} (x_i^2 - 1)^2,$

on a grid of shape $(8,8)$ and $\beta=100$. Sampling from $\pi$ is complex as this requires sampling so-called ordered states. In this context, the teacher model is the Langevin diffusion sampling $\pi(x)$ with $100,000$ iterations and stepsize $10^{-3}$ while our speculative sampling algorithm uses the frozen prediction draft model and a window size of $20$. The results are averaged over $500$ runs. The mean energy is the average of $U(x_i)$ computed on the $500$ samples $x_i$, while the standard deviation energy is the standard deviation of $U(x_i)$.

[1] Guth et al. (2022) – Wavelet Score-Based Generative Modeling

[2] Milchev et al. (1986) – Finite-size scaling analysis of the $\phi_4$ field theory on the square lattice