基于强化学习的快速权重与下一序列预测 Reinforced Fast Weights with Next-Sequence Prediction
REFINE：用RL序列级反馈优化快速权重模型的长上下文能力
前置知识
快速权重架构
快速权重架构是一类替代传统注意力机制的神经网络设计，用固定大小的权重矩阵存储上下文信息。与标准Transformer维护不断增长的KV缓存不同，快速权重模型在处理每个新token时通过梯度更新来动态修改权重矩阵，从而实现常数内存开销的推理。代表模型包括DeltaNet、GatedDeltaNet和LaCT等。权重更新规则为 $W_{t+1} \leftarrow W_t - \eta \nabla_{W_t} \ell(W_t k_t, v_t)$，其中 $W$ 是快速权重，$\eta$ 是学习率，$k_t, v_t$ 是输入token的键值表示。
快速权重架构是本文研究的核心对象，理解其工作原理是理解REFINE如何改进这些模型训练方式的基础
下一Token预测(NTP)
下一Token预测是语言模型预训练的标准目标，通过最小化交叉熵损失来训练模型预测序列中的下一个token。给定输入序列 $S = (x_1, ..., x_T)$，NTP损失为 $\mathcal{L}_{NTP} = \sum_t -\log p(x_{t+1} | x_{\leq t})$。这个目标在每个token位置计算交叉熵，只考虑单个token的预测，而忽略多个后续token之间的语义关系。
NTP是当前快速权重模型训练的标准方法，本文正是针对NTP的局限性提出了新的训练目标
下一序列预测(NSP)
下一序列预测是本文提出的新训练目标，旨在让模型预测给定前缀后语义连贯的多个token序列。与NTP逐token优化不同，NSP在选定位置 $T^* \subseteq \{1, ..., T\}$ 上优化多token序列对齐：$\mathcal{L}_{NSP} = \sum_{t \in T^*} \mathcal{L}_{seq}(\hat{x}_{t+1:t+k}, x_{t+1:t+k})$，其中 $k > 1$。这个目标直接反映快速权重中存储的信息是否能支持多步准确续写。
NSP是REFINE框架的核心创新，它提供了比NTP更合适的训练信号来优化快速权重模型的长上下文记忆能力
GRPO算法
GRPO(Group Relative Policy Optimization)是一种强化学习算法，用于根据rollout及其相对优势计算策略梯度。在REFINE中，同一序列S的所有奖励会被标准化以计算优势值，然后使用GRPO更新策略模型。该算法结合了组内相对比较的优势，使得训练更加稳定有效。
GRPO是REFINE框架中用于优化NSP目标的具体RL算法，理解它有助于把握REFINE的技术实现细节
研究动机
快速权重架构如DeltaNet和LaCT通过固定大小的记忆替代全局注意力机制，实现了常数内存开销的长上下文建模，是Transformer的有力替代方案。然而，这些模型目前都使用标准的下一Token预测(NTP)目标进行预训练，这在快速权重模型上存在根本性缺陷。NTP目标只对下一个token有直接影响，忽略了依赖相同内部状态的后续预测质量。具体而言，NTP的token级反馈只鼓励优化短期似然，限制了快速权重的自适应能力和模型在更长时间范围内的行为表现。实验表明，纯NTP训练的快速权重模型在长上下文检索任务(如Needle-in-a-Haystack)上仍然表现不佳，这源于快速权重更新的不稳定或长期监督的不足。
本文的目标是本文的核心目标是提出一个名为下一序列预测(NSP)的新型训练目标，专门针对快速权重模型的特性进行优化。NSP鼓励模型在给定前缀条件下预测语义连贯的未来token序列，而不是孤立的单个token。同时，作者提出REFINE(Reinforced Fast weIghts with Next sEquence prediction)框架，将NSP形式化为强化学习问题，通过熵驱动的token选择、多token rollout生成、自监督序列级奖励和GRPO优化来高效训练快速权重模型。该框架旨在贯穿语言模型训练的全生命周期：中期训练、后训练和测试时训练。
与已有工作不同的是，虽然先前工作已经探索了标准Transformer上的多token预测(Gloeckle等人, 2024; Liu等人, 2025)，但本文是首个研究基于强化学习的下一序列预测在快速权重模型上应用的工作。REFINE的独特切入角度在于：它将NSP形式化为RL问题而非监督学习问题，解决了标准交叉熵损失无法自然扩展到多token预测的挑战；通过熵驱动的token选择策略，在计算可控的条件下聚焦于信息量大的上下文区域；使用隐状态的余弦相似度作为奖励信号，而非直接的token匹配，从而提高了泛化能力和训练稳定性。这种方法避免了对每个位置生成完整续写的计算开销，同时允许多个合理续写获得非零奖励。
核心方法
REFINE框架的核心思路是将下一序列预测(NSP)形式化为强化学习问题，让快速权重模型在序列级别获得监督信号，而不仅仅依赖逐token的交叉熵损失。直觉上，如果一个快速权重模型能够正确预测多个连续token，说明它有效地将上下文信息压缩到了固定大小的记忆中。REFINE通过四个关键步骤实现这一目标：首先，通过计算NTP熵来识别序列中信息量大的位置；然后，在这些位置生成多token的续写rollout；接着，通过比较预测和真实续写的隐状态相似度来分配奖励；最后，使用GRPO算法更新策略模型。这种方法的优势在于：(1)优化k步续写利用了比单token更高的信息内容；(2)可以对多个合理续写基于语义相似度分配奖励。
REFINE的核心创新在于三个方面。第一，首次将基于RL的下一序列预测应用于快速权重模型，这是一个全新的研究方向。不同于直接在每个位置生成完整续写的朴素方法，REFINE通过熵驱动的token选择策略，只在预测不确定性高的位置进行序列级优化，大幅降低了计算成本。第二，使用隐状态的余弦相似度而非token精确匹配作为奖励信号，这解决了交叉熵损失在多token预测中的过惩罚问题——例如对"cars are fast"和"automobiles move quickly"这样语义等价但表述不同的序列，REFINE能够给予相似的奖励。第三，框架的阶段无关性设计使得同一个REFINE框架可以应用于中期训练、后训练和测试时训练三个不同阶段，只需调整奖励函数和超参数即可适配不同场景。
方法步骤详情
REFINE框架包含四个关键步骤。第一步是熵驱动的token选择：给定输入序列 $S = (x_1, ..., x_T)$，通过策略模型前向传播计算每个token位置的NTP熵值 $H_t = H(\pi_\theta(\cdot | x_{\leq t-1}))$，然后用1D平均池化平滑熵分布，将序列分成 $c$ 个等长chunk，在每个chunk中按softmax熵值采样一个token位置。具体地，$p_i(t) = \frac{e^{H_t/\tau}}{\sum_{t' \in T_i} e^{H'_{t}/\tau}}$。第二步是rollout生成：对每个采样的高熵位置 $t_i$，构建截断前缀 $x_{\leq t_i}$，生成 $k$ 个token的续写 $\hat{x}_{t_i+1:t_i+k}$，并提取预测和真实续写的最后一层隐状态 $h^{pred}_k(t_i)$ 和 $h^{gt}_k(t_i)$。第三步是奖励分配：计算预测和真实隐状态的余弦相似度 $R^k_\phi(t_i) = \frac{1}{k} \sum_{j=1}^k \phi(h^{pred}(t_i+j), h^{gt}(t_i+j))$，对TTT场景还引入二值精确匹配奖励 $R^k_{binary}$ 的混合奖励 $R^k_{hybrid}$。第四步是GRPO优化：将同一序列的奖励标准化计算优势值，使用GRPO算法更新策略，最终损失是NSP损失和标准NTP损失的加权和，系数分别为 $\lambda_{RL}$ 和 $\lambda_{SFT}$。
技术新颖性
REFINE的技术新颖性体现在多个层面。首先，它解决了标准交叉熵损失无法自然扩展到多token预测的问题——朴素应用CE损失需要为所有前缀生成完整续写，计算代价极高；直接匹配单一参考答案会过度惩罚语义等价但表述不同的续写。REFINE通过RL框架和相似度奖励解决了这两个问题。其次，熵驱动的token选择策略是一个精巧的设计：作者发现熵加权采样优于均匀采样(+6.9%)、最大熵选择(+1.8%)和最小熵选择(+1.2%)，这表明NSP训练在不确定性水平平衡的区域最有效。第三，框架的灵活性允许通过调整奖励函数和RL系数来适配不同训练阶段：中期训练使用余弦相似度奖励($R_\phi$)和较低的 $\lambda_{RL}=0.2$；后训练使用混合奖励($R_{hybrid}$)；测试时训练使用二值精确匹配奖励($R_{binary}$)和更高的 $\lambda_{RL}=0.4$以增强上下文记忆能力。
实验结果
REFINE在两个快速权重模型(LaCT-760M和DeltaNet-1.3B)上的实验展示了显著且一致的性能提升。在中期训练阶段，REFINE在RULER NIAH任务上的表现尤为突出：DeltaNet-1.3B在Multi-key NIAH上比无中期训练提升23.5%，比SFT中期训练提升8.8%。LaCT-760M在Multi-key NIAH上也比SFT提升2.4%。值得注意的是，即使SFT直接优化NTP目标，REFINE在NTP准确率上的提升反而更高(图4)，这表明NSP目标的序列级监督提供了NTP无法提供的学习信号。在后训练阶段，REFINE的嵌套学习策略表现优异：LaCT-760M在SQuADQA上比嵌套SFT提升17%(25.5 vs 21.8)，DeltaNet-1.3B提升24.1%(10.3 vs 8.3)。在测试时训练阶段，REFINE在LongBench的12个任务上平均得分LaCT-760M为18.0(对比SFT的17.0)，DeltaNet-1.3B为17.9(对比SFT的16.5)。消融研究证实了余弦相似度奖励优于二值精确匹配奖励(中期训练分别提升1.8%和3.0%)，以及 $k=5$ 是最优的rollout长度。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| RULER S-NIAH (4K-16K平均) | 准确率 | LaCT: 96.3%, DeltaNet: 99.9% | LaCT SFT: 95.6%, DeltaNet SFT: 100.0% | LaCT: +0.7%, DeltaNet: -0.1% |
| RULER Multi-key NIAH (4K-16K平均) | 准确率 | LaCT: 47.7%, DeltaNet: 18.4% | LaCT SFT: 46.6%, DeltaNet SFT: 16.9% | LaCT: +2.4%, DeltaNet: +8.9% |
| RULER SQuADQA (4K-16K平均，REFINE后训练) | F1分数 | LaCT: 25.5%, DeltaNet: 10.3% | LaCT嵌套SFT: 21.8%, DeltaNet嵌套SFT: 8.3% | LaCT: +17.0%, DeltaNet: +24.1% |
| RULER HotpotQA (4K-16K平均，REFINE后训练) | F1分数 | LaCT: 19.8%, DeltaNet: 11.5% | LaCT嵌套SFT: 17.5%, DeltaNet嵌套SFT: 11.3% | LaCT: +13.1%, DeltaNet: +1.8% |
| LongBench 12任务平均 (TTT) | 平均分数 | LaCT: 18.0, DeltaNet: 17.9 | LaCT SFT TTT: 17.0, DeltaNet SFT TTT: 16.5 | LaCT: +5.9%, DeltaNet: +8.5% |
| Booksum NTP准确率 (中期训练) | 验证准确率 | LaCT: ~30.5%, DeltaNet: ~33.0% | LaCT SFT: ~29.5%, DeltaNet SFT: ~32.8% | LaCT: +3.4%, DeltaNet: +0.6% |
局限与改进
REFINE存在几个明显的局限性。首先，作者承认余弦相似度奖励在rollout长度过长时会开始退化，当k从5增加到7时性能下降(图5左)。这是因为奖励在更长序列上平均化后变得过于平滑，无法提供足够精确的信号。其次，最优的rollout长度取决于具体上下文，当前的固定k值可能无法对所有前缀都找到语义上有意义的区域。第三，生成多token rollout需要额外的前向传播计算，增加了训练开销。第四，REFINE目前只在760M和1.3B参数量的模型上验证，在更大规模模型上的效果尚不确定。此外，REFINE的性能高度依赖于熵驱动的token选择策略的有效性，如果熵值无法准确反映模型的不确定性，训练效果可能受到影响。最后，虽然REFINE展示了在三个训练阶段的应用，但实验设置相对受限，评估基准主要集中在长上下文检索和问答任务上。
独立分析的弱点
REFINE存在几个值得关注的弱点。第一，固定rollout长度k=5可能不是所有场景的最优选择，不同任务和上下文可能需要不同的预测长度。改进方向是引入动态rollout长度调整机制，根据前缀的复杂度和不确定性自适应选择k值。第二，余弦相似度奖励虽然提供了平滑的训练信号，但对语义相似性的捕捉仍然有限，例如对词序变化不敏感。可以考虑引入更丰富的语义相似度度量，如编辑距离或学习到的嵌入空间距离。第三，熵驱动的token选择策略假设熵值能够准确反映模型的不确定性，但在训练早期或分布偏移时这一假设可能不成立。第四，REFINE的计算开销主要来自rollout生成阶段，对于长上下文场景可能成为瓶颈。第五，当前实验只验证了LaCT和DeltaNet两种快速权重架构，对其他类型(如GatedDeltaNet)的泛化能力尚待验证。
未来方向
作者在论文中提出了几个重要的未来研究方向。首先，将NSP完全整合到快速权重模型的标准训练框架中需要架构层面的改变，特别是快速权重在截断前缀间的高效传递将显著加速rollout生成，使REFINE能够扩展到更大规模的数据和计算。其次，引入能够捕捉更丰富语义相似度的奖励函数(如编辑距离)可以解决当前余弦相似度奖励在长rollout上的退化问题。第三，动态rollout长度调整可能更有效地隔离语义上有意义的区域。基于本文成果还可以延伸的方向包括：(1)将REFINE应用到更大规模的快速权重模型(如7B+参数)上验证其可扩展性；(2)探索NSP目标在标准Transformer上的效果，而非仅限于快速权重模型；(3)将REFINE与其他训练技巧(如课程学习、数据增强)结合；(4)研究REFINE对快速权重内部记忆机制的影响，深入理解NSP如何改善上下文压缩。
复现评估
从复现角度来看，REFINE具有较好的可复现性。作者声称代码已公开发布(虽然论文中未给出具体链接)，这有助于其他研究者复现实验结果。实验设置相对清晰：使用了两个公开的快速权重模型(LaCT-760M和DeltaNet-1.3B)，训练数据使用Long-Data-Collections等公开数据集，评估基准使用RULER和LongBench等标准测试集。训练配置在表2中明确给出，包括chunk数量c=8、rollout长度k=5、batch size等关键参数。然而，复现存在一些挑战：(1)中期训练需要约200M tokens的训练数据和100步训练，对计算资源有一定要求；(2)GRPO算法的实现细节需要参考Shao等人(2024)的工作；(3)实验结果存在方差(图4显示了三次独立实验的最小最大值)，完全复现可能需要多次运行。总体而言，复现难度中等，有经验的研究者应该能够在合理时间内完成。
论文图表
表格比较了四种token选择策略：均匀采样、最大熵选择、最小熵选择和熵加权采样(REFINE的方法)。熵加权采样在LaCT-760M上比均匀采样高4.3%，比最大熵高3.0%，比最小熵高1.8%；在DeltaNet-1.3B上分别高6.9%、1.8%和1.2%。
这个消融研究证明了熵驱动token选择策略的有效性，表明在不确定性水平平衡的区域进行NSP训练最为有效