💻 computer science

Dual-Flow Transformers: Decoupling the Primary Prefill Path from Additional Decode Computation

本文介绍了双流 Transformer（Dual-Flow Transformer），这是一种通过采用一个用于提示词编码和 KV 缓存生成的初级流，以及一个仅在解码期间激活的辅助流，从而将提示词处理与自回归解码解耦的新颖架构，进而实现各阶段计算资源的独立缩放，以在降低推理成本的同时提高预测质量。

原作者： Liming Liu, Mingze Wang, Tuo Zhao

发布于 2026-08-14

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Liming Liu, Mingze Wang, Tuo Zhao

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

想象一下你正在经营着一座规模宏大、超级聪明的图书馆，其中一位图书管理员可以回答你的任何问题。这位图书管理员是一个大型语言模型（LLM），一种读过了几乎整个互联网内容的智能人工智能。但问题在于，这位图书管理员有两种截然不同的工作模式。首先是“阅读模式”：当递给图书管理员一本长书（提示词/prompt）时，他们会非常快速地一次性读完，利用脑力来理解上下文。这很快，且消耗大量的思考能力。然后是“写作模式”：一旦理解了这本书，他们就必须逐字逐句地写出答案，每写一个词都要查阅一次笔记。这很慢，而且非常消耗内存，因为他们必须不断翻阅那叠越来越厚的笔记。

长期以来，科学家们认为让这位图书管理员变得更聪明的唯一方法就是让图书管理员变得更大——增加更多的书架、更多的书籍和更大的大脑。但让图书管理员变大，会让“阅读”和“写作”两个阶段都变得更慢、更昂贵。现在的核心问题是：我们能否在不让“阅读”部分变慢或变贵的前提下，专门让图书管理员在“写作”时变得更聪明？这就像是在问，我们能否在图书管理员写下每个词之前，赋予他们一个神奇的“思考停顿”，让他们在写字前先复核一下逻辑，而不需要强迫他们每次都重新阅读整本书。

这篇题为《Dual-Flow Transformers》（双流变换器）的论文提出了一个巧妙的新设计，叫做 Dual-Flow Transformer 来解决这个问题。请把标准的 AI 模型想象成一条单车道的公路，汽车（数据）从提示词的开头一直开到结尾，然后开始逐字写作。Dual-Flow 设计在旁边紧挨着又建了一条平行的第二车道。其运作方式如下：

“主车道”（Primary Lane）是原始的标准路径。它像普通 AI 一样阅读整个提示词，创建一个完美的剧情地图（称为键值缓存/Key-Value cache）。这条车道快速、高效，且保持不变。“辅助车道”（Auxiliary Lane）是全新的、秘密的通道。它完全跳过了阅读阶段。相反，它会在主车道完成提示词阅读后才启动。它直接从提示词的最后一个位置开始，观察主车道制作的地图，进行额外的“思考”以优化预测，然后协助写作。

神奇之处在于，这两条车道共享相同的重型机械（大型数学矩阵），但拥有各自的小型“思考垫”（嵌入/embeddings）。因为它们共享重型机械，所以计算机不需要为了进行额外的思考而加载新的巨大文件。这就像是在厨房里有一个第二位厨师，他使用与第一位厨师相同的炉灶和刀具，但他只在第一位厨师准备好食材后才开始烹饪。第一位厨师（主）完成了繁重的预处理工作；第二位厨师（辅）只在最后上菜前出现，负责添加特殊的酱汁。

研究人员发现这种设置效果非常好。在实验中，他们在不同的模型规模和数据集上进行了测试。他们发现，通过在写作阶段加入这个额外的“思考车道”，AI 比同等规模的标准模型犯错更少（验证损失更低）。这就像图书管理员在读完书后，在开口说话前先想了一下“嗯，这是正确的词吗？”，而这短短一瞬的思考让整个故事变得更好了。

