Resumen Técnico: AMOR (Enrutador de Salida Metacognitivo Adaptativo)

Enunciado del Problema

Los modelos híbridos recurrentes-atención buscan combinar la eficiencia de inferencia de los modelos de espacio de estados recurrentes lineales (SSM) con el poder expresivo de los mecanismos de atención. Sin embargo, las arquitecturas híbridas existentes suelen depender de estrategias de programación fija, donde la atención se aplica uniformemente a todos los tokens o en capas predeterminadas (por ejemplo, cada $N$-ésima capa). Este enfoque incurre en el costo cuadrático completo de la atención ($O(N^2)$) incluso cuando el estado recurrente por sí solo es suficiente para predecir el siguiente token con alta confianza. En consecuencia, estos modelos no logran capitalizar plenamente las ganancias de eficiencia de los SSM, particularmente en escenarios de contexto largo donde la sobrecarga computacional de los bloques de atención innecesarios se convierte en un cuello de botella. Además, los híbridos fijos a menudo se degradan bajo desplazamientos de distribución (por ejemplo, extrapolación de longitud) porque fuerzan mecanismos de atención que pueden no estar calibrados para la dificultad específica de la entrada.

Metodología: Arquitectura AMOR

Los autores proponen AMOR (Enrutador de Salida Metacognitivo Adaptativo), una arquitectura híbrida post-hoc que asigna la atención dinámicamente basándose en la incertidumbre predictiva en lugar de posiciones fijas.

Arquitectura Central

AMOR se construye sobre una columna vertebral recurrente estándar (probada con Mamba2 y Gated DeltaNet) y añade $K$ (típicamente 3) bloques de atención con puerta de entropía después de las capas recurrentes pero antes de la cabeza final del LM.

• Columna Vertebral: El mezclador recurrente procesa toda la secuencia, manteniendo un estado de tamaño fijo.
• Bloques Post-Hoc: A diferencia de los híbridos en serie que reemplazan capas recurrentes, los bloques de AMOR son aditivos. Operan sobre la salida normalizada de la columna vertebral.
• Asimetría: El primer bloque de AMOR (Bloque 0) opera directamente sobre el residuo normalizado final de la columna vertebral (post-norm), mientras que los bloques subsiguientes (1 y 2) utilizan una formulación estándar pre-norm. Este diseño preserva la columna vertebral como un modelo de lenguaje válido independiente, tratando los bloques de AMOR como refinamientos condicionales.

El Mecanismo de Puerta de Entropía

La innovación central es un mecanismo de puerta duro y libre de gradientes inspirado en la distinción entre Sistema 1 y Sistema 2 en la ciencia cognitiva (Kahneman, 2011).

1. Estimación de Incertidumbre: Para cada token, la cabeza compartida del LM calcula los logits. La entropía de salida normalizada ($H_t$) se calcula a partir de la distribución softmax.
2. Umbral Adaptativo: Un umbral dinámico $\tau$ se calcula como una media móvil exponencial (EMA) de la mediana de entropía del lote ($\mu$) más una desviación estándar escalada ($\kappa\sigma$):
   $$\tau = \mu + \kappa\sigma$$
   Esto permite que el umbral rastree los desplazamientos en la distribución de incertidumbre del modelo durante el entrenamiento sin requerir parámetros de enrutador aprendidos.
3. Ejecución Condicional:
   • Si $H_t > \tau$, la puerta se activa ($g_t=1$): El token sufre atención causal (SDPA) con RoPE, y la salida se suma al residuo con un factor de escalado por canal aprendido $\alpha$.
   • Si $H_t \le \tau$, la puerta permanece inactiva ($g_t=0$): El cálculo de atención se omite y el token pasa sin cambios.
4. Flujo de Gradientes: La puerta es no diferenciable. Los gradientes fluyen a través de los pesos de atención ($W_Q, W_K, W_V, W_O$) solo en las posiciones de activación. Crucialmente, $W_K$ y $W_V$ reciben gradientes de futuras posiciones de activación para asegurar que la caché KV permanezca válida para la recuperación, incluso si el token actual omitió la atención.

Entrenamiento e Inferencia

