Resumen Técnico: Raven - Modelado de Secuencias de Alta Recuperación con Enrutamiento de Memoria Dispersa

1. Planteamiento del Problema

Los modelos de secuencias de tiempo lineal, como los Modelos de Espacio de Estados (SSMs) y los Transformers lineales, enfrentan un compromiso fundamental en la gestión de su memoria para la recuperación de contexto largo. Las arquitecturas existentes generalmente caen en dos extremos, ambos de los cuales sufren limitaciones específicas para recuperar tokens pasados específicos:

• Actualizaciones de Estado Densas (SSMs): Modelos como Mamba-2, GLA y GDN mantienen un estado de tamaño fijo que se actualiza densamente en cada paso. Cada ranura de memoria se actualiza para cada nuevo token, confiando en mecanismos de decaimiento aprendidos para olvidar. Si bien esto permite una persistencia indefinida en principio, las actualizaciones densas causan interferencia, lo que dificulta la preservación de tokens individuales específicos sin que sean sobrescritos o diluidos por la información subsecuente.
• Escrituras Dispersas Basadas en Posición (Atención de Ventana Deslizante - SWA): SWA escribe de forma dispersa almacenando representaciones de tokens explícitos, pero los restringe a una ventana reciente fija. Utiliza una estrategia de Primero en Entrar - Primero en Salir (FIFO) donde los tokens antiguos son expulsados de forma estricta basándose en su posición. Si bien esto permite una recuperación fiable dentro de la ventana, la capacidad de recuperación cae casi a cero una vez que el token relevante es expulsado, independientemente de su importancia.

El problema central es que los modelos existentes confunden dónde se escribe la información con cuánto tiempo persiste. Los SSMs escriben en todas partes y confían en el decaimiento para olvidar; SWA escribe en ninguna parte (excepto en la ventana actual) y confía en la posición para olvidar. Ninguno combina eficazmente la escritura dispersa dependiente del contenido con un decaimiento explícito y controlable.

2. Metodología: Routing Slot Memories (RSMs)

Los autores proponen las Routing Slot Memories (RSMs), un marco unificado que desacopla la ubicación de las escrituras de la duración de la persistencia. Las RSMs mantienen un conjunto fijo de $M$ ranuras de memoria y, en cada paso de tiempo $t$, utilizan un enrutador para determinar qué subconjunto de ranuras recibe nueva información y aplica decaimiento.

2.1 El Marco RSM

La recurrencia general para una RSM se define como:
$$S_t = (1 - r_t) \odot S_{t-1} + r_t \odot (D_t S_{t-1} A_t + U_t)$$
Donde:

• $S_t \in \mathbb{R}^{M \times d}$ es el estado de memoria con $M$ ranuras independientes.
• $r_t \in \mathbb{R}^M$ es el vector de enrutamiento que determina la intensidad de escritura para cada ranura.
• $D_t$ y $A_t$ controlan el decaimiento y la mezcla de características.
• $U_t$ es el contenido de escritura derivado del input $x_t$.

El marco posiciona a los modelos existentes en un espectro:

• SSMs: Utilizan un enrutador denso ($r_t = \mathbf{1}_M$), actualizando todas las ranuras con un decaimiento gradual.
• SWA: Utiliza un enrutador dependiente del tiempo y de un solo elemento ($r_t = e_t$), actualizando exactamente una ranura con expulsión estricta (sin decaimiento).
• GSA/ABC: Utilizan enrutadores dependientes del input pero típicamente actualizan todas las ranuras (escrituras densas), lo que limita su capacidad para aislar memorias.

2.2 La Arquitectura Raven

Raven es una instanciación específica de las RSMs diseñada para ocupar el "punto medio", combinando el enrutamiento disperso dependiente del input con un decaimiento explícito.

• Estructura de Estado Dual: Al igual que SWA, Raven mantiene estados de Clave ($S^k_t$) y Valor ($S^v_t$) separados.
• Enrutamiento Disperso Dependiente del Input: En lugar de una ventana fija o una actualización densa, Raven utiliza un enrutador $g_t$ derivado del input $x_t$. Emplea una estrategia de selección Top-K (adaptada del enrutamiento de DeepSeek MoE) para seleccionar un pequeño subconjunto de $K$ ranuras para actualizar.
  • Los puntajes se computan mediante una función sigmoide: $m_t = \sigma(W x_t)$.
  • Solo se mantienen las $K$ entradas superiores; las demás se establecen en cero.
  • El enrutador se normaliza para asegurar que las ranuras seleccionadas reciban la intensidad de escritura apropiada.
• Decaimiento Selectivo: El decaimiento se aplica solo a las ranuras seleccionadas. La recurrencia es:
  $$S^k_t = \exp(a_t r_t) \odot S^k_{t-1} + (1 - \exp(a_t r_t)) \odot k^\top_t$$
  Aquí, $a_t$ es un escalar de decaimiento dependiente del input (similar a Mamba-2). Si una ranura no es seleccionada ($r_t[i] \approx 0$), el término $\exp(a_t r_t[i]) \approx 1$, lo que significa que la ranura está congelada y retiene su contenido perfectamente hasta que sea seleccionada nuevamente.
