The State-Prediction Separation Hypothesis¶

会议: NeurIPS 2026
arXiv: 2607.01218
OpenReview: https://openreview.net/forum?id=nzCLgrnrNe
代码: https://github.com/lil-lab/sps
领域: 其他
关键词: 状态预测解耦、自回归语言模型、KV缓存、注意力机制、预训练效率

一句话总结¶

论文提出状态-预测分离假设（State-Prediction Separation Hypothesis），通过在输入序列中交替插入专用的 <predict> 标记，将自回归 Transformer 的前向计算与表示解耦为负责持久历史存储的状态流与负责即时预测的预测流，在保持推理期 KV 缓存显存不变和仅牺牲 6% 吞吐的前提下，跨 53M 至 1.68B 参数规模一致取得 2-3 个百分点的下游任务零样本准确率提升与数倍的数据利用效率。

研究背景与动机¶

标准自回归 Transformer 架构中，每个序列位置的前向计算与隐状态表示天然承受着双重职责：它既要汇聚上下文并经过输出头即时预测下一个 token，又要在因果注意力中生成键值对（Key-Value）存入 KV 缓存，作为后续所有 token 可检索的历史状态信息。在现有的单流网络中，这两种截然不同的功能被强行塞入同一套隐藏层激活与参数路径中。这种职责纠缠在反向传播的反向梯度流中体现得尤为明显：针对当前步损失 $\ell_i$ 的即时预测梯度，与来自后续所有位置 $j > i$ 经由 KV 缓存反传回来的状态表示梯度，必须在同一个隐藏状态 $h_i^{(l)}$ 上线性相加并优化同一组参数。

这种双重身份的强行绑定导致了深层次的优化张力。为了做好即时预测，表示倾向于捕捉特定候选词的局部语义与细粒度表层特征；而为了充当持久上下文状态，表示则需要保留全局语义、句法骨架与长期依赖线索。现有工作虽然观察到语言模型中存在“预缓存（pre-caching）”与“记忆痕迹（breadcrumbs）”现象，或尝试通过插入 Pause Token、填充符来增加推理容量，但都未从架构底层将“记忆状态维护”与“即时输出预测”在时序和表示空间上解耦开来。

本文的切入角度是直接质疑 Transformer 将状态与预测混为一谈的设计根基，并正式确立“状态-预测分离假设”：将这两个角色在表征与计算流上彻底解耦，能够从根本上消除梯度冲突，大幅提升语言模型的泛化能力与数据/计算效率。核心 idea：在每个输入 token 后插入专用的学习标记 <predict> 构建双流交织架构，由输入流承担永久 KV 状态存储，由预测流挂载交叉熵损失并在滑动窗口后淘汰其临时 KV 缓存，以极简的注意力掩码实现职责解耦。

方法详解¶

整体框架¶

SPS Transformer 的核心在于将传统的单一因果计算流重构成“输入流（Input Stream）”与“预测流（Prediction Stream）”交错推进的双流结构。给定输入 token 序列 $x = (x_1, x_2, \ldots, x_T)$，模型通过交织插入虚拟预测标记 $\rho_i = \text{<predict>}$，将其扩充为两倍长度的序列 $x \to (x_1, \rho_1, x_2, \rho_2, \ldots, x_T, \rho_T)$。在任意索引 $i$ 处，真实 token $x_i$ 与其配对的预测标记 $\rho_i$ 共享相同的位置编码。

前向传播时，因果注意力掩码被精心设计为不对称结构：所有后续查询都可以访问历史输入 token 构成的永久 KV 缓存，而由 $\rho_i$ 生成的临时 KV 条目则仅在长度为 $w$（默认 $w=64$）的局部滑动窗口内可见，超出窗口后即被逐出。语言建模的目标损失仅在 $\rho_i$ 对应的时间步施加，用于预测下一个实际 token $x_{i+1}$。在推理阶段，解码器仅保留真实 token 的持久 KV 缓存，辅以大小为 $w$ 的极小预测环形缓冲区，单步同时计算 $(x_i, \rho_i)$，彻底消除显存膨胀。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列：真实 Token 与 predict 标记交织"] --> B["双流序列交织与位置编码共享<br/>共享位置嵌入，拆分输入流与预测流"]
    B --> C["不对称因果注意力与滑动窗口淘汰<br/>输入流永久缓存，预测流仅保留 w 步"]
    C --> D["单向解耦损失反向传播<br/>仅在 predict 位置施加下一词预测损失"]
    D --> E["推理期双 Token 联合前向<br/>环形缓冲复用，持久显存开销仅 1.01x"]

