Balance-Aware Pivot Selection for Parallel Diffusion LLM Decoding

Yejin Lee; Hyundong Jin; Yo-Sub Han

cs.CLOct 19, 2026

Balance-Aware Pivot Selection for Parallel Diffusion LLM Decoding

Authors: Yejin Lee, Hyundong Jin, Yo-Sub Han

Organizations: Yonsei University, Seoul, Republic of Korea

Abstract

Diffusion language models enable parallel token generation, offering a promising alternative to sequential autoregressive decoding. Existing acceleration methods use predictive confidence to enable parallel token commitment while preserving generation quality. However, confidence-based selection prioritizes reliable token predictions without explicitly considering how token positions shape the remaining decoding schedule. Moreover, such ordering does not always reflect inter-token dependencies, suggesting that confidence need not be the sole criterion for token selection. Motivated by this limitation, we propose Balance-Aware Pivot Decoding (BPD), which favors central pivots within contiguous masked intervals and decodes the selected positions in parallel. Balanced pivots yield logarithmic decoding depth in the interval length. Our practical method accounts for model uncertainty by combining confidence-based prefilling with confidence-thresholded center selection. We evaluate the resulting quality-efficiency tradeoff on LLaDA and Dream across five mathematical reasoning, code generation, and instruction-following benchmarks. Our implementation is available at: https://anonymous.4open.science/r/DLM_random_pivot-EE62/

Figures & tables

Figure 1: Motivation for position-aware unmasking strategy. Center-based selection promotes balanced partitions, unlike confidence-based selection that may leave long masked intervals.

Figure 2: Overview of Balance-Aware Pivot Decoding (BPD). BPD first commits high-confidence tokens in parallel, then selects one pivot per remaining masked interval. It favors the most central confidence-eligible position to promote balanced splits and reduce decoding depth.

Dataset	Method	LLaDA-8B-Instruct			LLaDA-1.5			Dream-7B
Dataset	Method	Score $\uparrow$	TPS $\uparrow$	QTG $\uparrow$	Score $\uparrow$	TPS $\uparrow$	QTG $\uparrow$	Score $\uparrow$	TPS $\uparrow$	QTG $\uparrow$
GSM8K	Vanilla	81.65	12.02	1.00 $\times$	81.88	11.92	1.00 $\times$	82.87	14.62	1.00 $\times$
	Hierarchy-dLLM	77.94	52.87	4.20 $\times$	79.45	54.51	4.44 $\times$	73.54	73.43	4.46 $\times$
	WINO	80.29	46.24	3.78 $\times$	79.53	46.46	3.79 $\times$	80.21	55.24	3.66 $\times$
	Fast-dLLM	81.50	52.83	4.39 $\times$	81.58	54.96	4.59 $\times$	80.52	74.60	4.96 $\times$
	DAPD	80.44	36.01	2.95 $\times$	82.03	37.69	3.17 $\times$	82.56	41.84	2.85 $\times$

Table 1: Main results on five benchmarks using LLaDA-8B-Instruct, LLaDA-1.5, and Dream-7B. We report generation quality (Score), throughput measured in tokens per second (TPS), and Quality-Throughput Gain (QTG). QTG is the product of relative throughput and relative task performance with respect to vanilla.

Method	GSM8K	MATH	H.Eval	MBPP	IFEval
Vanilla	11.41	12.69	10.99	11.11	12.28
Vanilla	(81.65)	(38.60)	(35.98)	(39.60)	(56.01)
Hierarchy	51.36	48.74	85.36	53.96	37.69
Hierarchy	(77.94)	(36.20)	(32.93)	(36.80)	(46.40)
WINO	48.67	33.84	50.23	29.55	21.49
WINO	(80.29)	(37.20)	(39.63)	(38.80)	(55.82)

Table 2: Generation quality and throughput on five benchmarks using LLaDA-8B-Instruct. Each method reports TPS in the first row and task performance (Score) in the second row.

Variant	Pass@1	TPS	# Steps	QTG
Vanilla	35.98	10.44	256.00	1.00 $\times$
w/o KV caching	37.20	54.01	49.21	5.35 $\times$
w/o Prefilling	33.54	59.15	73.77	5.28 $\times$
w/ Random Pivot	34.15	68.72	64.20	6.25 $\times$
w/ Confident Pivot	35.37	62.34	71.79	5.87 $\times$
Ours	35.98	78.29	55.73	7.50 $\times$

