Parallel Language Models via Variational Decoding

Felix Draxler; Robin C. Geyer; Daniel Geyfman; Farrin Marouf Sofian; Justus Will; Philipp Hildebrandt; Giosue Migliorini; Prakhar Srivastava

cs.LGOct 19, 2026

Parallel Language Models via Variational Decoding

Authors: Felix Draxler, Robin C. Geyer, Daniel Geyfman, Farrin Marouf Sofian, Justus Will, Philipp Hildebrandt, Giosue Migliorini, Prakhar Srivastava,

Organizations: University of California, Irvine, CA, USA · Department of Computer Science, ETH Zurich, Switzerland · Institute for Artificial Intelligence in Medicine, Charité Berlin, Germany · Hasso Plattner Institute, Potsdam, Germany

Abstract

Autoregressive language models emit one token per forward pass. This causes high inference latency, in particular as models are scaled up for quality. In this paper, we introduce Parallel Language Models (PLMs), which generate mutually consistent tokens in one forward pass by routing dependencies through latent auxiliary variables. PLMs are trained using a variational bound that reduces to an autoregressive loss plus a parallel decoding loss. Empirically, PLMs are on a faster and better Pareto frontier than autoregressive and discrete diffusion baselines. Pretrained PLMs achieve near-autoregressive pretraining loss on model scales of 80M to 3B parameters, and speedup follows a power law in model size. When post-trained, our 1.3B model matches or exceeds an identically trained autoregressive model on math, coding, and instruction-following benchmarks, speeding up inference $2.3\text{-}5.7\times$ . Together, we show that parallel generation can be trained from a single training objective at autoregressive quality. We will make our code available at https://github.com/mandt-lab/plm.

Figures & tables

Figure 1: Parallel Language Models (PLMs, orange ) shift the Pareto frontier of speed vs quality for language models trained from scratch. (Right) At the same inference speed, PLMs outperform the quality of publicly available autoregressive models (AR, blue ), despite the latter being trained on trillions of tokens and our models only trained on 25B (PLM 350M) to 300B tokens (PLM 3B). (Left) PLMs outperform discrete diffusion approaches at various NFEs. Inference speeds using vLLM on NVIDIA RTX 6000 Ada. PLM circle and triangle refer to exact respectively relaxed inference, see Section 3.4 . Error bars are 95% confidence intervals.

Figure 2: Parallel Language Models jointly train an encoder that maps training sequences

{\bm{x}}

to auxiliaries

{\bm{u}}

, and a decoder that reconstructs them back to sequences

\hat{\bm{x}}

in parallel. The encoder leverages a standard autoregressive (AR) model predicting next-token probabilities

Q_{i}=q_{\theta}(\,\cdot\mid{\bm{x}}_{<i})

, which we then turn into auxiliaries

u_{i}\sim q_{\theta}(u_{i}\mid x_{i},Q_{i})

. See Fig. 3 for KV cache details.

Figure 3: (Top) The encoder

q_{\theta}

has a standard causal attention mask ( blue ). (Bottom) The decoder

p_{\theta}

attends to the input tokens before its block. For example, the first block attends to the first input token

x_{1}

via the encoder ( orange ), and predicts

{\bm{x}}_{2:4}

at the positions receiving auxiliaries

{\bm{u}}_{2:4}

( green ).

Figure 4: Auxiliaries are grouped around token embeddings

{\bm{e}}_{v}

, with group sizes matching token probabilities.

Figure 5: PLMs become better and more parallel with scale. (Left) PLMs (colored) achieve comparable quality to an AR model of the same size (black). (Right) PLMs achieve more correct tokens with greater model size. Varying the encoder-to-decoder ratio

\beta

reduces quality and increases speed. Markers show measurements at token to non-embedding parameter ratio

D/N=90

, solid lines are Eq. 14 .

