Spherical Brownian Diffusion for Language Modeling

Andrea Miele; Ilija Bogunovic

cs.LGOct 26, 2026

Spherical Brownian Diffusion for Language Modeling

Authors: Andrea Miele, Ilija Bogunovic

Organizations: University of Basel

Abstract

Continuous diffusion language models involve two coupled design choices: how to corrupt continuous token representations and how to invert that corruption at generation time. Recent work has explored language generation on the unit sphere through geodesic and von Mises--Fisher (vMF) paths, raising the question of which forward process to use on this geometry and how much that choice matters relative to inference. We introduce Spherical Brownian Diffusion, using Brownian motion as the canonical isotropic diffusion on the sphere. Its heat semigroup provides a principled forward process, an additive dimension-free noise clock, and an exact ancestral bridge, while a clean-token posterior learned with cross-entropy supports multiple samplers. To separate corruption from inference, we derive high-dimensional approximation rates between Brownian and moment-matched vMF corruption and use this relationship for controlled substitutions. On Sudoku, Brownian and moment-matched vMF give nearly identical results across four samplers and five budgets under otherwise identical settings. In contrast, with the corruption process fixed, sampler choice produces a 25.0-point spread at a sampler budget of 180 calls. Using the ancestral bridge, we split this gain into a memory effect from re-noising the state, present on Sudoku and GSM8K, and a token-selection effect on GSM8K and largest at small budgets. Across reasoning tasks, categorical refinement improves performance, whereas unconditional text generation benefits from stochastic updates that preserve diversity.

Figures & tables

Figure 1: Overview of SBD. Training (top): token embeddings are corrupted with the spherical heat kernel, and a Transformer predicts the clean-token posterior with cross-entropy. Sampling (bottom): the posterior either builds the score for reverse SDE/ODE updates or drives categorical refinement, which selects tokens and re-noises their codewords.

$d=128$
Kernel	Residual
Heat	$\leq.001$
moment vMF	$.001$ – $.025$
Mix $.5$	$2.0$ – $16$
Mix $.9$	$9.2$ – $98$
$d=768$

Table 1: Shape-residual ranges.

Figure 2: GSM8K consistency restoration at

B=256

(257 NFEs, left) and mismatch by program length (right), both with posterior decoding. Error bars show three inference seeds; fine-tunes use one training seed per arm.

Task ( $B$ )	Memory	Token
Sudoku (16)	$+6.4\pm 1.0$	$-1.3\pm 1.0$
Sudoku (180)	$+18.4\pm 1.8$	$-1.4\pm 1.0$
GSM8K (16)	$+3.6\pm 0.3$	$+11.4\pm 1.0$
GSM8K (256)	$+9.5\pm 1.4$	$+4.7\pm 2.0$

Table 2: Paired q-argmax gains from removing state memory or using argmax token selection

Model	Easy	Medium	Hard
AR (greedy)	14.6	5.1	1.0
MDLM	92.0	77.1	30.2
Duo	96.3	84.7	58.4
CANDI	79.3	45.9	16.7
FLM	94.2	82.7	44.5
S-FLM	94.8	85.2	45.0

Table 3: Sudoku whole-grid accuracy (%) at

B=180

sampler calls (gated SBD: 181 NFEs). FMLM+ and SBD: mean

\pm

std. over five sampler seeds.

Figure 3: GSM8K accuracy vs. NFE. (a)

T=1

. (b) q-argmax for SBD/S-FLM,

T=0.1

for MDLM/Duo, endpoint argmax for FMLM+.

\dagger

: published vMF PC. SBD shows mean

\pm

std. over three seeds; baselines are single runs.

Figure 4: OpenWebText frontier at 1,024 NFEs. Lower generative perplexity is better at matched unigram entropy.

Appendix figures & tables59 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Fixed-noise heat-vMF diagnostic at

r=1

, which is distinct from the critical scaling in Proposition 1 . Left: at fixed

r

, the moment-matched vMF radial score becomes a better local approximation to the heat score as dimension increases. Right: both discrepancies are normalized by their values at

d=16

. The RMS relative radial-score error decreases approximately as

d^{-1/2}

, while the conditional channel KL remains nearly constant. This does not contradict Equation 29 : at fixed

r

the retained cosine is constant, whereas the proposition follows

r_{d}=(\log d+\xi)/2

, for which the retained cosine is

m/\sqrt{d}

Method	Forward family	Noise parameterization	Semigroup / path
S-FLM	Geodesic interpolation	Interpolation schedule $\alpha_{t}$	Designed path
vMF	vMF family	Concentration $\kappa_{t}$	Chosen concentration schedule
Ours	Spherical heat kernel	Retained-cosine clock $r=-\log\rho$	Heat semigroup

Appendix

Table 4: Comparison with S-FLM [ 17 ] and vMF spherical flows [ 10 ] . Brownian motion defines a Markov heat semigroup, while the other methods specify a path or concentration schedule.

Figure 6: Numerical validation of the Brownian clock and forward sampler. (a) For

d\in\{32,128,768\}

, forward samples recover

\mathbb{E}\langle\mathbf{Z}_{r},\mathbf{e}\rangle=e^{-r}

with relative error from

0.002\%

0.1\%

wherever the target is above the Monte Carlo resolution. Because

r=-\log\rho

, this checks the heat tables and sampler. (b) At

d=128

, direct corruption at

r_{1}+r_{2}

and sequential corruption at

(r_{1},r_{2})=(0.5,0.5)

differ by KS distance

D=0.0069

, compared with a null floor of

0.0011

from two independent direct draws. Across all tested dimensions and seven radius pairs,

D\leq 0.023

. The exact Brownian semigroup identity is therefore only approximate in the numerical sampler.

Task	Heat model	Evaluation set	Restart protocol
Hard Sudoku	8-layer DiT, width 512; frozen $d=128$ codebook; 20k updates, batch 256	1,000 tuning and 2,000 disjoint evaluation puzzles	$B=128$ ; seed regimes specified per table
GSM8K	Heat-trained online 250k checkpoint; context length 512	All 1,319 test problems	$B=64$ ; three paired inference seeds

Appendix

Table 5: Model and evaluation configuration for the Heat restart experiments. Sudoku uses

B=128