这一发现中最令人兴奋的部分之一是它如何处理“混合专家模型”（Mixture of Experts, MoE）。想象一下，图书管理员拥有一支由 100 名专家组成的团队，但对于任何给定的句子，只调用其中的 5 名专家。在普通模型中，如果你想让更多专家来帮忙，你也必须在“阅读”阶段调用更多专家，这会拖慢速度。有了 Dual-Flow，你可以保持阅读阶段简单（只调用 5 名专家），但在写作阶段调用 10 或 15 名专家。这给了你一种新的权衡方式：你可以保持阅读过程既快又便宜，但只在 AI 真正生成答案时投入更多的“思考预算”。

该论文并不声称它是一个解决了所有 AI 问题的万能灵药，但其结果非常有力。他们证明了在各种场景下，这种双车道方法都能持续提升性能，而无需增加模型的“阅读”负担。他们甚至测试了一个两个车道通过特殊“耦合向量”（信息小桥）进行交流的版本，并发现这能让第二条车道更好地理解第一条车道的想法。

简而言之，Dual-Flow Transformer 是一种聪明的架构微调，它将“阅读”成本与“写作”成本解耦。它证明了你不需要把整个 AI 做得更大来让它变得更聪明；你只需要在它读完问题之后、开始回答之前，给它一点额外的思考时间。这有点像意识到，提高一个学生的作文水平，最好的方式不是让他们把教科书再读两遍，而是让他们在写下第一个句子前，先深呼吸并思考两次。

技术摘要：双流 Transformer (Dual-Flow Transformers)

问题陈述

大型语言模型（LLM）的推理成本正日益由推理阶段而非训练阶段主导。然而，推理的两个主要阶段——预填充（prefill）（处理提示词）与自回归解码（autoregressive decode）（生成 Token）——表现出截然不同的硬件瓶颈：

预填充： 提示词 Token 的并行处理通常是计算密集型（compute-bound），受限于算术吞吐量。
解码： 序列化生成通常是内存带宽密集型（memory-bandwidth-bound），受限于为每个新 Token 必须重复读取不断增长的键值（KV）缓存和模型权重。

传统的缩放策略（增加模型宽度或深度）会同时提高这两个阶段的成本，因为每增加一层，在预填充和解码阶段都需要进行评估。这造成了一种僵化性：即提高生成质量（这通常受益于更多的计算）不可避免地会增加处理提示词的成本。本文探讨了是否可能将额外的学习计算专门分配给续写预测（continuation prediction，即解码阶段），同时保留完整的提示词计算并维持单一且持久的 KV 缓存。

方法论：双流 Transformer

作者提出了 Dual-Flow Transformer，这是一种在结构上将主要的提示词处理路径与辅助的续写预测路径解耦的架构。

核心架构

该模型由两条残差状态轨迹组成，它们共享大部分模型参数，但在状态管理和激活范围上有所不同：

主流（Primary Flow）： 一个完整的、标准的因果 Transformer 路径。它处理整个提示词，写入持久的 KV 缓存，并确定主要状态。它是自包含的，不依赖于辅助流。
辅助流（Auxiliary Flow）： 一个残差轨迹，它读取主状态和缓存，但从不向持久 KV 缓存中写入数据，也从不反馈回主流程。
- 激活机制： 在初始提示词处理（预填充）期间，辅助流会被省略。它仅在最后一个提示词位置（以预测第一个续写 Token）以及随后的每一个生成位置被激活。
- 权重共享： 两个流共享所有的主要注意力（Attention）、MLP 和输出矩阵。
- 分离性： 它们使用独立的 Token 嵌入和轻量级的学习耦合向量。

关键机制

非对称共享 KV 注意力（Asymmetric Shared-KV Attention）： 主流执行标准的因果自注意力。辅助流生成一个新的查询（ $Q_2$ ），并通过学习向量（ $a_\ell$ ）与主查询（ $Q_1$ ）耦合。至关重要的是，辅助流注意力机制指向的是主流的键和值（ $K_1, V_1$ ）。这使得辅助流能够利用缓存的状态，而无需写入自己的状态，从而实现分组执行，使权重和 KV 区域被加载一次即可同时服务于两个流。
轻量级耦合： 信息通过学习向量（ $a_\ell, b_\ell, c_\ell$ ）在查询层、MLP 中间层和 MLP 输出层，从主流单向流动到辅助流。这丰富了辅助流的表示，而不会产生改变主路径依赖关系的副作用。
混合目标（Mixture Objective）： 最终预测是主分布与辅助分布的混合： $p_{mix} = \alpha_t p_{primary} + (1-\alpha_t) p_{aux}$ 。混合权重 $\alpha_t$ 根据主分布的置信度（碰撞概率）进行动态调整，确保在推理过程中主流保持主导地位（至少 50% 的权重）。
路由重放（Router Replay，针对 MoE）： 在混合专家（MoE）配置中，辅助流重用与主流相同的选定专家集（由主路由决定），但以其自身的混合权重应用这些专家。这保留了主流的专家选择逻辑，同时增加了在专家上执行的算术运算。

