💬 NLP

On Fine-Grained I/O Complexity of Attention Backward Passes

Este artículo establece límites ajustados de complejidad de E/S para los retropasses de atención en todos los tamaños de caché utilizando el marco del juego de guijarros rojo-azul, valida la optimalidad de FlashAttention en escenarios de caché grande y propone un nuevo algoritmo que logra la optimalidad teórica para entornos de caché pequeña mientras extiende estos resultados a la atención dispersa.

Autores originales: Xiaoyu Li, Yingyu Liang, Zhenmei Shi, Zhao Song, Song Yue, Jiahao Zhang

Publicado 2026-01-26

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Xiaoyu Li, Yingyu Liang, Zhenmei Shi, Zhao Song, Song Yue, Jiahao Zhang

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que eres un maestro chef (el modelo de IA) intentando cocinar un banquete masivo para una lista muy larga de invitados (el "contexto" o la secuencia de palabras). Para que el plato sea perfecto, necesitas contrastar las preferencias de cada invitado con las de todos los demás para decidir cuánto usar de cada ingrediente. Este es el mecanismo de "Atención" en los Grandes Modelos de Lenguaje.

¿El problema? A medida que la lista de invitados crece, el número de comprobaciones que necesitas realizar explota. Si tienes 1.000 invitados, haces un millón de comprobaciones. Si tienes 10.000, haces 100 millones. Este es el cuello de botella de la "escalabilidad cuadrática" mencionado en el artículo.

Ahora, imagina que tu cocina tiene dos tipos de almacenamiento:

La Encimera (Caché): Un espacio pequeño, rápido y costoso justo al lado de la estufa donde puedes agarrar los ingredientes instantáneamente.
La Despensa (Memoria): Una sala de almacenamiento enorme, lenta y profunda donde se guardan todos tus ingredientes.

Cada vez que tienes que caminar desde la despensa hasta la encimera para agarrar un ingrediente, te cuesta tiempo y energía. Este ir y venir es lo que los científicos de la computación llaman Complejidad de I/O (Entrada/Salida). El objetivo es minimizar estos viajes.

El Problema Principal: El "Paso hacia Atrás" (Backward Pass)

Cuando el chef está aprendiendo (entrenando), no solo cocina el plato; también necesita averiguar qué salió mal para poder ajustar la receta la próxima vez. Esto se llama Paso hacia Atrás (Backward Pass).

Durante mucho tiempo, el estándar de la industria para cocinar de manera eficiente fue un método llamado FlashAttention. Fue brillante organizando los viajes a la despensa para el paso hacia adelante (cocinar el plato). Pero los autores de este artículo se preguntaron: "¿Es FlashAttention también la forma más eficiente de organizar los viajes a la despensa para el paso hacia atrás (aprender de los errores), especialmente cuando nuestra encimera es pequeña?"

El Descubrimiento: Depende del Tamaño de la Encimera

Los autores se dieron cuenta de que la respuesta depende enteramente de qué tan grande es tu encimera (Caché) en comparación con el tamaño de tu receta (la dimensión oculta, $d$ ). Encontraron un "punto de inflexión" en un tamaño específico ( $d^2$ ).

1. El Escenario de la "Encimera Grande" ( $M \ge d^2$ )

Si tu encimera es lo suficientemente grande como para contener una parte significativa de tus ingredientes a la vez, FlashAttention es perfecto.

La Analogía: Tienes una isla enorme en tu cocina. Puedes disponer todos los ingredientes que necesitas para una sección completa de la receta allí mismo. Cocinas, aprendes y limpias sin necesidad de correr de vuelta a la despensa.
El Resultado: El artículo demuestra matemáticamente que FlashAttention no puede ser superado aquí. Es el método más eficiente posible tanto para cocinar como para aprender.

2. El Escenario de la "Encimera Pequeña" ( $M < d^2$ )

Si tu encimera es diminuta (como en computadoras más viejas o baratas), FlashAttention empieza a tambalearse. Intenta usar una estrategia que funciona para encimeras grandes, lo que le obliga a realizar viajes innecesarios a la despensa.

La Analogía: Imagina intentar cocinar un estofado complejo en una encimera pequeña. FlashAttention sigue trayendo ollas enormes de ingredientes, solo para darse cuenta de que la encimera es demasiado pequeña, por lo que tiene que devolverlas a la despensa y traer lotes más pequeños. Es ineficiente.
La Solución: Los autores inventaron un nuevo algoritmo (Algoritmo 6). En lugar de traer grandes trozos, este nuevo método descompone la receta en piezas diminutas y manejables que encajan perfectamente en la pequeña encimera. Lee y escribe los datos de una manera que coincide exactamente con el tamaño de la encimera.
El Resultado: Este nuevo método es estrictamente mejor que FlashAttention para encimeras pequeñas. Demuestra que FlashAttention no es la mejor opción cuando la memoria es escasa, y los autores encontraron el "límite de velocidad" teórico de qué tan rápido se puede hacer esto.

