The State-Prediction Separation Hypothesis

论文解决了什么问题

标准 Transformer 在每个时间步使用同一组隐藏状态既预测下一个 token，又为后续步骤提供 KV cache，两种梯度在同一表示上竞争，可能限制语言建模效率。论文想验证：如果将“状态保存”和“下一 token 预测”拆分到两条独立计算流，是否能在相同参数规模下获得更低的损失和更好的下游泛化。

核心方法

作者在每两个输入 token 之间插入一个可学习的 <predict> token，形成输入流与预测流交替的序列。输入 token 的 KV 条目永久保留在 cache 中，而 <predict> 条目仅在一个固定大小的滑动窗口内可见，随后被丢弃；损失只在 <predict> 位置计算。这样状态准备梯度主要回传到输入流，预测梯度主要回传到预测流，实现角色分离。

实验结论

• 在 53M 到 1.678B 参数五个规模上，分离变体的验证 NLL 始终低于标准 Transformer，下游零样本准确率平均提升 2.3–3.1 个百分点，且只需约一半训练 token 即可匹配标准模型的损失。推理时持久 KV cache 大小与标准模型相同，峰值显存增加约 1%，端到端吞吐下降不超过 10%。优势随规模扩大而放大，但将角色反转（让预测流持久、输入流窗口化）会导致性能下降，表明特定角色分配是关键。

适合谁阅读