关键设计¶

1. 双流序列交织与位置编码共享：解耦状态持久化与即时预测流
针对单流 Transformer 中隐藏表示被迫同时承担即时预测与长期记忆的优化冲突，SPS 在序列维度将每个时间步 $i$ 显式拆解为两个计算步：输入标记 $x_i$ 与预测标记 $\rho_i$。为了保持时序语义的严格一致，二者共享完全相同的旋转位置编码（RoPE），即 $\mathcal{R}(x_i) = \mathcal{R}(\rho_i)$。这种设计的巧妙之处在于，输入流 $\{x_i\}_{i=1}^T$ 不直接产生当前步的预测概率分布，其前向激活完全解放出来专职提炼供后续上下文检索的高质量键值表示；而预测流 $\{\rho_i\}_{i=1}^T$ 则专注于结合已有上下文和当前输入 $x_i$，对即时候选 token $x_{i+1}$ 进行最精准的概率建模。

2. 不对称因果注意力与滑动窗口淘汰：阻断未来梯度的过度污染
为了让分离在计算图和梯度流上真正生效，SPS 通过定制的因果注意力掩码施加严格的不对称访问权限。对于任意查询 $q$ 在步数 $i$： $$\mathcal{A}_{\mathrm{SPS}}(i, q) = \{x_k : k \le i\} \cup \begin{cases} \{\rho_k : i-w \le k < i\} & \text{若 } q = x_i \\ \{\rho_k : i-w \le k \le i\} & \text{若 } q = \rho_i \end{cases}$$ 输入流 $x_i$ 的 KV 条目永久驻留在持久缓存中，供未来任意 $j > i$ 随时查阅，因此输入表示 $h_{x_i}^{(l)}$ 累积了纯粹的状态准备梯度 $\sum_{j>i} \nabla_{h_{x_i}^{(l)}} \ell_j$。相反，预测标记 $\rho_i$ 的 KV 仅允许在局部滑动窗口 $w$（实验证明 $w=64$ 最优）内被查阅，超出 $w$ 步后即被物理废弃。这使得预测标记的梯度几乎全部由即时预测损失 $\nabla_{\theta_{\rho_i}} \ell_i$ 主导，切断了遥远未来对当前预测表征的反向干扰，使两类梯度的分布流向实现彻底分流。

3. 推理期双 Token 联合前向与有界环形缓冲区：零显存膨胀与极低延迟开销
在自回归生成阶段，许多双流或前向倍增方案往往会面临显存成倍增加或延迟翻倍的缺陷。SPS 充分利用了现代 GPU 计算受限于内存带宽（Memory-bound）而非算力的特性，将每次自回归步转化为对对偶元组 $(x_i, \rho_i)$ 的并行批处理前向计算：真实输入 token $x_i$ 写入持久 KV 缓存，而 $\rho_i$ 的键值条目写入长度为 $w$ 的小尺寸环形覆盖缓冲区（Ring Buffer）。下一个 token 的 logits 直接由 $\rho_i$ 的隐藏状态投影得出。由于持久 KV 缓存仅包含真实 token，SPS 的推理期显存占用与标准 Transformer 几乎严格为 1.00:1.01，端到端生成吞吐仅微降 5-10%，打破了结构改动必损推理效率的传统顾虑。

损失函数 / 训练策略¶

SPS 的自回归语言建模损失严格限制在预测标记位置施加，输入 token 位置不计入交叉熵损失： $$\mathcal{L} = \frac{1}{T-1} \sum_{i=1}^{T-1} -\log p(x_{i+1} \mid x_1, \rho_1, \ldots, x_i, \rho_i)$$ 模型骨干采用 RMSNorm、SwiGLU 激活函数（隐藏层倍率 3d）、RoPE 旋转位置嵌入与权重绑定的输出投影头，训练上下文长度为 4,096。为了高效支持非对称交织与滑动窗口掩码，作者基于 Triton 定制了 FlashAttention 核函数，使得 SPS 在训练吞吐上与标准 FlashAttention 保持同等效率，避免了 FlexAttention 的显存低效与计算减速。优化器采用 AdamW，并配合短线性预热、常数峰值与最后 10% 步数余弦线性衰减的学习率调度方案。

