Reinforced Fast Weights with Next-Sequence Prediction

Abstract: 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 LLMs: 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.

• The paper introduces ReFINE, a reinforcement learning-based four-stage framework that shifts from token-level to sequence-level supervision to enhance long-context reasoning.
• It employs entropy-based token selection and rollout generation with cosine similarity rewards, achieving significant performance gains on retrieval and QA tasks.
• Empirical results on LaCT-760M and DeltaNet-1.3B demonstrate improvements up to +23.5% in needle-in-a-haystack retrieval and increased recall in long-context settings.

Reinforced Fast Weights with Next-Sequence Prediction: A Technical Review

Motivation: Fast Weights and the Limitation of Next-Token Prediction

Standard transformer architectures, while effective for long-context applications in NLP, suffer from quadratic scaling of memory with respect to context length due to their attention mechanisms. Fast weight models address this by replacing global attention with fixed-size, dynamically updated memories, yielding constant memory overhead as context grows. Recent works, such as DeltaNet and LaCT, instantiate this approach by modifying the transformer block to employ these fast memory updates.

Despite this architectural innovation, the prevalent training objective borrowed from transformers—Next-Token Prediction (NTP) via cross-entropy—remains unsatisfactory for optimizing the full capabilities of fast weight models. NTP offers only token-level supervision, ignoring whether the neural memory representations can support coherent, contextually relevant multi-token generation. This inherently short-sighted feedback restricts fast weight models’ ability to capture and utilize long-range dependencies.

Figure 1: Comparison of standard NTP and ReFINE. Standard NTP delivers token-level supervision, whereas ReFINE enables sequence-level supervision using multi-token rollouts and RL-based optimization.

The Next-Sequence Prediction Objective and the ReFINE Framework

To rectify these deficiencies, the paper introduces Next-Sequence Prediction (NSP): a training objective that shifts supervision from individual token probabilities to rollout-level, sequence-based alignment. NSP evaluates the model on its ability to generate multi-token continuations conditioned on a given prefix—more faithfully matching the deployment scenario of long-context LLMs.

Crucially, applying NSP as a pure supervised objective is computationally prohibitive; multi-token rollouts for all positions are infeasible for long inputs, and cross-entropy loss fails to robustly capture semantic equivalence among different valid continuations. The proposed solution is to recast NSP as a reinforcement learning (RL) problem.

The paper presents ReFINE (Reinforced Fast Weights with Next-Sequence Prediction), a four-stage RL training framework for fast weight models:

1. Entropy-Based Token Selection: The sequence is split into chunks, with target rollout positions sampled non-uniformly according to the local entropy of the NTP distribution, ensuring focus on high-uncertainty regions.
2. Rollout Generation: For each selected position, the model generates a rollout of $k$ tokens conditioned on the prefix.
3. Reward Assignment: Sequence-level rewards are computed by comparing the representations of generated continuations with ground truth, using cosine similarity in hidden space for smooth semantic alignment.
4. RL Optimization: Using Group Relative Policy Optimization (GRPO), the policy is updated to maximize rollout-level sequence rewards, while optionally regularizing with the standard NTP objective.

   Figure 3: The ReFINE pipeline: from entropy-driven rollout targeting, through multi-token generation, to self-supervised reward assignment and RL optimization.

Notably, ReFINE is agnostic to the training phase and can be applied during mid-training (continued pre-training), post-training (instruction or task tuning), and even at test-time (TTT adaptation).

Empirical Evaluation: Long-Context Retrieval and QA

ReFINE is evaluated on both LaCT-760M and DeltaNet-1.3B, two competitive fast weight LLMs. Across the RULER long-context retrieval suite and LongBench multi-domain tasks, ReFINE-trained models consistently outperform those using pure SFT with NTP, with particularly substantial improvements in:

• Needle-in-a-Haystack (RULER MK-NIAH): DeltaNet-1.3B demonstrates a +23.5% improvement over the non-mid-trained baseline and +8.8% over SFT mid-trained.
• Long-context multi-document QA: LaCT-760M with nested ReFINE in the post-training loop achieves up to 43.5 average recall on SQuADQA at long contexts, a marked increase over all SFT counterparts.

In all phases—mid-training, post-training, and test-time training—ReFINE yields non-trivial gains across context lengths up to 16K tokens.

Figure 2: DeltaNet-1.3B performance visualization across tasks, highlighting ReFINE’s improvements in the long-context regime.

Ablation and Analysis

Key analyses support the robustness and design of ReFINE:

• Reward Function: Cosine similarity in hidden space yields better generalization than pure binary token-level exact match for mid-training and sequence adaptation, though for more memorization-driven test-time adaptation, hybrid reward formulations are preferred for strong transfer.
• Token Selection: Entropy-weighted sampling outperforms uniform, max-entropy, and min-entropy selection, leading to better downstream results by matching learning focus to model uncertainty.
• Rollout Length and Chunk Count: There exists an optimal rollout length $k$ (typically 5), with performance degrading for longer rollouts due to diluted reward sharpness. Increasing the number of rollout chunks $c$ improves the density of sequence-level supervision and overall performance.

  Figure 4: Ablation on rollout length ($k$) and number of chunks ($c$): demonstration of performance saturation and monotonic gain, respectively.

Theoretical and Practical Implications

ReFINE’s core contribution is its demonstration that RL-based sequence-level supervision robustly enhances the memory and adaptation capacities of fast weight LLMs. From a theoretical perspective, it provides evidence against the prevailing assumption that advantage-weighted single-token objectives suffice for training dynamic neural memory modules. The use of latent-space, self-supervised rewards further suggests a path toward more semantically aligned optimization for generative models.

Practically, these findings motivate broader integration of RL and sequence-level objectives in next-generation efficient LLM architectures, particularly for resource-constrained or extremely long-context scenarios. The strong phase-agnostic effect of ReFINE—across pre-, post-, and test-time—positions it as a unifying methodology for fast weight model adaptation and transfer.

Future Directions

The framework highlights bottlenecks in current fast weight architectures, including inefficiencies in prefix truncation and limitations of static rollout horizon selection. Future work could incorporate adaptive $k$ selection, richer reward formulations (e.g., edit distance or semantic retrieval), and architectural modifications for rapid fast weight reuse across rollouts.

By demonstrating clear gains in long-context retrieval and QA, the paper motivates both further architectural innovation in neural memory design and continued development of RL-based sequence-level objectives for language modeling.

Conclusion

ReFINE fundamentally advances the training methodology for fast weight LLMs by introducing RL-based Next-Sequence Prediction, replacing token-level objectives with rollout-centric, representation-anchored, entropy-targeted optimization. Extensive evidence shows that this improves both immediate and generalizable long-context reasoning performance, creating productive avenues for sequence-level learning in efficient neuro-symbolic architectures.