• Entrenamiento: AMOR se entrena en modo "denso con máscara" utilizando FlashAttention-2 para la eficiencia, matemáticamente equivalente a una ejecución dispersa real. El umbral $\tau$ se congela en la inferencia.
• Inferencia: El modelo realiza un salto condicional. Cuando la puerta está apagada, la proyección $Q$ costosa, el cálculo de atención y la proyección $O$ se eluden, ahorrando cómputo. Sin embargo, las proyecciones $K$ y $V$ y las actualizaciones de caché siempre ocurren para mantener la continuidad del contexto.

Contribuciones Clave

1. Enrutamiento Impulsado por Incertidumbre: AMOR introduce un mecanismo simple y post-hoc que invoca la atención selectivamente solo cuando la columna vertebral recurrente exhibe alta incertidumbre predictiva (alta entropía), en lugar de depender de programas de capas fijos.
2. Hibridación Eficiente en Parámetros: El método añade una sobrecarga mínima de parámetros (aprox. <5% para los bloques de puerta y atención) mientras logra un rendimiento comparable o superior a los híbridos de programación fija.
3. Robustez ante Desplazamientos de Distribución: Al confiar en la columna vertebral recurrente para tokens de baja incertidumbre, AMOR mitiga la degradación del rendimiento observada en híbridos fijos durante la extrapolación de longitud, donde los mecanismos de atención a menudo sufren de deriva posicional.
4. Puerta Libre de Gradientes: El enfoque evita la complejidad de enrutadores aprendidos o objetivos de aprendizaje por refuerzo, utilizando un umbral estadístico derivado de la propia distribución de salida del modelo.

Resultados Experimentales

Los autores evaluaron AMOR en tres escalas (180M, 440M, 1.5B de parámetros) utilizando columnas vertebrales Mamba2 y Gated DeltaNet, comparándolas con modelos recurrentes puros, Transformers puros e híbridos en serie/fusionados de programación fija.

• Razonamiento de Sentido Común: AMOR coincidió consistentemente o superó tanto a las columnas vertebrales recurrentes puras como a los híbridos de programación fija en todas las escalas y puntos de referencia (LAMBADA, HellaSwag, PIQA, ARC, etc.).
• Recuperación en Contexto: Mientras que los modelos recurrentes puros luchan con la recuperación de copias exactas, AMOR mejoró sustancialmente sobre su columna vertebral, logrando resultados competitivos con híbridos fijos mientras activaba la atención en solo ~22% de los tokens.
• Rendimiento en Contexto Largo: En el punto de referencia LongBench, AMOR mantuvo un rendimiento estable a medida que aumentaba la longitud del contexto. Por el contrario, los híbridos de programación fija y los Transformers puros se degradaron significativamente debido a desplazamientos de distribución RoPE y restricciones de memoria. La capacidad de AMOR para omitir la atención en tokens predecibles le permitió escalar de manera más fluida.
• Eficiencia de Inferencia: AMOR demostró una mejor escalabilidad del rendimiento de decodificación en regímenes de contexto largo. El análisis teórico y las pruebas empíricas (hasta 448K de contexto) mostraron que los ahorros por omitir la atención superan la sobrecarga del cálculo de entropía una vez que la longitud del contexto excede cierto punto de cruce (aprox. 256K para modelos de 440M).

Significado y Afirmaciones

El artículo argumenta que cuándo se aplica la atención es tan crítico como cuánto se aplica. Los autores afirman que:

1. Eficiencia y Expresividad no son un intercambio: Al tratar la asignación de cómputo como un problema de enrutamiento basado en la dificultad de la entrada, los modelos pueden lograr alta expresividad sin incurrir en el costo completo de la atención.
2. La simplicidad es efectiva: No son necesarios enrutadores aprendidos complejos ni objetivos auxiliares; una puerta de entropía simple e impulsada por la salida puede coordinar eficazmente arquitecturas híbridas.
3. Cómputo Adaptativo: AMOR apunta hacia una clase de modelos que igualan dinámicamente el cómputo a la dificultad de la entrada, preservando las fortalezas de las columnas vertebrales de espacio de estados ligeras mientras despliegan selectivamente mecanismos de atención pesados solo cuando el modelo está incierto.

Los resultados sugieren que la asignación selectiva de atención basada en la incertidumbre predictiva ofrece una alternativa robusta a las estrategias de enrutamiento uniformes o fijas, particularmente para tareas de generación y recuperación en contexto largo.