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

Tianyi Wu

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 $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.

Figures & tables

Figure 1: The four-bit example. Columns: the six valid strings (two ones) and, pooled, the ten invalid ones; disc areas are unconditional probabilities, and nothing is renormalized. The example implies no rate.

design, criterion, cost	balanced slice	full support	bounded weights
fixed schedule, KL, length	$\Theta((\log n)^{2})$	$\Theta((\log n)^{2})$	$\Theta((\log n)^{2})$
mixture, $H$ , mean	$\Theta((\log n)^{2})$	$\Theta((\log n)^{2})$	$\Theta((\log n)^{2})$
relabelled profile, $F$ , length	$\ell+O(1)$	$\ell+O(1)$	$\ell+O(1)^{\ddagger}$
causal policy, $F$ , cap	$\ell+O(1)$	$\Theta(\ell)^{\dagger}$	$\ell+O(1)$
mixture, $F$ , mean	$(e^{-\varepsilon}+o(1))\ell$	$(\frac{1}{2}e^{-\varepsilon}+o(1))\ell$	$(e^{-\varepsilon}+o(1))\ell$
causal policy, $F$ , $\mathbb{E}\tau$	$(e^{-\varepsilon}+o(1))\ell$	$(\frac{1}{2}e^{-\varepsilon}+o(1))\ell$	$(e^{-\varepsilon}+o(1))\ell$

Table 1: Least number of rounds at a fixed tolerance

\varepsilon>0

, by design, criterion and cost;

\ell=\log_{2}\log_{2}n=\log_{2}\ln n-\log_{2}\ln 2

Figure 2: Balanced slice,

\varepsilon=\ln 2

: round counts of one fixed labelled schedule (circles) and one relabelled profile (squares). Each has one length per run, so circles and squares bound

L_{H}

and

L_{F}

from above; open squares (best-found witnesses) and the kernel-checked

L_{F}\leq 11

n=2^{16}

are upper bounds not shown optimal. No rate is fitted.

Figure 3: Exact grid points and bars; relations schematic. Top: level

c=8

; middle: levels

6

–

11

; bottom: their union, with gaps of order

c_{\min}^{-2}

away from small denominators (shaded: larger gaps near

1/3

1/2

2/3

; not the converse’s set

S^{\star}

). Dots: support, not probability.

Figure 4: One level of the bounded-weight sampler (top, executed) and of the reference law (bottom, analysis only). Both share the kept set

C

and the block values

x_{B}

;

W

compares the two completions of

C

(

W=0

j

is infeasible).

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Statement	Proof	Key step
Lemma 1	Appendix B	invariant laws are uniform on the single orbit $V$
The survival envelope ( Env )	Appendix D , Proposition 9 (its weaker $\kappa_{r}$ form $\Gamma_{r}$ , at the end of that appendix, is used in Appendices E and L )	one-round negative-moment bound, then Cauchy–Schwarz exponent halving
Lemma 3 , Theorem 4 , Corollary 6	Appendix E (total-variation versions in its paragraph Total variation ; the window in Appendix E.1 )	count-state supermartingale bounds success; tail sums charge rounds
Theorem 2 (a),(b)	Appendix B (sketches in Section 3 )	(a) per-reveal entropy decrement; (b) square-root ladder
Theorem 2 (c), fixed-schedule upper bound on $L_{H}$	Appendix C	per-block bound (Step 1) on a geometric schedule
Corollary 5	Appendix E , end, from Theorem 2 , Lemmas 1 (a),(c),(e) and 3 (a), and ( Env )	point masses in Theorem 2 ; ( Env ) caps a profile’s success

Appendix

Table 6

class, criterion, cost	order	status
Balanced slice
fixed schedule, KL, length	$\Theta(\log^{2}n)$	K (Cor. 5 )
one relabelled profile, $F$ , length	$\log_{2}\ln n+O_{\varepsilon}(1)$	K (Cor. 5 )
mixture, $H$ , mean	$\Theta(\log^{2}n)$	K (Thm 2 a,c)
mixture, $F$ , mean	$(e^{-\varepsilon}+o(1))\log_{2}\ln n$	K (Thm 2 b; Cor. 6 b,c)
$\mathcal{A}_{r}$ , $F$ , cap	$\log_{2}\ln n+O_{\varepsilon}(1)$	K (Thm 4 a; Cor. 6 a)

Appendix

Table 2: Classes, criteria and costs (

\varepsilon>0

fixed; for tilts and bounded weights also

\Lambda,\kappa_{0}

;

\mathcal{A}_{r}

is the class of Section 2 , which for two-value tilts is

\mathrm{RAND}_{r}

of Appendix G ). “Mean”: scheduled mean

