🤖 machine learning

Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

本文介绍了学习中继表示（Relay），这是一种使掩码扩散模型能够通过可微分的每令牌通道在去噪步骤间传播潜在信息的方法，从而在规划与编码任务上提升性能，同时显著降低推理延迟。

原作者： Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

发布于 2026-05-25

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是用通俗语言和日常类比对论文《用于前瞻式离散扩散模型的学习型中继表示》的解释。

问题：“硬重置”带来的失忆

想象你正在解一个复杂的谜题，比如数独，或者编写一段代码。你正使用一位非常聪明的助手（一个 AI 模型）来帮助你。

目前，这位助手就像一个患有失忆症的人，每写完一句话，瞬间就会忘记之前发生的一切。

助手看着一张布满问号的空白页。
它苦思冥想，进行复杂的心理运算，决定填入一个或几个字母。
硬重置： 一旦它写下这些字母，它立即丢弃了关于页面其余部分的所有复杂思考、计算和“感觉”。
在下一步中，它从头开始，只看着它刚刚写下的字母，完全忘记了它刚才所做的丰富内部工作。

这被称为**“硬重置”问题**。论文认为这是低效的。这就像一位厨师尝了一口汤，写下“需要加盐”，然后在加盐之前立刻忘记了汤底的味道。他们每次添加食材时，都必须从头重新品尝整锅汤。

解决方案：“中继”接力棒

作者提出了一种名为**中继（Relay）**的新方法。

想象一场接力赛。跑者不是停下来、忘记比赛并重新开始，而是将接力棒传递给下一位跑者。这根接力棒承载着上一位跑者的动量、策略和疲劳，以便下一位跑者能精确地从他停下的地方继续。

在 AI 世界里，这根“接力棒”是一条连续的信息流（隐藏状态），由模型向前传递。

中继之前： 模型进行计算，写下一个 token（词元），然后删除该计算。
中继之后： 模型进行计算，写下一个 token，并且为下一步给自己传递一张“便条”。这张便条写着：“嘿，我正在思考这个特定的模式，我对接下来的这部分有 80% 的把握。”

工作原理：学习传递接力棒

论文引入了一些关键技巧来实现这一目标：

“前瞻式”训练：
通常，AI 模型只被训练去确保当前答案正确。作者训练模型具备“前瞻性”。他们教导模型：“不要只写对今天的字母；写一张便条（中继），帮助你在明天写对字母。”
截断反向传播（BPTT）：
这是一个听起来很复杂的数学术语，意为“为了学习而向前看”。想象你在练习一套舞蹈动作。与其只练习一个动作，不如练习连续三个动作。如果你在第三个动作上绊倒了，你就会意识到：“哦，我在第一个动作上搞砸了，因为我没有为第三个动作做准备。”
论文使用一种称为截断 BPTT的方法，让模型在训练期间能向前看几步。它学习专门向前传递信息，以使未来的步骤更容易。
“软”通道：
向前传递的信息不仅仅是一个单词或字母；它是一个“软”数值（连续值）。这就像调光开关，而不是简单的开/关灯。它允许模型携带那些尚未转化为最终决定的微妙线索和概率。

结果：更快、更智能

作者在两件事上测试了这种方法：

数独谜题： 他们将数独视为一项规划任务。“中继”模型比标准模型用更少的尝试（更少的“前向传递”）解开了谜题，并且犯了更少的错误。它在填充网格时，更能保持对“大局”的把握。
编程（编写软件）： 他们选取了一个最先进的编程 AI（Fast-dLLM v2）并加入了中继系统。
- 更高的准确率： 它写出了更好的代码。
- 更快的速度： 为了获得相同的结果，它需要运行其“思考引擎”的次数减少了32%。

为什么这很重要

论文声称，通过停止“硬重置”并让模型像传递接力棒一样将自己的“思想”向前携带，我们可以使 AI 模型：

更聪明： 它们能在长序列（如复杂代码或逻辑谜题）上进行更好的推理。
更快： 它们不必每次都从头重新计算一切，从而节省时间和能量。

简而言之，这篇论文教导 AI 模型停止忘记自己的工作，并开始像人类解决难题时那样，一步步地在此基础上构建。

技术摘要：面向前瞻性离散扩散模型的学习型中继表示

1. 问题陈述：掩码扩散模型中的“硬重置”

掩码扩散模型（MDMs）通过迭代去噪生成离散序列。从完全掩码的序列开始，模型迭代地解除部分位置的掩码。然而，标准 MDM 推理中存在一个关键的低效问题：在每次前向传播后，模型会丢弃为掩码位置生成的丰富内部连续计算（隐藏状态）。下一步仅从新提交的离散 token 开始。

作者将这种现象称为硬重置问题。虽然 MDM 在生成过程中执行多次前向传播，但由于无法重用中间的连续计算，模型无法有效地在步骤间分摊推理成本。这与循环计算形成对比，后者已证明将固定参数模型在步骤间展开可以扩展模型可逼近的函数类，从而有助于复杂推理任务的性能。本文解决的核心问题是：如何使 MDM 支持循环计算，在不丢弃中间潜在状态的情况下，在连续的解掩码步骤间传递更丰富的信息。

2. 方法论：学习型中继表示（Relay）

提出的解决方案Relay引入了一种连续的、可微的潜在状态，该状态在解码步骤间向前传递，使模型具备“前瞻性”。

核心机制

