A Few More Tokens Ahead: Reinforcing Parallelism in Diffusion LLMs

Jiahao Kuang; Yangfan Gao; Tao Gui; Yuanbin Wu; Qi Zhang; Xuanjing Huang; Tao Ji

cs.LGOct 13, 2026

A Few More Tokens Ahead: Reinforcing Parallelism in Diffusion LLMs

Authors: Jiahao Kuang, Yangfan Gao, Tao Gui, Yuanbin Wu, Qi Zhang, Xuanjing Huang, Tao Ji

Organizations: East China Normal University · Fudan University

Abstract

Parallel generation is widely regarded as an inherent advantage of diffusion large language models. Yet our fixed-quota evaluation of four representative models reveals severe degradation across mathematics and code, with all scores below 15% at 16 tokens per step. Token-level probes identify local predictive potential near the observed prefix, especially in mathematical text. We propose RePara, which reinforces short consecutive token blocks through group-relative policy optimization and reduces rollout and log-probability recomputation costs. On LLaDA and SDAR, \ours improves GSM8K and MATH-500 accuracy over JustGRPO under reduced denoising budgets, with a 1.96x training-step speedup in our LLaDA/MATH setup.

Figures & tables

Figure 1: Generation order and decoding parallelism in dLLMs. Arbitrary-order decoding commits potentially nonconsecutive positions; causal-order decoding commits consecutive blocks from left to right. Token numbers indicate commitment steps; equal numbers denote parallel commitment, and darker shades denote later steps.

Figure 2: Task performance across decoding granularities. Four representative dLLMs are evaluated on mathematical reasoning (GSM8K and MATH-500) and code generation (HumanEval and MBPP) with a 512-token generation budget.

Figure 3: Local predictive potential in mathematical text and code. Dashed lines mark

h_{80}

, the 80% retention boundary relative to the first position for Probability, Confidence, or inverse perplexity

e^{-H}

. Entropy is plotted in nats; boundary definitions are in Appendix A.2 .

Figure 4: Local prediction across clean-prefix lengths. Each cell averages 2,000 sequences at a fixed target position and prefix length. Red outlines apply the $h_{80}$ criterion from Figure 3 relative to each row’s first position. Entropy colors show raw nats.

Figure 5: Consecutive vs. confidence-based decoding at matched token quotas. Filled/hollow markers denote confidence-based/consecutive decoding. Scores and generation budget follow Figure 2 .

Figure 6: Overview of RePara: Reinforcement Learning with Multi-Token Rollouts. (a) JustGRPO generates one token per forward pass. (b) RePara samples consecutive token blocks from a shared masked state in a single forward pass and optimizes their joint probabilities with GRPO. For an

N

-token completion, committing an average of

M

tokens per block reduces the forward-pass count from

N

N/M

for both rollout generation and training-time log-probability recomputation.

Model	GSM8K					MATH-500
Diff# $k$	1	2	4	8	16	1	2	4	8	16
LLaDA-8B-Instruct	81.50	81.12	75.89	45.94	5.08	44.40	41.20	31.60	17.60	4.20
LLaDA-1.5	82.64	81.05	76.72	48.14	6.82	43.20	40.20	36.40	20.00	6.80
ESPO	85.14	82.64	76.42	46.70	7.43	45.20	42.20	34.20	20.00	5.60
d-TreeRPO	83.85	82.03	77.10	47.84	6.52	41.20	40.60	36.20	19.00	7.60
wd1-Math	82.03	81.88	75.59	44.96	5.23	44.40	38.80	33.60	16.20	6.00

Table 1: LLaDA accuracy (%,

\uparrow

) at

L=512

. Colors mark Best , Second best , and Third best per column; ties share a rank at the reported precision.

Figure 7: LLaDA accuracy (top) and gains over the backbone in pp. (bottom) at

L=512

. Arrows show nominal TPD ratios to the backbone at matched accuracy, interpolated on the log-rate axis.

Model	GSM8K			MATH-500
Diff# $k$	1	2	4	1	2	4
SDAR-8B-Chat	90.75	55.12	74.45	50.00	36.20	28.20
JustGRPO	82.11	55.12	46.63	54.20	34.20	35.40
RePara	85.37	72.63	81.80	52.40	39.80	40.40

Table 2: SDAR accuracy (%,

\uparrow

) at

L=512

under dynamic threshold decoding. Bold marks the best score per column.

Training setting	Core	Step	Core speedup
JustGRPO	422.64	427.89	$1.00\times$
Fix2	213.39	218.60	$\mathbf{1.98\times}$
Adaptive2-0.95	372.26	377.66	$1.14\times$

Table 3: LLaDA mathematical training efficiency on the same 127 steps. Core time includes rollout, reward processing, learner computation, and optimizer updates; step time also includes checkpoint and metric-preparation costs. Times are seconds per optimizer step; speedup is relative to JustGRPO.

