Parallel Sampling via Counting

TL;DR

This paper introduces a parallel sampling algorithm that achieves sublinear runtime for arbitrary distributions on product spaces using counting queries, with implications for faster sampling in autoregressive models.

Contribution

It presents the first sublinear-in-n parallel sampling algorithm for arbitrary distributions using counting oracles, and establishes a matching lower bound.

Findings

Achieves $O(n^{2/3} ext{polylog}(n,q))$ parallel time for sampling.

Shows a lower bound of $ ilde{ ext{O}}(n^{1/3})$ for any such sampling algorithm.

Implications for improving sampling speed in autoregressive neural network models.

Abstract

We show how to use parallelization to speed up sampling from an arbitrary distribution $\mu$ on a product space $[q]^n$, given oracle access to counting queries: $\mathbb{P}_{X\sim \mu}[X_S=\sigma_S]$ for any $S\subseteq [n]$ and $\sigma_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 \mu}[X_i=\sigma_i\;\vert\; X_S=\sigma_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{\Omega}(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$.