KL Before and After Mixing: A Round-Complexity Separation in Parallel Sampling
Organizations: Amazon
Abstract
A parallel sampler can choose at random which coordinates to generate together. For a target and the output law of a schedule drawn afresh for each output, we compare the schedule-averaged divergence with the output-law divergence , 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 with ones, the least expected number of rounds at a fixed tolerance is under the first criterion but 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--Bernoulli target, where the sharp expected cost under the output-law criterion is , attained with rare -round runs. Prior work implies the order of the schedule-averaged cost on the balanced target and its upper bounds on both extensions.
Figures & tables
| design, criterion, cost | balanced slice | full support | bounded weights |
|---|---|---|---|
| fixed schedule, KL, length | |||
| mixture, , mean | |||
| relabelled profile, , length | |||
| causal policy, , cap | |||
| mixture, , mean | |||
| causal policy, , |
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 |
| The survival envelope ( Env ) | Appendix D , Proposition 9 (its weaker form , 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 | 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 |
| class, criterion, cost | order | status |
| Balanced slice | ||
| fixed schedule, KL, length | K (Cor. 5 ) | |
| one relabelled profile, , length | K (Cor. 5 ) | |
| mixture, , mean | K (Thm 2 a,c) | |
| mixture, , mean | K (Thm 2 b; Cor. 6 b,c) | |
| , , cap | K (Thm 4 a; Cor. 6 a) | |
| design | one round | rounds |
|---|---|---|
| balanced, ladder | unrevealed | |
| full support, one profile | level | |
| full support, mixture | band |
| Symbol | Meaning | Defined in |
|---|---|---|
| Interface, criteria and costs | ||
| , | target law on ; the number of coordinates | § 1 ; § 2 |
| , | balanced slice , ; the required number of ones (also: ) | § 2 |
| count state: coordinates unrevealed, of them must still be one (also: rung index ); with a subscript, an unrevealed count at a stage: before round (§ 3 ), the rungs of Theorem 2 (b)’s ladder and of Theorem 4 (b)’s witness (rows and ). In § 4.2 the nominal sizes are deterministic and the free count is , random around ; there the of is the reservoir size of part (b)’s construction (Theorem 8 ’s keeps the count-state meaning) | § 2 | |
| ; (length); (profile) | labelled schedule , an ordered partition of ; its length (rounds); its profile (block sizes) | § 2 |
| a permutation of , as in the uniform relabelling ; is also the -algebra generated, as in . In the appendices also denotes, each defined where used, a standard deviation ( , , in Apps. F , K , G ), the logistic function (App. I ) and a state (App. H ) | § 2 | |
| shape | TV fixed / mixed | |||||
|---|---|---|---|---|---|---|
| 2.66 | 1.999 | – | – | 2 | – | |
| 8.19 | 3.28–3.39 | 9 | 89% | 4 | 0.412 / 0.005 | |
| 12.56 | 3.71–3.86 | 13 | 95% | 4 | 0.574 / 0.014 | |
| 15.14 | 3.92–4.04 | – | – | – | ||
| 18.06 | – | 19 | 98% | 0.673 / 0.015 | ||
| 26.20–33 | 1.99–4.69 | 33 | 99% | 0.831 / 0.022 |
| ratio | ratio | |||||
|---|---|---|---|---|---|---|
| 14.55 | 4.36 | 3.3 | 4.35 | 2.29 | 1.9 | |
| 29.48 | 5.59 | 5.2 | 8.03 | 2.91 | 2.7 | |
| 51.61 | 6.50 | 7.9 | 13.10 | 3.43 | 3.8 | |
| 81.45 | 7.12 | 11.4 | 19.29 | 3.82 | 5.0 | |
| 119.34 | 7.68 | 15.5 | 26.75 | 4.04 | 6.6 | |
| quantity | value | |
|---|---|---|
| K | ||
| (one schedule) | Fin | |
| (one relabelled profile) | K | |
| K | ||
| (finite brackets) | Fin |
| 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 | 26 | 3.49 | 7 |
| Result | Criterion; samplers or policies | Round count; target | Bearing on the least rounds under |
|---|---|---|---|
| Bounds on , which bound only from above, through | |||
| Chen et al. (2026, Def. 15, Thms. 4, 9, 16) | , the expected per-schedule KL: exact formula, optimal block sizes, a bound via total correlation; . Samplers: random unmasking: prescribed block sizes, uniformly random positions | number of blocks; any target | upper bounds on through |
| Secchi and Zanella (2026, Prop. 4) | factorization error: the of the exact-conditional sampler; output KL at most it plus a learning error. Samplers: -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 , random block sizes cost at least as much in the limit; enters only through (exact conditionals) or (learned ones) |
| Lower bounds on the output law | |||
| Feng et al. (2025, Thm. 4.4) | output mass off the support; above it gives . Samplers: masked-diffusion sampling with any masking schedule, each masked position revealed independently; conditionals exact or with per-token KL strictly below | sampling steps; an explicit hidden Markov target | steps for , also for mixtures of schedules and for ; does not separate from |
| Chen et al. (2026, Def. 22, Thms. 24–25) | output TV. Samplers: randomized adaptive conditional-marginal queries, unrestricted output | at most queries with probability 1; worst case over a known family, the target in it unknown | family-level; does not constrain a fixed known target |
| 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 , 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 | — |