Setting	Diff# $k$
	1	2	4	8	16
JustGRPO	87.95	87.49	85.60	75.06	28.20
Adaptive2-0.80	86.81	86.73	84.31	81.27	66.34
Adaptive2-0.90	87.11	86.50	83.17	76.65	51.78
Adaptive2-0.95	88.48	87.19	86.35	81.35	47.76

Table 4: LLaDA/GSM8K accuracy (%) by adaptive training threshold at

L=512

and

B=2

. Bold marks the best score per column.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Optimizer	AdamW ( Loshchilov and Hutter, 2019 )
Learning rate	$5\times 10^{-6}$ , constant
AdamW $(\beta_{1},\beta_{2})$	$(0.9,0.999)$
AdamW $\epsilon$	$10^{-8}$
Weight decay	$0$
Gradient norm limit	$1.0$

Appendix

Table 5: Shared training configuration for JustGRPO and RePara on GSM8K, MATH, and AceCode.

Study	Changes
Basic comparison	Fixed/adaptive two-token coding rollouts.
Medium data	AceCode-Medium replaces AceCode-Hard.
JustGRPO initialization	Multi-token RL starts from a JustGRPO checkpoint.

Appendix

Table 6: Coding studies and their changes.

$L=128$
Model / training setting	GSM8K					MATH-500
Diff# $k$	1	2	4	8	16	1	2	4	8	16
LLaDA-8B-Instruct	75.97	74.00	58.53	40.71	16.60	37.40	32.60	25.40	14.20	8.60
LLaDA-1.5	77.48	75.06	61.41	43.97	26.46	36.40	33.60	28.60	17.60	12.60
ESPO	80.06	74.60	60.96	39.65	21.23	38.60	32.20	26.20	19.40	10.40
d-TreeRPO	77.33	75.21	60.96	45.26	19.11	36.40	35.00	28.40	19.20	14.00

Appendix

Table 7: Detailed LLaDA accuracy (%,

\uparrow

) on GSM8K and MATH-500 at generation lengths

L\in\{128,256,512\}

. Nominal TPD is

k=L/S

, where

S

is the total number of denoising steps. RePara uses Fix2 . Best , Second best , and Third best are marked within each length, benchmark, and nominal TPD; ranks use the reported two-decimal scores, and ties share a rank.

LLaDA
Setting	Steps	Diff#1	Diff#2	Diff#4
Fix2	128	42.60	43.60	42.40
Fix4	128	39.80	38.40	37.40
Fix2 $\rightarrow$ Fix4	$128+128$	43.00	42.00	41.20
SDAR
Fix2	128	49.20	39.00	40.00

Appendix

Table 8: Training block-size ablation on MATH-500 at

L=256

and sampling temperature

0

. Entries are accuracy (%) under the same Diff#

k

settings as the main comparisons. Steps count optimizer updates; the staged setting adds 128 updates with

B=4

to Fix2 . Bold marks the highest accuracy within each model and decoding setting.

$L$	Model	GSM8K			MATH-500
	Diff# $k$	1	2	4	1	2	4
128	SDAR-8B-Chat	39.80	30.17	35.71	3.80	3.20	4.60
	JustGRPO	52.16	40.64	35.25	39.20	30.20	31.60
	RePara ( Fix2 )	60.27	53.98	59.51	39.20	33.60	33.80
256	SDAR-8B-Chat	86.66	53.83	72.40	20.40	15.60	14.00
	JustGRPO	81.05	53.90	45.64	53.00	34.20	35.20

Appendix

Table 9: Accuracy (%,

\uparrow

) with SDAR at generation lengths

L=128,256,512

and temperature

0

. Higher nominal TPD

k

uses fewer scheduled denoising rounds. Trained methods use step-128 checkpoints; RePara uses fixed two-token training blocks. Bold marks the best score within each generation length, benchmark, and nominal TPD; ties are bolded.

Figure 8: Effect of the adaptive training threshold on LLaDA GSM8K accuracy (%) with sampling temperature

0

. The maximum proposal length is two tokens. Panels show generation lengths

L=128,256,512

; the horizontal axis is nominal TPD

k=L/S

. JustGRPO is included as a reference.

Panel of Figure 9 — Figure 9: LLaDA mathematical training efficiency. Adaptive denotes Adaptive2-0.95 . Panel (a) shows the 128-step training traces; panels (b) and (c) summarize the 127 common recorded steps. Logical blocks are averaged per completion and executed rounds per optimizer step.

Metric	JustGRPO	Fix2	Adaptive
Rollout generation	101.22	51.04	89.54
Log-probability	115.33	56.23	100.83
Backward	202.33	102.13	177.99

Appendix