增强状态轨迹：在推理步骤 $k$ ，模型维护一个增强状态 $s_k = (x_{t_k}, h_k)$ ，其中 $x_{t_k}$ 是当前部分掩码的序列， $h_k$ 是学习到的中继状态（连续潜在向量）。
架构：模型由骨干网络 $f_\theta$ $f_{θ}$ 、token 嵌入 $Emb_\theta$ $E m b_{θ}$ 、解嵌入头 $UnEmb_\theta$ $U n E m b_{θ}$ 以及特定的中继模块 $R_\theta$ $R_{θ}$ 组成。
- 下一个隐藏状态计算为： $h_{k+1} = f_\theta(Emb_\theta(x_{t_k}) + R_\theta(h_k))$ 。
- 中继状态 $h_k$ 被投影并添加到嵌入 token 中，然后输入骨干网络。
- 输出 logits $\ell_k$ 通过 $UnEmb_\theta$ 从 $h_{k+1}$ 导出。
通过截断 BPTT 进行训练：与使用独立步骤训练的标准 MDM 不同，Relay 使用**截断时间反向传播（BPTT）**进行训练。
- 模型从完全掩码的序列开始，执行 $K$ 步（例如 $K=2$ ）的展开。
- 在展开过程中，根据策略（例如置信度阈值）解除位置掩码，并将提交的 token“教师强制”为真实值，以防止展开期间的误差累积。
- 损失是 $K$ 步上交叉熵损失的总和。
- 梯度通过中继路径（ $h_k \to h_{k1}$ ）反向传播，使模型不仅能学习预测当前步骤，还能在 $h_k$ 中编码对未来步骤最有用的信息。

推理

在推理时，解掩码计划或采样策略基本保持不变。唯一的增加是模型将中继状态 $h_k$ 与提交的 token 一起向前携带。中继状态充当一种可微的内存通道，用于存储尚未提交为离散 token 的信息。

3. 主要贡献

Relay 框架：一种将循环计算纳入 MDM 的通用方法，通过截断 BPTT 训练模型在解码步骤间传递学习到的潜在中继状态。该方法与架构无关，并兼容现有的 MDM 加速技术，如块扩散和 KV 缓存。
可扩展性与适应性：通过对 Fast-dLLM v2 (1.5B) 进行全参数适配，验证了 Relay 在大语言模型（LLM）规模下的有效性。该方法以最小的额外训练 FLOPs 成功适配了最先进的离散语言模型（DLM）。
性能 - 延迟改进：实证表明，Relay 推进了性能 - 延迟帕累托前沿，在编码任务上优于标准监督微调（SFT），同时显著降低了推理延迟。
设计空间分析：在具有挑战性的基于数独的规划任务上进行了广泛的消融研究，隔离了展开过程、中继通道和 BPTT 训练目标的贡献。

4. 实验结果

数独规划任务

使用 Sudoku-Extreme 数据集，作者比较了四种训练目标：

MLM：标准均匀掩码扩散（单步，无中继）。
Rollout：多步展开，带教师强制但无中继通道（ $R_\theta \equiv 0$ ）。
Relay (sg)：包含中继路径，但使用停止梯度（中继不经过 BPTT）。
Relay：完整方法，中继经过 $K=2$ 的 BPTT。

发现：

准确率 - 前向评估次数（NFE）前沿：Relay 实现了准确率与 NFE 之间的最佳权衡。
组件影响：从 MLM 过渡到 Rollout 提升了性能。添加中继通道（Relay sg）带来了显著的飞跃。最后，启用 BPTT（Relay）取得了最佳结果，表明模型学会了在更激进的解掩码下保持部分棋盘的自洽性。
效率：在固定的置信度阈值下，Relay 在每次前向传播中提交了更多的单元格，同时保持合法性（无行/列/宫违规），与基线相比，以更少的总前向传播次数达到相同的准确率。

预训练模型适配（Fast-dLLM v2）

作者在代码和数学数据集的混合体上使用 Relay 适配了 Fast-dLLM v2 (1.5B)。

编码性能：在 HumanEval 和 MBPP 基准测试上，Relay 优于标准 SFT。
- HumanEval：Relay 达到 42.1% 的准确率（SFT 为 38.4%），NFE 为 88.3（SFT 为 130.7）。
- MBPP：Relay 达到 46.6% 的准确率（SFT 为 43.9%），NFE 为 78.8（SFT 为 84.8）。
延迟：与标准 SFT 相比，Relay 将推理延迟降低了高达 32%，同时提高了准确率。
内存开销：在 A100 GPU 上的分析显示，虽然 Relay 在第二次前向传播期间增加了活跃激活内存，但峰值内存使用量与 vanilla SFT 几乎相同（约 20.1 GiB 对比 21.2 GiB），因为瓶颈在于词汇投影的反向传播，而不是中继状态的存储。

5. 意义与主张

本文主张 Relay 通过允许 MDM 累积中间连续计算，从根本上改变了迭代非自回归生成的性质。

前瞻性：该方法使模型能够显式优化中间表示以用于未来的解掩码决策，而不是将每一步视为全新的预测问题。
帕累托前沿推进：通过减少达到目标准确率所需的前向传播次数，Relay 推进了性能 - 延迟前沿，使离散扩散模型在复杂推理和生成任务中更加高效。
兼容性：该方法被提出为一种即插即用的增强方案，可与现有的最先进 DLM 和加速技术（如块扩散、KV 缓存）配合使用，无需进行架构 overhaul。

作者总结道，Relay 代表了迈向非贪婪、前瞻性生成的有意义的一步，为潜在推理提供了连续基底，可进一步探索其在可解释性及其他模态（如图像、分子）中的应用。