Early Decisions Matter: Proximity Bias and Initial Trajectory Shaping in Non-Autoregressive Diffusion Language Models

Jiyeon Kim; Sungik Choi; Yongrae Jo; Moontae Lee; Minjoon Seo

cs.CLApr 12, 2026

Early Decisions Matter: Proximity Bias and Initial Trajectory Shaping in Non-Autoregressive Diffusion Language Models

Authors: Jiyeon Kim, Sungik Choi, Yongrae Jo, Moontae Lee, Minjoon Seo

Organizations: KAIST AI · LG AI Research · University of Illinois Chicago

Abstract

Diffusion-based language models (dLLMs) have emerged as a promising alternative to autoregressive language models, offering the potential for parallel token generation and bidirectional context modeling. However, harnessing this flexibility for fully non-autoregressive decoding remains an open question, particularly for reasoning and planning tasks. In this work, we investigate non-autoregressive decoding in dLLMs by systematically analyzing its inference dynamics along the temporal axis. Specifically, we uncover an inherent failure mode in confidence-based non-autoregressive generation stemming from a strong proximity bias-the tendency for the denoising order to concentrate on spatially adjacent tokens. This local dependency leads to spatial error propagation, rendering the entire trajectory critically contingent on the initial unmasking position. Leveraging this insight, we present a minimal-intervention approach that guides early token selection, employing a lightweight planner and end-of-sequence temperature annealing. We thoroughly evaluate our method on various reasoning and planning tasks and observe substantial overall improvement over existing heuristic baselines without significant computational overhead.

Figures & tables

Panel of Figure 1 — Figure 1: (a) Confidence-based sampling in non-autoregressive decoding suffers from premature constriction of the generation window due to proximity bias , where neighboring tokens accumulate confidence sequentially. This leads to spatial error propagation, making early unmasking decisions decisive for the entire trajectory. (b) We introduce a lightweight planner and EOS temperature annealing to strategically guide initial position selections toward the target trajectory

Figure 2: Performance in confidence-based non-autoregressive decoding across different diffusion timesteps(

T

) when the generation length(

L

) is fixed at 256.

Figure 3: Unmasked token position index and average ratio of predicting EOS token across diffusion timesteps (x-axis) when evaluating on GSM8K testset with

T=128,\;L=256

. Since two tokens are unmasked simultaneously at each step, we plot both positions, distinguishing between the earlier(front) and later(back) tokens in the sequence. Token position index and ratio of EOS is averaged over all test instances.

Figure 4: Pass@k accuracy on GSM8K with a fixed budget of

T=32

across different k(x-axis). The impact of uniform sampling in Position Selection injected at different decoding steps is compared with token-level Temperature Sampling . The solid red line represents randomness applied only at the initial step, while dashed red lines indicate delayed randomness introduced at intermediate steps. The dashed black line denotes the deterministic confidence-based baseline.

Figure 5: Accuracy for randomness introduced via Position Sampling and Temperature Sampling, categorized into Correct Paths, Incorrect Paths, and their union. Error bars represent 95% bootstrapped confidence intervals.

Sampling Method	Randomness			Dataset				Avg
Sampling Method	T.S.	P.S.	Step	GSM	MATH	CTD	SDK	Avg
Top1 Prob	-	-	-	46.6	19.2	42.2	71.2	44.8
+ Planner	-	L	1st	55.0	22.4	44.1	65.2	46.7
+ EOS Temp.	-	E	all	50.9	22.4	46.1	63.6	45.7
+ Both	-	L&E	all	56.8	22.8	43.8	67.0	47.6
Prob Margin	-	-	-	47.2	19.6	41.4	71.7	45.0

Table 1: Accuracy of basic sampling methods and our strategy on GSM8K (GSM), MATH, Countdown (CTD), and Sudoku (SDK) under a constrained compute budget (

T=32

). All results are obtained under the non-autoregressive decoding regime. T.S. indicates randomness in Token Selection, while P.S. denotes randomness in Position Selection. T represents Temperature sampling, L denotes Learned position selection, E means EOS token logit annealing, and U indicates Uniform sampling.The highest score within each column is marked in bold , and the second best is underlined . Both includes Planner and EOS temperature annealing.

