跳转至

几乎零开销的状态-预测解耦架构

Almost Free State Prediction Separation

  • arXiv ID: 2609.03807 [cs.LG]
  • Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI)
  • Submission Date: Submitted on 3 Sep 2026; last revised 11 Sep 2026 (v4)
  • Authors: John Langford, Nathan Godey, Giovanni Monea, Yoav Artzi, Harry Dong, Ying Fan, Gustavo de Rosa, Zheng Zhan

文章背景与核心概要

在传统自回归大语言模型 (Large Language Model, LLM) 中，隐藏状态往往要同时承担两项相互冲突的繁重任务：一方面需要高度压缩并总结历史上下文，另一方面又必须竭力预测下一个 Token。状态-预测解耦 (State–Prediction Separation, SPS) 技术虽然通过将单次前向传播拆分为“状态流”与“预测流”成功化解了这一矛盾，但以往方案却带来了约 1.9 倍的惊人预训练计算开销。本文提出了一套革新机制，借助巧妙的“免费暂停 Token” (Free Pause Token) 设计，使预测流在推理过程中不写入任何键值缓存 (KV Cache) 且不占用额外上下文位置，真正实现了推理阶段的近乎零开销。同时，配合双阶段拆分、共享门控 FFN 等多项训练期优化，在几乎不增加端到端耗时的情况下全面提升了模型预测精度，为大模型基础架构设计开辟了兼顾高收益与低成本的新路径。

📌 内容概要

Summary

状态-预测解耦 (State–Prediction Separation, SPS) 是一项旨在减轻语言模型双重竞争负担的架构技术——通过将模型的前向传播过程拆分为“状态流”与“预测流”，使模型无需同时兼顾“总结历史上下文”与“预测下一个 Token”这两个相互掣肘的目标。然而，传统的 SPS 方法计算代价高昂，通常需要消耗标准预训练浮点运算量 (FLOPs) 的约 \(1.9\times\)。

State–prediction separation (SPS) is a technique that relieves language models of competing burdens—summarizing the context versus predicting the next token—by splitting the forward pass into a state stream and a prediction stream. However, traditional SPS is computationally expensive, typically costing roughly \(1.9\times\) the standard pretraining FLOPs.

本文提出了一系列优化机制，不仅让状态-预测解耦在推理阶段变得近乎零开销，同时还大幅削减了训练阶段的额外负担：

This paper introduces mechanisms to make state–prediction separation almost free during inference while drastically reducing training overhead:

  • 免费暂停 Token (Free Pause Token) ： 构建一种完全不向显存写入任何键值 (Keys / Values) 的预测流，使其直接复用序列中现有的位置，既不会膨胀上下文长度，也无需额外的 KV 缓存 (KV Cache) 或解码步数。
  • 训练期加速优化 (Training Optimizations) ： 引入两阶段拆分算法 (完整保留高效的 FlashAttention 内核算子) 、\(w=0\) 预测窗口、共享门控前馈网络 (Shared Gated FFN) 以及仅在训练尾声分阶段引入解耦策略。
  • Free Pause Token: A prediction stream that writes no keys or values, allowing it to leverage existing sequence positions without increasing context length, requiring a KV cache, or adding decoding steps.
  • Training Optimizations: Incorporates a two-pass split (preserving FlashAttention kernels), a \(w=0\) prediction window, a shared gated FFN, and phasing the separation toward the tail of training.

这些优化方案协同发力，在等浮点算力 (Isoflop) 、等参数量 (Isoparameter) 与等训练 Token 量 (Isotoken) 的公平对比下，相比标准按次词预测训练的 Transformer 实现了显著性能飞跃，同时将实际物理耗时与硬件开销压至最低。

Together, these adjustments achieve an isoflop, isoparameter, and isotoken improvement over standard next-token-trained transformers while minimizing wall-clock and hardware overhead.

🔍 论文摘要

Abstract