C(\mu)

; “

\mathbb{E}\tau

”: executed mean; “cap”: hard cap; “length”: rounds of one schedule or relabelled profile, the same on every path. K: kernel-checked (declarations in Appendix Q , whose table also marks the results whose class identification uses one ordinary step; for an asymptotic order, the asymptotic statement); P: prior work (preprint). ‡ One fixed labelled schedule, relabelled uniformly within an equal-weight group of linear size; without one, a common length is open. Finite regime: on the balanced slice, explicit bounds (Table 8 ), kernel-checked bounds and finite computations at

n=2^{16}

(Table 7 ), and an exact least profile length at

n=2^{4}

, computed optima (float64), linear-program brackets, best-found upper bounds and one seeded Monte Carlo estimate (Table 5 , Figure 2 ); on full support, explicit bounds on

L_{H}

(Proposition 14 ) but astronomical constants for

F

, with computed

L_{H}

optima and best-found upper bounds on

L_{F}

(both float64; the latter interval-certified at

n\leq 2^{10}

; Table 6 ); for two-value tilts, explicit but impractical constants (e.g., caps of hundreds to thousands of rounds; Appendix H , §6.9) and, at

n\leq 24

and

r\leq 4

, rigorous enclosures of the optimal output KL, some kernel-checked, beside float64 values (Appendix K ); for bounded weights, existence-only constants in the reservoir-free

F

upper bounds (Appendix I ).

design	one round	rounds
balanced, ladder	unrevealed $M\to\sqrt{M}$	$\log_{2}\log_{2}n$
full support, one profile	level $c\to c^{2}$	$\log_{2}\log_{2}n$
full support, mixture	band $c_{\min}\to c_{\min}^{4}$	$\frac{1}{2}\log_{2}\log_{2}n$

Appendix

Table 3: One round’s growth, up to lower-order factors (

M

: number of unrevealed coordinates;

c

: level;

c_{\min}

: start of a band of levels), and the leading round count: the length of every path for the two profiles, and the scheduled mean

C(\mu)

for the mixture, which has rare paths of length

n

. Mixing in the one-round schedule multiplies the leading term of the scheduled mean by

e^{-\varepsilon}

Symbol	Meaning	Defined in
Interface, criteria and costs
$p$ , $n$	target law on $\{0,1\}^{n}$ ; $n$ the number of coordinates	§ 1 ; § 2
$V$ , $k$	balanced slice $\{x:\sum_{i}x_{i}=k\}$ , $n=2k$ ; $k$ the required number of ones (also: $V(s)$ )	§ 2
$(M,j)$	count state: $M$ coordinates unrevealed, $j$ of them must still be one (also: rung index $j$ ); with a subscript, an unrevealed count at a stage: $M_{t-1}$ before round $t$ (§ 3 ), the rungs $M_{j}$ of Theorem 2 (b)’s ladder and $M_{t}$ of Theorem 4 (b)’s witness (rows $j$ and $L^{\ast}$ ). In § 4.2 the nominal sizes $M_{t}$ are deterministic and the free count is $N=\|\mathcal{C}_{t}\|$ , random around $M_{t}$ ; there the $M$ of $M\geq\rho n$ is the reservoir size of part (b)’s construction (Theorem 8 ’s $j/M$ keeps the count-state meaning)	§ 2
$U$ ; $\|U\|$ (length); $b(U)$ (profile)	labelled schedule $(B_{1},\dots,B_{r})$ , an ordered partition of $[n]$ ; its length (rounds); its profile (block sizes)	§ 2
$\sigma$	a permutation of $[n]$ , as in the uniform relabelling $\sigma(U)$ ; $\sigma(\cdot)$ is also the $\sigma$ -algebra generated, as in $\mathcal{F}_{t-1}$ . In the appendices $\sigma$ also denotes, each defined where used, a standard deviation ( $\sigma(y)$ , $\sigma_{n}$ , $\sigma^{2}$ in Apps. F , K , G ), the logistic function (App. I ) and a state $\sigma_{t}$ (App. H )	§ 2

Appendix

Table 4: Notation index. The first locator is the first place the paper defines the symbol; a second locator gives the formal or per-round definition. “(also: …)” marks a symbol that has another meaning in another row.

