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.