扩散语言模型用 Register Tokens 实现有界状态推理

HuggingFace Daily Papers(社区热门论文)·2026-09-14 08:00·3天前
AI 导读

研究者为掩码扩散语言模型(dLLM)提出 register tokens:用少量固定位置 token 的连续隐状态跨生成块携带推理进度,从而在清空已生成文本后仍能继续推理。

HuggingFace Daily Papers(社区热门论文)
41AI 编辑部评分,满分 100

扩散语言模型用 Register Tokens 实现有界状态推理

2026-09-14 08:00· 3天前
AI 导读

研究者为掩码扩散语言模型(dLLM)提出 register tokens:用少量固定位置 token 的连续隐状态跨生成块携带推理进度,从而在清空已生成文本后仍能继续推理。

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.

来源:HuggingFace Daily Papers(社区热门论文)· arxiv.org