$n$	$L_{H}$	$L_{F}$	$A$	shape	$R_{F}^{\mathrm{PROF}}$	TV fixed / mixed
$2^{4}$	2.66	1.999	–	–	2	–
$2^{8}$	8.19	3.28–3.39	9	89%	4	0.412 / 0.005
$2^{10}$	12.56	3.71–3.86	13	95%	4	0.574 / 0.014
$2^{11}$	15.14	3.92–4.04	–	–	$\geq 4$	–
$2^{12}$	18.06	–	19	98%	$\leq 5$	0.673 / 0.015
$2^{16}$	26.20–33	1.99–4.69	33	99%	$\leq 5$	0.831 / 0.022

Appendix

Table 5: Finite comparison at

\varepsilon=\ln 2

(finite computation; details in Appendix K ).

L_{H},L_{F}

: mean-cost optima of the same class (two-length linear programs;

L_{F}

bracketed below by the lower program over the computed

G_{r}

and above by best-found relabelled profiles; at

2^{16}

, where the program’s

d^{*}_{r}

were not computed, the

L_{H}

entry brackets the optimum between the kernel-checked lower bound and the length

A

of one schedule);

A

: least length of one fixed schedule with KL

\leq\ln 2

(computed optimum: float64 evaluation of an exact dynamic program) and the shape share of its divergence;

R_{F}^{\mathrm{PROF}}

: least length of one relabelled profile with

F\leq\ln 2

(exact at

2^{4}

; at

2^{8}

and

2^{10}

a computed optimum, float64 with margins of at least

0.02

from

\frac{1}{2}

; best-found upper bound at

2^{12},2^{16}

;

\geq 4

2^{11}

, since

G_{3}(2048,1024)=0.3145<\frac{1}{2}

in float64 bounds the success of every relabelled profile with at most three rounds); TV: total-variation distance between the target and the output law of the

\delta=\frac{1}{2}

ladder of Theorem 2 (b)’s proof with the Hoeffding–Ville choice

K=\lceil 2\ln(2/\delta)\rceil=3

(the stated bound takes the smaller

K_{\delta}=2

of the Kolmogorov branch; both ladders fail with probability at most

\delta

), whose bulk blocks are

(160,72)

(256,672,72)

(3328,672,72)

and

(40960,23808,672,72)

2^{8}

2^{10}

2^{12}

and

2^{16}

, each followed by

24

singletons; fixed labelling versus relabelled (

=1-s

, with

s

the ladder’s success probability; exact at

2^{8}

–

2^{12}

, within

2\cdot 10^{-57}

2^{16}

); the fixed-labelling TV is a

30

-digit enumeration at

2^{8}

–

2^{12}

(pruned mass below

10^{-31}

) and at

2^{16}

a seeded Monte Carlo estimate of

\mathbb{E}_{P}[(1-Q/P)_{+}]

over

400{,}000

target draws (standard error

0.0004

). Optima, shares and TV values are rounded to nearest, bounds outward; –: not computed.

	$\varepsilon=0.1$			$\varepsilon=\ln 2$
$n$	$L_{H}$	$L_{F}\leq$	ratio $\geq$	$L_{H}$	$L_{F}\leq$	ratio $\geq$
$2^{6}$	14.55	4.36	3.3	4.35	2.29	1.9
$2^{8}$	29.48	5.59	5.2	8.03	2.91	2.7
$2^{10}$	51.61	6.50	7.9	13.10	3.43	3.8
$2^{12}$	81.45	7.12	11.4	19.29	3.82	5.0
$2^{14}$	119.34	7.68	15.5	26.75	4.04	6.6

Appendix

Table 6: Full-support target Beta

(1,1)

–Bernoulli: computed

L_{H}

optima (float64 evaluation of the exact optimization) and best-found upper bounds on

L_{F}

(float64; interval-certified only where stated). Each

L_{F}

bound mixes the best-found profiles with

R

and

R+1

rounds whose

F

values bracket

\varepsilon

, weighted so that the average

F

\varepsilon

; as

\mathrm{KL}(p\|\cdot)

is convex, the mixture has

F\leq\varepsilon

, and the

(R+1)

-round profile alone has

F\leq\varepsilon

, so each bound, rounded up to an integer, is the length of the shortest best-found single profile with

F\leq\varepsilon

. All

L_{F}

entries come from one evaluation pipeline applied to the candidate profiles of two independent optimizers; at

n\leq 2^{10}

a separate implementation, which recomputes each profile’s output count law exactly and encloses

F

in interval arithmetic, certifies these

L_{F}

entries as upper bounds; a second implementation reproduces

L_{H}

n=2^{8}

–

2^{12}

