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

Mar 5, 2026

Masked diffusion LLMs can use EoS tokens for hidden reasoning

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.

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

Ask PDF

Jul 16, 2026

Mask-Aware Policy Gradients for Diffusion Language Models

Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging due to the intractability of the log-likelihood estimation. Existing approaches approximate this log-likelihood by modeling only the token predictions, ignoring the order in which positions are unmasked during generation. We observe that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask. We formalize this as a two-stage action MDP, showing that the policy gradient naturally decomposes into a token term and a masking term. Combining optimization of both terms leads to state-of-the-art outcomes on mathematical reasoning and coding benchmarks, with scores of 87.1% on GSM8K and 53.4% on MBPP.

Haran Raajesh, Kulin Shah, Adam Klivans +1

Department of Computer Science The University of Texas at Austin

Masked Diffusion Models RL for Language Model Reasoning Policy Gradient Methods RL for Diffusion Models

Ask PDF

Mar 16, 2026

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

Discrete diffusion models offer global context awareness and flexible parallel generation. However, uniform random noise schedulers in standard DLLM training overlook the highly non-uniform information density inherent in real-world sequences. This wastes optimization resources on low-density structural glues while leaving high-density logical pivot points severely under-optimized. To address this, we propose an Information Density Driven Smart Noise Scheduler. By extracting information-dense hubs and applying Complementary Priority Masking, our method decouples a single training instance into mutually reinforcing reasoning and syntax samples, forcing the model to master both logical deduction and foundational sequence structure. Experiments demonstrate that our approach improves average accuracy by ~4% across four Code and Math reasoning benchmarks, significantly outperforming uniform baselines. Mechanistic analyses further reveal that probabilistic priority masking effectively mitigates contextual collapse during block diffusion training. Overall, this density-aware strategy efficiently unlocks the reasoning potential of diffusion language models at minimal annotation cost, emerging as a promising new masked data training paradigm for Diffusion LLMs. Our processed dataset can be found at https://huggingface.co/datasets/malr07/opc-sft-stage2-dense-extracted.

Linrui Ma, Yufei Cui, Kai Han +1

Noah’s Ark Lab, Huawei Montreal, Canada & Beijing, China

Diffusion Language Models Blockwise Diffusion Language Models LLM Reasoning

Ask PDF