Register Tokens for Bounded-State Reasoning in Diffusion Language Models

The paper investigates whether masked diffusion language models (dLLMs) can continue reasoning after earlier generated text is cleared by using a fixed-size carried state implemented as register tokens. These are dedicated fixed-position tokens whose continuous hidden states are trained to carry reasoning progress across chunks. After post-training dLLMs to decode a chunk, clear the text while preserving register values, and resume decoding, evaluations on LLaDA and Dream show that registers outperform discrete-text carry on every benchmark. Gains reach up to 8.5 points on math and 19.5 points on code, with effectiveness particularly prominent in bounded code generation.

In our main comparisons on LLaDA and Dream, registers outperform discrete-text carry on every benchmark.

Registers achieve gains of up to 8.5 points on math.

Registers achieve gains of up to 19.5 points on code.