(

\varepsilon=0.1

L_{H}

rounded down, upper bounds up, ratios (computed from the unrounded values) down.

quantity	value
$L_{H}$	$>26.20$	K
$L_{H}$	$\leq 33$ (one schedule)	Fin
$L_{F}$	$\leq 11$ (one relabelled profile)	K
$L_{H}/L_{F}$	$>2$	K
$L_{F}$ (finite brackets)	$[1.996,\,4.681]$	Fin

Appendix

Table 7: Balanced slice,

n=65536

\varepsilon=\ln 2

. K: kernel-checked; Fin: finite computation, of two kinds:

L_{H}\leq 33

is one schedule found by a float64 dynamic program, its KL recomputed at

60

digits; the

L_{F}

bracket comes from exact-rational linear programs over inputs selected by a float64 search. The lower bound on

L_{H}

uses the three-term form of Theorem 2 (a) (the relaxed form gives

23.02

); the

L_{F}\leq 11

certificate relabels one explicit profile (numbers of unrevealed coordinates

65536\to 1735\to 201\to 50\to 19\to 9\to 5

, then five singletons). The kernel check bounds that profile’s failure probability by the second-moment budget

\sum_{t}b_{t}/m_{t}^{2}

over its shrinking blocks (

m_{t}

the coordinates left after block

t

; the singleton tail adds nothing), here

0.4889<\frac{1}{2}

, so its success exceeds

\frac{1}{2}

and

F<\ln 2

. The five-round profile behind the finite computation

F=0.650

has a budget above

7

, which is why the certificate uses eleven rounds.

Figure 5: The mechanism at finite

n

(

\varepsilon=\ln 2

; finite computation, Table 5 ): total variation to

p

of the

K=\lceil 2\ln(2/\delta)\rceil=3

ladder of Theorem 2 (b)’s proof (

\delta=\frac{1}{2}

; the Hoeffding–Ville choice, where the stated bound takes

K_{\delta}=2

; Table 5 lists its blocks) under the fixed labelling (circles) and under its relabelling (squares;

=1-s

, the quota-violation probability). Filled: computed; open: a seeded Monte Carlo estimate (standard error

0.0004

). The least lengths of the same table are plotted in Figure 2 .

$\log_{2}n$	$L_{H}\geq$	$L_{F}\leq$	$L_{F}\geq$	$R_{F}^{\mathrm{adapt}}\geq$
8	5.09	19	1.32	2
16	26.20	20	1.48	2
32	122.43	21	1.88	3
64	530.82	22	1.99	4
128	2211.37	23	2.49	4
1024	$1.46\cdot 10^{5}$	26	3.49	7

Appendix

Table 8: Proved bounds versus

n

(

\varepsilon=\ln 2

). Formula evaluations; lower bounds rounded down, upper bounds up (script hash below). The second and third columns evaluate Theorem 2 (a) and (b), and the two right-hand columns the kernel-checked envelope forms (Appendix Q ), with

\Gamma_{t}(n)

of Appendix D .

Figure 6: Proved bounds as functions of

n

\varepsilon=\ln 2

and the kernel-checked finite certificates at

n=2^{16}

. Curves are formula evaluations, not measurements.

Figure 7: Controlled test with learned conditionals (144 fresh trajectories, 48 independent clusters, six held-out regimes and fresh seeds on two discovery targets; output KL of the complete output law, computed in float64 by dynamic programming over revealed sets (no sampling),

n\leq 16

; predictions frozen before evaluation). (a) Learned regret of the path-optimal profile PATH versus its oracle loss

g_{P}

(identity dashed). (b) Learned versus oracle gap between a fixed labelled order and its relabelling. (c) Learned regret of the exact-oracle terminal-optimal profile TERM against the head’s development-loss excess (“DEV excess”), for cells with oracle margin

0.01

–

0.03

and the EARLY, ES and END checkpoints of each run (descriptive, not a frozen pass/fail item; dotted line at

0.05

, the top boundary of the freeze’s descriptive head-error bands

\leq 0.02

0.02

–

0.05

and

>0.05

; DEV excess, which can be negative, is drawn clipped below at

10^{-5}

). Controlled prediction only.

Result	Criterion; samplers or policies	Round count; target	Bearing on the least rounds under $F\leq\varepsilon$
Bounds on $H$ , which bound $F$ only from above, through $F\leq H$
Chen et al. (2026, Def. 15, Thms. 4, 9, 16)	$H$ , the expected per-schedule KL: exact formula, optimal block sizes, a bound via total correlation; $F\leq H$ . Samplers: random unmasking: prescribed block sizes, uniformly random positions	number of blocks; any target	upper bounds on $F$ through $F\leq H$
Secchi and Zanella (2026, Prop. 4)	factorization error: the $H$ of the exact-conditional sampler; output KL at most it plus a learning error. Samplers: $\tau$ -leaping (random block sizes) versus deterministic block sizes, both with random positions	steps; limit at a fixed ratio of length to steps; under all assumptions of their Prop. 4	on $H$ , random block sizes cost at least as much in the limit; $F$ enters only through $F\leq H$ (exact conditionals) or $F\leq H+\varepsilon_{\mathrm{learn}}$ (learned ones)
Lower bounds on the output law
Feng et al. (2025, Thm. 4.4)	output mass off the support; above $1/2$ it gives $F>\ln 2$ . Samplers: masked-diffusion sampling with any masking schedule, each masked position revealed independently; conditionals exact or with per-token KL strictly below $1/128$	sampling steps; an explicit hidden Markov target	$\Omega(n)$ steps for $F\leq\ln 2$ , also for mixtures of schedules and for $H$ ; does not separate $H$ from $F$
Chen et al. (2026, Def. 22, Thms. 24–25)	output TV. Samplers: randomized adaptive conditional-marginal queries, unrestricted output	at most $T$ queries with probability 1; worst case over a known family, the target in it unknown	family-level; does not constrain a fixed known target

Appendix

Table 17

result	ordinary proof	kernel-checked declarations	not kernel-checked
Lemma 1 (a)–(e)	App. B	(a) A36RW.Relabel.KLx_eq_neg_log_add and, for $q(V)=0$ , A36RW.Lemma1aZero.lemma1a_zero ; (b) A36RW.W11.Fext_le_Hext , A36RW.W11.Hext_eq_add ; (c) AGF.agfv_succMass_eq_snOf ; (d) A36RW.Relabel.KLx_schedLawE_profile ; (e) A36RW.Relabel.lemma1e_balanced	—
( Env )	App. D	A36Sharp.sharp_envelope , A36Sharp.sharpKap_lt_36	—
Lemma 3	App. E	(a) A36RW.Balanced.lemma2a_printed , A36.GrRec_mono ; for fixed schedules: A36.succMass_qSched_le_GrRec ; (b) A36RW.ExecMean.lemma2b_tail , A36RW.Converse.deadline_le_GrRec , A36RW.Converse.margUnit_balanced , A36RW.Converse.policy_outBy_le_succ , A36RW.Converse.policy_outBy_nonneg , A36RW.Converse.policy_doneBy_le_one	†
Theorem 4 (a)	App. E	A36RW.Balanced.hardcap_optimum_outLim ; attained by the randomly ordered rule: A36RW.Attain.sigmaStar_KL , A36RW.Attain.sigmaStar_capAS , A36RW.Converse.outLim_eq_outBy_of_capAS , A36RW.Converse.margUnit_balanced	†
Theorem 4 (b)	Apps. D , E.1	A36RW.Window.window_two_sided , A36RW.Window.thm2b_upper	†
Theorem 2 (a)	App. B	AGFH.Hopt_ge_three_unfolded , AGFCert.Hopt_family_lb	—

Appendix

Table 18

Explore similar work

Lower Bounds for Parallel Diffusion Sampling
Oct 6, 2026Yiwen Kou, Yimeng WangDiffusion Sampling Query Complexity
UCLA
Conditional Total Correlation and the Serial Depth of Adaptive Parallel Sampling
Aug 26, 2026Chuling Wen, Weijie Liang, Jian LuAdaptive Sampling Parallel Decoding
Shenzhen Key Laboratory of Advanced Machine Learning and Applications, School of Mathematical Sciences, Shenzhen University, Shenzhen 518060, China · National Center for Applied Mathematics Shenzhen (NCAMS), Shenzhen 518055, China
The Cost of Conservation: Coordination-Memory Laws for Exact-Support Generation
Aug 17, 2026Zhen Zhang, Amr Alanwar
School of Computation, Information and Technology Technical University of Munich, Germany

Oct 6, 2026, cs.DSJ/K move · Enter open · S save

Lower Bounds for Parallel Diffusion Sampling

Yiwen Kou, Yimeng Wang

UCLA

Diffusion Sampling Query Complexity

Open paper AskPDF arXiv

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

\widetildeΩ(d^{1/3})

-round lower bound for sampling smooth, near-isotropic Gaussian mixtures in

R^d

, and (2) an

Ω(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.

Abstract

Explore similar work

Lower Bounds for Parallel Diffusion Sampling

Conditional Total Correlation and the Serial Depth of Adaptive Parallel Sampling

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