状态-预测分离假设 — HF Daily Papers

通过分离Transformer的状态表示和token预测任务，显著提升语言模型的数据和计算效率

前置知识

Transformer自回归语言模型

Transformer通过自注意力机制处理序列，在语言建模中，每个位置的hidden state既要计算当前位置的下一个token预测，又要作为键值对被未来位置访问。这种双重角色通过同一个计算流实现，参数在序列位置上重复使用。

本文的核心假设正是基于这种双重角色的冲突，不理解Transformer的前向传播和反向传播机制就无法理解SPS的设计动机。

KV缓存

在自回归生成中，为了避免重复计算，每个位置的注意力键和值会被缓存起来供后续位置访问。缓存的KV对数量等于已生成的token数量，是推理时显存占用的主要来源。

SPS的关键设计就是控制哪些位置的KV进入持久化缓存，只有理解KV缓存的作用才能理解SPS如何分离状态流和预测流。

梯度分解与归因

通过分析参数对不同位置损失的贡献，可以分解出每个位置i的梯度由两部分组成：预测当前token的梯度和为未来预测提供状态的梯度。这种分解揭示了模型学习的两个不同目标。

本文的实验核心就是证明SPS能更好地分离这两部分梯度流，这是其性能提升的根本机制。

研究动机

标准Transformer中每个隐藏状态承担两个相互竞争的角色：一方面要计算当前位置的下一个token预测（预测任务），另一方面要准备键值对供未来所有位置访问（状态表示任务）。这两个任务通过同一组参数和计算流实现，导致梯度冲突。具体来说，位置i的损失对参数theta的梯度包含两部分：即时预测梯度和未来状态梯度。当模型需要同时优化这两个目标时，可能出现优化冲突——为了预测下一个token，模型可能需要某些表示特征；而为了服务于未来的预测，又需要完全不同的特征。这种张力在高容量模型中尤为明显，相关研究表明Pythia等模型中存在预缓存现象。

本文的目标是本文提出状态-预测分离假设：通过架构上的分离，让状态表示和token预测各自使用独立的计算流，可以提升语言建模性能。具体目标是设计一种Transformer变体，在保持相同参数量和推理开销的前提下，实现更好的数据效率和计算效率，用更少的训练token达到更低的验证损失和更好的下游任务性能。

与已有工作不同的是，现有工作主要关注如何增加计算容量（如插入额外的pause token）或增强未来预测信号（如多token预测），但没有从根本上解决两种任务的架构级冲突。XLNet的双流注意力虽然区分了内容和预测，但服务于排列语言建模而非标准的左到右训练。本文的独特切入点是：不是通过更复杂的损失函数或更多计算，而是通过简单的架构分离——让一个流专门负责状态持久化，另一个流专门负责即时预测，从而结构性地解消这两种任务的竞争关系。

核心方法

SPS Transformer的核心直觉是将Transformer的双流功能拆分到两个并行的计算路径上：输入流专门负责生成持久化的键值对，供未来所有位置访问；预测流专门负责计算下一个token的分布，其键值对仅在一个小的滑动窗口内保留。这种分离通过在每个输入token后插入一个特殊的predict token实现，形成交织序列，其中xi是输入token，rho i是predict token。两者共享位置编码，但扮演完全不同的角色。输入流xi的键值对进入持久化KV缓存，预测流rho i的键值对在窗口大小w后就被丢弃。

SPS的核心创新是精心设计的注意力掩码模式，精确控制哪些位置的键值可以被哪些位置访问。对于查询q在步骤i（可以是xi或rho i），其注意力掩码为：如果q等于xi（输入流），可以访问所有因果的输入位置和窗口内的预测位置；如果q等于rho i（预测流），则额外可以访问其对应的输入位置xi。这个设计确保了梯度流的分离：输入位置的梯度主要来自未来的损失（状态准备），预测位置的梯度主要来自即时的损失（token预测）。关键区别在于持久化策略——只有输入流进入持久化缓存，预测流是短暂的。这与基线模型2X MEMORY（两个流都持久化）和DELAYED STATE（状态延迟提交）形成本质对比。

方法步骤详情

