arXiv:2602.16704cs.CL2026-02

用强化学习提升快速权重模型的长文本理解能力

Reinforced Fast Weights with Next-Sequence Prediction

  • 引入基于预测熵选位的多词序列奖励机制
  • 在多个任务上超越传统单标记训练,长上下文表现显著
  • 适配预训练模型全生命周期,可实时优化

快速权重架构通过恒定内存开销实现长上下文建模,但受限于下一词预测(NTP)训练范式。NTP仅优化单个词预测,忽略前缀后多个词的语义连贯性,导致快速权重模型学习到次优表示,难以捕捉长程依赖。本文提出REFINE(基于下一序列预测的强化快速权重),采用强化学习框架,在下一序列预测(NSP)目标下训练快速权重模型。REFINE基于预测熵选择关键词位置,生成多词回溯轨迹,赋予自监督序列级奖励,并使用组相对策略优化(GRPO)进行训练。该方法适用于预训练语言模型的中段训练、后训练及测试时训练阶段。在LaCT-760M和DeltaNet-1.3B上的实验表明,REFINE在针堆中找针检索、长上下文问答以及LongBench多样化任务中持续优于基于NTP的监督微调。REFINE为快速权重架构提供了高效且通用的长上下文建模改进方案。

原文摘要 · Abstract (English)

Fast weight architectures offer a promising alternative to attention-based transformers for long-context modeling by maintaining constant memory overhead regardless of context length. However, their potential is limited by the next-token prediction (NTP) training paradigm. NTP optimizes single-token predictions and ignores semantic coherence across multiple tokens following a prefix. Consequently, fast weight models, which dynamically update their parameters to store contextual information, learn suboptimal representations that fail to capture long-range dependencies. We introduce REFINE (Reinforced Fast weIghts with Next sEquence prediction), a reinforcement learning framework that trains fast weight models under the next-sequence prediction (NSP) objective. REFINE selects informative token positions based on prediction entropy, generates multi-token rollouts, assigns self-supervised sequence-level rewards, and optimizes the model with group relative policy optimization (GRPO). REFINE is applicable throughout the training lifecycle of pre-trained language models: mid-training, post-training, and test-time training. Our experiments on LaCT-760M and DeltaNet-1.3B demonstrate that REFINE consistently outperforms supervised fine-tuning with NTP across needle-in-a-haystack retrieval, long-context question answering, and diverse tasks in LongBench. REFINE provides an effective and versatile framework for improving long-context modeling in fast weight architectures.

快速权重长文本强化学习

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。