Parallel Sampling via Counting

Abstract

We show how to use parallelization to speed up sampling from an arbitrary distribution $μ$ on a product space $[q]^n$, given oracle access to counting queries: $\mathbb{P}_{X\sim μ}[X_S=σ_S]$ for any $S\subseteq [n]$ and $σ_S \in [q]^S$. Our algorithm takes $O({n^{2/3}\cdot \operatorname{polylog}(n,q)})$ parallel time, to the best of our knowledge, the first sublinear in $n$ runtime for arbitrary distributions. Our results have implications for sampling in autoregressive models. Our algorithm directly works with an equivalent oracle that answers conditional marginal queries $\mathbb{P}_{X\sim μ}[X_i=σ_i\;\vert\; X_S=σ_S]$, whose role is played by a trained neural network in autoregressive models. This suggests a roughly $n^{1/3}$-factor speedup is possible for sampling in any-order autoregressive models. We complement our positive result by showing a lower bound of $\widetildeΩ(n^{1/3})$ for the runtime of any parallel sampling algorithm making at most $\operatorname{poly}(n)$ queries to the counting oracle, even for $q=2$.