双流Transformer:将主预填充路径与额外解码计算解耦
Dual-Flow Transformer 将主流程与辅助流程解耦,辅助流仅在解码阶段从最后一个提示位置开始激活,在不写入持久 KV 缓存的情况下增加续写计算。实验显示在匹配 token 数量下验证损失更低,且在 MoE 模型中可按阶段独立分配专家计算。
随着大语言模型(LLM)在实际服务中面对越来越多的请求,推理成本正在从一次性训练成本之外的次要因素,逐渐变成影响总体拥有成本的关键变量。论文指出,推理的两个阶段对底层硬件提出了差异很大的要求。提示预填充(prefill)阶段需要处理整个输入上下文,拥有很高的并行度,通常受限于计算吞吐;自回归解码阶段则是逐个 token 生成,本质上是一个串行过程,在 KV 缓存变大后常常受制于内存带宽。传统的模型扩展方式,无论是增加宽度还是深度,都会让每一层新增的参数同时参与预填充和解码,因此两个阶段的成本会被同步推高。
在这篇论文中,Liming Liu、Mingze Wang 和 Tuo Zhao 提出了 Dual-Flow Transformer,目的是在保持既有预填充计算和单一持久 KV 缓存的同时,把额外的可学习计算集中在续写预测阶段。该架构将模型分成两条流。主流程(primary flow)是一个完整的因果语言模型,负责读取整个提示并像标准 Transformer 一样写入 KV 缓存。辅助流程(auxiliary flow)在提示处理阶段完全不参与计算,只有生成到最后一个提示位置之后才被激活。这样,辅助流程可以在不修改持久状态的情况下增加后续 token 预测的建模能力,也不会对主流程的表示产生干扰。
两条流程共享大部分注意力层、MLP 和输出矩阵,但同时使用独立的 token 嵌入和轻量级的耦合方式。这种设计既控制了新增参数量,也带来了执行层面的复用好处:在批量或分组解码时,权重已经被加载到显存/内存中,主流程 KV 缓存里的键和值也可以被辅助流程重复使用。因此,增加辅助计算不一定意味着成比例地增加访存压力。
为了验证思路,作者进行了匹配 token 数量的对比实验。也就是说,在比较 Dual-Flow 与基线模型时,让它们处理相同数量的 token,而不是简单比较相同参数量。结果显示,Dual-Flow 在不同架构和数据配置下都取得了更低的验证损失,说明额外计算被更有效地分配到了续写预测上。
论文还进一步探讨了混合专家(MoE)模型中的情况。在 MoE 中,每条 token 只会激活一部分专家网络,专家扇出(expert fan-out)决定了使用的专家数量。Dual-Flow 的分离让主流程和辅助流程的专家扇出变成两个相互独立的控制旋钮,从而可以分别调节预填充成本、续写成本和预测质量。作者研究了两种配置:一是在固定预填充专家计算量的前提下增加解码阶段的计算;二是在固定的解码专家预算下,把资源在两条流程之间重新分配。实验揭示了预填充—解码—质量之间的权衡关系,并展示了按推理阶段单独分配专家计算的潜力。
总的来说,Dual-Flow Transformer 为推理成本敏感的 LLM 部署提供了一个新的架构方向:不是简单地把更多计算平摊到所有阶段,而是把额外容量精确投放到最需要“思考”的续写阶段。对于长上下文、高并发服务、边缘设备上的生成任务等场景,这种面向阶段的解耦设计可能会成为模型选型和推理优化的重要参考。