Register Tokens for Bounded-State Reasoning in Diffusion Language Models
Register tokens let diffusion language models like LLaDA and Dream carry reasoning state across cleared chunks, gaining up to 19.5 points on code.
Researchers propose register tokens: dedicated fixed-position tokens whose continuous hidden states are trained to carry reasoning progress across generation chunks in masked diffusion language models. After decoding and clearing a chunk, the model continues from the prompt and the carried register state instead of retaining earlier text. On LLaDA and Dream, registers outperform discrete-text carry on every benchmark, with gains up to 8.5 points on math and 19.5 points on code. Registers are especially effective for bounded code generation and can be further refined with reinforcement learning on long-horizon reasoning tasks.
- Registers are dedicated fixed-position tokens trained to carry reasoning progress across generation chunks.
- Tested on LLaDA and Dream, registers outperform discrete-text carry on every benchmark.
- Gains reach 8.5 points on math and 19.5 points on code benchmarks.
- Registers are especially effective for bounded code generation spanning multiple chunks.
- Register-based reasoning can be refined with reinforcement learning on long-horizon tasks.
Full article157 words · extracted from huggingface.co · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://huggingface.co/papers/2609.16372