状态-预测解耦 (SPS) 通过将前向传播划分为状态流与预测流，有效卸下了语言模型隐藏状态肩负的两大对立重担——即压缩上下文语义与预测下一个 Token。这种解耦固然收效显著，但代价也极其高昂：预测流本质上是对整个主干网络的二次前向遍历，使得预训练所需的 FLOPs 膨胀至约 \(1.9\times\)；而在采用灵活的注意力掩码时，实际训练物理耗时更是进一步恶化。

State–prediction separation (SPS) relieves a language model's hidden state of two competing burdens—summarizing the context and predicting the next token—by splitting the forward pass into a state stream and a prediction stream. The separation works, but it is expensive: the prediction stream is a second pass over the whole backbone, costing \(\sim\)\(1.9\times\) the pretraining FLOPs, and even more in terms of wall-clock time when using a flexible attention mask.

本篇论文致力于让状态-预测解耦变得近乎免费。我们将解耦设计推向极致，提出了“免费暂停 Token”机制：构建一种完全不写入任何键 (Keys) 或值 (Values) 的预测流，从而顺理成章地“搭便车”复用序列现有的位置坐标。在 10 亿 (1B) 参数规模的模型实测中，该机制为标准 Transformer 的下一个 Token 预测精度带来了 2 到 3 centinats 的切实提升；而且因为它没有引入任何额外的位置开销，因而在推理端完全没有代价——既不增加上下文长度，无需消耗 KV 缓存，也无需多余的解码步骤，延迟几乎毫无增加；至于推理阶段轻微增加的理论浮点计算量，由于当前推理吞吐的核心瓶颈并不在计算本身，因此基本可以忽略不计。

This paper makes state–prediction separation almost free. We take the separation to its limit with a free pause token: a prediction stream that writes no keys or values at all and so rides the sequence's existing positions. It improves next-token prediction of a standard Transformer by 2-3 centinats in practice on a 1B parameter model, and because it adds no position it costs nothing at inference—no added context length, no KV cache, no decode steps, and essentially no latency, with the growth in inference flops typically irrelevant as it is not the active bottleneck on throughput.

如此一来，所有的计算额外开销便被完整转移并浓缩在训练阶段。为此，我们设计了四重核心机制来极限压缩训练成本：

The cost is therefore entirely in training where we use four mechanisms to drive it down:

  1. 两阶段拆分机制 (Two-Pass Split) ： 确保高效的 FlashAttention 内核算子始终可用；
  2. \(w{=}0\) 预测窗口 (\(w{=}0\) Prediction Window) ： 极限裁剪预测流所关注的上下文窗口；
  3. 共享门控 FFN (Shared Gated FFN) ： 每个物理位置仅计算一次前馈网络，而非每个数据流各自重复计算；
  4. 尾程分阶段解耦 (Phasing Separation) ： 将解耦机制巧妙安排在预训练运行周期的尾声阶段切入。
  1. A two-pass split that keeps FlashAttention kernels viable,
  2. The \(w{=}0\) prediction window,
  3. A shared gated FFN that evaluates one FFN per position rather than one per stream, and
  4. Phasing the separation onto the tail of the run.

多项优化齐头并进，使得相对于高度优化的预训练基线管线，整体物理训练时间开销被大幅压低至 \(1.33\times\)，同时成功挽回了经典 SPS 约 94% 的性能增益；在更温和平滑的“质量/算力”权衡曲线下，训练耗时增幅甚至可以低至 \(1.09\times\)。此外，共享 FFN 优化还同步削减了推理所需的原始浮点运算量。最终，该研究达成了在等算力 (Isoflop) 、等参数 (Isoparameter) 以及等 Token (Isotoken) 条件下，相对于标准单次预测 Transformer 的全方位实质突破。

Together these bring the overhead versus an optimized pretraining pipeline to \(1.33\times\) wall-clock while recovering ~94% of the gain compared to SPS, and to as low as \(1.09\times\) along a graceful quality/compute tradeoff. Furthermore, the FFN optimization reduces the raw flops required at inference time. The result is an isoflop, isoparameter, and isotoken improvement over standard next token trained transformers.

🔗 全文阅读与相关资源