SPS的训练和推理过程分为三个关键步骤。首先，序列构造：给定输入序列x，在每个位置插入predict token，形成交织的序列，其中所有rho i被设置为同一个特殊token。其次，注意力掩码应用：通过精心设计的掩码控制注意力模式，输入流xi可以访问所有之前的输入和窗口内的预测，预测流rho i还可以访问对应的xi，其键值对仅在窗口w等于64内保留。第三，损失计算：语言建模损失只在预测位置计算。推理时，每生成一个token需要前向传播这一对，从rho i的hidden state读取next-token logits，持久化KV缓存只包含输入流，预测流使用固定大小的环形缓冲区。虽然训练时上下文长度加倍，但推理时的开销几乎可以忽略。

技术新颖性

SPS的技术新颖性体现在三个方面：一是架构级的任务分离，不同于以往通过损失函数或数据增强来缓解问题，SPS从计算图结构上解消了两个梯度流的竞争；二是梯度路由的精确控制，通过注意力掩码和KV缓存策略，让预测梯度和状态梯度分别路由到不同的hidden state，这在论文的梯度分析实验中得到验证；三是效率与性能的平衡，虽然训练时上下文长度加倍导致计算增加，但推理时KV缓存大小与标准Transformer相同（1.01倍），吞吐量也保持在90-94%之间，这与2X MEMORY的显存doubling形成鲜明对比。论文通过系统的消融实验（2X MEMORY、DELAYED STATE、REVERSE SPS）证明了分离本身而非额外的计算或内存是性能提升的关键。

**Figure 1**: Standard versus State-Prediction Separation Transformer

实验结果

SPS在5个规模（53M到1.678B参数）的实验中均展现出优异的性能。在FineWeb-Edu验证集上，SPS达到的NLL从XS规模的3.065到XL规模的2.390，相对于STANDARD的改善从-0.042扩大到-0.068，呈现随规模单调增长的趋势。数据效率方面，在1.678B参数规模下，SPS用18B tokens训练就能达到STANDARD用47B tokens才能达到的验证损失，数据效率提升约2.6倍。即使将STANDARD的预衰减训练预算从18B翻倍到36B tokens，仍无法弥补与SPS的差距。泛化性能同样强劲，在4个OOB语料库（WikiText、C4、Pile-Books3、GovReport）上的Corpus NLL平均降低0.09-0.11，在5个零样本基准（ARC-Easy、HellaSwag、PIQA、SciQ、LAMBADA）上的平均准确率提升2.3-3.1个百分点。推理效率方面，SPS的峰值GPU内存仅为STANDARD的1.01倍，吞吐量保持在STANDARD的90-94%之间。梯度分析显示，SPS的输入流在所有偏移量上都承载了比STANDARD更多的未来损失梯度，而预测流的即时预测梯度较低，实现了理想的梯度分离。状态重要性实验表明，限制SPS的持久化状态到64窗口会使NLL增加更多（1.4-2.2倍），证明SPS的持久化状态在推理时确实更重要。

**Table 1**: Model configurations across the five scales

**Table 2**: SPS outperforms all baselines on quality while remaining comparable to STANDARD in memory and throughput

**Figure 2**: SPS trains faster and reaches lower loss at every scale

**Figure 3**: SPS gives consistent downstream accuracy gains

**Figure 4**: SPS works best at small but non-zero w, and outperforms REVERSE SPS at every window

查看结构化数据

任务	指标	本文	基线	提升
FineWeb-Edu验证NLL	负对数似然（越低越好）	2.390 (XL, 20B tokens)	2.458 (STANDARD, 20B tokens)	-0.068 (-2.8%)
Corpus NLL（4个OOB语料平均）	负对数似然（越低越好）	3.338 (XL, 20B tokens)	3.487 (STANDARD, 20B tokens)	-0.149 (-4.3%)
零样本任务平均准确率	准确率百分比（越高越好）	66.3% (XL, 20B tokens)	63.2% (STANDARD, 20B tokens)	+3.1个百分点
数据效率	达到STANDARD质量所需的训练tokens	18B tokens (XL, 1.678B参数)	47B tokens (STANDARD, 1.678B参数)	减少61.7% (2.6×效率)
推理内存开销	峰值GPU内存倍数（相对STANDARD）	1.01×	1.00× (STANDARD), 1.81× (2X MEMORY)	仅增加1%
推理吞吐量	tokens/sec倍数（相对STANDARD）	0.94× (XL)	1.00× (STANDARD), 0.64× (2X MEMORY)	保持94%吞吐量

