Resumen Técnico: Búsqueda de Atención Neuronal Lineal (NAtS-L)

Planteamiento del Problema

La complejidad computacional cuadrática ($O(L^2)$) de la autoatención basada en softmax estándar en los Transformers se ha convertido en un cuello de botella significativo para escenarios de contexto largo. Si bien los modelos de atención lineal ofrecen una alternativa prometedora al comprimir los valores de Clave-Valor (KV) pasados en un estado oculto de tamaño fijo para lograr una complejidad lineal ($O(L)$), a menudo sufren de una expresividad limitada debido a las restricciones de este tamaño de estado fijo. Por el contrario, los Transformers estándar preservan toda la información del contexto KV, lo que permite una mejor recuperación a largo plazo, pero incurre en altos costos de memoria y computación.

Los enfoques híbridos existentes intentan mitigar esto intercalando capas de atención softmax y lineal o utilizando esquemas fijos para alternar entre ellas. Sin embargo, estos métodos suelen retener estructuras fijas que pueden no capturar con flexibilidad toda la información requerida a través de diversos contextos, y su eficiencia sigue estando limitada por la presencia de capas completas de softmax. Además, las variaciones de atención dispersa (sparse attention) a menudo dependen de heurísticas humanas predefinidas, lo que podría ignorar correlaciones no cubiertas por sus mapas dispersos.

Metodología

El artículo propone Neural Attention Search Linear (NAtS-L), un marco que aplica operaciones de atención tanto lineales como de softmax dentro de la misma capa pero en diferentes tokens. En lugar de una arquitectura fija, NAtS-L determina automáticamente la operación de atención óptima para cada bloque de tokens basándose en el contexto de entrada.

Mecanismos Centrales

1. Hibridación a Nivel de Token: La secuencia de entrada se divide en fragmentos (chunks). Una Capa de Puntuación de Atención (consistente en un pooling de media seguido de una capa lineal) mapea cada fragmento a un conjunto de puntuaciones que representan diferentes tipos de operación (Atención Softmax vs. Gated DeltaNet, una variante de atención lineal). La operación con la puntuación más alta se asigna a ese fragmento.
2. Espacio de Búsqueda Unificado: El marco unifica las operaciones de atención de softmax no lineal y de atención lineal, permitiendo que sean procesadas dentro de una sola capa. El modelo aprende la importancia de cada bloque de tokens, aplicando atención softmax a los tokens con alto impacto a largo plazo (que requieren recuperación a largo plazo) y atención lineal a los tokens con impacto a corto plazo.
3. Computación por Fragmentos (Chunk-wise): Para mantener la eficiencia del hardware, NAtS-L utiliza un enfoque paralelo por fragmentos. Dentro de un fragmento, se computan tanto las correlaciones lineales como las de softmax. El modelo mantiene estados ocultos para ambos tipos de operación. Cuando llega un nuevo fragmento, el modelo realiza una pseudo-actualización de ambos estados; sin embargo, para los fragmentos que no fueron seleccionados como la operación objetivo, los estados ocultos se reinician o se revierten al inicio del fragmento, asegurando que solo se retengan los valores KV relevantes para el tipo de operación seleccionado.
4. Fusión de Salida: Las salidas de las operaciones de softmax y lineal seleccionadas se normalizan individualmente (usando normalización RMS) y luego se combinan mediante una suma ponderada. Los pesos para esta combinación se generan dinámicamente a partir de las proyecciones de la consulta (Query, $Q$), asegurando que el modelo aprenda a equilibrar las contribuciones de ambos tipos de atención.
5. Optimización Basada en Gradientes: La selección de los tipos de atención se aprende de extremo a extremo (end-to-end). Los gradientes se calculan para la capa de puntuación de atención, permitiendo que el modelo optimice conjuntamente la selección de la operación y los pesos del modelo sin requerir pérdidas auxiliares o modelos pre-entrenados.

Contribuciones Clave

1. Marco Unificado: Los autores proporcionan una descripción unificada para la atención de softmax no lineal y la atención lineal, permitiendo que sean procesadas y seleccionadas dentro de una sola capa.
2. Arquitectura Adaptativa: Se propone NAtS-L como un marco que determina automáticamente la operación de atención óptima para la información del contexto de entrada en una base de nivel de token (nivel de fragmento), en lugar de depender de ratios fijos por capa.
3. Eficiencia y Expresividad: El artículo demuestra que NAtS-L logra un sólido equilibrio entre rendimiento y eficiencia, ofreciendo capacidades de modelado de contexto largo superiores a los modelos puramente lineales, mientras mantiene una latencia computacional menor que los Transformers de softmax completo.

Resultados Experimentales

Los autores evaluaron NAtS-L en tareas de pre-entrenamiento de escala académica utilizando el conjunto de datos Fineweb-Edu con tamaños de modelo de 380M y 800M de parámetros.

• Modelado de Lenguaje y Razonamiento: En benchmarks de razonamiento de sentido común de cero disparos (zero-shot) (ej. LAMBADA, PIQA, Hellaswag), NAtS-L y su variante híbrida se desempeñaron de manera comparable o ligeramente mejor que los baselines de Gated DeltaNet (GDN), Mamba2 y Transformer puro.
• Tareas de Recuperación: En tareas de recuperación del mundo real (ej. SWDE, SQD, DROP), NAtS-L Hybrid logró los mejores puntajes promedio, superando a los modelos lineales puros (que tuvieron dificultades con tareas específicas como FDA) y a los Transformers puros (que tuvieron dificultades con SQD). Esto resalta la robustez del modelo en la preservación de la información de contexto necesaria.
• Extrapolación de Longitud: Evaluado en contextos de hasta 65,536 tokens (PG19, NarrativeQA, CodeParrot), NAtS-L Hybrid mantuvo los niveles de perplejidad incluso cuando el contexto de entrada excedía la longitud de entrenamiento (4096 tokens), mientras que los híbridos lineales puros (GDN Hybrid) no lograron extrapolar eficazmente.
• Benchmarks de Contexto Largo: En RULER y LongBench, NAtS-L Hybrid alcanzó un rendimiento de vanguardia (state-of-the-art) entre los modelos comparados, particularmente en longitudes de contexto de 8k y 16k.
• Latencia de Inferencia: NAtS-L Hybrid mostró una aceleración de 5.4x en el pre-llenado (pre-filling) en comparación con los Transformers y una aceleración de 2.3x en la decodificación a un contexto de 128k, siendo solo 1.66x más lento que los modelos GDN puros.

Significancia y Reivindicaciones

El artículo afirma que NAtS-L proporciona una arquitectura híbrida flexible que asigna adaptativamente los tokens ya sea a la atención lineal o a la de softmax basándose en sus necesidades informativas específicas. Al aprender esta distribución de extremo a extremo, el modelo evita la rigidez de las estructuras híbridas fijas y la ineficiencia de la atención de softmax completa.

Los autores enfatizan que este enfoque permite al modelo:

• Preservar la información relacionada con el contexto largo mediante la atención softmax cuando es necesario.
• Reducir los costos computacionales generales mediante la atención lineal para tokens con impacto a corto plazo.
• Lograr un mejor equilibrio entre precisión y eficiencia sin requerir modelos pre-entrenados ni restricciones auxiliares sobre el ratio de los tipos de atención.

El trabajo sugiere que una mezcla de atención lineal y de softmax a nivel de token es un camino viable hacia modelos de lenguaje de contexto largo escalables que no comprometan sus capacidades de recuperación ni su velocidad de inferencia.