Register Tokens for Bounded-State Reasoning in Diffusion Language Models

2026-09-14

AI Summary

研究通过在扩散语言模型中使用固定大小的寄存器令牌来解决跨生成块推理的问题，提高了数学和代码生成任务的表现。

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.

Problem

• diffusion language models
• bounded-state reasoning
• register tokens

Innovation

• register tokens
• bounded-state reasoning
• diffusion language models
• continuous hidden states