Size	$\beta$	Decoding	GSM8K		MATH-500		HumanEval		IFEval
Size	$\beta$	Decoding	Score	Q / s	Score	Q / s	Score	Q / s	Score	Q / s
1.3B	AR	AR	27.9	0.9 (1.0×)	7.4	0.4 (1.0×)	12.8	0.7 (1.0×)	48.2	0.8 (1.0×)
	$0.5$	Relaxed	32.6	5.2 (5.7×)	8.4	1.4 (3.7×)	17.7	3.0 (4.5×)	45.1	1.9 (2.4×)
	$1$	Relaxed	29.4	4.9 (5.4×)	7.2	1.2 (3.1×)	10.4	5.2 (7.8×)	45.7	1.9 (2.5×)
	$2$	Relaxed	29.5	4.8 (5.2×)	8.4	1.0 (2.8×)	15.8	3.0 (4.5×)	48.1	1.8 (2.3×)
3B	$1$	AR †	39.0	0.4 (1.0×)	12.2	0.2 (1.0×)	23.2	0.4 (1.0×)	54.0	0.4 (1.0×)

Table 1: Post-trained 1.3B PLMs are comparable to or better than an autoregressive model at

2.3\times

5.7\times

the inference speed. The 3B PLM-

\beta=1

model achieves better quality than all 1.3B variants. “Q / s” is questions per second. † We do not train an autoregressive model, so we use the 3B PLM in autoregressive inference mode for comparison.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Layers	$d$	FFN	Heads	KV	$N$	Total	$D_{\mathrm{max}}$	GBS	$\eta$	Nodes
80M	12	512	1408	8	2	33.8M	85.3M	3B	32	$6.22\cdot 10^{-3}$	1
150M	12	768	2048	12	3	74.3M	151.6M	6B	64	$4.39\cdot 10^{-3}$	2
350M	24	1024	2560	16	4	251.7M	354.7M	20B	128	$2.66\cdot 10^{-3}$	4
1.3B	24	2048	5632	32	8	1.08B	1.29B	80B	256	$1.44\cdot 10^{-3}$	8
3B	28	3072	8192	48	12	2.77B	3.08B	300B †	512	$1.10\cdot 10^{-3}$	32

Appendix

Table 2: Model sizes and per-run settings.

N

counts non-embedding parameters. “Total” adds the untied input and output embeddings. The planned token budget

D_{\mathrm{max}}

sets the learning rate

\eta

and the global batch size (GBS, in sequences of 4096 tokens) through Eq. 35 . † The 3B constant-learning-rate run stopped at 240B tokens; its

\eta

and GBS are those for the planned 300B.

Model	GPUs	AR	PLM- $\beta$ =2 ( $s=32$ )	PLM- $\beta$ =1 ( $s=16$ )	PLM- $\beta$ =0.5 ( $s=8$ )
350M	16	98	172	201	285
1.3B	32	761	1,436	1,884	2,701
3B	64–128	–	–	11,872	–

Appendix

Table 3: GPU-hours of the final pretrained models , including the final annealing branch. The 3B model was trained only as PLM-

\beta

=1; its constant-learning-rate run moved from 64 to 128 GPUs after 31 hours.

Auxiliaries	Alignment	AR NLL	$R$
AR model		2.788	1
None	no	2.821	1.443
None	yes	2.807	1.464
Quantile	no	2.849	1.705
Quantile	yes	2.837	1.739
Voronoi	no	2.844	1.785

Appendix

Table 4: Building the parallel decoder up from an AR model. 150M models at stride

s=32

(

\beta=2

), trained for 7.5B tokens including the learning rate decay. The last row is our recipe. The AR model uses the same schedule. “None” uses a constant auxiliary; “Quantile” uses quantiles as in PTP; “Voronoi” uses our geometry-aware auxiliaries. “Alignment” denotes the representation-alignment loss. AR NLL is measured in nats;

R

is the mean reconstructed prefix length.

Figure 6: Alignment reaches the same number of reconstructed tokens in fewer steps. Reconstructed tokens per model call

R

