Parallel Sampling via Counting

We show how to use parallelization to speed up sampling from an arbitrary distribution µ on a product space [q]ⁿ, given oracle access to counting queries: ℙ_{X∼ µ}[X_S=σ_S] for any S⊆ [n] and σ_S ∈ [q]^S. Our algorithm takes O(n^2/3· 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 ℙ_{X∼ µ}[X_i=σ_i |  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 Ω(n^1/3) for the runtime of any parallel sampling algorithm making at most poly(n) queries to the counting oracle, even for q=2.