sampler calls; posterior decoding in the bridge comparison gives 129 total NFEs.

Sampler	Ungated	$\Delta$	Gated (tuned)	$\Delta$
q-argmax	$77.1\pm 2.1$	–	$82.1\pm 1.3$	–
Bridge, argmax token	$59.4\pm 3.1$	$-17.8\pm 1.6$	$79.1\pm 1.9$	$-3.0\pm 1.1$
Bridge, sampled token	$60.4\pm 2.2$	$-16.8\pm 1.3$	$79.1\pm 1.3$	$-3.0\pm 0.6$
Memoryless, $T=1$	$77.5\pm 2.5$	$+0.4\pm 0.8$ (n.s.)	–	–

Appendix

Table 6: Hard-Sudoku exact-match accuracy (%) on the 2,000-puzzle evaluation split. Each sampler uses

B=128

sampling calls plus one final posterior-decoding call (129 network calls total). Four paired training seeds, with sampler outputs paired within seed;

\Delta

columns are paired differences from q-argmax. The main-text memory gain uses the completed five-seed comparison instead.

$\Delta$	Retained cosine $e^{-\Delta}$	Argmax token	Sampled token ( $T=1$ )
$0$ (bridge)	$1.00$	$59.0$	$60.0$
$0.5$	$0.61$	$75.9$	$76.7$
$1$	$0.37$	$76.5$	$77.8$
$2.43\;(\frac{1}{2}\log d)$	$0.09$	$76.9$	$78.0$
$8$ (near-memoryless)	$\approx 0$	$76.6$	$78.8$

Appendix

Table 7: Hard-Sudoku restart sweep on the 1,000-puzzle tuning split at

B=128

sampler calls, four sampler seeds, without gating.

Diagnostic	Memoryless	Bridge	Restart $\Delta=0.5$
KS versus forward cosine marginal	0.004	0.006	0.005
Codeword margin	0.115	0.115	0.116
Calibration error on trajectory states	0.050	0.070	0.053
Cell accuracy after $r\approx 1.9$ (%)	95.0	92.7	94.6

Appendix

Table 8: Hard-Sudoku trajectory diagnostics on 1,000 tuning puzzles at

B=128

, averaged over the informative noise band and three Heat training seeds. KS uses cosine to the nearest codeword against forward samples of the true solution; the reference codeword margin is

0.115

Diagnostic	Memoryless	Bridge	Bridge-fresh	Restart $\Delta=0.5$
Valid grids (%)	74.8	58.6	73.1	74.8
Wrong $\to$ right	0.291	0.168	1.390	0.273
Right $\to$ wrong	0.252	0.148	1.351	0.233
Net revision	0.040	0.020	0.039	0.040
Wrong-cell confidence, trajectory / redrawn	.976/.976	.968/.969	.975/.975	.976/.976
Calibration error, trajectory / redrawn	.052/.052	.068/.067	.055/.055	.052/.052

Appendix

Table 9: Hard-Sudoku memory intervention on the same 1,000 tuning puzzles, three Heat training seeds and

B=128

. Bridge-fresh retains the previous token choice but redraws the noise. Flip counts are per cell over the informative band; net revision is wrong-to-right minus right-to-wrong.

$B$ (Sudoku / GSM8K)	Sudoku			GSM8K
	Bridge	q-argmax	Difference	Bridge	q-argmax	Difference
16	58.8	65.1	$-6.4\pm 1.0$	14.6	18.2	$-3.6\pm 0.3$
32	59.1	69.2	$-10.1$	15.3	20.9	$-5.5$
64	58.9	73.7	$-14.8\pm 0.8$	15.5	23.6	$-8.1\pm 0.6$
128	59.3	77.0	$-17.7\pm 1.4$	16.6	25.6	$-9.0$
180 / 256	59.5	78.0	$-18.4\pm 1.8$	17.6	27.1	$-9.5\pm 1.4$

Appendix

Table 10: Bridge versus q-argmax across sampler budgets

B

. Sudoku uses 1,000 hard tuning puzzles and five sampler seeds; GSM8K uses all 1,319 test questions, the Heat 250k checkpoint, three paired inference seeds and common prior draws. Values are mean accuracy (%); differences are bridge minus q-argmax. The last row uses

B=180

for Sudoku and

B=256

for GSM8K.

Task ( $B$ )	q-argmax $-$ bridge	q-argmax $-$ $T=1$
Sudoku (16)	$+6.4\pm 1.0$	$-1.3\pm 1.0$
Sudoku (64)	$+14.8\pm 0.8$	$-0.6\pm 0.8$
Sudoku (128)	$+17.7\pm 1.4$	$-0.5\pm 0.7$
Sudoku (180)	$+18.4\pm 1.8$	$-1.4\pm 1.0$
GSM8K (16)	$+3.6\pm 0.3$	$+11.4\pm 1.0$
GSM8K (64)	$+8.1\pm 0.6$	$+8.5\pm 1.3$

Appendix

Table 11: Paired token-selection effect, measured as q-argmax minus memoryless sampling at

T=1

(percentage points; mean

\pm

sd). Sudoku uses 2,000 evaluation puzzles, five training seeds, and posterior decoding (NFE

=B+1

). GSM8K uses all 1,319 test questions, three paired inference seeds, and common prior draws. Negative values favor

T=1

sampling.

Sampler	Accuracy (%)	Difference (points)
q-argmax	23.63	–
Exact bridge	15.49	$-8.1\pm 0.6$
Bridge-fresh	13.34	$-10.3\pm 1.3$

Appendix

Table 12: GSM8K memory intervention at sampler budget

B=64

on all 1,319 test questions: one Heat 250k checkpoint, three paired inference seeds and common prior draws. Differences are paired against q-argmax.

$\Delta$	GSM8K accuracy (%)	Difference from q-argmax (points)
$0$ (exact ancestral)	$15.49$	$-8.14\pm 0.61$
$0.5$	$21.89$	$-1.74\pm 0.39$
$8$	$23.60$	$-0.03\pm 0.95$
$\infty$ (q-argmax)	$23.63$	–

Appendix

Table 13: GSM8K accuracy for the Heat restart sweep at

B=64

. Results use the 250k Heat checkpoint and three paired inference seeds.