实验关键数据¶

主实验¶

实验基于 FineWeb-Edu 语料库预训练，覆盖从 XS (53M) 到 XL (1.68B) 共 5 个参数级别，默认预训练数据量为 20B tokens。下表总结了在 FineWeb-Edu 验证集、4 个域外测试语料（WikiText、C4、Pile-Books3、GovReport）平均 NLL、5 项标准零样本基准（ARC-Easy、HellaSwag、PIQA、SciQ、LAMBADA）平均准确率，以及单卡 H100 上的推理效率对比：

模型规模	方法	验证集 NLL (↓)	域外语料 NLL (↓)	零样本准确率 (%, ↑)	推理相对吞吐 (×, ↑)	峰值显存比 (×, ↓)
XS (53M)	Standard	3.107	4.335	44.9	1.00	1.00
	2x Memory	3.073 (-0.034)	4.311 (-0.024)	47.3 (+2.4)	0.93	1.81
	Delayed State	3.086 (-0.021)	4.312 (-0.023)	46.9 (+2.0)	0.94	1.01
	SPS (本文)	3.065 (-0.042)	4.243 (-0.092)	47.3 (+2.5)	0.94	1.01
S (131M)	Standard	2.858	4.018	49.5	1.00	1.00
	2x Memory	2.829 (-0.030)	3.981 (-0.036)	51.1 (+1.5)	0.93	1.81
	Delayed State	2.829 (-0.029)	3.992 (-0.026)	51.8 (+2.2)	0.94	1.01
	SPS (本文)	2.810 (-0.048)	3.913 (-0.105)	51.8 (+2.3)	0.94	1.01
M (379M)	Standard	2.648	3.732	55.8	1.00	1.00
	2x Memory	2.610 (-0.038)	3.687 (-0.045)	57.3 (+1.4)	0.93	1.81
	Delayed State	2.611 (-0.037)	3.701 (-0.031)	57.5 (+1.7)	0.95	1.01
	SPS (本文)	2.591 (-0.058)	3.627 (-0.106)	58.7 (+2.9)	0.95	1.01
L (831M)	Standard	2.537	3.598	60.1	1.00	1.00
	2x Memory	2.495 (-0.042)	3.542 (-0.055)	62.1 (+1.9)	0.83	1.78
	Delayed State	2.491 (-0.045)	3.532 (-0.066)	61.3 (+1.2)	0.90	1.01
	SPS (本文)	2.470 (-0.067)	3.484 (-0.113)	62.6 (+2.5)	0.90	1.01
XL (1.68B)	Standard	2.458	3.487	63.2	1.00	1.00
	2x Memory	2.411 (-0.047)	3.423 (-0.064)	64.3 (+1.0)	0.64	1.75
	Delayed State	2.410 (-0.048)	3.433 (-0.055)	64.8 (+1.6)	0.94	1.01
	SPS (本文)	2.390 (-0.068)	3.338 (-0.149)	66.3 (+3.1)	0.94	1.01

消融实验¶

为证实性能增益源于“状态与预测的职能分离”，而非额外增加的前向计算步数或缓存容量，作者在 S 尺度（131M）上设计了严格控制变量的架构变体消融与滑动窗口大小扫描：

配置/变体	核心机制差异	验证集 NLL	域外 NLL	零样本准确率 (%)	KV 显存开销	职能是否解耦
Standard	单流前向，状态与预测纠缠	2.858	4.018	49.5	1.00×	否
2x Memory	插入 predict 但永久保留全部 KV	2.829	3.981	51.1	1.81×	否（predict 仍兼具状态功能）
Delayed State	延迟一拍，在 predict 处生成持久状态	2.829	3.992	51.8	1.01×	否（predict 汇聚双重功能）
Reverse SPS	输入流滑动淘汰，predict 充当持久状态	2.845 (w=64)	3.968	50.8	1.01×	是（但反向割裂输入可见性）
SPS ($w=0$)	预测流完全无窗口，直接即发即弃	2.816	3.922	51.6	1.01×	是
SPS ($w=16$)	预测流极短局部窗口	2.812	3.918	51.7	1.01×	是
SPS ($w=64$)	默认设计（局部窗口为 64）	2.810	3.913	51.8	1.01×	是
SPS ($w=256$)	扩大预测流局部窗口	2.811	3.915	51.7	1.01×	是