El Giro "Esparso" (Sparse)

El artículo también analizó una variación llamada Atención Dispersa (Sparse Attention).

La Analogía: Imagina que para la mayoría de los invitados, en realidad no necesitas contrastar sus preferencias contra todo el mundo. Tal vez solo necesites contrastarlas contra sus vecinos. Estos son datos "dispersos".
El Resultado: Los autores crearon un nuevo conjunto de reglas (límites inferiores) para cuántos viajes a la despensa son inevitables incluso con estos datos dispersos. Mostraron que el punto de inflexión entre una "encimera pequeña" y una "encimera grande" cambia según cuántos ingredientes tengas que mover realmente, pero la lógica sigue siendo la misma.

Resumen de las Afirmaciones del Artículo

FlashAttention es un héroo para cocinas grandes: Cuando tienes mucha memoria rápida (caché), FlashAttention es la mejor forma de manejar la fase de "aprendizaje" (paso hacia atrás). No puedes hacerlo mejor.
FlashAttention es superado en cocinas pequeñas: Cuando tienes muy poca memoria rápida, FlashAttention es ineficiente. Los autores diseñaron un nuevo algoritmo especializado que es demostrablemente más rápido y alcanza el límite teórico de eficiencia para estos espacios pequeños.
Ahora tenemos el mapa completo: Antes de este artículo, conocíamos los límites para "cocinar" (paso hacia adelante) y teníamos una suposición para el "aprendizaje" (paso hacia atrás) en cocinas grandes. Este artículo completa las piezas faltantes, dándonos los límites matemáticos exactos para tanto cocinar como aprender en cualquier tamaño de cocina, ya sea que los datos sean densos (completos) o dispersos (vacíos).

En resumen, el artículo nos dice: "Si tienes una cocina grande, quédate con FlashAttention. Si tienes una cocina pequeña, cambia a nuestro nuevo método para ahorrar tiempo y energía".

Resumen Técnico: Sobre la Complejidad de E/S de Grano Fino en los Pasos Atrás de la Atención

Planteamiento del Problema

Los Grandes Modelos de Lenguaje (LLMs) dependen fuertemente del mecanismo de autoatención, el cual exhibe una complejidad computacional cuadrática ( $O(n^2)$ ) con respecto a la longitud de la secuencia $n$ . Si bien FlashAttention se ha convertido en el estándar de la industria para optimizar el paso hacia adelante (forward pass) aprovechando la conciencia de E/S (minimizando las transferencias de datos entre la caché rápida en chip y la memoria lenta de alto ancho de banda), la comprensión teórica del paso atrás (backward pass) sigue incompleta.

La literatura existente proporciona límites superiores para la complejidad de E/S del paso atrás de FlashAttention para tamaños de caché grandes ( $M \in [d, nd]$ ), pero carece de límites inferiores coincidentes. En consecuencia, se desconoce si FlashAttention es óptimo para las computaciones del paso atrás a través de todo el espectro de tamaños de caché (especialmente en escenarios de caché pequeña, $M < d^2$ , o para mecanismos de atención dispersa). Este artículo aborda la pregunta abierta: ¿Cuál es la complejidad de E/S óptima de las computaciones de la atención en el paso atrás para cualquier tamaño de caché?

Metodología

Los autores emplean un marco teórico rigurooso basado en el juego de los guijarros rojo-azul [HK81] para modelar la jerarquía de memoria de dos niveles (caché rápida de tamaño $M$ y memoria lenta). Este marco trata la computación como un grafo acíclico dirigido (DAG) donde los guijarros rojos representan los datos en la caché y los guijarros azules representan los datos en la memoria. La complejidad de E/S se define como el número mínimo de operaciones de entrada/salida requeridas para computar la salida dado un límite en el número de guijarros rojos (tamaño de la caché).

Pasos metodológicos clave incluyen:

