arxiv:2609.16372
Published on Sep 14
· Submitted by
Albert Ge on Sep 16
Upvote
7
Authors:
,
,
,
,
,
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.
View arXiv page View PDF GitHub 2 Add to collection
Community
Paper submitter 6 days ago
We give masked diffusion LMs a few register tokens: fixed positions the model writes its reasoning progress into and reads back from while generating. Bidirectional attention makes these positions both readable and writable during decoding, which causal models cannot do. The model denoises a fixed-size chunk, writes to the registers, the text is cleared, and the next chunk continues from the prompt and registers alone, so the active window never grows.
Main results on LLaDA-8B-Base and Dream-7B-Base: registers beat carrying the last few generated tokens on every math and code benchmark (up to +8.5 on GSM8K, +19.5 on MBPP), with the largest gains on code, where a correct program has to span several chunks. Registers also work with a chunked version of diffu-GRPO, and the final answer and running total are linearly decodable from them on a long-arithmetic task.
Everything is released:
-
Code and training recipes: https://github.com/SprocketLab/dllm-registers
-
Checkpoints (registers, discrete-text, memory-token, and full-SFT baselines for both backbones): https://hf-mirror.com/collections/albertge/dllm-registers-6a2e409ed8c60039981a229c
-
60K math + code training mixture: https://hf-mirror.com/datasets/albertge/mix60k-math-code-sft
6 days ago
This is an automated message from the Librarian Bot. I found the following papers similar to this paper.
The following papers were recommended by the Semantic Scholar API
-
Distilled Continuous Diffusion Language Models Can Write Code in Few Steps---or One (2026)
-
Reducing Pretraining-Generation Mismatch in Diffusion Language Models (2026)
-
RecurTrace: Adaptive Latent Reasoning with Loop-Time Memory (2026)
-
Penelope: Localized Latent Recurrence for Efficient Structured Reasoning (2026)
-
Line-Coupled Language Model (2026)
-
Speculative Correction: Draft-then-Refine Decoding for Diffusion Language Models (2026)
Please give a thumbs up to this comment if you found it helpful!
If you want recommendations for any Paper on HF Mirror checkout this Space
You can directly ask Librarian Bot for paper recommendations by tagging it in a comment: @librarian-bot recommend
Upload images, audio, and videos by dragging in the text input, pasting, or clicking here.
Tap or paste here to upload images
· Sign up or log in to comment
Upvote
7
Get this paper in your agent:
hf papers read 2609.16372
Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash
Models citing this paper 8
Browse 8 models citing this paper
Datasets citing this paper 1
Spaces citing this paper 0
No Space linking this paper
Cite arxiv.org/abs/2609.16372 in a Space README.md to link it from this page.