Verstärkte Schnellgewichte mit Next-Sequence-Prädiktion
Reinforced Fast Weights with Next-Sequence Prediction
February 18, 2026
Autoren: Hee Seung Hwang, Xindi Wu, Sanghyuk Chun, Olga Russakovsky
cs.AI
Zusammenfassung
Schnelle-Gewichte-Architekturen bieten eine vielversprechende Alternative zu aufmerksamkeitsbasierten Transformern für die Modellierung langer Kontexte, da sie konstanten Speicherbedarf unabhängig von der Kontextlänge aufrechterhalten. Ihr Potenzial wird jedoch durch das Next-Token-Prediction (NTP)-Trainingsparadigma eingeschränkt. NTP optimiert Einzel-Token-Vorhersagen und ignoriert semantische Kohärenz über mehrere Token nach einem Präfix hinweg. Folglich lernen schnelle Gewichtsmodelle, die ihre Parameter dynamisch zur Speicherung kontextueller Informationen aktualisieren, suboptimale Repräsentationen, die keine langreichweitigen Abhängigkeiten erfassen. Wir stellen REFINE (Reinforced Fast weIghts with Next sEquence prediction) vor, ein Reinforcement-Learning-Framework, das schnelle Gewichtsmodelle unter der Next-Sequence-Prediction (NSP)-Zielfunktion trainiert. REFINE wählt informative Token-Positionen basierend auf Vorhersageentropie aus, generiert Multi-Token-Rollouts, weist selbstüberwachte sequenzielle Belohnungen zu und optimiert das Modell mit Group Relative Policy Optimization (GRPO). REFINE ist während der gesamten Trainingslebensdauer vortrainierter Sprachmodelle anwendbar: während des Trainings, nach dem Training und während Test-Time-Training. Unsere Experimente mit LaCT-760M und DeltaNet-1.3B zeigen, dass REFINE durchgängig das überwachte Fine-Tuning mit NTP bei Needle-in-a-Haystack-Retrieval, Frage-Antwort-Aufgaben mit langem Kontext und diversen Aufgaben in LongBench übertrifft. REFINE bietet einen effektiven und vielseitigen Rahmen zur Verbesserung der Langzeitkontextmodellierung in schnellen Gewichtsarchitekturen.
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.