Reinforced Fast Weights with Next-Sequence Prediction

1. El Problema

Las arquitecturas de pesos rápidos (fast weights) son una alternativa prometedora a los transformadores basados en atención para el modelado de contextos largos. A diferencia de los transformadores tradicionales, que tienen un costo de memoria y computación cuadrático respecto a la longitud del contexto, las arquitecturas de pesos rápidos mantienen una sobrecarga de memoria constante al actualizar dinámicamente sus parámetros (matrices de peso) para almacenar información contextual.

Sin embargo, estos modelos sufren una limitación fundamental: se entrenan bajo el paradigma de predicción del siguiente token (NTP).

• Limitación del NTP: El NTP optimiza la probabilidad de un solo token a la vez, ignorando la coherencia semántica a través de múltiples tokens que siguen a un prefijo.
• Consecuencia: Esto lleva a que los modelos de pesos rápidos aprendan representaciones subóptimas que no capturan bien las dependencias de largo alcance, ya que la retroalimentación es demasiado local y no considera el impacto de las predicciones futuras en el estado interno del modelo.

2. Metodología: REFINE

Los autores proponen REFINE (Reinforced Fast weIghts with Next sEquence prediction), un marco de aprendizaje por refuerzo (RL) diseñado para entrenar modelos de pesos rápidos bajo un objetivo de predicción de la siguiente secuencia (NSP).

El objetivo NSP busca que el modelo prediga una secuencia coherente de múltiples tokens futuros, en lugar de un solo token. Para abordar los desafíos computacionales de generar continuaciones completas para cada prefijo, REFINE utiliza un enfoque de cuatro pasos:

1. Selección de Tokens Basada en Entropía:

   • En lugar de generar continuaciones para cada posición de la secuencia (lo cual es costoso), el sistema calcula la entropía de la predicción del NTP en cada token.
   • La secuencia se divide en chunks (fragmentos). De cada chunk, se selecciona una posición objetivo con alta probabilidad si tiene alta entropía (alta incertidumbre). Esto permite centrar el entrenamiento en las regiones más informativas y desafiantes del contexto.

2. Generación de Rollouts (Despliegues):

   • Para cada posición seleccionada, el modelo genera una continuación de $k$ tokens (rollout) basándose en el prefijo hasta ese punto.
   • Se extraen los estados ocultos tanto de la secuencia generada como de la secuencia de verdad fundamental (ground truth).

3. Asignación de Recompensas:

   • En lugar de usar una pérdida de entropía cruzada estricta, se asigna una recompensa basada en la similitud de los estados ocultos (usando similitud de coseno) entre la predicción y la verdad fundamental.
   • Esto permite recompensar continuaciones semánticamente equivalentes aunque no sean idénticas lexicalmente, mejorando la generalización.
   • Para escenarios de entrenamiento en tiempo de prueba (Test-Time Training), se introduce una recompensa híbrida que incluye coincidencia exacta binaria para priorizar la memorización.

4. Optimización con RL:

   • El modelo se optimiza utilizando GRPO (Group Relative Policy Optimization), maximizando la recompensa esperada a nivel de secuencia.
   • La pérdida final es una combinación ponderada de la pérdida NSP (RL) y la pérdida NTP estándar (SFT) para evitar el olvido catastrófico.

3. Contribuciones Clave

1. Objetivo NSP para Pesos Rápidos: Introducen la predicción de secuencia (NSP) como un objetivo de entrenamiento superior al NTP para arquitecturas de pesos rápidos, alineando el objetivo de entrenamiento con la función de memoria a largo plazo de estos modelos.
2. Marco de RL Eficiente: Proponen REFINE, un marco que combina la selección de tokens basada en entropía con recompensas a nivel de secuencia, haciendo viable el entrenamiento de NSP sin el costo computacional prohibitivo de generar continuaciones para cada token.
3. Aplicabilidad en Todo el Ciclo de Vida: Demuestran que REFINE es efectivo en tres etapas distintas del entrenamiento de modelos de lenguaje:
   • Mid-training: Adaptación de modelos preentrenados a la nueva objetivo.
   • Post-training: Ajuste fino para tareas específicas y alineación.
   • Test-Time Training (TTT): Adaptación en tiempo de inferencia sobre el prompt de entrada sin etiquetas adicionales.

4. Resultados Experimentales

Los experimentos se realizaron en dos modelos de pesos rápidos: LaCT-760M y DeltaNet-1.3B.

• Rendimiento General: REFINE superó consistentemente a los modelos ajustados finamente solo con NTP (SFT) en todas las etapas.
• Recuperación de Información (Needle-in-a-Haystack): En la tarea RULER NIAH, REFINE mostró mejoras significativas, especialmente en contextos largos (8K y 16K tokens). Por ejemplo, en DeltaNet-1.3B, hubo una mejora del 20.3% en mid-training y 15.0% en TTT comparado con SFT.
• Preguntas de Respuesta (QA) y LongBench:
  • En tareas de QA con múltiples documentos, REFINE mejoró el rendimiento promedio en un 15-25% sobre las líneas base de SFT.
  • En el conjunto de pruebas LongBench (12 tareas), REFINE logró un puntaje promedio superior en todas las categorías (QA, resumen, codificación, etc.).
• Análisis de Componentes:
  • La selección de tokens basada en entropía superó a la selección uniforme o aleatoria.
  • La longitud óptima del rollout ($k$) se encontró en 5 tokens; longitudes mayores diluyeron la señal de la recompensa.
  • Las recompensas basadas en similitud semántica (coseno) funcionaron mejor que la coincidencia exacta binaria durante el mid-training, aunque la binaria fue superior en TTT.

5. Significado e Impacto

El trabajo de REFINE es significativo porque:

• Resuelve la brecha de entrenamiento: Identifica que el entrenamiento estándar (NTP) es insuficiente para explotar el potencial de las arquitecturas de pesos rápidos en contextos largos.
• Eficiencia y Escalabilidad: Ofrece un método para mejorar la capacidad de modelado de contextos largos sin aumentar la complejidad de memoria durante la inferencia (manteniendo la ventaja de los pesos rápidos).
• Versatilidad: Proporciona un marco unificado que funciona desde la adaptación de modelos base hasta la inferencia en tiempo real, lo que sugiere que el aprendizaje por refuerzo con objetivos de secuencia es una vía práctica y flexible para la próxima generación de modelos de lenguaje eficientes.

En resumen, REFINE demuestra que al cambiar el objetivo de entrenamiento de "predecir el siguiente token" a "predecir la siguiente secuencia coherente" mediante aprendizaje por refuerzo, se pueden desbloquear capacidades de memoria a largo plazo significativamente superiores en arquitecturas de pesos rápidos.