Reinforced Fast Weights with Next-Sequence Prediction

1. Problem Statement

The Limitation of Next-Token Prediction (NTP) in Fast Weight Models:

• Context: Long-context modeling is essential for tasks like document understanding and multi-shot in-context learning. While attention-based Transformers excel here, their memory and compute costs scale quadratically with context length.
• The Alternative: Fast weight architectures (e.g., DeltaNet, LaCT) offer a solution by replacing global attention with a fixed-size memory matrix that is dynamically updated as tokens are processed. This ensures constant memory overhead regardless of context length.
• The Core Issue: Despite their architectural efficiency, fast weight models are currently trained using the standard Next-Token Prediction (NTP) objective (minimizing cross-entropy loss token-by-token).
  • NTP optimizes for immediate single-token likelihood, ignoring the semantic coherence of the sequence following a prefix.
  • This leads to suboptimal representations where the model fails to capture long-range dependencies effectively because the training signal does not encourage the model to maintain a consistent internal state over multiple steps.
  • Generating full multi-token continuations for every position during training is computationally prohibitive, making standard sequence-level supervision difficult to implement.

2. Methodology: REFINE

The authors propose REFINE (Reinforced Fast weIghts with Next sEquence prediction), a Reinforcement Learning (RL) framework designed to train fast weight models under a Next-Sequence Prediction (NSP) objective.

Core Concept: Next-Sequence Prediction (NSP)

Instead of optimizing for $P(x_{t+1} | x_{\le t})$, NSP aims to optimize for a coherent sequence of future tokens $P(x_{t+1:t+k} | x_{\le t})$. This forces the fast weights to store information that supports accurate prediction over a horizon of $k$ tokens, not just the next one.

The REFINE Framework

REFINE formulates NSP as an RL problem to overcome the computational cost of generating full continuations for every token. It consists of four key steps:

1. Entropy-Based Token Selection:

   • Instead of sampling tokens uniformly, the model computes the NTP entropy ($H_t$) at each position.
   • The input sequence is divided into $c$ contiguous chunks.
   • From each chunk, one target token position is sampled with probability proportional to its entropy (high entropy indicates high uncertainty/challenging regions). This focuses training on informative parts of the context.

2. Rollout Generation:

   • For each selected high-entropy position $t_i$, the model generates a $k$-token continuation (rollout) $\hat{x}_{t_i+1:t_i+k}$.
   • The model extracts hidden states for both the generated rollout and the ground-truth continuation.

3. Reward Assignment:

   • Mid-training & Post-training: A Cosine Similarity Reward ($R_\phi$) is calculated between the hidden states of the predicted rollout and the ground truth. This provides a smooth, semantic signal that rewards plausible continuations even if they don't match the ground truth character-for-character.
   • Test-Time Training (TTT): A Binary Exact Match Reward ($R_{binary}$) is used to enforce strict memorization of the context.
   • Hybrid Reward: For post-training, a mix of both rewards is used to balance generalization and memorization.

4. Optimization with RL (GRPO):

   • The model is optimized using Group Relative Policy Optimization (GRPO).
   • The loss function is a weighted sum of the NSP loss (derived from rewards) and the standard NTP loss to prevent catastrophic forgetting.
   • The objective maximizes the expected sequence-level reward: $J(\theta) = \mathbb{E}[R]$.

Applicability Across Training Phases

REFINE is designed to be "phase-agnostic" and is applied in three stages:

• Mid-training: Reinforcing fast weights on pre-training-like corpora to improve long-context adaptation.
• Post-training: Integrated via nested learning. The model updates its fast weights on the prompt using REFINE before generating the final response, which is then fine-tuned via standard SFT.
• Test-Time Training (TTT): Directly reinforces fast weights on the inference prompt using self-supervised signals (no external labels required).

3. Key Contributions

1. Novel Objective: Introduction of the Next-Sequence Prediction (NSP) objective specifically for fast weight models to address the limitations of token-level feedback.
2. RL Framework: Proposal of REFINE, an efficient RL framework that combines entropy-based token selection (to reduce compute) with sequence-level rewards (to improve semantic coherence).
3. Lifecycle Integration: Demonstration that this approach is effective across the entire model lifecycle (Mid-training, Post-training, and Test-Time Training), a versatility not previously explored for fast weights.

4. Experimental Results

The authors evaluated REFINE on two models: LaCT-760M and DeltaNet-1.3B.

• Long-Context Retrieval (RULER NIAH):

  • REFINE significantly outperformed standard Supervised Fine-Tuning (SFT) with NTP.
  • DeltaNet-1.3B: Showed a 23.5% improvement in Multi-key NIAH tasks compared to no mid-training and 8.8% over SFT mid-training.
  • LaCT-760M: Consistently improved performance across Single, Multi-key, Multi-query, and Multi-value NIAH tasks.

• Long-Context QA (LongBench & RULER):

  • Mid-training: REFINE improved average performance on LongBench by 16.9% (LaCT) and 17.0% (DeltaNet) compared to baselines.
  • Post-training: Nested REFINE improved SQuADQA scores by 17% (LaCT) and 24.1% (DeltaNet) over nested SFT.
  • Test-Time Training: REFINE TTT consistently outperformed SFT TTT, demonstrating superior adaptation to new contexts without additional labels.

• Ablation Studies:

  • Reward Function: Cosine similarity rewards ($R_\phi$) outperformed binary exact match rewards during mid-training, proving that semantic similarity is more effective for generalization than strict matching.
  • Token Selection: Entropy-weighted sampling outperformed uniform, max-entropy, and min-entropy sampling, confirming that focusing on uncertain regions yields better learning signals.
  • Rollout Length ($k$): Performance peaked at $k=5$; longer rollouts ($k=7$) led to diminishing returns due to reward signal degradation.

5. Significance

• Bridging the Gap: REFINE successfully aligns the training objective of fast weight models with their architectural strength (long-term memory), moving beyond the short-sighted NTP paradigm.
• Efficiency: By using entropy-based sampling and RL, REFINE achieves sequence-level supervision without the prohibitive cost of generating full continuations for every token.
• Versatility: The framework proves that RL can be effectively applied to pre-trained fast weight models even without prior instruction tuning, offering a practical path to scaling long-context capabilities in memory-efficient architectures.
• State-of-the-Art: The results establish a new benchmark for fast weight models, showing they can rival or exceed attention-based models in specific long-context tasks when trained with the correct objective.

In summary, REFINE demonstrates that by shifting from token-level to sequence-level reinforcement learning, fast weight architectures can overcome their historical limitations in long-context modeling, offering a scalable and efficient alternative to Transformers.