用于扩散语言模型中有界状态推理的注册令牌

# 为扩散语言模型中的有界状态推理注册令牌
来源：https://arxiv.org/html/2609.16372 \\tl\_set:Ne\\tcboxmath tcboxmath\\tl\_set:Ne\\tcbhighmathtcbhighmath
###### 摘要
掩码扩散语言模型（dLLMs）通过迭代地对掩码令牌进行去噪，并利用双向注意力来生成文本。要将推理过程扩展到多个生成块（generation chunks），通常需要将早期生成的文本保留在上下文中。我们探究的问题是：dLLM 能否在清除该文本后，仅使用一个固定大小的携带状态（carried state）继续推理？我们将此状态实现为少量**寄存器令牌**：这些是专用的、位置固定的令牌，其连续隐藏状态经过训练，能够在不同的生成块之间承载推理进度。我们对 dLLMs 进行后训练，使其能够解码一个文本块，在保留寄存器值的同时清除该块，并基于提示（prompt）和携带状态继续解码。在对 LLaDA 和 Dream 的主要对比实验中，寄存器在每个基准测试上都优于离散文本携带方式，在数学任务上最高提升 8.5 分，在代码任务上最高提升 19.5 分。寄存器在有界代码生成（bounded code generation）中尤其有效，因为正确的程序通常跨越多个块。最后，寄存器可以通过在长期推理任务上进行强化学习来进一步优化。
## 1 引言
寄存器令牌问题：一个球体的体积和表面积（分别以立方单位和平方单位计）在数值上相等，该球体的半径长度是多少？块 0 输出：设球体的半径为 r。球体的体积为 V = 4/3 π r^3。球体的表面积为 A = 4 π r^2。... 1/3 r^3 = r^2 ... 寄存器令牌：相同问题；之前的文本已清除；寄存器状态已携带。块 1 输出：重新整理得到 r：... => r^2 (r - 3) = 0。这给出了两个可能的 r 值：r = 0 或 r = 3。由于问题要求的是球体的半径，半径必须为正。寄存器令牌：相同问题；之前的文本已清除；寄存器状态已更新。块 2 输出：\boxed{3}。... dLLM 上下文窗口。
图 1：使用寄存器进行有界状态多块推理。一个 dLLMs 对固定大小的掩码令牌块进行去噪，然后在生成窗口被清除之前将其进展写入寄存器令牌插槽。下一个窗口接收任务和携带的插槽，并从这个有界的连续状态继续，而不是从先前的生成文本继续。
扩散大语言模型（dLLMs）[1 (https://arxiv.org/html/2609.16372#bib.bib34),2 (https://arxiv.org/html/2609.16372#bib.bib35)]是自回归模型的一种有前途的替代方案，通常在并行生成的同时达到其性能。然而，dLLMs 中的推理仍然具有挑战性：缺乏从左到右的自回归结构使得在多个步骤中保持连贯的思维链变得困难 [3 (https://arxiv.org/html/2609.16372#bib.bib49),4 (https://arxiv.org/html/2609.16372#bib.bib37)]。现有方法通过监督和强化学习目标来改进 dLLM 的推理 [3 (https://arxiv.org/html/2609.16372#bib.bib49),5 (https://arxiv.org/html/2609.16372#bib.bib52)]，但仍保持先前生成的文本可用。我们则研究在上下文重置后，推理能否仅使用固定大小的携带状态继续进行。
BERT [6 (https://arxiv.org/html/2609.16372#bib.bib26)] 和视觉 Transformer [7 (https://arxiv.org/html/2609.16372#bib.bib27),8 (https://arxiv.org/html/2609.16372#bib.bib25)] 中的专用令牌聚合全局序列信息，而自回归模型在早期位置表现出注意力汇聚（attention sinks）[9 (https://arxiv.org/html/2609.16372#bib.bib24)]。然而，在因果解码中，这些位置是只读的：后续令牌可以关注它们，但无法更新它们。在 dLLMs 中，双向注意力使得固定位置在解码过程中既可读又可写，为学习的携带状态提供了自然机制。我们利用这一特性构建了**寄存器令牌**：固定位置的令牌，模型在其中存储和更新其跨生成块的推理进度。
我们在**有界状态多块推理**（bounded-state multi-chunk reasoning）的设置下研究寄存器：模型在固定大小的窗口内生成文本，来自早期生成块的信息必须通过一个有界表示（而非可见文本）来持久化；参见图 1 [ (https://arxiv.org/html/2609.16372#S1.F1)] 中的示例。在实践中，无论生成多少个块，活动窗口和携带状态都保持相同大小。
块之间携带的状态可以采取几种形式。**离散文本**方法保留最后几个生成的令牌，如 *马尔可夫思维*（Markovian Thinking）[10 (https://arxiv.org/html/2609.16372#bib.bib56)]，或将先前生成的内容总结成一段压缩文本，即自动压缩（autocompaction）[11 (https://arxiv.org/html/2609.16372#bib.bib2),12 (https://arxiv.org/html/2609.16372#bib.bib3),13 (https://arxiv.org/html/2609.16372#bib.bib17)]。纯文本携带状态受限于其令牌所能表达的内容；Memento 还保留了连续的 KV 状态。自回归系统已经探索了其他连续替代方案，包括用于上下文压缩的记忆令牌 [14 (https://arxiv.org/html/2609.16372#bib.bib1),15 (https://arxiv.org/html/2609.16372#bib.bib19)] 和潜在空间推理 [16 (https://arxiv.org/html/2609.16372#bib.bib41)]。寄存器是另一种选择：一种专门为持续推理而训练的连续表示。
训练寄存器具有挑战性，因为新块可以直接关注提示或其自身的未掩码令牌来预测掩码令牌，而无需使用寄存器。因此，我们在块之间清除生成的文本，同时保留寄存器值，并且在训练过程中，有时会屏蔽新块对提示的注意力，并强制某些传递从完全掩码的状态预测该块（见图 2 [ (https://arxiv.org/html/2609.16372#S3.F2)]）。当两者都应用时，寄存器是与提示和先前推理的唯一链接。
我们的贡献包括：
- • 学习在寄存器中存储推理状态：我们训练 dLLMs 将连续状态写入固定寄存器位置，清除生成的文本，并将保存的状态重新插入下一个块。我们的分块 SFT（监督微调）方案支持两个目标：由下一区块损失训练的任务导向寄存器，以及训练用于重建前一个区块的记忆令牌。
- • 比较寄存器与替代方案：我们将寄存器与无携带的全序列 SFT、离散文本携带和经重建训练的记忆令牌进行比较。寄存器在每个主要比较中都优于离散文本，并在 12 个基准中的 10 个中领先于两个携带基线，其在代码方面的优势最为明显，因为大多数成功的寄存器程序跨越多个块。
- • 通过强化学习优化寄存器：我们将寄存器集成到分块 GRPO（Group Relative Policy Optimization）中，这进一步改善了在两个多块任务（倒计时和长算术）上的携带状态。
## 2 相关工作
#### 全局状态与上下文压缩
Vision Transformer 的寄存器为全局计算提供了专用位置 [8 (https://arxiv.org/html/2609.16372#bib.bib25)]；而我们的寄存器则存储跨生成块变化的状态。Gisting [14 (https://arxiv.org/html/2609.16372#bib.bib1)]、AutoCompressors [17 (https://arxiv.org/html/2609.16372#bib.bib18)]、ICAE [15 (https://arxiv.org/html/2609.16372#bib.bib19)] 和 Activation Beacon [18 (https://arxiv.org/html/2609.16372#bib.bib20)] 将输入上下文压缩为连续表示，而 LLMLingua 则裁剪文本令牌 [19 (https://arxiv.org/html/2609.16372#bib.bib21),20 (https://arxiv.org/html/2609.16372#bib.bib22)]。循环记忆 Transformer [21 (https://arxiv.org/html/2609.16372#bib.bib23)] 在循环结构上最接近，在文本段之间传递学习到的记忆。对于自回归推理，马尔可夫思维携带最后几个生成的令牌 [10 (https://arxiv.org/html/2609.16372#bib.bib56)]，而 Memento、Reasoning Cache 和测试时递归思维则总结先前生成的内容 [11 (https://arxiv.org/html/2609.16372#bib.bib2),12 (https://arxiv.org/html/2609.16372#bib.bib3),13 (https://arxiv.org/html/2609.16372#bib.bib17)]。我们在双向 dLLM 中研究连续携带：状态在固定位置被读取和重写，而先前的生成文本被清除。
#### 扩散语言模型与推理
离散扩散模型学习去噪损坏的令牌序列 [22 (https://arxiv.org/html/2609.16372#bib.bib28),23 (https://arxiv.org/html/2609.16372#bib.bib29),24 (https://arxiv.org/html/2609.16372#bib.bib30)]。最近的工作将这种方法扩展到语言建模 [1 (https://arxiv.org/html/2609.16372#bib.bib34),2 (https://arxiv.org/html/2609.16372#bib.bib35),25 (https://arxiv.org/html/2609.16372#bib.bib36)]，并通过监督和强化学习目标改进推理 [3 (https://arxiv.org/html/2609.16372#bib.bib49),4 (https://arxiv.org/html/2609.16372#bib.bib37),5 (https://arxiv.org/html/2609.16372#bib.bib52)]。这些方法保留先前生成的上下文；我们研究在固定上下文预算下如何存储推理。Block Diffusion [26 (https://arxiv.org/html/2609.16372#bib.bib33)] 也以块为单位生成，但保留了块之间的自回归历史。两个并发的方法特别接近：MetaState [27 (https://arxiv.org/html/2609.16372#bib.bib50)] 在冻结的扩散 LM 周围的去噪步骤中增加了循环连续记忆，而《使用潜在令牌的推理》[28 (https://arxiv.org/html/2609.16372#bib.bib51)] 使用仍然掩码的位置进行潜在计算。我们的寄存器则在完成的文本被移除后，在块之间持续存在。
#### 潜在推理
隐式链式思考方法 [29 (https://arxiv.org/html/2609.16372#bib.bib39),30 (https://arxiv.org/html/2609.16372#bib.bib40)]、压缩和软思考 [31 (https://arxiv.org/html/2609.16372#bib.bib42),32 (https://arxiv.org/html/2609.16372#bib.bib43),33 (https://arxiv.org/html/2609.16372#bib.bib45)] 以及暂停或填充令牌 [34 (https://arxiv.org/html/2609.16372#bib.bib46),35 (https://arxiv.org/html/2609.16372#bib.bib47)] 将推理移出显式文本。Coconut [16 (https://arxiv.org/html/2609.16372#bib.bib41)] 是最接近的概念类比，因为它将隐藏状态作为连续输入反馈回来。寄存器将此思想用作扩散块之间的有界记忆，而非一系列潜在推理步骤。更多关联讨论见附录 J [ (https://arxiv.org/html/2609.16372#A10)]。
## 3 方法：用于有界状态推理的寄存器令牌
### 3.1 背景
dLLMs 学习对其中一些令牌已被替换为 \[mask\] 占位符的序列进行去噪。前向过程由时间 t∈[0,1] 索引，其中 t=1 对应完全掩码的序列，t=0 对应原始未掩码的序列。训练目标是预测掩码位置的原始令牌。令 y_0 表示待预测的目标令牌，c 表示任何未损坏的干净条件上下文。我们采样 t∼U[0,1]，并以概率 t 独立地将 y_0 的每个令牌替换为 \[mask\]，产生 y_t∼q_{t|0}(⋅|y_0)。掩码令牌目标由下式给出：
L_mask(θ) = -E_{t,(c,y_0),y_t} [ 1/t ∑_{k∈M_t} log p_θ(y_0^k | c ⊕ y_t) ], (1)
其中 ⊕ 表示连接，M_t={k: y_t^k=[mask]}。预训练是 c=∅ 且 y_0=x_0 的特例。对于在提示-完成对上进行监督微调（SFT），干净提示是条件上下文 c=p，而补全是目标 y_0；只有完成令牌被损坏。在推理时，dLLM 通过模拟反向过程来生成响应。从一个完全掩码的包含 C 个令牌的块开始，掩码预测器 p_θ 迭代地并行预测所有当前掩码的位置；在每一步，选择预测的一个子集（例如，置信度最高的位置）保留，而其余的在下一步被重新掩码，直到所有 C 个位置都被解掩码。当期望的生成超过 C 个令牌时，标准方法是将一个全新的掩码块附加到不断增长的上下文中，使注意力成本随总生成长度呈二次方增长。这引出了我们的核心问题：寄存器令牌能否携带 dLLM 的解码状态跨越生成块？
### 3.2 使用寄存器令牌进行推理
(a) 训练过程：块 0 -> 块 1 -> ... -> 块 K
(b) 推理伪代码
要求：提示 p，寄存器位置 r_1,...,r_R，dLLM d_θ
registers ← MaskEmbeddings(R)
while 答案未生成 do
    chunk ← Denoise_{d_θ}(p with registers, [mask]^C)
    h ← d_θ(p with registers, chunk)
    registers ← h[r_1,...,r_R]
图 2：寄存器训练与推理。(a) 我们将一个长的推理轨迹分割成多个块，并使用掩码 SFT 目标训练每个块。在每个非终端块之后，对其干净文本进行一次未掩码的前向传播，将寄存器状态从模型的隐藏状态中写入；下一个块在先前的文本被移除后，在相同位置接收这些嵌入作为寄存器。在提示掩码的后续传递中，完成和寄存器查询无法关注提示键，因此跨块信息必须通过寄存器传递。橙色箭头显示每个去噪损失通过紧接其前的寄存器提取传递反向传播。(b) 用于在上下文重置时携带寄存器状态的推理伪代码。
我们首先描述寄存器在推理时如何运作，假设它们已经过适当训练。令 R 为寄存器位置的数量，固定为提示中的 r_1,...,r_R。在每个块被去噪后，我们对提示和完成的块运行一次额外的前向传播。我们保存模型在寄存器位置的最后一层隐藏状态，产生一个 R×d 的寄存器嵌入张量。对于下一个块，我们用保存的值替换这些位置的输入嵌入，并保持提示的其余部分不变。然后我们重复此过程。我们在图 F（F）中提供了推理的高层伪代码。