Table 3: Ablation study on HumanEval using LLaDA-8B-Instruct to examine the contribution of each component. We report generation quality, throughput, and quality-throughput gain (QTG).

Method	Pivoting			Prefilling
Method	Pass@1	TPS	QTG	Pass@1	TPS	QTG
Exact	23.17	76.65	4.73	22.56	108.57	6.52
Margin	25.00	64.96	4.32	29.27	97.41	7.59
Entropy	29.88	63.53	5.05	27.44	94.54	6.91
Random	20.12	69.25	3.71	21.34	102.17	5.80
Ours	33.54	59.15	5.28	35.98	78.29	7.50

Table 4: Comparison of pivot selection strategies with and without confidence-based prefilling on LLaDA-8B-Instruct and HumanEval.

Pivot Type	# Events	# Balanced	BPR (%)
Eligible	2,570	1,451	56.46
Fallback	2,113	423	20.02
Total	4,683	1,874	40.02

Table 5: Balance statistics of BPD pivot selections. The first three iterations of each block and intervals of length at most four are excluded.

Figure 3: Effect of the confidence threshold on generation quality and throughput using LLaDA-8B-Instruct. We report Pass@1 on HumanEval and Accuracy on MBPP, along with tokens per second (TPS). Dotted lines indicate vanilla decoding performance.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Description	Examples
GSM8K ( Cobbe et al., 2021 )	Grade-school mathematical word problems requiring multistep reasoning.	1,319
MATH500 ( Lightman et al., 2024 )	A 500-problem subset of MATH covering competition-level mathematics.	500
HumanEval ( Chen et al., 2021 )	Python function completion from signatures and docstrings, evaluated with unit tests.	164
MBPP ( Austin et al., 2021b )	Basic Python programming problems with natural-language specifications and tests.	500
IFEval ( Zhou et al., 2023 )	Instruction-following prompts with automatically verifiable constraints.	541

Appendix

Table 6: Evaluation benchmarks and split sizes.

Dataset	Method	LLaDA-8B-Instruct			LLaDA-1.5			Dream-7B
Dataset	Method	Score $\uparrow$	TPS $\uparrow$	QTG $\uparrow$	Score $\uparrow$	TPS $\uparrow$	QTG $\uparrow$	Score $\uparrow$	TPS $\uparrow$	QTG $\uparrow$
GSM8K	Vanilla	81.65	12.02	1.00 $\times$	81.88	11.92	1.00 $\times$	82.87	14.62	1.00 $\times$
	Fast-dLLM (prefix)	81.50	52.83	4.39 $\times$	81.58	54.96	4.59 $\times$	80.52	74.60	4.96 $\times$
	Fast-dLLM (dual)	80.36	60.37	4.94 $\times$	82.03	62.16	5.22 $\times$	82.11	81.74	5.54 $\times$
	Fast-dLLM (parallel)	82.18	35.50	2.97 $\times$	82.03	36.41	3.06 $\times$	34.19	26.46	0.75 $\times$
	DOS (top-1)	80.36	12.04	0.99 $\times$	82.64	12.04	1.02 $\times$	82.71	14.84	1.01 $\times$

Appendix

Table 7: Additional results on five benchmarks using LLaDA-8B-Instruct, LLaDA-1.5, and Dream-7B. We report generation quality (Score), throughput measured in tokens per second (TPS), and Quality-Throughput Gain (QTG). QTG is the product of relative throughput and relative task performance with respect to vanilla.

Category	#Inst.	Accuracy (%) $\uparrow$		Early Term. (%) $\downarrow$		TPS (tok/s) $\uparrow$
Category	#Inst.	Hier.	Ours	Hier.	Ours	Hier.	Ours
startend	67	77.6	83.6	60.0	9.1	32.42	42.11
detectable_content	53	45.3	67.9	34.5	0.0	26.83	37.20
keywords	163	57.7	57.7	29.0	7.2	31.82	41.88
combination	65	41.5	66.2	28.9	22.7	49.21	60.11
detectable_format	157	63.7	72.6	22.8	4.7	31.11	43.68

Appendix

Table 8: Instruction-level analysis of IFEval across instruction categories. We report instruction-following accuracy, the proportion of failures attributable to early termination, and decoding TPS (tokens/second) for Hierarchy-dLLM and our method. We define early termination as a failed instruction whose response length falls below 50% of the corresponding vanilla response length.

