aiqbench analysis
This research introduces "register tokens" as a mechanism for masked diffusion language models (dLLMs) to maintain reasoning continuity without retaining full prior text in the context window. By training dedicated fixed-position tokens to carry hidden state progress, the approach significantly outperforms discrete-text carrying methods on math and code generation benchmarks. The technique effectively addresses long-horizon reasoning tasks by enabling models to clear generated chunks while preserving critical computational state for subsequent decoding steps.
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.
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.
Computation and Language (cs.CL)
arXiv:2609.16372 [cs.CL]
(or arXiv:2609.16372v1 [cs.CL] for this version)
https://doi.org/10.48550/arXiv.2609.16372 Focus to learn more arXiv-issued DOI via DataCite (pending registration)
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 252 words
Related coverage
Runway’s WorldPrompt and the Engineering of Real-Time Worlds
Latent Space · Sep 25, 2026
Build a Physical AI model factory with NVIDIA Cosmos 3 on SageMaker HyperPod
AWS · Sep 4, 2026
Does DiffusionGemma do latent reasoning?
AI Alignment Forum · Aug 16, 2026
Hierarchical Continuous Diffusion Language Models
arXiv · Oct 2, 2026