Masked diffusion language models (dLLMs) generate text by iteratively denoising masked tokens using bidirectional attention. Normally, extending reasoning across multiple generation chunks means keeping earlier generated text in the context window. A new research paper asks whether a dLLM can instead keep reasoning after that text is cleared, relying only on a fixed-size carried state.
The proposed mechanism uses a small number of register tokens: dedicated fixed-position tokens whose continuous hidden states are trained to carry reasoning progress across generation chunks. Because bidirectional attention makes these positions both readable and writable during decoding — something causal models cannot do — the model can denoise a fixed-size chunk, write to the registers, have the text cleared, and continue the next chunk from the prompt and registers alone. The active window therefore never grows.
In main comparisons on LLaDA and Dream, registers outperformed discrete-text carry on every benchmark, with gains of up to 8.5 points on math and 19.5 points on code. The largest improvements appeared in bounded code generation, where a correct program typically spans several chunks. The authors also report that registers can be further refined with reinforcement learning on long-horizon reasoning tasks, and that they work with a chunked version of diffu-GRPO.
On a long-arithmetic task, the paper states that the final answer and running total are linearly decodable from the register states. Code and training recipes are available on GitHub, alongside checkpoints for registers, discrete-text, memory-token, and full-SFT baselines for both backbones, plus a 60K math-and-code training mixture on Hugging Face.