Pesos Rápidos Reforzados con Predicción de Siguiente Secuencia
Reinforced Fast Weights with Next-Sequence Prediction
February 18, 2026
Autores: Hee Seung Hwang, Xindi Wu, Sanghyuk Chun, Olga Russakovsky
cs.AI
Resumen
Las arquitecturas de pesos rápidos ofrecen una alternativa prometedora a los transformadores basados en atención para el modelado de contextos largos, ya que mantienen una sobrecarga de memoria constante independientemente de la longitud del contexto. Sin embargo, su potencial está limitado por el paradigma de entrenamiento de predicción del siguiente token (NTP). El NTP optimiza las predicciones de tokens individuales e ignora la coherencia semántica a través de múltiples tokens que siguen a un prefijo. En consecuencia, los modelos de pesos rápidos, que actualizan dinámicamente sus parámetros para almacenar información contextual, aprenden representaciones subóptimas que no logran capturar dependencias de largo alcance. Presentamos REFINE (Pesos Rápidos Reforzados con Predicción de la Siguiente Secuencia), un marco de aprendizaje por refuerzo que entrena modelos de pesos rápidos bajo el objetivo de predicción de la siguiente secuencia (NSP). REFINE selecciona posiciones de tokens informativas basándose en la entropía de la predicción, genera simulaciones de múltiples tokens, asigna recompensas autosupervisadas a nivel de secuencia y optimiza el modelo con optimización de políticas relativas por grupos (GRPO). REFINE es aplicable a lo largo de todo el ciclo de vida de entrenamiento de los modelos de lenguaje preentrenados: durante el entrenamiento, posterior al entrenamiento y en el entrenamiento en tiempo de prueba. Nuestros experimentos en LaCT-760M y DeltaNet-1.3B demuestran que REFINE supera consistentemente al ajuste fino supervisado con NTP en tareas de recuperación "aguja-en-un-pajar", respuesta a preguntas de contexto largo y diversas tareas en LongBench. REFINE proporciona un marco eficaz y versátil para mejorar el modelado de contextos largos en arquitecturas de pesos rápidos.
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.