Research map / Architectures and efficiency / Attention and transformers

Reinforced Fast Weights via Next-Sequence Prediction

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse efficient transformers for long context

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes ReFINE, an RL framework (GRPO with entropy-based token selection and cosine/binary sequence-level rewards) that trains fast weight language models under a next-sequence prediction objective across mid-training, post-training, and test-time training.

Problem

Fast weight architectures for long-context modeling are trained with next-token prediction, which only optimizes single-token predictions and ignores how the evolving fast-weight memory supports multi-token continuations, leading to suboptimal long-range representations.

Approach

Reformulate training as next-sequence prediction solved with RL: sample high-entropy positions per chunk, generate k-token rollouts, reward them with hidden-state cosine similarity (or exact-match binary reward during TTT), and optimize with GRPO plus an NTP regularizer.

Key results

• Mid-training with ReFINE improves LaCT-760M RULER NIAH average over SFT mid-training (e.g., 46.2 to 47.7 MK-NIAH avg) and DeltaNet-1.3B Multi-key NIAH by +8.8 over SFT mid-training.
• Nested ReFINE post-training on DeltaNet-1.3B raises RULER SQuADQA average from 8.3 (nested SFT) to 10.3, and with ReFINE TTT reaches 12.3.
• On 12 LongBench tasks, ReFINE TTT reaches 18.0 average for LaCT-760M vs 13.5 for SFT and 13.1 for no training.

Surprising findings

• Validation NTP accuracy on Booksum improves under ReFINE mid-training more than under direct NTP SFT, even though SFT optimizes NTP directly.

Limitations

• Cosine similarity reward degrades for overly long rollouts, limiting gains from larger k.
• Rollout generation requires re-running prefixes for each sampled position, which is computationally costly for long contexts.
• Optimal rollout length is context-dependent and fixed in the current method.

Experimental setup

Datasets: Long-Data-Collections Booksum RULER SQuADQA RULER HotpotQA LongBench

Benchmarks: RULER NIAH RULER SQuADQA RULER HotpotQA LongBench Booksum short-context tasks (PIQA, HellaSwag, WinoGrande, ARC, Wikitext, LAMBADA, FDA, SWDE)

Metrics: recall F1 rouge accuracy NTP accuracy perplexity

Baselines: SFT with next-token prediction nested SFT LaCT-760M DeltaNet-1.3B

Builds on: GRPO LaCT DeltaNet multi-token prediction nested learning

Closest papers by meaning

• Better & Faster Large Language Models via Multi-token Prediction (ICML 2024 · similarity 0.89)
• Implicit Optimization Bias of Next-token Prediction in Linear Models (NeurIPS 2024 · similarity 0.88)
• A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts (ICLR 2025 · similarity 0.88)
• Next-Latent Prediction Transformers Learn Compact World Models (NeurIPS 2026 · similarity 0.88)
• RAT: Bridging RNN Efficiency and Attention Accuracy via Chunk-based Sequence Modeling (NeurIPS 2025 · similarity 0.88)
• Logits are All We Need to Adapt Closed Models (ICML 2025 · similarity 0.88)