Análisis del Grafo Computacional: Los autores construyen un grafo computacional detallado para el paso atrás de la atención, modelando explícitamente la computación del gradiente para los pesos de atención ( $X$ e $Y$ ). Esto implica analizar los árboles de suma inherentes a la multiplicación de matrices estándar (específicamente $A_1 X A_2^\top$ ).
Diferenciación de Casos: El análisis divide el problema en dos regímenes basados en la relación entre el tamaño de la caché $M$ $M$ y el cuadrado de la dimensión de la cabeza $d^2$ $d^{2}$ :
- Caché Pequeña: $M = o(d^2)$
- Caché Grande: $M = \Omega(d^2)$
Diseño de Algoritmos: Para el régimen de caché pequeña, los autores proponen un nuevo algoritmo (Algoritmo 6) que lee y escribe explícitamente matrices de atención de $n \times n$ y matrices intermedias de $n \times d$ , divergiendo de la estrategia por bloques de FlashAttention.
Extensión a la Atención Dispersa: El marco se extiende a la atención dispersa, donde las matrices se almacenan por entradas no nulas, estableciendo límites inferiores basados en el número de elementos no nulos ( $Z_{input}$ , $Z_{QK}$ ).

Contribuciones Clave y Resultados

1. Límites de Complejidad de E/S Ajustados para la Atención Densa

El artículo establece límites superiores e inferiores coincidentes para la complejidad de E/S de los pasos atrás de la atención para todos los tamaños de caché $M$ . El resultado principal (Teorema 1.1) establece que la complejidad de E/S es:
$\Theta\left(\min\left\{ \frac{n^2d^2 + nd^3}{M}, \frac{n^2d + nd^2}{\sqrt{M}} \right\}\right)$

Este resultado revela un punto de transición crítico en $M = \Theta(d^2)$ :

Caché Grande ( $M = \Omega(d^2)$ ): La complejidad es $\Theta(\frac{n^2d^2 + nd^3}{M})$ . Los autores demuestran que FlashAttention es óptimo en este régimen, ya que su límite superior existente coincide con el límite inferior derivado.
Caché Pequeña ( $M = o(d^2)$ ): La complejidad es $\Theta(\frac{n^2d + nd^2}{\sqrt{M}})$ . Los autores demuestran que FlashAttention no es óptimo en este régimen. Proponen un nuevo algoritmo (Algoritmo 6) que logra el límite más ajustado $O(\frac{n^2d + nd^2}{\sqrt{M}})$ , superando el límite de FlashAttention de $O(\frac{n^2d^2 + nd^3}{M})$ .

2. Límites Inferiores de la Atención Dispersa

El artículo proporciona límites inferiores finos para los pasos hacia adelante y hacia atrás de la atención dispersa (Teorema 4.5). Para matrices dispersas con $Z_{input}$ entradas no nulas en la entrada y $Z_{QK}$ en la matriz clave-consulta, la complejidad de E/S está acotada por:
$\Omega\left(\min\left\{ \frac{Z_{input}^2}{M}, \frac{Z_{input}\sqrt{Z_{QK}}}{\sqrt{M}} \right\}\right)$
Esto generaliza los resultados del caso denso y confirma que el punto de división entre el comportamiento de caché grande y pequeña permanece consistente ( $M = Z_{input}^2 / Z_{QK}$ ).

3. Completitud Teórica

Al combinar estos resultados del paso atrás con el análisis del paso hacia adelante existente [SY24], el artículo completa el panorama teórico de la complejidad de E/S para el mecanismo de atención. Confirma que, si bien FlashAttention es óptimo para cachés grandes, se requieren estrategias algorítmicas diferentes para cachés pequeñas para lograr la optimalidad teórica.

Significancia

El artículo afirma consolidar el marco teórico relativo a la complejidad de E/S en los mecanismos de atención. Su importancia radica en:

Completar la Teoría: Proporcionar los primeros límites superiores e inferiores coincidentes para los pasos atrás de la atención en todos los tamaños de caché, resolviendo un vacío en la literatura existente.
Guía Algorítmica: Identificar que los estándares actuales de la industria (FlashAttention) son subóptimos para entornos de caché pequeña (por ejemplo, GPUs más antiguas o configuraciones de alta dimensión donde $d$ es grande en relación con $M$ ). El propuesto Algoritmo 6 ofrece una alternativa teóricamente superior para estas restricciones de hardware específicas.
Fundamentos de la Atención Dispersa: Ofrecer los primeros límites inferiores teóricos para la complejidad de E/S de la atención dispersa, lo cual es crucial para comprender los límites de las técnicas de dispersión en el modelado de contextos largos.

Los autores concluyen que sus hallazgos proporcionan una guía crítica para el desarrollo de sistemas eficientes de entrenamiento e inferencia de LLM, particularmente a medida que las restricciones de hardware y las dimensiones de los modelos evolucionan.