NFE	Duo $T=1$	MDLM $T=1$	SBD memoryless $T=1$	Duo $T=.1$	MDLM $T=.1$	SBD q-argmax
16	$10.0$	$1.8$	$6.80$	$22.7$	$14.3$	$18.15$
32	$13.9$	$5.7$	$9.70$	$28.2$	$23.0$	$20.85$
64	$16.6$	$10.0$	$15.14$	$34.3$	$27.5$	$23.63$
128	$15.4$	$13.5$	$19.99$	$32.2$	$31.5$	$25.60$
256	$18.5$	$15.1$	$22.39$	$35.3$	$33.1$	$27.12$

Appendix

Table 14: GSM8K accuracy (%) by NFE and decoding temperature. MDLM and Duo values are single runs; SBD values are means over three inference seeds.

NFE	Memoryless, $T=1$	q-argmax
1	$0.0$	$0.7$
4	$1.7$	$8.3$
16	$6.8$	$18.2$
64	$15.1$	$23.6$
256	$22.4$	$27.1$
512	$26.2$	$30.0$

Appendix

Table 15: GSM8K accuracy (%) for memoryless SBD at

T=1

and SBD q-argmax across the full evaluated ladder. Means are over three inference seeds.

	Updates/s		Duration (h)		GPU-hours
Model	TinyGSM	OWT	TinyGSM	OWT	TinyGSM	OWT
AR	5.4	6.0	12.8	46.6	102.1	745.7
MDLM	3.0	3.5	22.8	80.3	182.7	1,284.5
Duo	3.0	3.5	22.8	80.3	182.7	1,284.5
CANDI	3.0	2.5	23.5	108.9	188.3	1,742.9
FLM	3.2	3.3	21.7	84.2	173.6	1,346.8

Appendix

Table 16: Training cost on TinyGSM (250k updates) and OpenWebText (1M updates). Baseline values are from Table 3 of Deschenaux and Gulcehre [17] . Our duration and GPU-hours for TinyGSM refer to the complete run. OWT GPU-hours are implied by the reported 4.219 updates/s over 1M updates on 16 GPUs; the full-run wall time was not recorded. Throughput is measured over training segments, not averaged over the full run. Baselines use H100s and our runs use GH200s. NR indicates an unreported full-run duration.

Method	Latency (s/sample)	Throughput (tokens/s)	Peak memory (GiB)
MDLM	3.40	979	3.73
Duo	3.21	1,259	3.73
S-FLM	3.62	620	5.59
Ours: q-argmax	3.37	1,221	2.31
Ours: Heun SDE	7.54	257	2.73
Ours: OWT hybrid, $k=1$	7.03	277	5.80

Appendix

Table 17: Sampling cost on one NVIDIA GH200 GPU at 256 NFEs, sequence length 1,024, and

T=1

where applicable. All configurations use two warmups and seven timed repeats. Latency uses batch size 1; throughput and peak memory use batch size 4.

Setting	Sudoku	TinyGSM $\to$ GSM8K	OpenWebText
Backbone	8 layers, width 512, 8 heads	12 blocks, width 768, 12 heads; conditioning width 128	Same as GSM8K
Parameters	26.2M	130.8M + 37.7M codebook	131.6M + 38.6M codebook
Codebook $d$	128	768	768
Tokenizer	Sudoku digits	SmolLM-135M; 49,152 + PAD	GPT-2; 50,257 + 1
Sequence length	180	512	1,024, wrapped
Training updates	20k	250k	1M

Appendix

Table 18: Main-model training configuration. Sequence length includes the prompt in conditional tasks. Language parameter counts separate the backbone from the trainable codebook; all models omit self-conditioning.

Difficulty	Sampler / $B$	16	32	64	128	180
Easy	SDE	$95.32\pm.26$	$96.93\pm.28$	$97.20\pm.18$	$97.32\pm.33$	$97.10\pm.15$
	ODE	$96.33\pm.63$	$96.38\pm.20$	$96.24\pm.21$	$96.54\pm.15$	$96.42\pm.21$
	q-argmax	$97.87\pm.18$	$98.24\pm.10$	$98.50\pm.08$	$98.68\pm.11$	$98.69\pm.07$
	gated q-argmax	$98.85\pm.11$	$98.94\pm.09$	$98.96\pm.13$	$99.14\pm.11$	$99.11\pm.12$
Medium	SDE	$84.83\pm.13$	$88.70\pm.48$	$89.15\pm.57$	$89.35\pm.53$	$89.75\pm.40$
	ODE	$86.82\pm.38$	$87.32\pm.95$	$87.29\pm.25$	$87.92\pm.22$	$88.15\pm.57$

Appendix

Table 19: Sudoku exact-match accuracy (%). The solver grid is chosen per sampler on held-out puzzles; SDE/ODE use the log-radius grid below

B=64

. Low-budget SDE/ODE entries average three sampler seeds; all other entries average five sampler seeds on one checkpoint per difficulty. Error bars do not include training variance. Column headings are sampler calls

B

; gated posterior decoding uses

B+1

total NFEs, and the other rows use

B

Study	Accuracy	Training	Grid	Decode	Seeds
Benchmark	$79.70\pm.75$	online, 20k	information	none, 180 calls	5 sampler; 1 checkpoint
Heat–vMF contrast	$73.88\pm.41$	20k updates	log-radius	q-argmax	3 paired training seeds
Kernel Heat diagonal	$77.9\pm 1.8$	online, 20k	information-matched	posterior, 181 calls	5 training seeds
Frozen-codebook ablation	$79.68\pm.42$	online, 20k	information	posterior, 181 calls	3 sampler; 1 checkpoint

Appendix

Table 20: Protocol provenance for four Heat q-argmax results (%) on 2,000 hard puzzles at a sampler budget of

B=180

calls. Standard deviations reflect the seed type in the last column. Gated runs include a final posterior decode (+1 call).

Difficulty	SDE log	SDE information	ODE log	ODE information
Easy	$97.10\pm 0.41$	$97.20\pm 0.18$	$96.47\pm 0.23$	$96.24\pm 0.21$
Medium	$89.44\pm 0.31$	$89.15\pm 0.57$	$88.21\pm 0.48$	$87.29\pm 0.25$
Hard	$62.97\pm 0.76$	$62.27\pm 0.73$	$58.34\pm 1.02$	$58.50\pm 0.88$

