Parallel Sampling via Autospeculation

We present parallel algorithms to accelerate sampling via counting in two settings: any-order autoregressive models and denoising diffusion models. An any-order autoregressive model accesses a target distribution 𝜇 on [𝑞] 𝑛 through an oracle that provides conditional marginals, while a denoising diffusion model accesses a target distribution 𝜇 on ℝ 𝑛 through an oracle that provides conditional means under Gaussian noise. Standard sequential sampling algorithms require 𝑂(𝑛) time to produce a sample from 𝜇 in either setting. We show that, by issuing oracle calls in parallel, the expected sampling time can be reduced to 𝑂(𝑛 1/2 ). This improves the previous 𝑂(𝑛 2/3 ) bound for any-order autoregressive models and yields the first parallel speedup for diffusion models in the high-accuracy regime, under the relatively mild assumption that the support of 𝜇 is bounded.

We introduce a novel technique to obtain our results: speculative rejection sampling. This technique leverages an auxiliary "speculative" distribution 𝜈 that approximates 𝜇 to accelerate sampling. Our technique is inspired by the well-studied "speculative decoding" techniques popular in large language models, but differs in key ways. Firstly, we use "autospeculation," namely we build the speculation 𝜈 out of the same oracle that defines 𝜇. In contrast, speculative decoding typically requires a separate, faster, but potentially less accurate "draft" model 𝜈. Secondly, the key differentiating factor in our technique is that we make and accept speculations at a "sequence" level rather than at the level of single (or a few) steps. This last fact is key to unlocking our parallel runtime of 𝑂(𝑛 1/2 ).