推理与训练成本

推理： 预填充阶段与标准 Transformer 相比不产生额外成本；辅助流在此阶段被跳过。解码阶段会产生额外的计算（每个 Token 的算术运算量大约翻倍），但不会增加 KV 缓存的大小或唯一权重的集合，因为两个流是分组执行的。
训练： 两个流都会在全序列上进行评估，导致训练 FLOPs 约为原主模型的 2 倍。然而，参数量的增加非常微小（仅增加了独立的嵌入和耦合向量）。

核心贡献

架构： 引入了一种非对称共享 KV 设计，将提示词范围内的计算与续写计算解耦。主流负责状态持久化，而辅助流提供了一个用于增强预测的“残差”轨迹，且不会改变主状态。
经验质量： 在三个实验维度上展示了验证损失（validation loss）的一致性降低：
- 数据缩放（NanoGPT）： 在匹配的 Token 预算下，Dual-Flow 的表现优于标准 Transformer，具有更低的拟合损失底限。
- 模型规模缩放（Dense LLaMA）： 随着模型规模增大，这种增益依然存在。
- 稀疏 MoE： 在带有路由重放的稀疏 LLaMA 式 MoE 配置中得到了验证。
阶段特定专家分配（Phase-Specific MoE Allocation）： 与以往应用统一计算分配的并行流设计不同，Dual-Flow 允许独立配置主流和辅助流的专家扇出（fan-outs）。这实现了一个权衡空间，可以在预填充成本、解码成本和预测质量之间进行优化（例如，在固定预填充成本的同时增加解码计算）。

结果

验证损失： 在 NanoGPT 实验的匹配 Token 预算对比中，Dual-Flow 始终比标准 Transformer 获得更低的验证损失。拟合渐近损失底限从 2.9416（Transformer）降至 2.9013（Dual-Flow）。
计算效率： 在训练计算量匹配的比较中（即 Dual-Flow 模型在完成一半数据步数时，其共享骨干网络的 FLOPs 大约等于 Transformer 完成两倍数据步数时的量），Dual-Flow 仍然实现了更低的验证损失。这表明，将计算分配到相互作用的流中，比单纯增加单个流的数据量更为有效。
消融实验： 组件分析显示，即使是采用了最小化 Dual 结构的构造（独立的嵌入、共享 KV），其表现也已经优于最接近的前作（PHD-2）。加入跨流耦合和概率混合进一步提升了性能。

意义与主张

本文声称，Dual-Flow 建立了阶段特定专家分配的结构性机会。通过将提示词路径与续写路径解耦，该架构允许推理系统将生成时的计算视为一种可扩展的资源，而不会按比例增加处理提示词的成本。

具体而言，作者强调在 MoE 模型中，主流和辅助流可以拥有独立的专家扇出。这创造了一个三向权衡杠杆：

预填充成本： 由主流的专家使用情况决定。
解码成本： 由主流与辅助流的专家使用总和决定。
预测质量： 通过增加分配给续写阶段的总计算量来提升。

这提供了部署特性的灵活性：系统可以在保持提示词计算成本不变的同时，将可用的解码预算投入到更多的辅助专家中，或者反之亦然，而无需改变共享参数池或单一的主 KV 缓存结构。这项工作支持了这样一个前提：即下一 Token 的计算是一个有用的缩放资源，它可以被分配在模型内部，而不是通过外部更长的推理链或多次采样来实现。