LogicDiff: Logic-Guided Denoising Improves Zero-Shot Reasoning in Masked Diffusion Language Models

Shaik Aman

cs.CLMar 24, 2026

LogicDiff: Logic-Guided Denoising Improves Zero-Shot Reasoning in Masked Diffusion Language Models

Authors: Shaik Aman

Organizations: Independent Researcher, Nellore, Andhra Pradesh, India

Abstract

Masked diffusion language models (MDLMs) generate text by iteratively unmasking tokens from a fully masked sequence. Their standard confidence-based unmasking strategy systematically defers high-entropy logical connective tokens, degrading reasoning performance. We introduce LogicDiff, an inference-time method that replaces confidence-based unmasking with logic-role-guided unmasking. A lightweight classification head (4.2M parameters, 0.05% of the base model) predicts the logical role of each masked position (premise, connective, derived step, conclusion, or filler) from the base model's hidden states with 98.4% accuracy, and a dependency-ordered scheduler unmasks tokens in logical order. In zero-shot settings, LogicDiff improves LLaDA-8B-Instruct accuracy from 22.0% to 60.7% on GSM8K (+38.7 percentage points) and from 23.6% to 29.2% on MATH-500 (+5.6 pp), with less than 6% speed overhead. However, with 8-shot chain-of-thought prompting, the baseline reaches approximately 70% and LogicDiff provides no additional improvement. Analysis reveals that few-shot prompting implicitly resolves the same ordering problem that LogicDiff explicitly addresses, and that fixed role-based ordering can cause premature commitment to numerical values before sufficient context is available. Our results characterize the Flexibility Trap as primarily a zero-shot phenomenon and identify context-adaptive ordering as a key direction for future work.

Figures & tables

Figure 1: LogicDiff system architecture. The frozen LLaDA model produces hidden states and logits. The logic role head classifies each masked position. The dependency scheduler computes priority scores and unmasks tokens in logical dependency order.

Figure 2: Unmasking order comparison. Top: Default confidence-based unmasking generates numbers first and defers connectives to the last step. Bottom: LogicDiff unmasks premises first, then connectives, then derived results, then conclusions.

† 100-example subset; disabled due to catastrophic failure.
Method	GSM8K	MATH-500	Base Modified	Speed
LLaDA Baseline	22.0% (290/1319)	23.6% (118/500)	No	0.18 ex/s
LogicDiff + Consistency	3.0% (3/100) †	—	No	0.09 ex/s
LogicDiff	60.7% (800/1319)	29.2% (146/500)	No	0.17 ex/s
	+38.7 pp	+5.6 pp		${<}6\%$ slower

Table 1: Zero-shot results (0-shot, 256 max tokens, 256 steps). LogicDiff uses the same frozen model with no RL and no fine-tuning.

Figure 3: Zero-shot accuracy on GSM8K and MATH-500.

Setting	Method	GSM8K	Tokens / Steps
0-shot	Baseline	22.0%	256 / 256
0-shot	LogicDiff	60.7% (+38.7 pp )	256 / 256
8-shot	Baseline	70.0%	512 / 256
8-shot	LogicDiff	70.0% (+0.0 pp )	512 / 256

Table 2: LogicDiff under different prompting conditions. With 0-shot prompting, LogicDiff provides a large improvement. With 8-shot prompting, the baseline is already strong and LogicDiff provides no additional benefit.

Method	GSM8K	MATH	RL	Modifies Base	Cost
LLaDA Baseline (0-shot)	22.0%	23.6%	No	No	0
LLaDA Baseline (8-shot)	$\sim$ 78%	$\sim$ 27%	No	No	0
DOS [ 17 ]	—	—	No	No	0
d1 [ 16 ]	$\sim$ 84.5%	$\sim$ 41.0%	Yes	Yes	Days (8 $\times$ A100)
JustGRPO [ 9 ]	89.1%	45.1%	Yes	Yes	Days (8 $\times$ A100)
LogicDiff (0-shot)	60.7%	29.2%	No	No	30 min

Table 3: Method comparison (0-shot setting). LogicDiff is the only inference-time method requiring no base model modification and no RL. Note that with 8-shot prompting, the LLaDA baseline reaches

{\sim}

70–78% without any method modifications.

Role	Count	Distribution	Weight
Premise	7,512	0.8%	1.0
Connective	11,453	1.3%	10.0
Derived	834,290	93.6%	1.0
Conclusion	34,581	3.9%	2.0
Filler	3,596	0.4%	0.5
Total	891,432

Table 4: Training data and role head performance.

Explore similar work

Masked diffusion LLMs can use EoS tokens for hidden reasoning
Mar 5, 2026Sarah Breckner, Sebastian SchusterMasked Diffusion Models Implicit Reasoning in Language Models
Faculty of Computer Science, University of Vienna, Vienna, Austria · UniVie Doctoral School Computer Science, University of Vienna, Vienna, Austria
Mask-Aware Policy Gradients for Diffusion Language Models
Jul 16, 2026Haran Raajesh, Kulin Shah, Adam Klivans +1Masked Diffusion Models RL for Language Model Reasoning
Department of Computer Science The University of Texas at Austin
Mask Is What DLLM Needs: A Masked Data Training Paradigm for Diffusion LLMs
Mar 16, 2026Linrui Ma, Yufei Cui, Kai Han +1Diffusion Language Models Blockwise Diffusion Language Models
Noah’s Ark Lab, Huawei Montreal, Canada & Beijing, China

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

Masked diffusion LLMs can use EoS tokens for hidden reasoning

Sarah Breckner, Sebastian Schuster

Faculty of Computer Science, University of Vienna, Vienna, Austria · UniVie Doctoral School Computer Science, University of Vienna, Vienna, Austria

Masked Diffusion Models Implicit Reasoning in Language Models

Open paper AskPDF arXiv

Diffusion LLMs have been proposed as an alternative to autoregressive LLMs. Curiously, they are especially capable if the generation length, i.e., the number of tokens the model has to output, is set to a much higher value than the correct answer length, and the model pads its answer with end-of-sequence (EoS) tokens. We hypothesize that off-the-shelf masked diffusion LLMs use the representations of EoS tokens as additional computing capacity, which enhances their performance. We experiment with the diffusion models LLaDA1.5, LLaDA2.0-mini, and Dream-v0 on three reasoning tasks: Addition, Entity Tracking, and Sudoku. In a controlled prompting experiment, we confirm that adding EoS tokens improves the LLMs' performance. To further verify whether their representations are used for hidden computations, we perform a causal intervention and transfer the hidden states of the EoS tokens between generations, which increases the models' relative likelihood of outputting the counterfactual answer. The behavioral experiments and the causal interventions indicate that fully bidirectional masked diffusion LLMs can indeed perform latent reasoning in the representations of EoS tokens. Furthermore, we find that these results generalize beyond toy tasks and that providing the model with additional EoS tokens also improves performance on GSM8K and two-hop reasoning.

Abstract

Explore similar work

Masked diffusion LLMs can use EoS tokens for hidden reasoning

Mask-Aware Policy Gradients for Diffusion Language Models

Mask Is What DLLM Needs: A Masked Data Training Paradigm for Diffusion LLMs