基于下一段预测的强化快速权重

快速权重架构为长上下文建模提供了一种有前景的替代方案，其内存开销与上下文长度无关。然而，其潜力受到下一个标记预测（NTP）训练范式的限制。NTP优化单个标记预测，忽略了前缀后多个标记的语义一致性。因此，快速权重模型通过动态更新参数存储上下文信息，学习到的表示无法捕捉长期依赖关系。本文引入REFINE（基于下一段预测的强化快速权重），一种强化学习框架，在下一段预测（NSP）目标下训练快速权重模型。REFINE基于预测熵选择信息量大的标记位置，生成多标记展开，分配自监督序列级奖励，并使用组相对策略优化（GRPO）优化模型。REFINE适用于预训练语言模型的整个训练周期：中期训练、后期训练和测试时训练。我们在LaCT-760M和DeltaNet-1.3B上的实验表明，REFINE在针尖定位检索、长上下文问答以及LongBench中的多样化任务上始终优于基于NTP的监督微调。REFINE为改进快速权重架构中的长上下文建模提供了有效且通用的框架。