Masked diffusion LLMs can use EoS tokens for hidden reasoning
Authors: Sarah Breckner, Sebastian Schuster
Organizations: Faculty of Computer Science, University of Vienna, Vienna, Austria · UniVie Doctoral School Computer Science, University of Vienna, Vienna, Austria
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.
Figures & tables
Figure 1: LLMs use the hidden representations of the prompt and the output tokens for computations (top). When increasing the generation length of diffusion LLMs, they pad excess positions with EoS tokens and thereby increase their reasoning budget, which leads to better task performance (bottom).
Figure 2: Illustration of the tasks.
Figure 3: Aggregate results of the prompting experiment (Exp. 1). The accuracy of the diffusion models with standard errors for generation lengths higher than required to fit the answer. For comparison, we show the performance of the autoregressive LLMs as dashed lines.
Figure 4: Fine-grained results of the prompting experiment (Exp. 1). The accuracy of the diffusion models across task-specific difficulty levels and generation lengths. The autoregressive LLMs are marked as dashed lines.
Figure 5: Aggregate results of the controlled prompting experiment (Exp. 2) where we pad an increasing number of EoS tokens in the start state. The accuracy is shown with standard error, and the autoregressive LLMs are marked with dashed lines.
Figure 6: Results of the causal intervention (Exp. 3). The difference between the rank of the counterfactual tokens in the original generation (blue) and their rank after patching (orange) indicates how much the patched representations affect the output distribution. A decrease suggests that the hidden states contain latent computations.
Figure 7: Accuracy with standard error on more naturalistic datasets, two-hop and GSM8K (Exp. 4). As for Experiment 2, we pad the start state with an increasing number of EoS tokens. For comparison, we show the performance of the autoregressive LLMs as dashed lines.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Accuracies with standard error of appending dots, whitespaces, and random token chains as potential reasoning tokens in the setup of Experiment 2.
Figure 9: Results of the intervention experiment (Exp. 3) executed with dots, whitespaces, and random token chains as potential reasoning tokens.
Figure 10: The number of semantic tokens generated over the predefined generation length. The remaining positions are padded with EoS tokens by the model. The dashed lines indicate the average length of the target answer.
Figure 11: Detailed results of the controlled prompting experiment (Exp. 2). The accuracy of the diffusion models across task-specific difficulty levels and trailing EoS tokens. We show the performance of the autoregressive LLMs as dashed lines.
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.
Shaik Aman
Independent Researcher, Nellore, Andhra Pradesh, India
Masked diffusion language models (dLLMs) can commit tokens in any order -- a freedom marketed as their core advantage over autoregressive decoding. We show that on reasoning tasks this freedom is instead the axis of failure. Logging every commitment during decoding of LLaDA-8B on GSM8K, we find that unconstrained (pure) decoding commits the final answer at 15-24% of the trajectory while half the reasoning region is still masked, and collapses to answer-only outputs on up to 90% of problems as the canvas grows. The cause is not the model's termination beliefs -- EOS "pressure" is nearly identical across decoders -- but reachability: whether the sampler may act on those beliefs at distant positions. A 2x2 prompt-decoder design shows that chain-of-thought helps only under ordered commitment (interaction +34.8 percentage points, 95% CI [26.8, 42.8]; without reasoning text the decoders are indistinguishable), an interaction we decompose into a collapse channel and an order channel and replicate on Dream-7B and MATH-500. A single-knob intervention -- frontier-gated commitment -- causally recovers the full gap (0.528 to 0.852) while preserving up to 4x parallel decoding, along a measured frontier whose optimal window flips from w=1 at full refinement to unconstrained at 8 tokens/step. Our results reframe existing window-style samplers, previously motivated by efficiency, as the minimal fix for a reasoning pathology they were never designed to address.
Jewon Yeom, Jaewon Sok, Seonghyeon Park +3
Graduate School of Data Science, Seoul National University · Department of Rural Systems Engineering, Seoul National University · Department of Aerospace Engineering, Seoul National University
Masked diffusion language models (dLLMs) generate text by iteratively denoising masked tokens with bidirectional attention. Extending reasoning across generation chunks normally requires keeping earlier generated text in context. We ask whether a dLLM can instead continue reasoning after that text is cleared, using only a fixed-size carried state. We implement this state as a small number of register tokens: dedicated fixed-position tokens whose continuous hidden states are trained to carry reasoning progress across generation chunks. We post-train dLLMs to decode a chunk of text, clear it while preserving the register values, and continue decoding from the prompt and carried state. In our main comparisons on LLaDA and Dream, registers outperform discrete-text carry on every benchmark, with gains of up to 8.5 points on math and 19.5 points on code. Registers are especially effective for bounded code generation, where correct programs usually span several chunks. Finally, registers can be further refined with reinforcement learning on long-horizon reasoning tasks.
Albert Ge, Chandan Singh, Yufan Zhuang +3
University of Wisconsin–Madison · Microsoft Research · UC San Diego