Dense Local Dependencies Induce Attention-Logit Explosion and Training Instability During Long-Sequence Transformer Training

Resumen Técnico: Las dependencias locales densas inducen la explosión de los logits de atención y la inestabilidad en el entrenamiento

Planteamiento del problema
Los modelos de lenguaje transformadores autorregresivos presentan frecuentemente inestabilidad en el entrenamiento cuando se entrenan con secuencias largas, un fenómeno que es particularmente agudo bajo aritmética de baja precisión (p. ej., bfloat16). Esta inestabilidad suele manifestarse como picos repentinos o divergencia en la pérdida de entrenamiento, frecuentemente acompañados de una "explosión de los logits de atención": un crecimiento rápido e ilimitado en la magnitud de los logits de atención. Mientras que trabajos previos han atribuido estos problemas a la amplificación de las conexiones residuales, interacciones del optimizador o desafíos generales de secuencias largas, la causa subyacente específica de por qué el aumento de la longitud de la secuencia promueve el crecimiento de los logits de atención permanece insuficientemente comprendida. Los remedios existentes, como la normalización QK o los optimizadores especializados, a menudo abordan los síntomas en lugar de la causa raíz.

Metodología y marco analítico
Los autores proponen que las dependencias locales densas son el motor principal de esta inestabilidad. Su análisis se basa en el cuello de botella de bajo rango inherente a los mecanismos estándar de autoatención completa (full self-attention):

1. Base teórica: En un transformador estándar, la matriz de logits de atención $S = QK^T$ tiene un rango restringido por la dimensión de atención $d$. Sin embargo, las tareas que involucran dependencias locales densas (donde un token depende de muchos tokens cercanos) requieren un patrón de atención que es efectivamente de alto rango, con un rango que crece linealmente con la longitud de la secuencia $n$.
2. El desajuste: Cuando $n \gg d$, el modelo enfrenta un desajuste representacional. Para aproximar un patrón de atención objetivo de alto rango utilizando una factorización de bajo rango ($QK^T$), el proceso de optimización se ve forzado a inflar la magnitud de los logits de atención.
3. Validación experimental: Los autores validan esta hipótesis a través de dos vías principales:
   • Tareas sintéticas: Se construyó una tarea controlada donde la matriz de atención objetivo contenía dependencias locales densas (una matriz de banda) pero sin estructura de largo alcance. El modelo tenía la tarea de reconstruir esta matriz.
   • Modelado de lenguaje: Se realizaron experimentos extensos utilizando arquitecturas GPT-2 en el conjunto de datos PG-19. Los autores variaron las longitudes de secuencia, las dimensiones de las cabezas de atención y la arquitectura del mecanismo de atención (introduciendo "Atención Local Completa", que mezcla cabezas de atención completa y de ventana deslizante). También manipularon la densidad de las dependencias locales mediante enmascaramiento.

Resultos clave
La evidencia empírica respalda consistentemente el mecanismo propuesto:

• Longitud de secuencia vs. Crecimiento de los logits: Tanto en las tareas sintéticas como en las de modelado de lenguaje, el logit de atención máximo requerido para lograr una precisión de reconstrucción fija aumenta rápidamente a medida que la longitud de la secuencia $n$ crece, siempre que la dimensión de atención $d$ permanezca fija.
• Escalamiento de la dimensión: Aumentar la dimensión de la cabeza de atención $d$ mitiga sustancialmente el crecimiento de los logits, confirmando que el cuello de botella está relacionado con la restricción de rango respecto a la longitud de la secuencia.
• Eficacia de la atención local: Reemplazar un subconjunto de cabezas de atención completa con cabezas de atención local de ventana deslizante (que modelan naturalmente las dependencias locales densas sin el cuello de botella de bajo rango) reduce drásticamente el crecimiento de los logits. Esta reducción ocurre incluso en las cabezas de atención completa restantes, lo que sugiere que las cabezas locales descargan la carga representacional.
• Densidad vs. Localidad: La inestabilidad es impulsada por la densidad de las dependencias locales, no simplemente por su existencia. Cuando las dependencias locales fueron esparcidas (mediante el enmascaramiento del 90% de las conexiones locales), el crecimiento de los logits se suprimió y el entrenamiento se estabilizó.
• Sensibilidad a la precisión: Aunque la explosión de los logits ocurre tanto en el entrenamiento FP32 como en el de baja precisión (BF16), la inestabilidad de entrenamiento resultante (divergencia de la pérdida) se observa principalmente en configuraciones de baja precisión, donde los errores numéricos amplifican los efectos de los logits grandes.
• Escala del modelo: La gravedad de la inestabilidad aumenta con el tamaño del modelo. Los modelos más grandes (p. ej., GPT-2 Large) exhiben un crecimiento de los logits e inestabilidad más pronunciados que los modelos más pequeños (GPT-2 Small) bajo las mismas condiciones, requiriendo una mayor capacidad de atención local para permanecer estables.

Significado y afirmaciones
El artículo afirma identificar un principio de diseño fundamental para transformadores de contexto largo estables, eficientes y escalables: modelar explícitamente las dependencias locales densas.

• Identificación de la causa raíz: El trabajo establece que el cuello de botella de bajo rango de la autoatención completa, cuando se le obliga a representar estructuras locales densas, es la causa raíz de la explosión de los logits de atención.
• Principio de diseño: En lugar de tratar la explosión de los logits como un artefacto numérico que debe normalizarse, los autores argumentan que las arquitecturas deberían asignar capacidad dedicada (p. ej., cabezas de atención local) para manejar las interacciones locales densas. Este enfoque alivia naturalmente el cuello de botella de bajo rango.
• Eficiencia y estabilidad: Al descargar las dependencias locales hacia las cabezas de atención local, los modelos no solo logran una mayor estabilidad en el entrenamiento, sino que también obtienen una mayor eficiencia computacional (reduciendo la complejidad de $O(n^2)$ a $O(nw)$) y menores requisitos de memoria del caché KV durante la inferencia.
• Implicaciones más amplias: Los hallazgos sugieren que el éxito de varias arquitecturas con sesgo de localidad (como aquellas que utilizan RoPE, ALiBi o atención estructurada) puede atribuirse a su capacidad para modelar, implícita o explícitamente, estas dependencias locales densas, evitando así el cuello de botella de bajo rango que afecta a los modelos de pura atención completa.