Sampling Method	GSM8K		Math		Countdown		Sudoku		Average
Sampling Method	64	128	64	128	64	128	64	128	64	128
Top1 Prob	32.1	23.3	17.0	15.2	41.4	41.8	68.4	70.5	39.7	37.7
Prob Margin	32.7	25.2	17.6	18.2	46.1	44.5	71.3	73.8	41.9	40.4
Ancestral	46.5	52.3	15.4	18.2	25.4	23.8	19.1	18.5	26.6	28.2
Temperature	38.2	26.8	18.2	15.8	41.4	43.8	69.7	70.4	41.9	39.2
Init. Position	52.2	44.9	21.2	20.6	36.7	39.5	56.4	66.9	41.6	43.0

Table 2: Accuracy of baselines and our method on GSM8K, MATH, Countdown, and Sudoku under higher timestep budget(

T=64,\;128

). The highest score within each column is marked in bold , and the second best is underlined . Both includes Planner and EOS temperature.

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Unmasked token position index and average ratio of predicting EOS token across diffusion timesteps (x-axis) when evaluating with

T=32,\;L=256

Figure 7: Unmasked token position index and average ratio of predicting EOS token across diffusion timesteps (x-axis) with Semi-AR decoding.

Figure 8: Top 1 probability predicted at each diffusion timestep(x-axis) for each token position(y-axis) in Non-Autoregressive decoding under lenient budget( $T=128$ ).

Figure 9: Top 1 probability predicted at each diffusion timestep(x-axis) for each token position(y-axis) in Non-Autoregressive decoding under tight budget( $T=32$ ).

Figure 10: Top 1 probability predicted at each diffusion timestep(x-axis) for each token position(y-axis) in Semi-Autoregressive decoding under lenient budget( $T=128$ ).

$T$	GSM8K	MATH	Countdown	Sudoku
32	39.9 (101.8)	17.2 (78.6)	50.0 (121.8)	7.5 (156.2)
64	45.5 (78.9)	17.0 (35.9)	50.8 (121.2)	8.6 (147.5)
128	33.9 (47.8)	17.2 (22.9)	49.6 (120.4)	0.0 (5.8)

Appendix

Table 3: Performance of Dream 7B Instruct in non-autoregressive decoding with confidence criteria across different diffusion timesteps(

T

) when the generation length is fixed at 256. Effective token counts reported in parentheses.

Figure 11: Unmasked token position index and average ratio of predicting EOS token across diffusion timesteps (x-axis) when evaluating Dream 7B Instruct on GSM8K testset with

T=128,\;L=256

Figure 12: Unmasked token position index and average ratio of predicting EOS token across diffusion timesteps (x-axis) when evaluating Dream 7B Instruct on GSM8K testset with

T=32,\;L=256

Figure 13: Pass@k accuracy of Dream 7B Instruct on GSM8K with a fixed budget of

T=32

across different k(x-axis). The impact of uniform sampling in Position Selection injected at different decoding steps is compared with token-level Temperature Sampling .

Figure 14: Accuracy of Dream 7B Instruct for randomness introduced via Position Sampling and Temperature Sampling, categorized into Correct Paths, Incorrect Paths, and their union. Error bars represent 95% bootstrapped confidence intervals.

Figure 15: Unmasked token position index and average ratio of predicting EOS token across diffusion timesteps (x-axis) when evaluating LLaDA 8B Instruct on MATH with

T=128,\;L=256

Figure 16: Unmasked token position index and average ratio of predicting EOS token across diffusion timesteps (x-axis) when evaluating LLaDA 8B Instruct on Countdown with

T=128,\;L=256

Figure 17: Unmasked token position index and average ratio of predicting EOS token across diffusion timesteps (x-axis) when evaluating LLaDA 8B Instruct on Sudoku with

T=128,\;L=256

		pass@1	pass@4	pass@8	pass@16	pass@32
GSM8K	P.S.	52.6	77.7	84.8	89.8	92.5
	T.S.	45.0	66.7	74.1	79.9	84.2
	diff	7.6	11.0	10.8	9.9	8.3
MATH	P.S.	17.6	33.4	41.0	48.6	55.2
	T.S.	19.4	28.4	33.2	39.2	43.4
	diff	-1.8	5.0	7.8	9.4	11.8

Appendix

Table 4: Comparison of Pass@k accuracy between uniform Position Selection (P.S.) injected solely at the first step and continuous token-level Temperature Sampling (T.S.).

Figure 18: Accuracy of LLaDA 7B Instruct on MATH for randomness introduced via Position Sampling and Temperature Sampling, categorized into Correct Paths, Incorrect Paths, and their union. Error bars represent 95% bootstrapped confidence intervals.

