Reinforced Fast Weights with Next-Sequence Prediction
This addresses the challenge of improving long-context modeling for fast weight architectures, offering a versatile framework applicable throughout training, though it appears incremental as it builds on existing fast weight methods.
The paper tackles the problem of suboptimal long-context modeling in fast weight architectures due to the next-token prediction training paradigm, and introduces REFINE, a reinforcement learning framework using next-sequence prediction, which consistently outperforms supervised fine-tuning on tasks like needle-in-a-haystack retrieval and long-context question answering.
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.