Appendix

Table 21: NFE-64 transport control with nearest-codeword decoding: exact-match accuracy (%), mean

\pm

sd over five sampler seeds per grid on one fixed Heat checkpoint per difficulty and 2,000 puzzles.

NFE	Easy	Medium	Hard
4	$97.60\pm.09$	$91.53\pm.39$	$70.36\pm.78$
16	$97.87\pm.27$	$92.97\pm.54$	$81.35\pm.55$
64	$97.87\pm.29$	$92.99\pm.42$	$81.44\pm.71$
180	$97.77\pm.13$	$92.99\pm.33$	$\underline{81.65\pm.36}$

Appendix

Table 22: FMLM+ Sudoku exact-match accuracy (%), mean

\pm

standard deviation over five sampler seeds on fixed released checkpoints.

Figure 7: Complete Sudoku accuracy curves for easy, medium, and hard puzzles. Error bars are standard deviations over five sampler seeds on fixed checkpoints, not training seeds. Comparing gated SBD with FMLM+, FMLM+ leads on hard Sudoku at

B=16

, they are tied within variation at

B=64

, and SBD leads at

B=180

. The plotted budget labels are sampler calls

B

; gated SBD’s posterior decode adds one call.

Sampler	Valid	Legal	Wrong puzzle	Rows	Columns	Boxes
SDE	63.25	68.96	5.71	75.01	75.77	81.00
ODE	58.24	62.80	4.56	70.32	71.53	77.75
q-argmax	79.70	86.94	7.24	89.39	89.70	91.92
gated, nearest	82.59	82.66	0.07	85.87	86.27	89.01
gated, posterior	83.36	85.00	1.64	87.98	88.07	90.56

Appendix

Table 23: Hard-Sudoku failure decomposition (%) at

B=180

sampler-call budget, averaged over five sampler seeds. The posterior-decoded gated row uses 181 actual calls. “Wrong puzzle” is legal but inconsistent with a clue.

Training kernel	ODE	SDE	q-argmax	Gated q-argmax
Heat	$58.42\pm 0.44$	$62.83\pm 1.16$	$73.88\pm 0.41$	$83.42\pm 0.32$
Moment-matched vMF	$58.81\pm 0.47$	$63.11\pm 1.14$	$73.92\pm 0.50$	$82.61\pm 0.31$
$\Delta$ (vMF $-$ heat)	$+0.39$	$+0.28$	$+0.04$	$-0.81$

Appendix

Table 24: Hard-Sudoku accuracy (%) at

B=180

sampler calls across three independent training seeds. The Heat and moment-matched vMF arms use an identical paired protocol, with inference settings matched within each procedure. Entries are mean

\pm

standard deviation. Posterior-decoded gated rows use 181 actual calls.

SLERP (full): $\|\Delta u\|\leq 0.016$ relative to Heat;
$u$	Heat	vMF	Mix $.5$	Mix $.9$
.05	1.05	1.03	.47	.056
.25	1.31	1.31	.81	.31
.50	1.56	1.55	1.18	.73
.90	2.39	2.38	2.51	2.30
separate truncation ablation: $r_{\min}=0.942$ .

Appendix

Table 25: Radius at remaining token-entropy fraction

u

. A smaller radius reveals the token later during generation. The full SLERP path matches Heat within

|\Delta u|\leq 0.016

. This comparison uses the full-path clock; the truncated variant is evaluated separately. Section 4.2 reports SLERP truncated at

r_{\min}=0.942

; Table 33 compares both paths.

$r$	Heat	moment vMF	mix $.5$	mix $.9$
.05	$7.2\times 10^{-4}$	$3.6\times 10^{-3}$	1.00	6.49
1.0	$7.9\times 10^{-5}$	$1.9\times 10^{-2}$	.87	3.69
1.6	$4.6\times 10^{-5}$	$1.3\times 10^{-2}$	.84	4.33
3.0	$3.2\times 10^{-5}$	$9.8\times 10^{-3}$	.80	1.41

Appendix

Table 26: Raw

T_{0}^{(1)}

: the

L^{1}

estimator in Equation 55 , using 20,000 draws per radius at

d=128

. The mixture estimates are most informative at the level of order of magnitude because of estimator variability.

Kernel	$r$	$T_{0}^{(1)}$ , repeats	$T_{0}^{(2)}$ , repeats	Tail share
vMF	1.0	.019 / .019	.024 / .024	.001
Mix $.5$	1.0	.87 / .87	.99 / .99	.90–.94
Mix $.9$	.05	7.40 / 7.34	2.46 / 1.51	.84–.96
Mix $.9$	1.0	4.14 / 4.32	1.05 / 1.12	.97–.99

Appendix

Table 27: Compatibility estimators recomputed on identical samples, with two repeats per setting at

d=128

. The last column is the fraction of

\mathbb{E}[V^{2}w]

contributed by the largest 0.1% of samples.

$B$	Heat	vMF	SLERP (full)	Mix $.5$	Mix $.9$
180	$54.5\pm 3.2$	$54.7\pm 2.3$	$56.2\pm 2.2$	$53.3\pm 2.4$	$55.3\pm 2.3$
720	$54.4\pm 1.6$	$55.2\pm 2.1$	–	$53.2\pm 2.6$	$55.6\pm 2.2$

Appendix

Table 28: Hard-Sudoku accuracy (%) using each kernel’s valid quantile ODE, mean

\pm

sd over three training seeds for every kernel. SLERP uses its full path here; Table 33 also reports the truncated path used in Section 4.2 . Posterior decoding gives NFE

=B+1

. We did not run SLERP at

B=720

Kernel	$d$	Raw $T_{0}^{(2)}$ ( $L^{2}$ )	$\alpha^{*}$	Residual	KS
Heat	128	$\leq.001$	1.000	$\leq.001$	$\leq.012$
Heat	768	$\leq.007$	1.000	$\leq.007$	$\leq.030$
moment vMF	128	$.004$ – $.025$	$.993$ – $1.012$	$.001$ – $.025$	$.0001$ – $.017$
moment vMF	768	$.004$ – $.012$	$.993$ – $1.012$	$.000$ – $.010$	$.0001$ – $.022$
info-vMF	768	$.004$ – $.45$	$.993$ – $1.45$	$.000$ – $.010$	$.0001$ – $.022$
SLERP (S-FLM path)	128	–	$1.003$	$.047$	–

