cs.LGOct 9, 2026

KL Before and After Mixing: A Round-Complexity Separation in Parallel Sampling

Authors: Tianyi Wu

Organizations: Amazon

Abstract

A parallel sampler can choose at random which coordinates to generate together. For a target pp and the output law qUq_U of a schedule UU drawn afresh for each output, we compare the schedule-averaged divergence EUDKL(p∥qU)\mathbb{E}_U D_{\mathrm{KL}}(p\|q_U) with the output-law divergence DKL(p∥EUqU)D_{\mathrm{KL}}(p\|\mathbb{E}_U q_U), in nats. Rounds commit independent draws from exact conditional marginals frozen at the round's start, each output's schedule is drawn before sampling from a finite mixture of schedules, and every scheduled round is charged. For the uniform distribution on binary strings of even length nn with n/2n/2 ones, the least expected number of rounds at a fixed tolerance ε>0\varepsilon>0 is Θε((log⁡n)2)Θ_\varepsilon((\log n)^2) under the first criterion but (e−ε+o(1))log⁡2ln⁡n(e^{-\varepsilon}+o(1))\log_2\ln n under the second, and the double-logarithmic order is attainable with the same round count on every run. For this target, a count-state recursion gives the exact least output divergence at every round cap over all causal randomized policies. The separation extends to weighted fixed-count targets with bounded weight ratios and quota fractions bounded away from zero and one, without permutation symmetry, and to the full-support Beta(1,1)(1,1)--Bernoulli target, where the sharp expected cost under the output-law criterion is (e−ε/2+o(1))log⁡2log⁡2n(e^{-\varepsilon}/2+o(1))\log_2\log_2n, attained with rare nn-round runs. Prior work implies the order of the schedule-averaged cost on the balanced target and its Oε((log⁡n)2)O_\varepsilon((\log n)^2) upper bounds on both extensions.

Figures & tables

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

Oct 6, 2026cs.DS

Lower Bounds for Parallel Diffusion Sampling

Standard diffusion samplers generate samples through repeated evaluations of a learned score function. Parallel sampling methods seek to accelerate generation by trading additional evaluations for fewer sequential rounds. This raises the question of how much sequential dependence is unavoidable, even when many score queries can be made simultaneously. We establish the first polynomial parallel-round lower bounds for diffusion sampling with approximate scores. Specifically, we prove (1) a Ω~(d1/3)\widetildeΩ(d^{1/3})-round lower bound for sampling smooth, near-isotropic Gaussian mixtures in RdR^d, and (2) an Ω(d)Ω(d)-round lower bound for uniform sampling from anisotropic axis-aligned boxes contained in the unit ball. Both bounds hold for arbitrary randomized algorithms making polynomially many queries per round at arbitrary locations and noise levels, with inverse-polynomial score error and constant total variation accuracy. The linear bound is tight for our box family. Our constructions use fixed approximate score oracles that enforce sequential access to hidden information while satisfying the accuracy guarantee at every noise level.
Aug 26, 2026cs.IT

Conditional Total Correlation and the Serial Depth of Adaptive Parallel Sampling

Motivated by parallel decoding in masked diffusion models, we study adaptive parallel sampling of discrete vectors: in each round, a deterministic policy selects unrevealed coordinates on the basis of the values observed so far, and the selected coordinates are sampled independently from their exact conditional marginals. Approximation error is measured by forward Kullback-Leibler divergence, and serial depth is the minimum target-averaged number of rounds meeting a prescribed error budget. Our central result is an exact identity: the divergence of every policy equals the expected conditional total correlation accumulated over its reveal rounds, so conditional total correlation is the exact information cost of within-round parallelism. The identity yields zero-error schedules for finite-order Markov chains with round complexity proportional to the Markov order and logarithmic in sequence length, a matching logarithmic characterization of the Bernoulli walk at every fixed error budget, and a linear-versus-logarithmic separation between left-to-right and hierarchical reveal orders. Uniform random permutations require linearly many expected rounds at every fixed budget; their hard-cap round-error tradeoff is an exact integer-composition problem whose fixed-round asymptotics and joint-scaling frontier we determine. Uniform balanced binary strings have depth of order squared logarithm, and binary one-hot blocks have square-root depth, with rectangular versions realizing every polynomial exponent up to one half. These results separate serial depth from entropy and negative log-likelihood, and establish conditional-dependence structure as a fundamental determinant of parallelizability. Experiments with a masked diffusion language model show that the pseudo-cost distinguishes deployed decoding rules and that its policy rankings agree closely with the quality of self-sampled outputs.
Aug 17, 2026cs.AI

The Cost of Conservation: Coordination-Memory Laws for Exact-Support Generation

Many AI systems make decisions locally, even when every realized output must obey an additive conservation law, such as selecting exactly a fixed number of items. This constraint can be statistically invisible: small subsets of a balanced fixed-budget output look increasingly independent, yet communication-free coordinate-parallel generation needs exponentially many pre-shared plans, while a sequential exact sampler needs only logarithmic memory. We study product measures conditioned on additive conservation laws in the intermediate regime where one plan is selected before a fixed-order pass, every plan is a bounded-state stochastic executor whose support is entirely legal, and the mixture of plan laws approximates the target distribution in total variation. Our main result identifies the optimal asymptotic selector rate, up to constant factors, with the killed spectral profile of the conservation-difference walk. The resulting coordination cost decreases as an inverse power of live-state width, with an exponent determined by intrinsic conservation rank rather than alphabet size; the law extends to noncentral budgets and heterogeneous local scores. The converse is driven by a state-versus-resource-sum obstruction, while a rate-matching construction compiles discrepancy control into exact-support finite-state plans. Complementary results characterize block-parallel plan complexity and the benefit of programmable output order. Together, these results show precisely how online memory substitutes for front-loaded coordination in exact-support generation.