局限与改进

作者承认了三个主要局限性：一是计算约束，仅在FineWeb-Edu这一个语料库上进行预训练，虽然在OOB语料和零样本任务上验证了泛化性，但未测试不同的预训练混合数据；二是规模限制，最大模型仅为1.678B参数，虽然SPS相对于STANDARD的NLL差距在XS-XL规模上单调增长，暗示趋势会延续到更大模型，但需要进一步验证；三是理论分析的缺失，本文基于受控的消融实验和梯度/状态分析论证了混合预测和状态准备是次优的，但缺乏形式化的刻画——在什么容量、深度或数据条件下这种混合会损害多少性能。此外，SPS的训练计算大约翻倍（因为上下文长度加倍），虽然推理开销很小，但在算力受限的场景下这可能是一个实际限制。论文还提到，两条流目前共享所有参数，进一步分离（如为每条流使用独立的注意力/FFN参数）是否能带来额外收益是一个开放问题。

独立分析的弱点

SPS存在几个潜在的改进方向。首先，训练计算开销问题：虽然推理效率高，但训练时上下文长度翻倍导致计算增加约2倍。改进方向是探索是否可以用更浅或更窄的预测流实现同样的分离效果，或者使用稀疏化的持久化状态来降低训练成本。其次，超参数敏感性：论文主要测试了w等于64的窗口大小，虽然在w为0、16、64、256上显示性能稳定，但在不同规模或数据集上的最优窗口可能不同，需要更系统的超参数搜索。第三，参数共享的局限：当前设计两条流共享所有参数，这可能限制了它们完全专业化各自任务的能力。改进方向是引入流特定的参数（如独立的FFN或注意力参数），在保持总体参数量的同时允许更强的专业化。第四，与多token预测的结合：SPS与多token预测方法在理念上是正交的，但尚未探索两者的结合是否能产生协同效应，这需要设计合适的损失函数和训练策略。

未来方向

作者提出了几个有价值的未来研究方向。一是更大规模的验证：将SPS扩展到10B甚至100B参数规模，验证NLL差距的单调增长趋势是否延续，以及数据效率优势是否会在更大规模上更加显著。二是混合数据的预训练：在多样化的语料库组合上预训练SPS，验证其在真实世界场景中的泛化能力。三是理论分析：对状态-预测冲突的形式化刻画，建立数学模型预测在什么条件下分离会带来多大收益。四是架构探索：探索是否可以用更低的计算开销实现相同的分离效果，例如通过更浅的预测流、稀疏注意力或更高效的注意力内核实现。五是与现有方法的结合：将SPS与多token预测、belief-state objectives等方法结合，可能产生更强的性能提升。六是应用场景扩展：将状态-预测分离的思想应用到其他架构（如RWKV、Mamba等）和任务（如视觉-语言模型、代码生成等），验证其普适性。

复现评估

论文提供了很好的复现条件。代码已开源，基于nanoGPT实现，使用标准的超参数设置。训练使用FineWeb-Edu数据集和GPT-2分词器，序列长度4096，全局批大小96个序列（约40万tokens）。所有模型在bfloat16混合精度下训练，使用学习率调度器（短暂线性预热、峰值学习率常数阶段、最后10%线性衰减）。硬件方面，XS到L使用2张NVIDIA H100 80GB GPU，XL使用2张NVIDIA B200 GPU。论文还提供了详细的消融实验设置（2X MEMORY、DELAYED STATE、REVERSE SPS）和评估脚本。复现难度中等，主要挑战是GPU资源需求——XL规模训练需要B200 GPU，但小规模可以在H100上快速验证。论文还在附录中报告了3个种子的显著性检验，结果显示SPS相对于STANDARD、DELAYED STATE和2X MEMORY的差距在p小于0.005水平上显著，这增加了结果的可信度。

状态-预测分离假设 The State-Prediction Separation Hypothesis