PaperlayerView exact revision on arXiv

Source-linked AI summary

Reinforced Fast Weights with Next-Sequence Prediction

Hee Seung Hwang, Xindi Wu, Sanghyuk Chun, Olga Russakovsky

arXiv:2602.16704v1cs.CL

TL;DR

Fast weight models offer constant-memory long-context modeling, but NTP’s token-level feedback does not assess coherent multi-token continuations. REFINE trains them with entropy-guided NSP reinforcement learning and consistently improves long-context benchmarks across training stages and model families.

  • Problem

    NTP optimizes immediate next-token predictions rather than semantic coherence across multi-token continuations, limiting its suitability for fast weights used as long-context memory.

  • Method

    REFINE selects high-entropy token positions, generates multi-token rollouts, and optimizes self-supervised sequence-level rewards with reinforcement learning under NSP.

  • Results

    REFINE consistently improves long-context benchmarks across mid-training, post-training, and test-time training for LaCT-760M and DeltaNet-1.3B.

  • Takeaways & Limitations

    REFINE provides a flexible and practical pathway for applying reinforcement learning to NSP in fast weight architectures.

  • Takeaways & Limitations

    Cosine-similarity rewards deteriorate for overly long rollouts, and the optimal rollout length depends on the context.

Abstract

from arXiv · show

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.

1. Introduction

Long-context tasks expose the quadratic cost of attention, while fast weights provide constant-memory alternatives whose NTP training may fail to reinforce multi-token coherence. REFINE addresses this mismatch with NSP-focused reinforcement learning and improves long-context performance across training stages.

  • Motivation: Attention-based transformers incur computational and memory costs that scale quadratically with context length, motivating alternative long-context architectures.Fast weight models replace global attention with dynamically updated fixed-size memory.
  • Motivation: NTP provides only immediate next-token feedback, encouraging short-term likelihood optimization rather than representations supporting longer-horizon predictions.This can limit the adaptive capacity and long-context behavior of fast weights.
  • REFINE: NSP trains fast weight models to predict semantically coherent multi-token continuations from a prefix, aligning feedback with long-context memory use.The objective addresses NTP’s sequence-level feedback limitation.
  • REFINE: REFINE formulates NSP as reinforcement learning, focusing on informative context regions and maximizing sequence-level rewards from model predictions.The framework uses entropy-based token selection and policy-gradient updates.
  • Results: REFINE applies across mid-training, post-training, and test-time training, providing a phase-agnostic approach to improving fast weight models.The three settings cover pretraining-like adaptation, downstream supervision, and prompt-based label-free reinforcement.
  • Results: 8.5% (mid-training), 15.3% (post-training), and 9.5% (test-time training) are LaCT-760M’s RULER gains over pure NTP training with SFT.For DeltaNet-1.3B, the corresponding gains are 20.3%, 11.0%, and 15.0%.

2. Background

Fast weight models store context in continually updated parameters and support test-time adaptation, while the paper situates REFINE across mid-training, post-training, and test-time training.

  • Fast Weight Architectures: Fast weight architectures replace global attention and growing key-value caches with fixed-size weight matrices that are continually updated as tokens arrive.This design stores contextual information directly in model parameters.
  • Fast Weight Architectures: The fast-weight update learns an online key-to-value mapping, and outputs are retrieved by applying the weights to token queries.The cited formulation identifies η as the learning rate and k_t, v_t as key and value representations.
  • Training Phases: Mid-training adapts pretrained models to domains or capabilities; here it adapts them to the NSP objective and reward using pretraining data.It is treated as continued pretraining.
  • Training Phases: Post-training typically uses task-specific instruction-response pairs, while REFINE is integrated through nested learning within training loops.The supplied passage describes post-training as involving relatively fewer gradient steps than pretraining.
  • Training Phases: Test-time training adapts parameters with self-supervised objectives, and fast weights update on each input token to memorize and adapt to the current context.The passage frames this as an on-the-fly adaptation mechanism.
  • RL for Language Modeling: Reinforcement learning can improve long-context capabilities of pretrained fast weight models during mid-, post-, and test-time training, even without prior instruction tuning.This places RL for fast weights beyond standard transformer applications described in the related work.