Table 10: Mean phase times (seconds per optimizer step) for LLaDA mathematical training over the 127 common steps. Adaptive denotes Adaptive2-0.95 ; rollout denotes generation only. Overall core and measured step times are reported in Table 3 .

Figure 10: Mean block length during LLaDA training on GSM8K (left) and MATH (right). The adaptive policy is shown alongside JustGRPO and Fix2 . The horizontal axis is the training step.

Metric	JustGRPO	Fix2	Adaptive
Mean block length	1.00	2.00	1.67
Logical blocks	256.00	128.00	153.23
Executed rounds	256.00	128.00	222.23
Inactive row–round slots	0.0%	0.0%	31.1%

Appendix

Table 11: Logical blocks and executed rounds over the same 127 steps. Adaptive denotes Adaptive2-0.95 . Mean block length pools committed tokens and logical blocks; the other counts are averaged per completion or training step.

Training setting	HumanEval	MBPP
Backbone	38.41	40.20
JustGRPO	50.00	47.60
Fix2	10.98	14.60
Adaptive2-0.95	15.24	29.40

Appendix

Table 12: Coding Pass@1 (%) under a common diffusion decoder (

L=S=256

, window size 32, temperature zero). Trained methods use their fixed step-128 endpoints.

Figure 11: Coding evaluation after 128 additional multi-token training updates from a JustGRPO checkpoint, at

L=512

and temperature zero. Gains at higher diffusion quotas coexist with accuracy losses at quota one.

Parallel Decoding for Diffusion Language Models

See all 37 papers

LightningRL: Breaking the Accuracy-Parallelism Trade-off of Block-wise dLLMs via Reinforcement Learning
Mar 4, 2026Yanzhe Hu, Yijie Jin, Pengfei Liu +2Parallel Decoding for Diffusion Language Models Blockwise Diffusion Language Models
Shanghai Jiao Tong University · Huazhong University of Science and Technology
Parallelism, critical windows, and separations among diffusion language models
Sep 17, 2026Sitan Chen, Liye WangMasked Diffusion Models Diffusion Sampling
SEAS, Harvard University · Tsinghua University
PSD: Pushing the Pareto Frontier of Diffusion LLMs via Parallel Speculative Decoding
May 15, 2026Shengyin Sun, Yiming Li, Renxi Liu +7Parallel Speculative Decoding Masked Diffusion Models
Department of Computer Science City University of Hong Kong · Huawei Technologies · City University of Hong Kong

Mar 4, 2026J/K move · Enter open · S save

LightningRL: Breaking the Accuracy-Parallelism Trade-off of Block-wise dLLMs via Reinforcement Learning

Yanzhe Hu, Yijie Jin, Pengfei Liu, Kai Yu+1

Shanghai Jiao Tong University · Huazhong University of Science and Technology

Parallel Decoding for Diffusion Language Models Blockwise Diffusion Language Models Diffusion Language Model Decoding Diffusion Model Inference Acceleration RL for Diffusion Models

Open paper AskPDF arXiv

Diffusion Large Language Models (dLLMs) have emerged as a promising paradigm for parallel token generation, with block-wise variants garnering significant research interest. Despite their potential, existing dLLMs typically suffer from a rigid accuracy-parallelism trade-off: increasing the number of tokens per forward (TPF) via aggressive parallel decoding often leads to performance degradation and increased generation instability. We identify that this limitation stems from the model's inability to navigate high-parallelism regimes where approximation errors and local corruptions accumulate, ultimately undermining the reliability of parallel generation. To address this, we propose LightningRL, a post-training framework designed to directly optimize the speed-quality Pareto frontier of pre-trained dLLMs. Instead of forcing uniform parallelization, our approach leverages reinforcement learning to identify and reinforce high-parallelism trajectories that maintain generation accuracy. Built upon the Group Relative Policy Optimization (GRPO) framework, LightningRL introduces several enhancements tailored for dLLMs: (1) stabilized training via per-reward decoupled normalization; (2) token-level negative log-likelihood (NLL) regularization on correct trajectories to anchor model performance; and (3) a dynamic sampling strategy with TPF-aware filtering to enhance training efficiency. Experimental results across mathematical and coding benchmarks demonstrate that LightningRL consistently advances the Pareto frontier, achieving competitive task accuracy while significantly increasing parallelism, reaching an average TPF of 7.32 (with a peak of 11.10 on the MBPP dataset). Our code is available at https://github.com/SJTU-DENG-Lab/LightningRL.

Abstract

Parallel Decoding for Diffusion Language Models

Explore similar work

LightningRL: Breaking the Accuracy-Parallelism Trade-off of Block-wise dLLMs via Reinforcement Learning

Parallelism, critical windows, and separations among diffusion language models

PSD: Pushing the Pareto Frontier of Diffusion LLMs via Parallel Speculative Decoding