KL Before and After Mixing: A Round-Complexity Separation in Parallel Sampling
Tianyi Wu
cs.LG
Oct 9, 2026 · v1
cs.IT
TL;DR
A Lean 4/Mathlib package (369 modules) kernel-checks the paper's key bounds, numerical brackets, and finite exhibits on parallel-sampling round complexity.
Abstract
A parallel sampler can choose at random which coordinates to generate together. For a target $p$ and the output law $q_U$ of a schedule $U$ drawn afresh for each output, we compare the schedule-averaged divergence $\mathbb{E}_U D_{\mathrm{KL}}(p\|q_U)$ with the output-law divergence $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 $n$ with $n/2$ ones, the least expected number of rounds at a fixed tolerance $\varepsilon>0$ is $Θ_\varepsilon((\log n)^2)$ under the first criterion but $(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)$–Bernoulli target, where the sharp expected cost under the output-law criterion is $(e^{-\varepsilon}/2+o(1))\log_2\log_2n$, attained with rare $n$-round runs. Prior work implies the order of the schedule-averaged cost on the balanced target and its $O_\varepsilon((\log n)^2)$ upper bounds on both extensions.
Problem
Parallel samplers for masked diffusion and any-order autoregressive models fill several coordinates per round with independent draws, and random schedules can be judged by two criteria. One is the schedule-averaged KL divergence; the other is the KL divergence of the mixed output law. The question is how many rounds each criterion requires.
Approach
The analysis uses an exact-oracle interface on the uniform distribution over balanced binary strings, and on its extensions to weighted fixed-count targets and the Beta(1,1)–Bernoulli target. Output divergence is reduced to a survival probability via uniform relabelling, and a count-state recursion gives exact optimal divergence at every round cap. Upper bounds come from square-root ladder schedules, and lower bounds from entropy-decrement arguments. Key bounds, numerical brackets and finite exhibits are kernel-checked in a Lean 4/Mathlib package.
Results
On the balanced slice the schedule-averaged criterion needs Θ((log n)^2) rounds, while the output-law criterion needs (e^{-ε}+o(1)) log2 ln n rounds. The separation extends to bounded-weight slices and to the full-support target, where the sharp cost is (e^{-ε}/2+o(1)) log2 log2 n. Lean checks include the bound L_H>26.20 and an 11-round profile with F≤ln 2 at n=65536.
| design, criterion, cost | balanced slice |
|---|
| fixed schedule, KL, length | Θ((log n)^2) |
| mixture, H, mean | Θ((log n)^2) |
| relabelled profile, F, length | ℓ+O(1) |
| mixture, F, mean | (e^{-ε}+o(1))ℓ |
Round complexity on the balanced slice (ℓ ≈ log2 ln n)