3. Method

REFINE replaces single-token training with next-sequence prediction optimized through reinforcement learning, selecting uncertain positions, generating continuations, and assigning sequence-level rewards. The framework combines entropy-based sampling, rollout evaluation, semantic or exact-match rewards, and GRPO updates for fast weight models.

  • Next-Sequence Prediction: NSP addresses NTP’s neglect of semantic relationships among future tokens and its uniform treatment of potentially informative long-context regions.These limitations motivate sequence-level feedback and selective training positions.
  • Next-Sequence Prediction: NSP optimizes multi-token sequence alignment at selected positions instead of token-by-token predictions.The sequence loss measures discrepancy between a predicted k-token continuation and its ground-truth continuation.
  • Framework Scope: The framework is presented as the first RL-based NSP approach for fast weight models and is designed for mid-training, post-training, and test-time training.Its four stages are entropy-based token selection, rollout generation, reward assignment, and RL optimization.
  • Entropy-Based Token Selection: REFINE samples one entropy-weighted token position from each sequence chunk, focusing training on locally uncertain predictions while distributing signals across the context.Entropy values are smoothed before chunk-wise sampling, and the resulting positions form the training set T*.
  • Rollout Generation: For each selected position, REFINE copies the preceding prefix and generates a k-token continuation, retaining predicted and ground-truth hidden states for reward computation.The generated and reference continuations are compared through their final-layer representations.
  • Reward Assignment and Optimization: REFINE uses sequence-level rewards and GRPO policy updates, with cosine similarity supporting semantically similar continuations and a hybrid reward balancing generalization with memorization.The final loss also combines NSP and standard NTP losses with phase-dependent weights to reduce catastrophic forgetting.

4. Experiments

Experiments evaluate REFINE across mid-training, post-training, and test-time training on two fast weight models and multiple long-context benchmarks. REFINE consistently outperforms SFT across retrieval, question answering, and LongBench tasks.

  • Setup: Experiments cover LaCT-760M and DeltaNet-1.3B across mid-training, post-training, test-time training, and long-context benchmarks.Evaluations include RULER NIAH, Booksum, multi-document QA, and 12 LongBench tasks.
  • Mid-training: REFINE mid-training consistently outperforms the original pre-trained and SFT mid-trained models on RULER needle-in-a-haystack retrieval tasks.For DeltaNet, Multi-key NIAH improves by +23.5% over no mid-training and +8.8% over SFT mid-training.
  • Mid-training: REFINE improves validation NTP accuracy during mid-training, whereas SFT is stagnant on Booksum.The figure reports minimum and maximum values across three independent trials.
  • Mid-training: 73.1%: REFINE mid-training improves pre-trained DeltaNet’s average RULER HotpotQA performance over no mid-training, and 22.0% over SFT mid-training.REFINE also consistently outperforms SFT on the reported long-context tasks.
  • Post-training: 17%: nested REFINE improves LaCT-760M’s average SQuADQA score over nested SFT, with corresponding gains for DeltaNet-1.3B.The reported scores are 25.5 versus 21.8 for LaCT-760M and 10.3 versus 8.3 for DeltaNet-1.3B.
  • Analysis: Entropy-based token selection performs best when sampled regions combine different uncertainty levels, while average performance increases through rollout length k = 5 and decreases at k = 7.Reported selection gains are +6.9% over uniform, +1.8% over max entropy, and +1.2% over min entropy for DeltaNet-1.3B.

5. Discussion

REFINE is presented as an NSP-based reinforcement-learning framework that improves long-context modeling throughout the fast weight training lifecycle. Its main boundary is that overly long rollouts degrade cosine-similarity rewards, while broader integration requires architectural changes.

  • Contributions: REFINE combines NSP, entropy-based token selection, and sequence-level rewards to train fast weight models across the training lifecycle.The framework is applied during mid-training, post-training, and test-time training.
  • Implications: REFINE shows consistent improvements on long-context benchmarks across the fast weight training lifecycle.The paper characterizes RL for NSP as a flexible and practical pathway for long-context modeling.
  • Limitations: Overly long rollouts deteriorate REFINE’s cosine-similarity reward, and optimal rollout length depends on the prefix context.The paper suggests richer semantic rewards and dynamic rollout adjustment as future directions.
  • Future work: Fully incorporating NSP into standard fast weight training requires architectural changes, while transferring fast weights across truncated prefixes could improve rollout efficiency.These changes are identified as future work for scaling data and compute.