Case 1: Keyword and Repetition Constraints
Prompt: Write a funny haiku about moms, containing keywords “mom” and “mother” in your response. First repeat the request word for word without change, then give your answer (…)
Hierarchy (16 tokens, Fail)	Ours (42 tokens, Pass)
Write a funny haiku about moms, containing keywords "mom" and "mother	Write a funny haiku about moms, containing keywords "mom" and "mother" in your response.<br><br>Mom’s smile is loud,<br>Mother’s laugh is loud,<br>Laughter all around.
Case 2: Explicit Ending Constraint
Prompt: How to write a good Chinese poem? At the end of your response, please explicitly add a note starting with “P.S.”
Hierarchy (59 tokens, Fail)	Ours (254 tokens, Pass)

Appendix

Table 9: Qualitative comparison between Hierarchy-dLLM and our method on IFEval. Our method satisfies the specified instructions, whereas Hierarchy-dLLM produces incomplete responses. Ellipses indicate omitted text.

Figure 4: Decoding trajectory on HumanEval. Gray bars represent masked intervals, and blue dots indicate committed tokens. Our method recursively partitions masked intervals through center-preferred pivot selection, completing a 32-token block in eight decoding steps.

Diffusion Language Model Inference

Diffusion language model inference studies how to generate, guide, and serve text efficiently through iterative denoising and parallel token updates.

See all 165 papers

Ripple-Pivot Search: Active Parallel Decoding for Diffusion Large Language Models
Aug 12, 2026Yushi Ye, Xu Chen, Haoyun Jiang +7Language Model Decoding Diffusion Language Model Inference
Cooperative Medianet Innovation Center, Shanghai Jiao Tong University · Alibaba Group
DAPD: Dependency-Aware Parallel Decoding via Attention for Diffusion LLMs
Mar 13, 2026Bumjun Kim, Dongjae Jeon, Moongyu Jeon +1Parallel Decoding for Diffusion Language Models Parallel Decoding
Department of Artificial Intelligence, Yonsei University, Seoul, Korea · Department of Computer Science, Yonsei University, Seoul, Korea
Mean-Field Parallel Decoding for Discrete Diffusion Language Models
Jun 14, 2026Tamim Zoabi, Ameen Ali, Liran Ringel +1Parallel Decoding Diffusion Language Model Decoding
School of Electrical & Computer Engineering Tel Aviv University, Israel · School of Computer Science and AI Tel Aviv University, Israel · Department of Computer Science Technion, Israel Institute of Technology, Israel

Aug 12, 2026J/K move · Enter open · S save

Ripple-Pivot Search: Active Parallel Decoding for Diffusion Large Language Models

Yushi Ye, Xu Chen, Haoyun Jiang, Jinsong Lan+6

Cooperative Medianet Innovation Center, Shanghai Jiao Tong University · Alibaba Group

Language Model Decoding Diffusion Language Model Inference Parallel Decoding Diffusion Language Model Decoding

Open paper AskPDF arXiv

Diffusion Large Language Models (dLLMs) have emerged as a competitive alternative to autoregressive language models, offering the potential for substantially faster inference through parallel decoding. Existing parallel decoding schedulers typically commit positions only after they meet a per-position criterion, overlooking how early commitments may benefit subsequent decoding. We identify a ripple effect in dLLM decoding: proactively committing a mid-entropy pivot position can induce a pronounced reduction in uncertainty across the remaining masked positions. This uncertainty reduction allows subsequent steps to unmask more tokens in parallel, thereby accelerating the overall decoding process. To exploit the ripple effect, we propose Ripple-Pivot Search (RPS), a novel training-free decoding method that seeks mid-entropy positions as promising candidate pivots (where to decode), and determines their token assignment that yields the greatest downstream benefit via lookahead evaluation (what to decode). Across 3 dLLMs and 4 reasoning and code-generation benchmarks, RPS achieves 4-10

\times

wall-clock speedup over the standard decoder while preserving generation quality, and improves accuracy over the previous lookahead baseline by up to 5.49% while delivering higher throughput in most settings. When integrated with KV caching, RPS further achieves up to 18

\times

wall-clock speedup over the standard decoder.

Abstract

Diffusion Language Model Inference

Explore similar work

Ripple-Pivot Search: Active Parallel Decoding for Diffusion Large Language Models

DAPD: Dependency-Aware Parallel Decoding via Attention for Diffusion LLMs

Mean-Field Parallel Decoding for Discrete Diffusion Language Models