Poids Rapides Renforcés avec Prédiction de Séquences Suivantes

Reinforced Fast Weights with Next-Sequence Prediction

February 18, 2026

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

cs.AI

Résumé

Les architectures à poids rapides offrent une alternative prometteuse aux transformeurs basés sur l'attention pour la modélisation de contextes longs, en maintenant une surcharge mémoire constante quelle que soit la longueur du contexte. Cependant, leur potentiel est limité par le paradigme d'entraînement par prédiction du token suivant (NTP). Le NTP optimise les prédictions token par token et ignore la cohérence sémantique sur plusieurs tokens suivant un préfixe. Par conséquent, les modèles à poids rapides, qui mettent à jour dynamiquement leurs paramètres pour stocker l'information contextuelle, apprennent des représentations sous-optimales qui échouent à capturer les dépendances à longue portée. Nous présentons REFINE (Reinforced Fast weIghts with Next sEquence prediction), un cadre d'apprentissage par renforcement qui entraîne les modèles à poids rapides avec l'objectif de prédiction de séquence suivante (NSP). REFINE sélectionne des positions de tokens informatifs basées sur l'entropie de prédiction, génère des déploiements multi-tokens, attribue des récompenses auto-supervisées au niveau de la séquence, et optimise le modèle avec l'optimisation de politique relative par groupe (GRPO). REFINE est applicable tout au long du cycle de vie d'entraînement des modèles de langage pré-entraînés : en cours d'entraînement, après l'entraînement, et pendant l'entraînement au moment du test. Nos expériences sur LaCT-760M et DeltaNet-1.3B démontrent que REFINE surpasse constamment le fine-tuning supervisé avec NTP sur des tâches de recherche d'aiguille dans une botte de foin, de question-réponse en contexte long, et diverses tâches de LongBench. REFINE fournit un cadre efficace et polyvalent pour améliorer la modélisation de contextes longs dans les architectures à poids rapides.

English

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.