A. Notation

The appendix provides glossary and notation tables for the paper.

  • Table A.1 contains the glossary and notation.
  • Table A.2 continues the glossary and notation.

B. Related Work

Prior work explores multi-token prediction, RL for next-token training, test-time training, and efficient attention, while REFINE targets fast-weight models with sequence-level adaptation.

  • Multi-Token Prediction: Multi-token prediction methods improve throughput or generate masked tokens, but parallel predictions may fail to capture dependencies among predicted tokens.These approaches primarily modify or target standard transformer architectures rather than fast-weight parameter updates.
  • Continued Pre-Training with RL: RL-based continued pre-training has been applied to next-token prediction using sampled reasoning traces and similarity-based rewards.
  • TTT in Language Models: Test-time training methods adapt transformer models using generated examples, pseudo-labels, or context-based updates before producing task responses.REFINE instead adapts fast-weight memory for long-context modeling, a setting described as unexplored by prior work.
  • Efficient Attention Variants: Sparse and linear attention variants reduce the computational or memory costs of standard attention through sparsity, low-rank operations, or kernel approximation.

C. Datasets and Benchmarks

The study uses a large pre-training corpus subset for mid-training and evaluates selected English LongBench subtasks for long-context performance.

  • Training Data: Long-Data-Collections contains 68.8B tokens assembled from several corpora, with a 200M-token subset used for mid-training.
  • Evaluation Benchmarks: LongBench evaluation uses 12 selected English-based subtasks.

D. Training Configuration

Training configurations vary compute and batch size by lifecycle stage while keeping most REFINE hyperparameters fixed; mid-training uses 8 L40 GPUs for approximately 24 hours.

  • Hyperparameters: REFINE changes train batch size, reward function, and RL loss coefficient across training phases while keeping other hyperparameters constant.
  • Compute: Mid-training and post-training use 8 L40 GPUs, whereas test-time training uses 4 because its batch size is smaller.TTT uses 8 samples per batch, compared with 128 for mid-training and 64 for post-training.
  • Compute: 24 hours are required to mid-train LaCT-760M and DeltaNet-1.3B with REFINE on 200M tokens at 16K context.

E. Additional Analysis

Additional analyses examine validation loss, token selection, reward behavior, short-context retention, and test-time reward choices across REFINE configurations.

  • Validation Loss: REFINE decreases Booksum validation loss for LaCT, while SFT remains constant because the mid-training data matches LaCT’s pre-training data.The authors interpret the decrease as evidence that NSP supplies a learning signal distinct from standard NTP training.
  • Entropy Distribution: Entropy-based token selection shows no index-dependent distribution pattern, supporting per-chunk target sampling weighted by entropy.
  • Short-Context Tasks: REFINE improves long-context modeling without causing catastrophic forgetting on the evaluated short-context tasks.
  • Reward Functions in TTT: Binary exact-match reward is optimal for test-time training, while cosine similarity also exceeds pure SFT.
  • Reward Distribution: Cosine-similarity reward mean and standard deviation both decrease as rollout length increases.The reward distribution remains stable throughout training, but larger k may produce a less sharp learning signal.

F. Qualitative Examples

Table F.1 gives qualitative examples of cosine-similarity rewards for model-generated continuations during mid-training. Each example compares predicted continuations with a ground-truth continuation, highlighting the highest rewards.

  • Four k = 5 continuations are generated for each randomly sampled prefix using the pre-trained LaCT-760M model.
  • Cosine similarity rewards measure semantic similarity between predicted continuations and the ground-truth continuation.
  • The highest reward value for each example is highlighted in bold, while P1–P4 identify the four predicted continuations.
Loading 2602.16704v1…