💬 NLP

Triplet-Block Diffusion RWKV

本文介绍了 $B^3D-RWKV$ ，这是一种新颖的架构，通过三元组块布局，将因果 RWKV 模型的 $O(L)$ 推理效率与并行双向离散扩散相结合，在实现与现有模型相当的准确性的同时，将解码吞吐量提升了 1.6 倍。

原作者： Ke Lin, Yiyang Luo, Zhaolong Su, Yunya Song, Anyi Rao

发布于 2026-05-26

📖 1 分钟阅读☕ 轻松阅读

原作者： Ke Lin, Yiyang Luo, Zhaolong Su, Yunya Song, Anyi Rao

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

以下是论文《Triplet-Block Diffusion RWKV》（B3D-RWKV）的解读，将其拆解为简单概念并辅以富有创意的类比。

核心难题：“单行道”与“小组项目”之争

想象两种截然不同的写故事方式：

严格因果模型（“单行道”）：
把传统 AI 想象成一位只能查看已输入文字的作家。他们从左到右，逐字书写。他们无法看到接下来的内容。
- 优点： 阅读长文档时速度极快，因为每添加一个新词时无需重读所有内容。
- 缺点： 生成文本时速度较慢，因为他们无法同时写出两个词。必须写完第 1 个词，才能开始写第 2 个词。
扩散模型（“小组项目”）：
想象另一种 AI，它从一张布满“乱码”或“掩码”（如 [MASK]）的空白页开始。它不是逐字书写，而是同时审视整页内容，猜测缺失的单词应该是什么，修正几个，然后重复此过程，直到故事通顺。
- 优点： 它可以同时修正多个单词（并行处理），因此潜在速度可能快得多。
- 缺点： 要做到这一点，它需要同时看到整个上下文（之前的内容和之后的内容）。这通常需要一个非常昂贵且缓慢的计算机架构。

冲突所在： 你很难将这两者混合。“单行道”作家无法向前看，但“小组项目”需要同时审视一切。

解决方案：“三元组块”技巧

作者 Ke Lin 及其同事提出了一种巧妙的训练技巧，称为三元组块布局（Triplet-Block Layout）。他们发现如何教导这位“单行道”作家表现得像一个“小组项目”团队，而无需改变作家的“大脑”。

以下是该技巧的运作方式，使用排练类比：

想象你是一位演员（AI），只能从左到右阅读剧本。你需要学习一个场景，其中必须猜测缺失的台词，但为了猜对，你需要知道当前位置之后的台词。

作者为每个场景设置了一个三部分排练：

第一部分（掩码副本）： 你阅读场景，但一半的台词被黑胶带（[MASK]）遮盖。你从左到右阅读此部分。
第二部分（可损失掩码副本）： 你再次阅读完全相同的场景，黑胶带位置不变。这是接受测试的部分。你必须猜测缺失的台词。
- 魔法时刻： 因为你刚刚阅读了第一部分，你的大脑（AI 的内部记忆）已经吸收了一部分中所有可见的台词。尽管你在第二部分严格地从左到右阅读，但你的大脑已经“知道”场景右侧的上下文，因为那部分信息已在第一部分中被存储。
- 结果： 你得以在保持从左到右阅读的同时，利用“伪双向”知识（既知晓过去也知晓未来）来猜测缺失的台词。
第三部分（清洁副本）： 你最后一次阅读完整、完美的场景。这将重置你的大脑，让你准备好进入下一个场景。

通过重复这种三元组模式（掩码 -> 掩码/测试 -> 清洁），AI 学会了利用来自“未来”的信息（实际上只是训练序列中的前一个块）来预测缺失的单词，同时保持其原始的“单行道”速度。

结果：既快又准

团队利用这种方法构建了一个名为 B3D-RWKV-7.2B 的模型。以下是他们的发现：

速度： 由于他们保留了“单行道”架构（RWKV），该模型在解码时速度极快。他们声称其速度比同一模型的标准版本快 1.6 倍。
智能： 在通用任务（如回答问题或写故事）上，其表现与其他顶级模型一样出色。
权衡： 论文指出，对于需要完美、逐步逻辑的非常复杂的数学问题，扩散的“猜测”性质有时会导致小错误。然而，对于大多数任务，它都能胜任。

一句话总结

这篇论文解决了一个悖论：如何让一位快速、从左到右的作家表现得像一位聪明、全知全能的编辑？

他们的做法是教导作家连续排练三次场景。第一次排练用上下文填充作家的记忆，第二次让他们利用该记忆练习猜测缺失部分，第三次则清空 slate 以准备下一个场景。这使得他们既能保持线性作家的速度，又能获得扩散模型的并行能力。

他们未声称的内容：

他们不声称这适用于医疗诊断或临床用途。
他们不声称这是解决所有 AI 问题的灵丹妙药；他们承认其在处理复杂数学结构时略有困难。
他们明确说明没有更改模型的安全功能，因此它继承了原始模型的任何偏见。

技术摘要：三元组块扩散 RWKV（B3D-RWKV）

问题陈述