Appendix

Table 29: Compatibility estimates across 13 radii from

.05

5

. Raw

T_{0}^{(2)}

is the

L^{2}

estimator with 50,000 draws per radius; Table 26 reports the corresponding

L^{1}

-based estimate. Mixture

L^{2}

ratios and residuals are tail-sensitive order-of-magnitude estimates. SLERP’s residual is a band average, not a range.

Figure 8: Mismatch cost and the local shape residual. Horizontal bars span the residual range over tested noise levels. Sudoku uses the valid quantile ODE at

B=180

with posterior decoding (181 NFEs); only the mixture points use information relabelling. vMF and SLERP use their native clocks. GSM8K uses q-argmax. The residual is local in

r

, so the figure shows ordering instead of a fitted relationship. Negative values mean better accuracy than Heat. Differences as small as the vMF/SLERP costs are not resolved with three training seeds; this does not apply to the larger mixture losses. Mixture residuals are tail-sensitive and indicate orders of magnitude. SLERP uses Heat-trained weights with the matched geodesic quantile ODE; its native exact velocity and radial transport agree numerically (Appendix E.4 ).

Kernel	Figure residual	Cost vs. Heat (points)
Moment-matched vMF	$.017$	$+.08\pm.08$
SLERP (S-FLM path)	$.047$	$+.38\pm.20$
Mixture $\lambda=.5$	$5.5$	$+4.8\pm.3$
Mixture $\lambda=.9$	$26.9$	$+9.6\pm.8$

Appendix

Table 30: Kernel mismatch costs shown in Figure 8 on hard Sudoku. Positive values mean worse accuracy than Heat. Costs are paired mean

\pm

sd over three Heat training seeds at

B=180

with posterior decoding (181 NFEs). Only mixtures use time relabelling; their residuals are order-of-magnitude estimates.

Kernel	Dynamics	Nearest	Posterior
Mix $.5$	Brownian ODE	$44.8\pm 1.4$	$55.0\pm 2.4$
	quantile ODE	$52.1\pm 2.3$	$53.3\pm 2.4$
Mix $.9$	Brownian ODE	$0.0\pm 0.0$	$15.3\pm 3.6$
	quantile ODE	$11.2\pm 0.7$	$55.3\pm 2.3$

Appendix

Table 31: Hard-Sudoku accuracy (%) at a sampler budget of

B=180

network calls (Brownian ODE: Heun, 90 steps; quantile ODE: 180 steps), mean

\pm

sd over three training seeds. Nearest decoding spends 180 NFEs; posterior decoding adds one call (181 NFEs).

$B$	Moment-matched vMF	SLERP	Mix $\lambda=.5$	Mix $\lambda=.9$
16	$+0.3\pm 0.8$	–	$-4.7\pm 2.3\;(t=-4.5)$	$-5.0\pm 4.1\;(t=-2.7)$
64	$+0.2\pm 1.1$	$+1.1\pm 3.6$	$-5.2\pm 2.4\;(t=-4.8)$	$-5.3\pm 3.6\;(t=-3.3)$
180	$-0.1\pm 1.0$	$+0.1\pm 3.4$	$-5.8\pm 1.8\;(t=-7.4)$	$-3.8\pm 3.0\;(t=-2.8)$

Appendix

Table 32: Effect of the training kernel on tuned confidence-gated refinement on hard Sudoku. Entries are accuracy differences in percentage points relative to Heat on the information grid, mean

\pm

standard deviation over five independent training seeds per kernel. Parentheses give the paired

t

statistic. SLERP is a separate three-training-seed comparison with paired spreads; see Table 33 . Budgets are sampler calls

B

, with one further call for posterior decoding.

Sampler	$B$	Heat	SLERP	SLERP truncated
ODE	16	$53.1\pm 2.2$	$55.0\pm 1.6\;(+2.0\pm 2.0)$	$51.2\pm 2.7\;(-1.8\pm 4.8)$
ODE	180	$54.5\pm 3.2$	$56.2\pm 2.2\;(+1.8\pm 3.3)$	$51.4\pm 1.9\;(-3.1\pm 4.8)$
q-argmax	64	$65.3\pm 1.5$	$67.4\pm 1.3\;(+2.1\pm 2.6)$	$66.8\pm 1.7\;(+1.5\pm 3.2)$
q-argmax	180	$71.5\pm 2.5$	$73.1\pm 1.3\;(+1.6\pm 2.5)$	$71.9\pm 1.6\;(+0.3\pm 3.9)$
Gating, tuned	64	$78.3\pm 2.0$	$79.5\pm 2.1\;(+1.1\pm 3.6)$	$77.7\pm 0.9\;(-0.7\pm 2.7)$
Gating, tuned	180	$81.7\pm 1.1$	$81.8\pm 3.0\;(+0.1\pm 3.4)$	$78.4\pm 1.4\;(-3.2\pm 2.5)$

Appendix

Table 33: SLERP kernel comparison on hard Sudoku. Values are accuracies in percent, mean

\pm

sd over three independent training seeds; parenthesized values are paired differences from Heat, also mean

\pm

sd. All runs use the same eight-layer backbone, 20k updates, batch 256, 2,000 evaluation puzzles, and posterior decoding, so total NFE is

B+1

. The truncated gating difference at

B=180

has paired

t=-2.3

and is not resolved with three seeds.

Autoregressive
Model	Easy	Medium	Hard
AR (sample)	13.9	5.1	0.6
AR (greedy)	14.6	5.1	1.0
Discrete diffusion
MDLM	92.0	77.1	30.2
Duo	96.3	84.7	58.4

Appendix

Table 34: Sudoku exact-match accuracy (%) at

B=180

sampler-call budget. AR, MDLM, Duo, CANDI, FLM, and S-FLM are published values. FMLM+ and SBD use five sampler seeds on fixed checkpoints; gated SBD uses 181 calls. Bold/underline mark the best/second-best non-AR value per column.

