Attention Sinks and Outliers in Attention Residuals

Resumen Técnico: OASIS – Enrutamiento Consciente de Outliers y Sinks para Arquitecturas AttnResidual

Declaración del Problema
El artículo aborda problemas críticos de estabilidad y robustez que surgen en las arquitecturas AttnResidual, un diseño reciente que reemplaza la normalización estándar de Softmax de un solo nivel con un esquema de doble normalización que gobierna tanto el enrutamiento de atención a nivel de token como a nivel de profundidad. Aunque AttnResidual mejora la expresividad representacional, los autores identifican un modo de fallo estructural: el diseño de doble normalización exacerba los sinks de atención y los outliers de activación.

En los Transformers estándar, los sinks de atención (por ejemplo, el token <|begin_of_text|>) acumulan masa de probabilidad para estabilizar las representaciones. Sin embargo, en AttnResidual, la restricción adicional del simplex en el enrutamiento de profundidad obliga a que las actualizaciones de "no-op" (identidad) se expresen a través de tokens reales y ramas de profundidad reales, en lugar de un canal nulo explícito. Esta restricción requiere logits pre-Softmax extremos para aproximar actualizaciones de identidad, lo que conduce a:

1. Sinks de Atención Intensificados: Mayor concentración en los tokens iniciales.
2. Outliers de Activación Amplificados: Los estados ocultos exhiben distribuciones de colas pesadas (alta curtosis) y normas infinito grandes.
3. Fragilidad de Cuantización: Estos outliers degradan la estabilidad de la inferencia y hacen que la cuantización de bajo bit (por ejemplo, W4A4, W8A8) sea significativamente menos robusta.
4. Colapso de Profundidad: Los pesos de enrutamiento colapsan hacia una única rama dominante, revertiendo efectivamente la arquitectura a un comportamiento residual tipo vanilla.

Metodología: OASIS
Los autores proponen OASIS (Enrutamiento Consciente de Outliers y Sinks mediante Señalización Nula Intercapa), una intervención ligera que modifica la capa de normalización sin alterar el objetivo de entrenamiento ni el esqueleto arquitectónico central. OASIS opera en dos frentes:

1. Enrutamiento Consciente de Nulo (Softmax1):
   El método reemplaza los operadores Softmax estándar tanto a nivel de token como a nivel de profundidad con Softmax1. Esta función introduce un estado nulo explícito (denotado como $\emptyset$) en el denominador de la normalización.

   • Nivel de Token: Permite que la masa de atención se enrute a un estado "nulo" en lugar de forzarla sobre tokens reales.
   • Nivel de Profundidad: Permite que los pesos de enrutamiento de profundidad absorban la presión de "no-op" a través de una capa nula, evitando la concentración de masa en ramas de profundidad específicas.
   • Mecanismo: Al desacoplar la masa no informativa de la competencia entre tokens/ramas reales, Softmax1 reduce la necesidad de amplificación extrema de logits.

2. Acoplamiento Token-a-Profundidad Nulo:
   Reconociendo que el enrutamiento a nivel de token y a nivel de profundidad permanece desacoplado en Softmax1 estándar, OASIS introduce un mecanismo de acoplamiento. Agrega la evidencia nula a nivel de token (la masa de probabilidad enrutada al estado nulo a través de las cabezas de atención) en una estadística escalar para cada rama fuente. Esta estadística se inyecta luego en los logits de enrutamiento de profundidad.

   • Efecto: Las ramas que exhiben un comportamiento alto de "no-op" (alta masa nula) se ponderan a la baja en el enrutamiento de profundidad, mientras que las ramas informativas se ponderan al alza. Esto evita que el enrutador de profundidad asigne un peso significativo a capas no informativas.

Contribuciones Clave

• Identificación de Fallo Estructural: El artículo identifica formalmente que el diseño de doble normalización de AttnResidual amplifica estructuralmente los outliers y los sinks debido a la falta de una vía nula explícita para actualizaciones cercanas a la identidad.
• Análisis Teórico: Los autores proporcionan pruebas teóricas (Lemas 5.1, 5.2; Teoremas 5.1, 5.2) que demuestran que:
  • Sin un canal nulo, las restricciones cercanas a "no-op" fuerzan la masa de atención sobre tokens de baja contribución, creando outliers.
  • La doble normalización conduce a un colapso de profundidad donde el enrutamiento se concentra en ramas de actualización mínima.
  • Softmax1 introduce una ruta nula viable que desacopla el comportamiento de no-op de la asignación de tokens reales, reduciendo teóricamente la presión estructural.
• Marco OASIS: Una modificación práctica y mínima (reemplazar Softmax con Softmax1 y añadir un término de acoplamiento) que mitiga estas patologías.

Resultados Experimentales
Los autores evaluaron OASIS en las arquitecturas base Llama-3.2-1B y Qwen3-0.6B a través de tres conjuntos de datos del mundo real (BookCorpus, Wiki40B, WikiText-2, GSM8K).

• Supresión de Outliers: En relación con la línea base Vanilla AttnResidual, OASIS logró una reducción promedio del 84.14% en la norma infinito máxima y del 96.77% en la curtosis promedio.
• Robustez de Cuantización:
  • Bajo cuantización W8A8, OASIS redujo la perplejidad en un 73.55% en comparación con Vanilla.
  • Bajo cuantización W4A4, OASIS mejoró la precisión Pass@1 en GSM8K en un 23.47%.
• Mitigación de Sinks de Atención: El análisis cualitativo de los mapas de atención mostró que OASIS redistribuyó eficazmente la atención alejándola del sink dominante del primer token, generando patrones de atención causal más suaves y locales en comparación tanto con Vanilla como con líneas base intermedias (como AoS o OutEffHop).
• Comparación: OASIS superó consistentemente a las líneas base, incluyendo Clip Attention, Gated Attention, OutEffHop y normalizadores dispersos (Sparsemax, Entmax15), tanto en métricas de outliers como en el rendimiento de tareas posteriores bajo cuantización agresiva.

Significado y Afirmaciones
El artículo afirma que OASIS proporciona una solución fundamentada a las limitaciones de robustez de las arquitecturas doble-normalizadas. Al introducir capacidad nula explícita, el método estabiliza el enrutamiento residual, permitiendo que los modelos manejen la compresión agresiva (cuantización de bajo bit) y la inferencia de contexto largo sin la degradación típicamente causada por los sinks de atención y los outliers de activación. Los autores enfatizan que estas mejoras se logran sin cambiar el objetivo de entrenamiento, lo que sugiere que el enrutamiento nulo explícito es un mecanismo fundamental para estabilizar arquitecturas residuales complejas.

Limitaciones
Los autores señalan que su evaluación está actualmente limitada a modelos base específicos (1B y 0.6B parámetros) y que se requiere una validación más amplia en modelos más grandes. También advierten que una asignación nula excesiva podría potencialmente suprimir señales débiles útiles o alterar la interpretabilidad de los mapas de atención, sugiriendo trabajo futuro sobre el control adaptativo de la fuerza nula.