Round Complexity Separation in Parallel Sampling: KL Divergence Before vs. After Mixing

Starting with a Simple Sampling Problem

Parallel sampling is one of the core techniques for accelerating inference in modern generative models. The basic idea is this: rather than generating samples one coordinate at a time, a sampler randomly selects which coordinates to generate together, compressing what would otherwise be sequential multi-step computation into fewer "rounds." The question is — how far can this compression go? The answer turns out to depend critically on how you measure error.

This is exactly what this paper investigates. It considers a target distribution $p$ and the output distribution $q_U$ produced by a schedule $U$, and compares two ways of measuring error: the schedule-averaged divergence $\mathbb{E}U D{\mathrm{KL}}(p|q_U)$, which computes the KL divergence for each schedule and then takes the expectation; and the output-law divergence $D_{\mathrm{KL}}(p|\mathbb{E}_U q_U)$, which first mixes the output distributions across all schedules and then computes the KL divergence. The difference in ordering — "before mixing" versus "after mixing" — gives rise to a striking complexity gap.

KL divergence (Kullback-Leibler Divergence) is the standard tool for measuring the difference between two probability distributions. For distributions $p$ and $q$, $D_{\mathrm{KL}}(p|q) = \sum_x p(x)\log\frac{p(x)}{q(x)}$ measures the information lost when approximating $p$ with $q$. KL divergence is non-negative, equals zero if and only if the two distributions are identical, and is asymmetric — $D_{\mathrm{KL}}(p|q) \neq D_{\mathrm{KL}}(q|p)$. In the context of sampling quality evaluation, choosing $D_{\mathrm{KL}}(p|q_U)$ means measuring the sampler's deviation relative to the target distribution $p$. Jensen's inequality guarantees that $D_{\mathrm{KL}}(p|\mathbb{E}_U q_U) \leq \mathbb{E}U D{\mathrm{KL}}(p|q_U)$, meaning "measure after mixing" is always no greater than "measure then average" — this is the mathematical root of the order-of-magnitude gap between the two metrics.

Two Metrics, Two Fates

It's worth first clarifying the paper's core model setup. In each round, samples are drawn independently from the exact conditional marginal distribution frozen at the start of that round. Each output schedule is drawn from a finite mixture of schedules prior to sampling. Every scheduled round counts toward the cost. This constitutes a fairly strict framework of "causal randomized schedules."

Within this framework, the authors focus on a canonical target: the uniform distribution over binary strings of even length $n$ containing exactly $n/2$ ones (the balanced target). The conclusion, for fixed tolerance $\varepsilon > 0$, is:

• Under schedule-averaged divergence, the minimum expected number of rounds required is $\Theta_\varepsilon((\log n)^2)$;
• Under output-law divergence, this drops to $(e^{-\varepsilon}+o(1))\log_2\ln n$.

The former is on the order of the logarithm squared; the latter is doubly logarithmic. This means that simply changing the order in which the error metric is defined causes the required number of rounds to collapse from $(\log n)^2$ down to $\log\ln n$ — an enormous asymptotic separation. More crucially, the doubly logarithmic bound can be achieved in the same number of rounds per run, not merely in expectation.

Causal randomized schedules form the central analytical framework of this paper, and understanding them is essential for grasping the scope of the results. "Causal" means that when deciding which coordinates to sample in round $t$, the strategy can only depend on information generated in the previous $t-1$ rounds — it cannot anticipate future randomness. "Randomized" means the schedule itself is allowed to be random, i.e., drawn from a probability distribution over candidate schedules. This contrasts with deterministic schedules, where the same coordinate ordering is fixed across every run. Randomization gives the strategy the ability to make different runs produce complementary biases — and it is precisely this ability that is fully exploited under the output-law criterion, driving the doubly logarithmic round savings. The framework excludes stronger adaptive strategies, such as dynamically adjusting sampling precision based on previously observed coordinate values, so all results hold under a relatively constrained but reasonable model.

Why the Order of Mixing Matters So Much

Intuitively, schedule-averaged divergence demands that each individual schedule's output $q_U$ closely approximates the target $p$ — a stringent, per-instance requirement. Output-law divergence only requires that the mixture $\mathbb{E}_U q_U$, averaged over all schedules, approximates $p$. Individual schedules are free to deviate, as long as their deviations cancel out in the mixture.

This is precisely where randomized strategies show their power. By letting different schedules produce complementary biases, the mixed distribution can achieve the same approximation accuracy with far fewer rounds. In other words, the randomness in "which coordinates to generate together" itself encodes additional approximation capability under the output-law metric — capability that is completely wasted under the per-schedule metric.

Methodologically, the authors construct a count-state recursion for the balanced target that exactly characterizes the minimum output divergence achievable by any causal randomized strategy, for any round budget. This is not a loose sandwich argument between upper and lower bounds — it is an exact characterization, which significantly strengthens the rigor of the separation result.

Generality of the Separation Result

A natural question is whether this separation is an artifact of the balanced target's combinatorial symmetry. The paper answers in the negative.

The separation is extended to two broader classes of targets:

• Weighted fixed-count targets: as long as the weight ratios are bounded and the quota proportions stay bounded away from 0 and 1, the separation holds even without permutation symmetry;
• Full-support Beta(1,1)–Bernoulli targets: under the output-law criterion, the exact expected cost is $(e^{-\varepsilon}/2+o(1))\log_2\log_2 n$, achieved via rare $n$-round runs.

This last point is particularly noteworthy: the doubly logarithmic expected cost is possible because the strategy is allowed to execute a full $n$ rounds with very small probability. These rare "heavy-tailed" runs are what pull down the overall expected round count — revealing a subtle tension between expected cost and worst-case cost in randomized strategies.

Beta(1,1)-Bernoulli targets are defined as follows: $n$ independent Bernoulli random variables whose success probability $\theta$ is itself drawn from a Beta(1,1) distribution (i.e., the uniform distribution on $[0,1]$). After marginalizing out $\theta$, the resulting joint distribution assigns positive probability to every binary string, and its structure is fundamentally different from the balanced target — there is no fixed Hamming weight constraint, and no permutation symmetry (although coordinates are conditionally i.i.d. given $\theta$, the randomness in $\theta$ introduces inter-coordinate dependence). Extending the separation to this target demonstrates that the doubly logarithmic round advantage does not rely on the combinatorial constraint of "exactly $n/2$ ones," and holds equally for a broader class of full-support distributions — strengthening the generality of the theoretical results.

Relation to Prior Work

The authors are also forthright about the boundaries with prior research. Earlier work had already implicitly established the order of schedule-averaged cost for the balanced target, as well as $O_\varepsilon((\log n)^2)$ upper bounds for the two generalized targets. The new contribution of this paper is to provide the first exact doubly logarithmic characterization under the output-law criterion, and thereby establish a strict round complexity separation between the two metrics.

Implications for Parallel Sampling in Practice

Despite being a primarily theoretical work, this paper carries potential implications for engineering practice in parallel sampling. When evaluating the quality of a parallel sampler, the choice of metric is not neutral — if what ultimately matters is the overall fidelity of the sampler's output distribution (i.e., the distribution after mixing), then theory tells us there exists a strategy space far more efficient than per-schedule approximation.

This is also a reminder to practitioners: designing an accelerated sampling scheme that demands every schedule path closely track the target may be an overly conservative goal. Allowing per-path deviation and relying on randomized scheduling to correct biases in aggregate may yield order-of-magnitude savings in round count. Of course, translating these asymptotic theoretical advantages into practically realizable algorithms remains an open engineering challenge.