Sampling Regime	Time Schedule	Dataset				Avg
Sampling Regime	Time Schedule	GSM	MATH	CTD	SDK	Avg
SemiAR	Linear	46.5	19.8	25.4	16.2	27.0
NAR	Linear	46.6	19.2	42.2	71.2	44.8
NAR	Progressive	45.0	16.4	39.5	67.4	42.1
NAR	+ Ours	56.8	22.8	43.8	67.0	47.6

Appendix

Table 5: Accuracy on GSM8K(GSM), MATH, Countdown (CTD), and Sudoku (SDK) under a constrained compute budget (

T=32

) in Semi-autoregressive and non-autoregressive regime. The highest score within each column is marked in bold .

Figure 19: Performance on each task across different numbers of candidates

P

Sampling Method	Dataset				Avg
Sampling Method	GSM8K	MATH	Countdown	Sudoku	Avg
Top1 Prob	46.6	19.2	42.2	71.2	44.8
+ Ours	56.8	22.8	43.8	67.0	47.6
Margin	47.2	19.6	41.4	71.7	45.0
+ Ours	58.6	23.0	45.3	69.5	49.1
Ancestral	42.1	13.2	22.7	17.8	23.9

Appendix

Table 6: Accuracy of basic sampling methods with and without our learned planner on GSM8K, MATH, Countdown, and Sudoku under a constrained compute budget (

T=32

). All results are obtained under the non-autoregressive decoding regime. The highest score within each column is marked in bold , and the second best is underlined . +Ours indicates the integration of our learned planner and EOS temperature.

Sampling Method	Dataset		Avg
Sampling Method	GSM8K	MATH	Avg
Top1 Prob	39.9	17.2	28.5
+ Ours	52.4	21.4	36.9
Prob Margin	40.9	16.2	28.5
+ Ours	54.7	19.6	37.1
Ancestral	28.7	12.8	20.7

Appendix

Table 7: Accuracy of basic sampling methods and our learned planner on GSM8K and MATH using Dream 7B Instruct under a constrained compute budget (

T=32

). The highest score within each column is marked in bold , and the second best is underlined . +Ours indicates the integration of our learned planner and EOS temperature.

Non-Autoregressive Text Generation

Non-autoregressive text generation produces multiple tokens in parallel to reduce decoding time while maintaining coherent, useful outputs.

See all 30 papers

Discrete Diffusion Models Exploit Asymmetry to Solve Lookahead Planning Tasks
Feb 23, 2026Itamar Trainin, Shauli Ravfogel, Omri Abend +1Diffusion Language Models Autoregressive Language Modeling
Hebrew University of Jerusalem · New York University
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
Efficient Diffusion LLMs via Temporal-Spatial Parallel Decoding and Confidence Extrapolation
May 29, 2026Zekai Li, Ji Liu, Yiqing Huang +3Adaptive Computation for Diffusion Models Diffusion Language Model Decoding
Advanced Micro Devices, Inc. (AMD)

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

Discrete Diffusion Models Exploit Asymmetry to Solve Lookahead Planning Tasks

Itamar Trainin, Shauli Ravfogel, Omri Abend, Amir Feder

Hebrew University of Jerusalem · New York University

Diffusion Language Models Autoregressive Language Modeling

Open paper AskPDF arXiv

While Autoregressive (AR) Transformer-based Generative Language Models are frequently employed for lookahead tasks, recent research suggests a potential discrepancy in their ability to perform planning tasks that require multi-step lookahead. In this work, we investigate the distinct emergent mechanisms that arise when training AR versus Non-Autoregressive (NAR) models, such as Discrete Diffusion Models (dLLMs), on lookahead tasks. By requiring the models to plan ahead to reach the correct conclusion, we analyze how these two paradigms fundamentally differ in their approach to the problem. We identify a critical asymmetry in planning problems: while forward generation requires complex lookahead at branching junctions, reverse generation is often deterministic. This asymmetry creates an opportunity for NAR models. Through mechanistic analysis of training and inference dynamics, we demonstrate that NAR models learn to solve planning tasks by utilizing future tokens to decode backwards, avoiding the need to learn complex traversal mechanisms entirely. Consequently, we report that both AR and NAR models are able to achieve perfect accuracy on the lookahead task. However, NAR models require exponentially fewer training examples and shallower architectures compared to AR models, which often fail to converge without specific curriculum adjustments.

Abstract

Non-Autoregressive Text Generation

Explore similar work

Discrete Diffusion Models Exploit Asymmetry to Solve Lookahead Planning Tasks

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

Efficient Diffusion LLMs via Temporal-Spatial Parallel Decoding and Confidence Extrapolation