Figure 9: Enhanced GSM8K comparisons using teacher initialization or self-conditioning. Our model uses neither; the controlled comparison is in Figure 3 .

Model	Accuracy
Autoregressive
AR (sample)	53.9
AR (greedy)	63.3
Discrete diffusion
MDLM	16.91
Duo	18.95

Appendix

Table 35: GSM8K accuracy (%) at the full reported budget. Diffusion/flow models use the reported 1,024-call budget; decoding regimes differ: MDLM/Duo use

T=1

, SBD uses q-argmax, and S-FLM top-1 and FMLM+ use deterministic token selection. vMF uses its published PC setting. SBD is the mean over three inference seeds; AR uses 512 sequential passes. Bold/underline mark the best/second-best non-AR result.

\dagger

: pretrained assistance;

\ddagger

: self-conditioning with predictor-corrector sampling. See Appendix E.6 for provenance and comparability.

Autoregressive
Model	1	2	4	8	16	32	64	128	256	512	1024
AR (sample)	NR	NR	NR	NR	NR	NR	NR	NR	NR	53.9	NR
AR (greedy)	NR	NR	NR	NR	NR	NR	NR	NR	NR	63.3	NR
Discrete diffusion and distillation
MDLM	0.00	0.00	0.00	0.23	1.82	5.69	10.01	13.57	15.09	14.94	16.91
Duo	0.00	0.38	1.67	3.64	10.01	13.87	16.60	15.39	18.50	18.27	18.95

Appendix

Table 36: GSM8K final-answer accuracy (%) across network-evaluation budgets. Dense rows are evaluations of released 250k-update checkpoints, excluding teacher/init pretraining costs. vMF PC (with/without SC) and IDLM values are taken from their papers. SBD q-argmax reports a mean and a separate standard-deviation row over three inference seeds; other dense rows are single runs.

\dagger

: pretrained teacher or initialization;

\ddagger

: self-conditioning and PC. NR means that a value is not reported. AR uses 512 sequential passes. Bold marks the best non-autoregressive result at each budget.

NFE	1	2	4	8	16	32	64	128	256	512	1024
Accuracy (%)	0.8	3.3	8.2	14.9	20.5	23.2	26.2	28.7	29.3	30.5	31.2

Appendix

Table 37: GSM8K accuracy of q-argmax on the released S-FLM S-arch checkpoint, using SLERP re-noising at the model’s own time labels. One evaluation seed.

Weights $\to$ kernel	q-argmax		quantile ODE
	$B=64$	$B=256$	$B=64$	$B=256$
Heat $\to$ Heat	$23.63\pm.46$	$27.12\pm 1.41$	$6.92\pm.61$	$7.56\pm.61$
Heat $\to$ moment vMF	$22.74\pm.08$	$27.32\pm.24$	$6.85\pm.69$	$7.48\pm.69$
Heat $\to$ mixture $.5$	$0.05\pm.09$	$0.03\pm.04$	$0.40\pm.16$	$0.45\pm.13$
Heat $\to$ mixture $.9$	$0.00$	$0.00$	$0.00$	$0.00$
Info-vMF $\to$ info-vMF	$23.05\pm.85$	$27.24\pm.76$	$7.15\pm.16$	$7.43\pm.27$

Appendix

Table 38: GSM8K accuracy (%), mean

\pm

standard deviation over three inference seeds. Each row names training weights and inference kernel. Budgets are sampler calls

B

; posterior decoding adds one call, giving 65 and 257 total NFEs.

$B$	Heat $\to$ Heat	info-vMF $\to$ info-vMF	Difference
1	0.66	0.96	$+0.30$
2	3.94	4.20	$+0.25$
4	8.34	8.82	$+0.48$
8	14.99	14.25	$-0.73$
16	18.15	18.12	$-0.03$
32	20.85	19.18	$-1.67\pm 0.92$ ( $t=-3.1$ , $p\simeq 0.09$ )

Appendix

Table 39: GSM8K kernel/training comparison: Heat-trained Heat inference versus information-vMF-trained information-vMF inference. Values are means over three inference seeds; differences are info-vMF minus Heat. Budgets are sampler calls

B

, before a separate posterior decode.

$r$	Heat	moment vMF	mix $.5$	mix $.5$ , best	mix $.9$	mix $.9$ , best
.2	1.000	1.000	.996	.999	.828	.971
.5	.999	.999	.994	.998	.620	.945
1.0	.996	.996	.949	.988	.381	.880
1.6	.984	.983	.860	.960	.213	.758
2.5	.936	.936	.814	.883	.153	.481

Appendix

Table 40: Code-token probe: one-step top-1 accuracy of the Heat posterior on the 9,345 program tokens in 128 TinyGSM validation answers, excluding PAD. Unqualified columns use the raw time label; “best” chooses the best tested Heat-network time label.

Length range (mean)	$n$	Heat	mix $.5$ , relabelled	ratio
67–108 (96)	259	44.3	21.0	.47
108–125 (116)	258	40.8	13.3	.33
125–142 (132)	265	26.7	6.2	.23
142–164 (152)	266	17.8	3.5	.20
164–290 (189)	271	7.3	.25	.03

Appendix

Table 41: GSM8K success by generated-program length at

B=256

with posterior decoding (257 NFEs). The same 1,319 problems define both columns; values average three inference seeds.

Posterior decoding, EMA	$B=256$ (257 NFEs)
Heat fine-tune $\to$ Heat	$27.32\pm.52$
Mix $.5$ fine-tune $\to$ mix $.5$	$26.23\pm 1.30$

Appendix

Table 42: Matched GSM8K fine-tunes (%), mean

\pm

sd over three inference seeds, with one training seed per arm. The main comparison uses posterior decoding at

B=256

(257 NFEs); nearest-codeword controls use

B

total NFEs. The posterior-decoded paired mixture-minus-Heat difference is

-1.09\pm 1.80

points.

Figure 10: Our GSM8K checkpoint and sampling ablations. All configurations use plain q-argmax without Gumbel perturbations, score updates, self-conditioning, or pretrained assistance. Missing NFE measurements are omitted; Table 43 gives the numerical results.