因果 Transformer 语言模型面临两个主要的结构性局限：严格的顺序解码阻碍了并行化，以及每步二次方注意力成本使得长上下文推理计算昂贵。虽然线性时间因果模型（例如 RWKV）和离散扩散模型各自独立地解决了这些弱点，但它们的整合在本质上是不一致的。扩散模型需要双向注意力以并行地对令牌块进行去噪，而因果模型则是严格单向的。现有的结合尝试通常需要架构修改，这会损害线性时间主干的效率，或者需要从零开始训练。

方法论

作者提出了B3D-RWKV，这是一种扩散风格的 RWKV 变体，它在不改变底层因果主干的情况下，统一了 $O(L)$ 推理效率与并行、双向的离散扩散。核心创新是三元组块布局（Triplet-Block Layout），这是一种训练和推理策略，旨在严格因果架构内模拟双向访问。

1. 三元组块布局（训练）

为了在因果模型上实现扩散，每个大小为 $B$ 的逻辑生成块在训练样本中被扩展为三个连续的物理块：

$b_1$ （掩码副本）： 逻辑块的掩码版本。
$b_2$ （可计算损失的掩码副本）： 一个相同的掩码副本，用于计算去噪损失。
$b_3$ （清洁真值副本）： 未掩码版本，用于为后续块刷新循环状态。

伪双向访问：
由于因果主干严格从左到右处理令牌，到达 $b_2$ 中任意掩码位置的隐藏状态已经处理了整个 $b_1$ 。由于 $b_1$ 和 $b_2$ 共享相同的掩码模式， $b_1$ 中的未掩码令牌在模型尝试预测 $b_2$ 中的掩码令牌之前，实际上已被“吸收”进隐藏状态中。这赋予了模型在其逻辑块内的未掩码上下文中伪双向访问能力，既满足了扩散的双向要求，又保持了 RNN/SSM 主干的单向性质。

训练目标：
模型使用掩码预测目标进行训练。损失由以下部分组成：

交叉熵损失（ $L_{CE}$ ）： 在令牌被掩码的监督位置计算。
置信度感知并行损失（ $L_{CAP}$ ）： 遵循 LLaDA-2.0，该术语在已正确预测的位置上锐化概率分布，鼓励推理采样器每步承诺更多令牌。

2. 推理：块级迭代去噪

在推理过程中，模型一次生成一个逻辑块。该过程涉及：

初始化一个所有令牌均被掩码的块。
运行迭代去噪步骤，模型预测掩码位置的 top-1 概率。
承诺那些 top-1 概率超过阈值 $\tau$ 的令牌。
如果少于 $k_{min}$ 个令牌通过阈值，则通过承诺置信度最高的 $k_{min}$ 个令牌来确保进展的后备机制。
一旦所有位置都被承诺，清洁块被附加到前缀，下一个逻辑块开始。

主要贡献

B3D-RWKV-7.2B： 首个在 7B 规模上以掩码预测目标训练的扩散风格、线性时间 RNN 语言模型。它将并行令牌选择集成到 RWKV-7 主干中，而无需修改其原始参数。
三元组块框架： 一种通用的训练配方，可将严格因果主干（满足严格因果性和前向传播状态属性）转换为块扩散模型。这种方法避免了将主干修改为双向块的架构复杂性。
效率与性能： 该方法在保持竞争性推理能力的同时，相比标准因果解码实现了显著的吞吐量提升。

实验结果

该模型在包含通用知识（MMLU、ARC、PIQA、RACE）和推理/数学任务（GSM8K、MATH、GPQA）的 8 项任务套件上进行了评估。

准确率： B3D-RWKV-7.2B 实现了与 RWKV-7 基线及其他同规模扩散模型（如 LLaDA-8B、Dream-7B）相当的准确率。它 notably 在 ARC-Challenge 和 RACE 等推理任务上优于基线，这归因于伪双向感知增强了推理能力。
吞吐量： 在 1K 到 512K 令牌的上下文长度范围内，该模型相比 RWKV-7 基线在解码吞吐量上实现了1.6 倍的平均加速。通过调整采样参数，在轻微质量权衡下可实现 2.02 倍的加速。
数学推理： 论文指出，与因果基线相比，在复杂数学任务（如 MATH）上的性能略有下降。作者将此归因于并行解码的性质，它难以处理符号数学问题所需的高度结构化、逐步验证过程，而此类问题通常不提供部分分数。

意义与主张

该论文声称提供了一种高效途径，可在不改变架构的情况下将预训练因果语言模型转化为扩散语言模型。通过利用三元组块布局，B3D-RWKV 成功统一了线性时间模型的 $O(L)$ 推理效率与扩散的并行解码能力。

作者将这项工作定位为因果范式与扩散范式之间不兼容性的结构性解决方案。他们承认了局限性，包括当前仅在单一主干（RWKV-7）上验证、训练期间 3 倍的物理序列长度成本，以及使用相对较小的 SFT 数据集且未进行 RL 对齐，这可能导致特定任务上的准确率回归。然而，核心贡献在于证明：通过数据布局而非架构重构，可以在严格因果、线性时间的主干上实现高吞吐量、并行扩散生成。