of the 150M PLM (block size

K=16

, Voronoi auxiliaries) with alignment (blue) and without (orange). Each curve is one training run during its constant learning rate phase of 6B tokens, before learning rate decay. Left: stride

s=16

(

\beta=1

). Right: stride

s=32

(

\beta=2

Stride $s$	$\beta$	Tokens per step	AR NLL	$R$
AR	–	$1\times$	3.023	1
64	4	$1.25\times$	3.041	1.598
32	2	$1.5\times$	3.065	1.719
16	1	$2\times$	3.110	1.844
8	0.5	$3\times$	3.174	1.936
4	0.25	$5\times$	3.266	1.987

Appendix

Table 5: Reconstructed tokens saturate at small strides while quality keeps degrading. 80M models trained for 3.75B tokens including the learning rate decay. “Tokens per step” is the number of positions processed per training step relative to an AR model,

1+K/s

K=16

			AR NLL			Reconstructed tokens $R$
Size	Tokens	Stride	from scratch	converted	AR model	from scratch	converted
80M	3.75B	8	3.174	3.181	3.023	1.936	1.734
80M	3.75B	16	3.110	3.124	3.023	1.844	1.649
80M	3.75B	32	3.065	3.085	3.023	1.719	1.566
150M	7.5B	8	2.930	2.926	2.788	2.081	1.835
150M	7.5B	16	2.863	2.878	2.788	1.993	1.747

Appendix

Table 6: Converting an AR model is worse than training from scratch. “Converted” trains an AR model at constant learning rate and switches to the PLM objective for the learning rate decay, the final 20% of training. “AR model” applies the same decay without the switch. Token counts include the decay.

Model	Train tokens	Decoding	Elo [95% CI]	Tokens/s	Tokens/call	$\times$ AR
AR 350M	25B	AR sampling	840 [803, 874]	746	1.00	1.00
PLM- $\beta$ =1 350M	25B	exact	860 [826, 890]	1323	1.95	1.77
PLM- $\beta$ =1 350M	25B	relaxed	764 [724, 802]	2001	2.96	2.68
AR 1.3B	100B	AR sampling	1000 (anchor)	310	1.00	1.00
PLM- $\beta$ =1 1.3B	100B	exact	1055 [1025, 1086]	579	2.14	1.87
PLM- $\beta$ =1 1.3B	100B	relaxed	920 [887, 956]	856	3.21	2.76

Appendix

Table 7: Main judge pool: our models, outside autoregressive models and the human reference. Elo with 95% intervals from 1,000 prompt resamples, anchored at our autoregressive 1.3B model. “Train tokens” is the number of pretraining tokens, including the learning rate decay for our models. Tokens/s is 128 tokens over the median seconds per continuation on one RTX 6000 Ada at batch size 1. “

\times

AR” is the speed relative to our autoregressive model of the same size. ∗ We did not train an autoregressive 3B model, so the 3B PLM decoded by AR sampling stands in for it.

		NFE
Model		1	2	4	8	16	32	64	128
LLaDA-8B	Elo				58	144	315	668	866
	tokens/s				511	250	140	73	36
Dream-7B	Elo				110	332	408	676	883
	tokens/s				456	244	128	65	32
DFM 1.3B	Elo		409	476	555	621	665	671	665

Appendix

Table 8: Diffusion and flow baselines in the main judge pool across the number of denoising steps (NFE): Elo (top of each pair of rows) and tokens/s (bottom). LLaDA and Dream reuse the prompt’s key–value cache. All models are timed on one RTX 6000 Ada. Elo intervals are about

\pm 45

–

90

for these rows, which played fewer matches than the rows of Table 7 .

		SFT		DPO
Model	Stride	Learning rates	Selected (lr, epoch)	Learning rate	Epoch
AR 1.3B	–	$\{0.4,0.8,1.6,3.2\}$	$1.6$ , 6	$10^{-6}$	5
PLM 1.3B	8	$\{0.4,0.8,1.6,3.2\}$	$3.2$ , 5	$10^{-6}$	3
PLM 1.3B	16	$1.6$	$1.6$ , 6	$10^{-6}$	5
PLM 1.3B	32	$\{0.4,0.8,1.6,3.2\}$	$1.6$ , 5	$10^{-6}$	2
PLM 3B	16	$\{0.4,0.8,1.6,3.2\}$	$1.6$ , 6	$5\cdot 10^{-7}$	6

Appendix

Table 9: Post-training sweeps. We first run SFT, then DPO from the selected SFT checkpoint. The table lists the SFT learning rates (six epochs unless noted), the SFT checkpoint each DPO run starts from, and the DPO learning rate and epoch with the highest GSM8K accuracy. SFT learning rates are in units of

10^{-4}

Benchmark	$n$	Prompt	Max. new tokens	Scoring
GSM8K	1,319	zero-shot instruction to reason step by step and end with “The answer is <number>.”	512	last number in the response
MATH-500	500	“Please reason step by step, and put your final answer within \boxed{} .”	2,048	symbolic equivalence of the boxed answer
HumanEval	164	implement the function in one Python code block	768	pass@1, reference tests executed
IFEval	541	the instruction as given	1,280	prompt-level strict accuracy

Appendix

Table 10: Post-training benchmarks.

n

is the number of questions, all of which we score.

Size	$\beta$	Decoding	GSM8K				MATH-500
Size	$\beta$	Decoding	Score	Q / s	Tok / s	Tok / call	Score	Q / s	Tok / s	Tok / call
1.3B	AR	AR	27.9	0.9 (1.0×)	291	1.0	7.4	0.4 (1.0×)	291	1.0
	$0.5$	AR	27.0	1.1 (1.2×)	308	1.0	7.6	0.4 (1.0×)	308	1.0
		Exact	31.5	4.4 (4.8×)	1112	4.2	8.0	1.0 (2.6×)	803	3.2
		Relaxed	32.6	5.2 (5.7×)	1337	6.1	8.4	1.4 (3.7×)	1107	5.1
	$1$	AR	27.7	1.1 (1.2×)	307	1.0	5.6	0.4 (1.1×)	309	1.0

Appendix

Table 11: Detailed Post-training results. Post-trained 1.3B PLMs are comparable to or better than an autoregressive model at

2.3\times

5.7\times

the inference speed. The PLM-

\beta=1

3B model achieves better quality than all 1.3B variants. “Q / s” is questions per second. † We do not train an autoregressive model, so we use the PLM 3B in autoregressive inference mode for comparison.

Law	Coefficient	Value
$\mathcal{L}_{\mathrm{enc}}$	$E$	1.105 $\pm$	0.12
	$A$	0.851 $\pm$	0.11
	$a$	0.194 $\pm$	0.023
	$B$	0.763 $\pm$	0.23
	$b$	0.127 $\pm$	0.032
	$c$	0.482 $\pm$	0.027

Appendix

Table 12: Fitted coefficients of the scaling laws. Values with one jackknife standard deviation from refits leaving out one model size at a time;

n=N/10^{8}

and

d=D/10^{10}

. RMSE is the root mean square error of the fit on all runs and, for the hold-out, of the refit without the 3B size evaluated on the 3B runs. The AR hold-out uses 1.3B runs instead, as noted below.

Figure 7: Parallel Language Models predict coordinated tokens in parallel. This schematic shows a simplified version of our inference algorithm. Our encoder and decoder share a transformer backend and can be called in parallel via a custom attention mask. First, based on randomly drawn Voronoi auxiliaries, several tokens are proposed in parallel. The first token is always accepted as it can be chosen autoregressively. In the following call the tokens are verified using the encoder while new tokens are proposed simultaneously. The KV cache is built from verified correct tokens.

Model		PLM 1.3B, SFT + DPO, $\beta=1$			PLM 3B, SFT + DPO, $\beta=1$
		tokens/s	tokens/call	ms/call	tokens/s	tokens/call	ms/call
AR		307	1.00	3.25	140	1.00	7.14
$K_{\mathrm{pred}}$	2	548	1.85	3.32	255	1.88	7.28
	4	876	3.06	3.40	425	3.25	7.46
	7	1090	4.05	3.59	558	4.53	7.86
	8	966	4.25	4.25	514	4.81	9.07

Appendix

Table 13: Quadratic Decoding of PLMs on GSM8K is optimal for

K_{\mathrm{pred}}=7

under exact acceptance . Higher widths

K_{\mathrm{pred}}

accept more tokens per call but are slower per call, with GPU-specific jumps after

K_{\mathrm{pred}}=7,10

. Across

K_{\mathrm{pred}}

, accuracy varies only within sampling noise.

Parallel Decoding

Parallel decoding generates multiple tokens or structured outputs at once to reduce language-model inference latency while controlling quality.

See all 66 papers

Why Diffusion Language Models Struggle with Truly Parallel (Non-Autoregressive) Decoding?
Feb 26, 2026Pengxiang Li, Dilxat Muhtar, Tianlong Chen +2Parallel Decoding for Diffusion Language Models Diffusion Language Model Decoding
The Hong Kong Polytechnic University, Hong Kong, China · ELLIS Institute Tübingen, Tübingen, Germany · Max Planck Institute for Intelligent Systems, Tübingen, Germany +3
Self-Orchestrating Language Models: Leveraging Semantic Dependence for Efficient Inference
Sep 13, 2026Tian JinAutoregressive Decoding Efficient Language Model Reasoning
Department of Electrical Engineering and Computer Science, MASSACHUSETTS INSTITUTE OF TECHNOLOGY
Gumbel Distillation for Parallel Text Generation
Mar 23, 2026Chi Zhang, Xixi Hu, Bo Liu +1Non-Autoregressive Text Generation Language Model Distillation
Department of Computer Science The University of Texas at Austin

Feb 26, 2026J/K move · Enter open · S save

Why Diffusion Language Models Struggle with Truly Parallel (Non-Autoregressive) Decoding?

Pengxiang Li, Dilxat Muhtar, Tianlong Chen, Lu Yin+1

The Hong Kong Polytechnic University, Hong Kong, China · ELLIS Institute Tübingen, Tübingen, Germany · Max Planck Institute for Intelligent Systems, Tübingen, Germany+3

Parallel Decoding for Diffusion Language Models Diffusion Language Model Decoding Diffusion Language Models Non-Autoregressive Text Generation Parallel Reasoning in Language Models LLM Reasoning

Open paper AskPDF arXiv

Diffusion Language Models (DLMs) are often advertised as enabling parallel token generation, yet practical fast DLMs frequently converge to left-to-right, autoregressive (AR)-like decoding dynamics. In contrast, genuinely non-AR generation is promising because it removes AR's sequential bottleneck, better exploiting parallel hardware to reduce synchronization/communication overhead and improve latency scaling with output length. We argue that a primary driver of AR-like decoding is a mismatch between DLM objectives and the highly sequential structure of widely used training data, including standard pretraining corpora and long chain-of-thought (CoT) supervision. Motivated by this diagnosis, we propose NAP (Non-Autoregressive Parallel DLMs), a proof-of-concept, data-centric approach that better aligns supervision with non-AR parallel decoding. NAP curates examples as multiple independent reasoning trajectories and couples them with a parallel-forced decoding strategy that encourages multi-token parallel updates. Across math reasoning benchmarks, NAP yields stronger performance under parallel decoding than DLMs trained on standard long CoT data, with gains growing as parallelism increases. Our results suggest that revisiting data and supervision is a principled direction for mitigating AR-like behavior and moving toward genuinely non-autoregressive parallel generation in DLMs. Our code is available at https://github.com/pixeli99/NAP.

Abstract

Parallel Decoding

Explore similar work

Why Diffusion Language Models Struggle with Truly Parallel (Non-Autoregressive) Decoding?

Self-Orchestrating Language Models: Leveraging Semantic Dependence for Efficient Inference

Gumbel Distillation for Parallel Text Generation