Reinforced Fast Weights via Next-Sequence Prediction¶
会议: NeurIPS 2026
arXiv: 2602.16704
OpenReview: https://openreview.net/forum?id=A1sw6FviPQ
代码: https://github.com/princetonvisualai/ReFINE
领域: 其他
关键词: 快速权重, 线性注意力, 强化学习, 序列预测, 长文本建模
一句话总结¶
针对快速权重架构在长文本自回归生成中因单步下一词预测（NTP）缺乏序列级反馈而导致的记忆漂移与误差累积问题，本文提出基于下一序列预测（NSP）的强化学习训练框架 ReFINE，通过熵加权采样关键位置生成多步 rollout 并以余弦/混合奖励驱动 GRPO 优化，在全生命周期训练中显著提升长程检索与问答性能。
研究背景与动机¶
标准 Transformer 依赖自注意力机制处理长序列，但随着上下文长度增加，其计算与显存开销呈二次方爆炸增长，且随序列线性膨胀的 KV Cache 为高吞吐推理带来了严重的显存墙瓶颈。为攻克这一计算瓶颈，以 DeltaNet、GatedDeltaNet 及 LaCT 为代表的快速权重（Fast Weight）与线性循环架构应运而生。这类模型通过将历史信息动态压缩进一个固定大小的循环状态（或权重矩阵），在推理阶段仅需常数级别的显存占用，展现出卓越的长文本建模效率与扩展潜力。
然而，将无限长上下文压缩进有界状态的循环更新机制对自回归生成的稳定性提出了严苛挑战。传统 Transformer 的 KV Cache 显式完整保留各位置历史表征，而快速权重则在每一步接收新 token 时递归覆盖和演化内部记忆矩阵；一旦自回归生成路径发生微小偏差，预测误差便会沿着循环步迅速级联累积，导致灾难性遗忘与表征严重漂移。令人遗憾的是，现有快速权重模型完全套用标准 Transformer 的逐词预测（NTP）交叉熵目标进行预训练。NTP 仅通过 Teacher-Forcing 提供单步 ground-truth 监督，根本无法评估模型内部快速权重在多步自由演化轨迹下的长期记忆质量，从而引发训练与自回归推理行为的根本脱节。
直接对全长序列密集进行多步序列预测不仅面临组合爆炸的计算代价，还会因单字符硬匹配对语义同义表达产生过度惩罚。为此，本文将下一序列预测转化为强化学习奖励最大化问题，提出通用训练框架 ReFINE。核心 idea：将下一词监督扩展为下一序列预测（NSP），通过局部熵加权自适应筛选不确定性关键位置展开多步 rollout，设计平滑表征相似度与精准匹配混合奖励，并利用 GRPO 强化学习端到端校准快速权重的长程演化轨迹。
方法详解¶
整体框架¶
ReFINE 旨在全生命周期训练（继续预训练/中训练、指令后训练、测试时训练 TTT）中强化快速权重模型的长程自适应能力。整个流水线包含四个核心步骤：首先将长文本输入策略模型计算逐 token 的 NTP 熵分布并平滑，在序列等长分块中按熵加权概率采样出不确定性最高的信息关键位置；随后截取对应前缀并让当前模型自回归生成 \(k\) 步 rollout 预测序列；接着提取预测序列与 ground truth 真实后续序列的最终隐藏层表征，计算平滑语义相似度奖励（或混合奖励）；最后基于组内优势标准化，采用 GRPO 算法反向传播优化模型参数，同时结合全序列 NTP 交叉熵损失防止语言建模基础能力发生灾难性遗忘。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["长序列输入 S<br/>计算逐 token NTP 熵分布"] --> B["熵加权 token 选择<br/>等长分块按 softmax 熵采样关键位置"]
B --> C["Rollout 序列生成<br/>截断前缀自由展开预测 k 步延续"]
C --> D["表征与多阶段奖励计算<br/>提取预测与真实隐藏层计算余弦/混合奖励"]
D --> E["GRPO 强化学习优化<br/>序列优势标准化更新参数 + NTP 联合正则"]关键设计¶
1. 熵加权 token 选择：兼顾全局覆盖与局部不确定性 在长达数万 token 的长文本中，在每个位置强行展开 \(k\) 步自回归生成的计算成本极其昂贵，且大量低熵确定性 token（如常见介词、标点）无法提供有价值的长程状态演化监督。为此，ReFINE 首先将长文本 \(S\) 等长划分为 \(c\) 个连续分块 \((S_1, \ldots, S_c)\)，对各 token 位置 \(t\) 计算 NTP 预测熵 \(H_t = H(\pi_\theta(\cdot \mid x_{\leq t-1}))\) 并通过一维均值池化核 \(w\) 进行局部平滑得到 \(\tilde{H}_t\)。在每个分块 \(S_i\) 内部，依据平滑熵的 Softmax 概率分布采样目标位置 \(t_i\)： $\(p_i(t) = \frac{\exp(\tilde{H}_t)}{\sum_{t' \in \mathcal{T}_i} \exp(\tilde{H}_{t'})}, \quad \mathcal{T}_i = \{t' \mid x_{t'} \in S_i\}\)$ 该设计使模型不仅能在整个长文本上下文空间中保持均匀全局跨度，又能自适应聚焦于最富语义信息量、最易产生认知歧义的局部决策边界，使后续的序列级强化信号产生最大梯度的有效校准。
2. 潜空间平滑语义奖励与阶段自适应混合机制：化解多词精确匹配惩罚 在 multi-token 预测序列中，纯交叉熵损失或严格字符级 Exact Match 具有严重的刚性缺陷——语义完全等价但字面不同的表述（例如“cars are fast”与“automobiles move quickly”）会被错误地施加高额惩罚，破坏训练稳定性。ReFINE 从策略模型最后一层输出中提取预测序列与 ground-truth 序列隐藏层表征 \(\mathbf{h}^{\text{pred}}(t_i+j)\) 与 \(\mathbf{h}^{\text{gt}}(t_i+j)\)，在潜空间中基于余弦相似度计算平滑连续奖励： $\(R^\varphi_k(t_i) = \frac{1}{k} \sum_{j=1}^k \cos\big(\mathbf{h}^{\text{pred}}(t_i+j),\, \mathbf{h}^{\text{gt}}(t_i+j)\big)\)$ 在不同训练阶段，长程记忆的目标发生转变：在继续预训练（Mid-training）阶段需要高度的上下文泛化能力，采用纯余弦奖励 \(R^\varphi\)；而在测试时训练（Test-Time Training, TTT）阶段，测试样本小批次且亟需精准记忆输入 Prompt 的事实细节，采用字符级指示函数 \(R^{\text{binary}}\)；在后训练（Post-training）阶段，采用平衡泛化与记忆的混合奖励 \(R^{\text{hybrid}} = R^\varphi + R^{\text{binary}}\)。
3. GRPO 策略梯度优化与全生命周期协同：端到端稳固快速权重状态 得到一组关键位置的 rollout 产物与其对应奖励后，同一序列内部的奖励集合进行均值与方差标准化计算优势值（Advantage），利用不需要额外 Critic 模型的 Group Relative Policy Optimization（GRPO）直接对策略网络进行参数更新： $\(\mathcal{J}(\theta) = \mathbb{E}_{x_{\leq t} \sim \mathcal{D},\, \hat{x}_{t+1:t+k} \sim \pi_\theta(\cdot \mid x_{\leq t})}[R_k(t)]\)$ 在多阶段落地中，ReFINE 展示出高度的生命周期适配性：在中训练阶段直接作为序列级自适应目标；在后训练阶段采用嵌套学习（Nested Learning）策略，先在 Prompt 段以 ReFINE 更新快速权重记忆，再在 Response 段进行标准 SFT 微调；在 TTT 阶段则在推理前利用测试 Prompt 快速前向更新状态。
损失函数 / 训练策略¶
为防止纯强化学习训练导致预训练语言模型基础生成能力与词汇概率分布退化（如语言流畅度下降与模式坍塌），ReFINE 在总损失中保留了覆盖整个序列 \(S\) 的标准 NTP 交叉熵损失，形成联合优化目标： $\(\mathcal{L}_{\text{total}} = \lambda_{\text{RL}} \mathcal{L}_{\text{NSP}} + \lambda_{\text{SFT}} \mathcal{L}_{\text{NTP}}\)$ 超参数设置上，中训练与后训练阶段设 \(\lambda_{\text{RL}} = 0.2, \lambda_{\text{SFT}} = 1.0\)；TTT 阶段为最大化上下文记忆容量，将 \(\lambda_{\text{RL}}\) 提升至 0.4。rollout 步长 \(k\) 默认设为 5，序列分块数 \(c\) 设为 8，采样温度 \(\tau = 1.0\)，采用 Adam 优化器配合梯度裁剪 0.2，以极小的计算开销即可取得显著的长上下文性能增益。
实验关键数据¶
主实验¶
论文在两种具有代表性的快速权重架构（带参数显式更新的 LaCT-760M 与固定参数更新内部矩阵状态的 DeltaNet-1.3B）上进行了系统评测，覆盖长上下文大海捞针检索基准 RULER（4K、8K、16K）以及多文档问答与 LongBench 长文本综合基准。下表展示了继续预训练（MidTr）阶段标准 SFT 与 ReFINE 在 RULER 各检索任务上的平均表现对比：
| 模型 | 方法 | S-NIAH (Avg) | MK-NIAH (Avg) | MQ-NIAH (Avg) | MV-NIAH (Avg) |
|---|---|---|---|---|---|
| LaCT-760M | 原始无MidTr | 95.3 | 46.2 | 27.1 | 28.4 |
| LaCT-760M | + SFT MidTr | 95.6 | 46.6 | 27.8 | 28.5 |
| LaCT-760M | + ReFINE MidTr | 96.3 | 47.7 | 28.1 | 29.8 |
| DeltaNet-1.3B | 原始无MidTr | 100.0 | 14.9 | 11.1 | 10.4 |
| DeltaNet-1.3B | + SFT MidTr | 100.0 | 16.9 | 17.8 | 17.1 |
| DeltaNet-1.3B | + ReFINE MidTr | 99.9 | 18.4 | 19.9 | 17.8 |
在长文本多文档问答（RULER SQuADQA 与 HotpotQA）全生命周期串联评测中，各阶段融合 ReFINE 的表现展现出显著累积收益：
| 架构 | Mid-Training | Post-Training | TTT | SQuADQA (Avg) | HotpotQA (Avg) |
|---|---|---|---|---|---|
| LaCT-760M | - | - | - | 12.7 | 15.7 |
| LaCT-760M | SFT | - | - | 13.2 | 16.2 |
| LaCT-760M | ReFINE | - | - | 14.7 | 17.2 |
| LaCT-760M | ReFINE | Nested SFT | - | 21.8 | 17.5 |
| LaCT-760M | ReFINE | Nested ReFINE | - | 25.5 | 19.8 |
| LaCT-760M | ReFINE | Nested ReFINE | ReFINE | 27.0 | 23.0 |
| DeltaNet-1.3B | - | - | - | 6.0 | 5.2 |
| DeltaNet-1.3B | SFT | - | - | 6.0 | 7.2 |
| DeltaNet-1.3B | ReFINE | - | - | 6.8 | 9.0 |
| DeltaNet-1.3B | ReFINE | Nested SFT | - | 8.3 | 11.3 |
| DeltaNet-1.3B | ReFINE | Nested ReFINE | - | 10.3 | 11.5 |
| DeltaNet-1.3B | ReFINE | Nested ReFINE | ReFINE | 12.3 | 13.2 |
消融实验¶
研究深入分析了奖励函数设计、token 采样策略以及 rollout 长度 \(k\) 和分块数量 \(c\) 对 LongBench 12 项长文本综合任务平均性能的影响：
| 实验维度 | 考察变量配置 | LaCT-760M Avg | DeltaNet-1.3B Avg | 说明 |
|---|---|---|---|---|
| 奖励函数对比 (MidTr) | 二值匹配 \(\mathcal{R}^{\text{binary}}\) | 16.6 | 16.5 | 字符刚性惩罚损失泛化性 |
| 奖励函数对比 (MidTr) | 混合奖励 \(\mathcal{R}^{\text{hybrid}}\) | 16.8 | 16.8 | 结合泛化与精确记忆 |
| 奖励函数对比 (MidTr) | 负欧氏距离 \(\mathcal{R}^{\text{NED}}\) | 16.8 | 17.0 | 与余弦相似度在优势标准化下单调等价 |
| 奖励函数对比 (MidTr) | 余弦表征奖励 \(\mathcal{R}^\varphi\) (本文) | 16.9 | 17.0 | 潜空间平滑语义匹配表现最优 |
| 采样策略对比 (MidTr) | 均匀随机采样 Uniform | 16.2 | 15.9 | 缺乏对信息稠密区域的针对性 |
| 采样策略对比 (MidTr) | 最大熵极值 \(\arg\max H\) | 16.4 | 16.7 | 极值易受极端噪点干扰 |
| 采样策略对比 (MidTr) | 最小熵极值 \(\arg\min H\) | 16.6 | 16.8 | 难以探索不确定性决策边界 |
| 采样策略对比 (MidTr) | 熵加权采样 (本文) | 16.9 | 17.0 | 兼顾全文本覆盖度与不确定性专注度 |
| 序列分块数 \(c\) 影响 | \(c=2 \rightarrow c=8\) | \(16.5 \rightarrow 16.9\) | \(16.3 \rightarrow 17.0\) | 增加序列级监督频次稳定提升长程质量 |
| Rollout 长度 \(k\) 影响 | \(k=1 \rightarrow k=5 \rightarrow k=7\) | 倒 U 型 (峰值 \(k=5\)) | 倒 U 型 (峰值 \(k=5\)) | \(k\) 过大导致表征余弦奖励锐度衰减 |
关键发现¶
- 快速权重内部状态漂移被大幅压制：通过 Frobenius 范数追踪生成过程中快速权重记忆矩阵与真实真值轨迹的漂移，发现 ReFINE 显著降低了状态偏离度，证明序列级反馈切实改善了长期状态的数值稳定性。
- 反哺单步 NTP 验证精度：在 Booksum 长篇叙事摘要数据集的验证集上，直接优化 NTP 的 SFT 基线精度早早停滞，而通过 ReFINE 优化 NSP 目标的模型在 NTP 准确率上持续上升，表明长视野反馈提供了单步 Teacher-Forcing 所缺乏的高阶认知信号。
- 全生命周期阶段收益可叠加：从 Mid-training 到 Nested Post-training 再到推理期 TTT，每个阶段引入 ReFINE 都能在上一阶段的基础上实现稳健叠加增益。
亮点与洞察¶
- 跳出注意力机制改进，直击训练目标失配：多数快速权重研究局限于设计更复杂的更新门控或硬件算子，本文敏锐指出根源在于 NTP 与自适应循环记忆的“单步监督 vs 长程演变”目标失配，以学习范式突破架构瓶颈。
- 轻量且无 Critic 的 GRPO 序列适配：巧妙运用同一序列内部多个分块 rollout 的相对优势标准化，摆脱了庞大 Critic 模型的显存开销，使快速权重模型无需预先经过庞大的 SFT/偏好对齐即可直接开展强化学习自进化。
- 生命周期普适的嵌套学习范式：在后训练与测试时微调中创造性地结合了 Prompt 级别的自监督序列强化与任务目标生成对齐，为长上下文在线自适应提供了通用方法论。
局限与展望¶
- 固定单步 rollout 长度的语义截断：当前超参数 \(k=5\) 在整个序列中保持固定，但自然语言的语义边界因词组、从句长度而异，缺乏根据句法结构自适应动态伸缩 rollout 视野的能力。
- 超长 rollout 下表征余弦奖励退化：实验表明当 \(k\) 增大到 7 时性能出现回落，揭示出深层隐藏状态表征在较长预测距离下余弦相似度区分度变钝，未来需探索基于分块自编码或信息瓶颈的更鲁棒语义奖励。
- 快速权重内部状态跨前缀迁移的系统开销：在长文本自回归生成截断前缀时，当前实现需要重复前向计算，未来若在架构底层支持快速权重状态的轻量化高效 checkpointing 与转移，将极大加速 rollout 并支持更大规模的强化学习伸缩。
相关工作与启发¶
- vs 多词预测（MTP, 如 Gloeckle et al., 2024）：MTP 依赖增加并行的 \(k\) 个独立输出预测头，强行在架构层修改网络且预测头间相互独立割裂；ReFINE 完全不改动原生模型架构，通过自回归 rollout 与序列整体强化学习维持语义内聚性。
- vs 纯文本 RL pre-training（如 RLP / Reinforcement Pre-training）：现有强化预训练工作多依赖思维链（CoT）推理以及在具备较强推理能力的大型 Transformer 上展开探索；ReFINE 首次证明在不具备显式 CoT 机制的线性快速权重预训练小模型上，自监督多步序列预测依然能激发出强大的长文本自稳定能力。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次系统揭示快速权重循环模型与 NTP 目标的失配矛盾，创新性构建序列级强化学习自适应训练框架。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖两种主流快速权重模型、贯穿中训练/后训练/TTT 三阶段，消融与机制分析（状态漂移量化）非常扎实。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑推进严谨自然，问题陈述一针见血，数学推导与实验论证结合紧密。
- 价值: ⭐⭐⭐⭐⭐ 为打破长文本线性模型（RNN/SSM/Fast Weights）在长序列自回归推理中的误差累积瓶颈提供了极其关键的方法学指导。