• Diseño de Bloque Simplificado: A diferencia de muchos modelos lineales (ej. Mamba, GLA), Raven omite las convoluciones de corto alcance. Se basa únicamente en los mecanismos de enrutamiento y decaimiento para la integración de contexto local y de largo alcance, resultando en una estructura de bloque más simple.

3. Contribuciones Clave

1. Routing Slot Memories (RSMs): Un marco formal que unifica los SSMs y SWA, separando explícitamente el "lugar de escritura" (enrutamiento) del "mecanismo de olvido" (decaimiento).
2. El Modelo Raven: Un modelo de secuencia de tiempo lineal que implementa las RSMs mediante el uso de un enrutamiento disperso dependiente del input para actualizar un subconjunto de ranuras de memoria mientras aplica decaimiento solo a aquellas ranuras actualizadas.
3. Mecanismo para la Recuperación de Largo Alcance: Al enrutar los tokens críticos para la recuperación (ej. contraseñas/passkeys) hacia ranuras dedicadas y congelarlos, Raven evita la interferencia destructiva común en los SSMs densos y la expulsión estricta de SWA.
4. Generalización de Longitud: El modelo logra un rendimiento sólido en contextos de hasta 16$\times$ su longitud de entrenamiento sin necesidad de un currículo de longitud explícito o técnicas de interpolación posicional (como RoPE), ya que las longitudes efectivas variables a través de las ranuras exponen naturalmente al modelo a diversas dinámicas temporales.

4. Resultados Experimentales

Los autores evaluaron Raven en benchmarks intensivos de recuperación y tareas de modelado de lenguaje general, comparándolo contra bases sólidas incluyendo Mamba-2, GLA, GDN, SWA y variantes de Transformer (FoX).

• Needle-in-a-Haystack (NIAH):
  • Escala de 400M: Raven mantuvo una precisión casi perfecta ($\ge 99\%$) hasta los 16K tokens y se mantuvo fuerte ($>91\%$) a los 32K tokens (16$\times$ la longitud de entrenamiento). En contraste, Mamba-2 y GDN se degradaron significativamente más allá de los 8K.
  • Escala de 800M: Raven alcanzó un 91.0% de precisión en NIAH-1 a 64K tokens, superando significativamente al siguiente mejor baseline (GDN con 45.2%).
• Tareas de Recuperación del Mundo Real:
  • En SWDE (extracción web), FDA (extracción de documentos) y SQuAD, Raven superó consistentemente a los baselines lineales. Por ejemplo, en FDA, fue el único modelo lineal en exceder el 22% de precisión, mejorando aproximadamente 8 puntos sobre el mejor baseline de SSM.
• Modelado de Lenguaje:
  • Raven igualó o superó el rendimiento de los SSMs densos (Mamba-2, GLA) en benchmarks estándar de zero-shot (Lambada, PIQA, HellaSwag, etc.), demostrando que el enrutamiento disperso no compromete las capacidades generales de modelado de lenguaje.
• Arquitecturas Híbridas:
  • Cuando se integra en modelos híbridos (emparejado con capas de atención completa), Raven superó significativamente a los híbridos que usan Mamba-2, GDN o SWA-RoPE en tareas de recuperación de contexto largo, particularmente en el rango de 16K–64K.

5. Significación y Reivindicaciones

El artículo afirma que Raven aborda una brecha arquitectónica fundamental en el modelado de secuencias lineales. Al separar explícitamente dónde se escribe la información de cuánto tiempo persiste, Raven logra un nivel de recuperación de largo alcance previamente inalcanzable para los modelos lineales sin recurrir a la atención cuadrática o configuraciones híbridas complejas.

Las reivindicaciones clave incluyen:

• Superioridad sobre SWA: Raven reemplaza la expulsión basada en la posición con enrutamiento dependiente del contenido y decaimiento selectivo, permitiendo una recuperación perfecta en tareas donde SWA falla completamente.
• Superioridad sobre SSMs Densos: Raven evita la interferencia de las actualizaciones de estado densas escribiendo selectivamente, lo que permite que tokens específicos persistan indefinidamente en ranuras dedicadas.
• Simplicidad: El modelo logra estas ganancias sin convoluciones de corto alcance, simplificando la arquitectura en comparación con modelos lineales previos.
• Especialización Emergente: El análisis muestra que Raven aprende naturalmente a asignar ranuras de memoria específicas para tokens críticos de recuperación (ej. contraseñas) mientras usa otras ranuras para el contenido general, creando una asignación de memoria estructurada y consciente del contenido sin supervisión explícita.

Los autores concluyen que el marco RSM sugiere que el enrutamiento y el olvido son ejes de diseño ortogonales que deben tratarse como elecciones de diseño de primer orden en lugar de subproductos de los mecanismos de decaimiento.