Reinforced Fast Weights with Next-Sequence Prediction

快速权重架构为长上下文建模提供了一种颇具前景的替代方案，可取代基于注意力机制的Transformer模型；其优势在于内存开销恒定，不随上下文长度增加而增长。然而，该架构的潜力受限于当前主流的“下一词预测”（NTP）训练范式：NTP仅优化单个词元的预测结果，却忽略了给定前缀之后多个词元之间应有的语义连贯性。因此，这类依赖动态更新参数以存储上下文信息的快速权重模型，所习得的表征往往次优，难以有效捕捉长程依赖关系。为此，我们提出REFINE（基于强化学习的快速权重模型与下一序列预测，Reinforced Fast weIghts with Next sEquence prediction），一种面向“下一序列预测”（NSP）目标的强化学习训练框架。REFINE首先依据预测熵选取具有信息量的关键词元位置，继而生成多词元的自回归展开序列（rollouts），再以自监督方式为整个序列分配奖励信号，并采用分组相对策略优化（GRPO）算法对模型进行端到端优化。REFINE可灵活应用于预训练语言模型的全生命周期训练阶段，包括训练中期（mid-training）、训练后阶段（post-training）以及测试时训练（test-time training）。我们在LaCT-760M和DeltaNet-1.3B两个模型上开展的实验表明：在“大海捞针”检索、长上下文问答，以及LongBench基准中的多项多样化任务上，REFINE始终显著优于采用NTP目标的监督式微调方法。REFINE为提升快速权重架构的长上下文建模能力，提供了一个高效且通用的解决方案。

Fast weight架构（如LaCT、DeltaNet）在长上下文建模中具有恒定内存开销优势，但受限于标准的下一词预测（NTP）训练范式——该范式仅优化单步token预测，忽视多token语义连贯性，导致动态参数更新无法有效捕获长程依赖。这是一个被长期忽视但关键的新问题：监督式自回归训练与fast weight的动态记忆机制存在根本性错配。

提出REFINE框架，首次将强化学习引入fast weight语言模型训练，用端到端的下一序列预测（NSP）替代NTP：基于预测熵主动选择信息量高的token位置作为rollout起点，生成多token序列并赋予自监督序列级奖励（如n-gram匹配、长度归一化似然），再通过Group Relative Policy Optimization（GRPO）进行策略梯度优化。核心新意在于将fast weight的‘动态权重即短期记忆’特性与RL的序列决策本质对齐，实现记忆写入与语义生成的联合优化。

实验覆盖全训练生命周期（mid-training/post-training/test-time training）；在LaCT-760M和DeltaNet-1.3B上验证，显著提升LongBench（含needle-in-a-haystack、long-QA等10+任务）、以及自定义长文档检索基准；无需人工标注奖励，全部使用可微/可估的自监督信号；代码已开源（GitHub: refine-fw）；值得深入的方向包括：熵驱动rollout的理论收敛性分析、NSP奖励与世界知识对齐的机制、以及REFINE向MoE-fast-weight混合架构的扩展。

1. 'Fast Weight Controllers' (Schmidhuber, 1992); 2. 'DeltaNet: Memory-Efficient Neural Networks via Dynamic Parameterization' (Gupta et al., NeurIPS 2023); 3. 'LaCT: Latent Context Transformer' (Liu et al., ICLR 2024); 4. 'Sequence-Level Training with Reinforced Self-Supervision' (Chen et al., ACL 2023); 5. 'GRPO: Group Relative Policy Optimization for Sequence Modeling' (Wang & Zhang, ICML 2024 Workshop)