NFE	$d=128$ frozen default	$d=768$ learned default	$d=768$ entropy + log sampling	$d=768$ entropy + entropy sampling	$d=768$ online adaptive
1	0.38 (5)	0.15 (2)	0.53 (7)	0.53 (7)	0.76 (10)
2	0.45 (6)	0.15 (2)	0.53 (7)	4.17 (55)	4.70 (62)
4	1.59 (21)	0.68 (9)	1.21 (16)	9.63 (127)	9.70 (128)
8	NR	NR	NR	NR	14.86
16	5.99 (79)	4.62 (61)	8.95 (118)	19.48 (257)	18.50 (244)
32	9.48 (125)	11.37 (150)	14.10 (186)	21.53 (284)	21.15 (279)

Appendix

Table 43: GSM8K final-answer accuracy (%) across NFE budgets for our TinyGSM-trained checkpoints. Parentheses give the number of correct examples out of 1,319. Bold marks the best result at each budget; NR means not reported.

NFE	ODE	Heun SDE
32	3.64	0.00
64	7.81	7.20
128	7.96	10.08
256	8.34	12.21
512	8.34	11.60
1024	8.26	11.98

Appendix

Table 44: GSM8K executed-answer accuracy (%) for continuous samplers on the full 1,319-question test set. Budgets are reported NFEs. Bold marks the better sampler at each budget. This table is separate from the checkpoint and generated-trajectory ablations.

$T$	$H$	PPL	$H$	PPL	$H$	PPL
	$k=1$		$k=2$		$k=4$
0.80	4.56	16.32	4.61	20.82	4.78	26.13
1.00	5.18	35.12	5.24	43.58	5.36	59.42
1.20	5.36	65.10	5.57	136.87	5.75	230.57

Appendix

Table 45: Temperature-matched OWT

k

ablation. Each entry uses 64 samples (two seeds, 32 each) at the same nominal 1,024-NFE budget; realised NFEs are 1,024 / 1,023 / 1,024 for

k=1/2/4

H

is mean unigram entropy in nats and PPL is mean generative perplexity.

Panel of Figure 11 — Figure 11: OpenWebText frontiers at all six budgets. Our 1M-update model uses $k=1$ unless noted. Curves show measured Pareto-filtered points from 15 temperatures, with 512 samples per point except $k=2/4$ (64). Reference markers use 512 samples; $k=2$ realises 1,023 evaluations. FMLM+ $\dagger$ denotes our budget-matched extrapolation of the inner flow-step count with four refinement rounds ( $256\times 4$ at 1,024); only $32\times 4=128$ is the released OWT configuration. Lower perplexity is better at matched entropy.

Figure 12: Evolution of the OWT generative frontier during training. Curves compare checkpoints at 50k, 100k, 500k, and 1M updates, with 256 NFEs (left) and 1,024 NFEs (right). This 128-sample checkpoint diagnostic uses GPT-2 Large to measure generative perplexity; lower perplexity is better at matched unigram entropy. The 500k and 1M curves are close, although both improve over the early checkpoints across much of the shared entropy range. Sample counts differ from the main 512-sample frontier evaluation.

Task	Checkpoint	Trajectory settings
Sudoku easy/medium	8-layer DiT, embedding 512; $d=768$ , learned-online, 20k updates	$T=1$ , log-radius grid
TinyGSM / GSM8K	Online 250k checkpoint, EMA	Length 512, $T=1$ , entropy grid
OWT	Free-online 1M checkpoint, EMA	Length 1,024, entropy grid; score-temperature sweep

Appendix

Table 46: Setup for the additional commitment diagnostics. These runs are separate from the matched-NFE Sudoku accuracy sweep.

Task	Posterior entropy	$\max q>0.99$	MAP stability
Sudoku easy	$\approx 0$	100	100
Sudoku medium	$\approx 0$	99.99	100
TinyGSM	0.20	79.4	97.2
OWT	5.61	0.51	55.3

Appendix

Table 47: Posterior diagnostics on forward-noised reference examples near

r=3.04

. Entropy is in nats; confidence and adjacent-noise MAP stability are percentages. Measurements are rounded.

Task	Metric	q-argmax	Heun SDE
TinyGSM	Posterior entropy	0.12	0.32
	Positions with $\max q>0.99$	83.6	70.9
	Executed-answer accuracy	39.1	20.3
GSM8K	Executed-answer accuracy	24.2	12.5
OWT	Sample unigram entropy	$\approx 0.014$	5.31

Appendix

Table 48: Additional generated-trajectory diagnostics with 256 sampling calls plus one separate final diagnostic call (257 total evaluations). Accuracy and high-confidence fractions are percentages; both entropy metrics use nats but measure different quantities. These runs are separate from the main checkpoint sweeps.

Figure 13: Posterior readiness along generated language trajectories. Rows: TinyGSM reference non-PAD answer positions, GSM8K answer slots including PAD, and OWT at score temperature

T=0.90

. Columns: posterior entropy, the fraction with

\max q>0.99

, and adjacent-state MAP stability. Blue: q-argmax; orange: Heun SDE; green: the native OWT hybrid. Each method uses 256 sampling calls plus one separate final diagnostic call. The horizontal axis is noise level

r

; generation proceeds from right to left. Stability uses different adjacent-radius increments across samplers and is not a common-interval comparison. GSM8K confidence includes PAD and should not be read as content-only certainty.

Figure 14: Geometry of the learned spherical codebook (

d=128

). (a) Empirical CDF of nearest-neighbor (NN) angles at 5k and 150k updates, compared with frozen max-min and random uniform controls. Only 1.2% fall below the separation margin

\arccos(0.7)\simeq 0.7954

radians. (b) GPT-2 cosine alignment increases while centered effective rank (participation ratio) decreases. The dashed line marks the log-to-entropy training-schedule switch at 100k; changes after it do not isolate the schedule’s effect. This is a single-run diagnostic.

Metric	Learned	Max-min	Random
NN angle: median	0.970	1.327	1.202
NN angle: first percentile	0.784	1.294	1.128
NN angle: minimum	0.237	1.232	1.035
Codewords with NN cosine $>0.7$	624	0	0
Centered effective rank	50.4	127.9	127.7
Mean pairwise cosine	0.041	0.00005	$-0.00006$

Appendix