关键发现¶

分离是核心驱动力而非容量红利：2x Memory 拥有双倍持久 KV 显存开销与相同的每步算力，但在所有模型尺度下 NLL 均显著落后于 SPS。Delayed State 同样提供了额外的单步计算缓冲，但由于状态持久化与预测仍在同一个 slot 执行，其性能显著劣于 SPS。双侧等价性检验（TOST）表明 Delayed State 与 2x Memory 在统计上完全等价，证明不解耦职责的算力堆砌会迅速遭遇瓶颈。
数据利用效率呈现倍数级跃升：在 1.68B 参数规模下，SPS 仅需约 18B tokens 的预退火训练即可达到标准 Transformer 训练 47B tokens 时的验证集损失，数据利用效率提升达 2.6 倍。即便将标准 Transformer 的训练预算翻倍至 36B-40B tokens，依然无法追平 SPS 的表现。
梯度分布的实证剖析：分析未来 $k$ 步损失对当前位置的梯度范数比率 $r(p, k) = \frac{\|\nabla_{\theta_p} \ell_{i+k}\|_2}{\|\nabla_{\theta_p} \ell_i\|_2}$，SPS 的输入流 $x_i$ 持续承载了高强度的未来状态优化梯度（甚至在 $k=512$ 处依然稳健），而预测流 $\rho_i$ 的未来梯度则微乎其微。相反，Standard 模型中两类梯度强行混杂，削弱了长程状态建模的专一性。
强制窗口截断下的敏感度对比：在推理时人为将持久状态截断为滑动窗口 $\omega=64$，SPS 遭受的 NLL 恶化幅度是 Delayed State 的 1.4 至 2.2 倍，这直接实证表明 SPS 的输入流确实承载了模型长期依赖不可或缺的高价值持久记忆。

亮点与洞察¶

直击经典 Transformer 隐式假设的思维盲区：标准 Transformer 将“生成当前字”与“记住历史以备后用”强行共用一套表示，看似优雅紧凑，实则造成了深层优化维度的负向迁移。SPS 通过形式化提出解耦假设，指出了自回归单流建模的内在缺陷。
优雅的不对称注意力与零额外持久显存架构：仅依靠简单的注意力掩码和 token 交错，既未引入复杂的额外网络模块，也未增加模型参数量，推理期持久 KV 缓存维持 1.01x，以极小的工程改造撬动了显著的性能收益。
为数据枯竭时代的 Scaling Law 提供全新抓手：高质量人类文本数据正面临耗尽危机，SPS 在不增加实际模型参数量的前提下，通过结构解耦将单位 token 的数据利用效率提升数倍，且收益随模型规模扩大呈单调放大趋势（从 XS 尺度的 -0.042 NLL 扩大到 XL 的 -0.068 NLL）。

局限与展望¶

预训练计算量翻倍：由于在训练阶段序列中插入了等量的 <predict> 标记，导致等效训练上下文长度翻倍，在同等 token 数量下单步训练算力开销约为标准模型的 2 倍。尽管按实际 GPU 训练耗时核算时 SPS 最终依然显著超越基线，但降低预测流的训练算力（如采用浅层预测流或稀疏注意力）仍是亟待探索的方向。
双流间仍完全共享网络参数：目前的 SPS 架构中输入流与预测流共享全部注意力和 FFN 权重，仅通过掩码和标记位置进行流区分；未来若针对状态准备与预测分派独立的专精参数子网络，或能带来更彻底的解耦增益。
超大规模（>10B）验证尚缺：论文最大实验规模止步于 1.68B 参数，尽管各项指标呈现明显的尺度单调递增规律，但在主流超大模型规模（7B/70B）以及多模态混合训练下的表现仍需进一步验证。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统性提出状态-预测分离假设，揭示了自回归模型长期被忽视的双重表示冲突，架构改动精炼深刻。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 53M 至 1.68B 5 个规模、多任务与多语料评估，精心构建了控制计算量与显存的变体消融，并辅以梯度分布和缓存破坏分析。
写作质量: ⭐⭐⭐⭐⭐ 理论阐述清晰透彻，数学定义严谨规范，图表详实且论证逻辑环环相扣。
价值: ⭐⭐⭐⭐⭐ 方案极易在现代大规模 LLM 预训练管线中落地实施，推理显存友好，数据利用效率提升显著，对架构演进具重要启发。