Register Tokens for Bounded-State Reasoning in Diffusion Language Models

arXiv:2609.16372v1 Announce Type: new Abstract: 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.

Related

• Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models
• LogicDiff: Logic-Guided Denoising Improves Zero-Shot Reasoning in Masked Diffusion Language Models
• DA-DLM: Explicitly Modeling Token Dependencies in Diffusion Language Models
• Knowledge Editing in Masked Diffusion Language Models

Source: arXiv cs.CL | 2026-09-16