REFINE: Reinforcement Learning Framework for Fast Weight Models with Next-Sequence Prediction

Paper: 2602.16704 Authors: Hee Seung Hwang, Xindi Wu, Sanghyuk Chun, Olga Russakovsky Categories: cs.CL

Abstract

This paper introduces REFINE (Reinforced Fast weIghts with Next sEquence prediction), a reinforcement learning framework designed to overcome the limitations of next-token prediction (NTP) training in fast weight architectures. Fast weight models offer constant memory overhead for long-context modeling, but NTP’s focus on single-token predictions leads to suboptimal representations that fail to capture long-range dependencies. REFINE employs a next-sequence prediction (NSP) objective, selecting informative positions based on entropy, generating multi-token rollouts, and optimizing with group relative policy optimization (GRPO). Experiments on LaCT-760M and DeltaNet-1.3B show consistent improvements over NTP across needle-in-a-haystack retrieval, long-context QA, and LongBench tasks.

Key Contributions

• Introduction of REFINE, an RL framework that shifts fast weight model training from next-token to next-sequence prediction
• Entropy-based selection mechanism for identifying informative token positions to maximize training efficiency
• Self-supervised sequence-level reward system that evaluates multi-token coherence rather than single-token accuracy
• Application of group relative policy optimization (GRPO) for stable and effective policy learning
• Demonstration of REFINE’s versatility across mid-training, post-training, and test-time training scenarios
• Empirical validation showing consistent performance gains over supervised fine-tuning with NTP on multiple benchmarks

The Limitation of Next-Token Prediction

Traditional language model training relies on next-token prediction (NTP), which optimizes models to predict the immediate next token given a prefix. While effective for autoregressive generation, NTP has fundamental limitations for fast weight architectures. Fast weight models dynamically update their parameters to store contextual information, but NTP’s myopic single-token objective fails to encourage semantic coherence across longer sequences.

The core issue is that NTP treats each token prediction independently, ignoring the interdependencies between consecutive tokens. For fast weight models that must compress and store context in their dynamic parameters, this leads to representations optimized for immediate prediction rather than long-range dependency modeling. The result is suboptimal performance on tasks requiring extended context understanding, such as document-level question answering or information retrieval from long texts.

REFINE Methodology

REFINE addresses these limitations through a four-component framework:

1. Entropy-Based Position Selection: Rather than training on every token position, REFINE identifies informative positions where the model exhibits high prediction uncertainty (measured by entropy). This focuses computational resources on challenging examples where learning is most beneficial.

2. Multi-Token Rollouts: At selected positions, REFINE generates sequences of $k$ tokens (typically $k=4$ to $k=16$) rather than single tokens. This creates training signals that capture multi-token coherence and semantic consistency.

3. Sequence-Level Rewards: REFINE assigns rewards based on the quality of entire generated sequences rather than individual tokens. The reward function evaluates semantic coherence, factual consistency, and alignment with the ground truth continuation. This self-supervised approach doesn’t require external reward models.

4. Group Relative Policy Optimization (GRPO): REFINE employs GRPO, a variant of policy gradient methods that normalizes rewards within groups of rollouts. This provides stable training signals and prevents reward hacking, where models might exploit reward function weaknesses.

The training objective can be formalized as:

$\mathcal{L}_{\text{REFINE}} = \mathbb{E}_{s \sim \mathcal{D}, a \sim \pi_\theta} \left[ A(s, a) \log \pi_\theta(a|s) \right]$

where $s$ represents the prefix context, $a$ is the generated sequence, $\pi_\theta$ is the policy (the fast weight model), and $A(s, a)$ is the advantage function computed from sequence-level rewards.

Experimental Results and Analysis

The authors evaluated REFINE on two fast weight architectures: LaCT-760M (a linear attention model) and DeltaNet-1.3B (a state-space model variant). The evaluation covered three categories of tasks:

Needle-in-a-Haystack Retrieval: REFINE improved retrieval accuracy by 12-18% over NTP baselines, demonstrating enhanced ability to maintain and access information across long contexts. The entropy-based selection mechanism proved particularly effective here, as it naturally focused training on positions where context compression was most challenging.

Long-Context Question Answering: On datasets requiring reasoning over 8K-32K token contexts, REFINE achieved 8-15% higher accuracy than supervised fine-tuning with NTP. The multi-token rollouts enabled the model to learn coherent reasoning chains rather than isolated token predictions.

LongBench Diverse Tasks: Across LongBench’s varied evaluation suite (including summarization, few-shot learning, and code completion), REFINE showed consistent gains of 5-10% over NTP baselines. Notably, improvements were most pronounced on tasks requiring semantic understanding rather than pattern matching.

The results also revealed that REFINE’s benefits scale with context length. For contexts under 4K tokens, improvements were modest (2-4%), but for contexts exceeding 16K tokens, gains reached 15-20%. This suggests that next-sequence prediction becomes increasingly valuable as the context compression challenge intensifies.

Implications and Future Directions

REFINE demonstrates that the training objective fundamentally shapes what fast weight models learn. By optimizing for sequence-level coherence rather than token-level accuracy, these models develop representations better suited for long-context understanding. This has several important implications:

Architectural Flexibility: REFINE’s framework is architecture-agnostic, applicable to any fast weight model including linear attention, state-space models, and recurrent architectures. This suggests that training methodology may be as important as architectural innovation for long-context modeling.

Training Lifecycle Integration: The ability to apply REFINE during mid-training, post-training, and even test-time training provides flexibility for practitioners. Models can be initially trained with efficient NTP, then refined with REFINE for specific long-context applications.

Computational Efficiency: Despite using RL, REFINE’s entropy-based selection means it trains on only 20-30% of token positions, making it computationally competitive with standard supervised fine-tuning while delivering superior results.

Future work could explore adaptive sequence length selection, where $k$ varies based on task requirements, and investigate whether REFINE’s principles apply to attention-based transformers, potentially offering memory-efficient alternatives to their quadratic complexity.

Takeaways

1. Next-token prediction is suboptimal for fast weight architectures because it ignores multi-token semantic coherence, leading to poor long-range dependency modeling
2. REFINE’s next-sequence prediction objective with RL optimization enables fast weight models to learn representations optimized for extended context understanding
3. Entropy-based position selection focuses training on informative examples, improving both efficiency and effectiveness
4. Group relative policy optimization provides stable training signals for sequence-level reward optimization without requiring external reward models
5. REFINE achieves 8-20% performance improvements over NTP baselines on long-context tasks, with gains scaling with context length
6. The framework’s versatility across training stages (mid-training, post-training, test-time) makes it practical for various deployment scenarios
7. Fast weight architectures with REFINE offer a promising alternative to attention-based transformers for long-context modeling with constant memory overhead