Table 49: OWT codebook diagnostics (

d=128

, 150k updates). Angles are in radians; rank is centered. Above-margin counts use NN cosine

>0.7

, or angle

<\arccos(0.7)\simeq 0.7954

. Top-10 overlap uses all 50,257 shared tokens.

Updates	NN median	Eff. rank	Pearson	GPT-2 top-10 (%)	Space agreement (%)
5k	1.279	127.0	$-0.017$	0.8	57.8
25k	1.160	103.5	0.050	6.7	79.1
50k	1.082	83.3	0.102	14.5	85.2
75k	1.055	76.2	0.130	18.9	86.0
100k	1.044	73.2	0.146	21.1	86.2
125k	0.970	52.6	0.167	21.3	94.8

Appendix

Table 50: Codebook organization through training in one OWT run. The noise draw changes from log to entropy at 100k; comparisons across that boundary do not isolate the effect of additional training. Angles are in radians, and overlap and form agreement are percentages.

Pair sampling	Pearson
Uniform, full vocabulary	0.176
Frequency-weighted by unigram prior	0.210
Uniform, 20,000 most frequent tokens	0.277
Uniform, 8,000 most frequent tokens	0.382
Uniform, 2,000 most frequent tokens	0.463

Appendix

Table 51: Sensitivity of GPT-2 pairwise cosine alignment to pair sampling for the 150k learned codebook. Uniform full-vocabulary sampling is the primary metric; restricted subsets are diagnostic comparisons.

Figure 15: Numerical validation of Proposition 1 along

r_{d}=(\log d+\xi)/2

for

\xi\in\{0,1,2\}

and

d=16

1024

. The score- and KL-based coefficient estimates approach the unfitted target

m=e^{-\xi/2}

; gray curves show

1/\log d

guides.

Diffusion Language Models

Diffusion language models generate and refine text through iterative denoising, enabling bidirectional context and parallel token updates.

See all 473 papers

May 11, 2026

Language Modeling with Hyperspherical Flows

Discrete Diffusion Language Models progressed rapidly as an alternative to autoregressive (AR) models, motivated by their parallel generation abilities. However, for tractability, discrete diffusion models sample from a factorized distribution, which is less expressive than AR. Recent Flow Language Models (FLMs) apply continuous flows to language, transporting noise to data with a deterministic ODE that avoids factorized sampling. FLMs operate on one-hot vectors whose dimension scales with the vocabulary size, making FLMs costly to train. Moreover, since all distinct one-hot embeddings are equidistant in

\ell_2

, adding Gaussian noise does not have a clear semantic interpretation (unlike images, where Gaussian noise progressively degrades structure). We introduce

\mathbb{S}

-FLM, a latent FLM in the hypersphere.

\mathbb{S}

-FLM generates sequences by rotating vectors in

\mathbb{S}^{d-1}

along a velocity field learned with cross-entropy, avoiding the overhead of materializing one-hot vectors. Previous FLMs match AR in Generative Perplexity (Gen.\ PPL), but samples with high likelihood are not necessarily correct in verifiable domains such as math and code.

\mathbb{S}

-FLM substantially improves continuous flow language models on large-vocabulary reasoning and closes the gap to masked diffusion under standard-temperature sampling (

T=1

), while a gap remains under optimized low-temperature (

T=0.1

) decoding.

Justin Deschenaux, Caglar Gulcehre

EPFL, Lausanne, Switzerland

Normalizing Flows Continuous Normalizing Flows Language Modeling Flow-Based Generative Modeling

Ask PDF

Oct 1, 2026

Clock Diffusion: Efficient Semi-Autoregressive Continuous Diffusion Language Models

Recent works on continuous diffusion for discrete data have demonstrated performance on par with comparable discrete diffusion models. However, these continuous counterparts lack key features that are essential to practical use as language models, namely variable-length generation and support for a key-value cache, and they still lag behind the frontier of autoregressive and discrete diffusion quality. In this work, we address these limitations. We do so by introducing a model parameterization that uses position-dependent noise schedules to define semi-autoregressive (SAR) continuous diffusion language models (DLMs). Together with efficient training and sampling algorithms, we call this framework Clock Diffusion, and we present two special cases of our method: block and sliding window generation. We then define ClockDLMs, a family of Gaussian DLMs based on sliding window Clock Diffusion that attain state-of-the-art diffusion likelihood bounds on OpenWebText, even beating the performant block SAR discrete diffusion models. ClockDLMs trained on TinyGSM also substantially outperform continuous baselines on the GSM8K benchmark and match and exceed comparable SAR discrete diffusion models. Finally, building on our parameterization, we propose more efficient samplers that we dub Cache Grab, which adapt techniques from accelerated inference in discrete diffusion, such as committing tokens whose probabilities exceed a confidence threshold and self-speculative decoding, further improving our models' quality and efficiency.

Yair Schiff, Omer Belhasin, Roy Uziel +7

NVIDIA · Cornell University

Diffusion Language Model Inference Language Modeling Diffusion Language Models Continuous Diffusion Language Models Diffusion Model Inference Acceleration Diffusion Language Model Training

Ask PDF

May 22, 2026

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

Diffusion language models intrinsically fail to capture correlations between decoded tokens, which leads to a harsh trade-off between sampling quality and throughput. To solve this issue, we propose DiLaDiff, a variant of masked diffusion language models with three components: (1) a continuous latent space with semantic capabilities, learned by an auto-encoder fine-tuned from an existing masked diffusion language model; (2) a latent diffusion model learning the prior over the encoder distribution; (3) a consistency model distilling the learned prior into a few-step latent generative model. We show that, even without distillation, our latent-guided diffusion model outperforms the masked diffusion baseline while significantly accelerating inference. Consistency distillation further lowers the computational overhead of continuous diffusion, such that the latent is generated in negligible time compared to discrete decoding.

Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3

NVIDIA

Diffusion Language Models Diffusion Language Model Distillation Diffusion Language Model Inference Masked Diffusion Models Continuous Diffusion Language Models Diffusion Model Inference Acceleration+1

Ask PDF

Abstract

Diffusion Language Models

Explore similar work

Language Modeling with Hyperspherical Flows

Clock Diffusion: Efficient Semi-Autoregressive Continuous Diffusion Language Models

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling