🤖 machine learning

AdaSplash-2: Faster Differentiable Sparse Attention

El artículo presenta AdaSplash-2, una implementación eficiente de atención dispersa diferenciable basada en $\alpha$ -entmax que utiliza una inicialización mediante histogramas para reducir drásticamente el costo computacional, logrando tiempos de entrenamiento comparables o superiores a FlashAttention-2 en contextos largos y mejorando el rendimiento en tareas de largo alcance.

Autores originales: Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

Publicado 2026-04-17

📖 4 min de lectura☕ Lectura para el café

CC BY 4.0

Autores originales: Nuno Gonçalves, Hugo Pitorro, Vlad Niculae, Edoardo Ponti, Lei Li, Andre Martins, Marcos Treviso

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

¡Hola! Imagina que las Inteligencias Artificiales (como los chatbots) son como bibliotecarios gigantes que leen millones de libros para responder tus preguntas.

El problema es que, cuando el libro es muy largo (un contexto largo), el bibliotecario se vuelve lento y costoso. ¿Por qué? Porque el bibliotecario tradicional (llamado "Softmax") tiene una regla estricta: debe leer y considerar cada palabra del libro, incluso las que son irrelevantes (como "el", "la", "y"), para decidir cuál es la más importante. Esto hace que, si el libro se duplica de tamaño, el trabajo se cuadruplique. ¡Es un caos!

Aquí es donde entra ADASPLASH-2. Es una nueva forma de hacer que el bibliotecario sea más rápido, más inteligente y capaz de ignorar lo que no importa.

Aquí te explico cómo funciona, usando analogías simples:

1. El Problema: El Bibliotecario Obsesivo

En el mundo de la IA, la "atención" es la capacidad de enfocarse en las partes importantes de un texto.

El método antiguo (Softmax): Es como un bibliotecario que, para encontrar la palabra clave, revisa todas las páginas del libro, una por una, y les da una puntuación a todas. Si el libro tiene 100.000 páginas, revisa 100.000 veces. Es lento y gasta mucha energía.
La solución deseada (Atención Escasa o "Sparse"): Queremos un bibliotecario que solo revise las páginas que realmente importan y descarte el resto (las que tienen puntuación cero). Pero... ¡aquí está el truco! Calcular cuáles son las páginas importantes es matemáticamente muy difícil y lento.

2. La Solución: ADASPLASH-2 (El Bibliotecario con un Mapa Rápido)

Los autores crearon ADASPLASH-2. Imagina que, en lugar de leer palabra por palabra para decidir qué ignorar, el bibliotecario hace algo muy inteligente:

A. El Histograma (El Mapa de Colores)

Imagina que el bibliotecario no lee el texto completo de inmediato. En su lugar, lanza un "proyector" rápido sobre el libro y crea un mapa de colores (un histograma) en su memoria rápida (SRAM).

Este mapa no dice qué dice cada palabra, sino cuántas palabras tienen un nivel de importancia "bajo", "medio" o "alto".
Es como si, en lugar de contar cada grano de arena de una playa, el bibliotecario hiciera una estimación rápida de cuántos granos hay en cada cubo de arena.

B. El Truco del "Ajuste Fino" (La Iteración Rápida)

Anteriormente, para encontrar la línea exacta de corte (qué palabras ignorar), el bibliotecario tenía que hacer muchas vueltas lentas (iteraciones) sobre el texto, como si estuviera adivinando y corrigiendo una y otra vez.

Con ADASPLASH-2: Gracias a ese "mapa de colores" (histograma) que hizo al vuelo, el bibliotecario tiene una aproximación casi perfecta desde el principio.
En lugar de dar 10 vueltas para acertar, ahora solo necesita 1 o 2 vueltas rápidas para ajustar el corte exacto. ¡Es como tener un GPS que ya sabe dónde estás y solo necesita un pequeño ajuste!

C. Saltar los Bloques Vacíos (El Atajo)

Una vez que el bibliotecario sabe qué partes del libro son irrelevantes (tienen puntuación cero), usa un atajo mágico.

En lugar de caminar por todo el pasillo de la biblioteca, salta directamente a las estanterías que tienen libros interesantes.
Esto se hace usando una técnica de "empaquetado de bits" (como comprimir una lista de tareas en un solo código de barras) que le permite a la computadora ignorar millones de ceros sin siquiera mirarlos.

3. ¿Por qué es tan importante esto?

Velocidad: En contextos largos (como leer un libro entero o analizar horas de video), ADASPLASH-2 es más rápido que la tecnología actual más avanzada (FlashAttention-2), especialmente cuando hay mucho "ruido" o información irredundante.
Calidad: No solo es más rápido; ¡es mejor! Al ignorar el ruido, la IA entiende mejor las relaciones complejas en textos largos. En pruebas, los modelos que usan esta tecnología entendieron mejor las historias largas y las preguntas difíciles que los modelos tradicionales.
Eficiencia: Ahorra mucha energía y dinero, lo que significa que en el futuro podríamos tener IAs que lean libros enteros en segundos en lugar de minutos.

En resumen

ADASPLASH-2 es como darle a un bibliotecario de IA un mapa rápido y un atajo. En lugar de revisar cada página de un libro gigante, hace una estimación rápida de qué páginas importan, ajusta el corte en un par de segundos y luego salta directamente a la información útil.

Esto permite que las IAs sean más rápidas, más baratas y más inteligentes cuando tienen que trabajar con textos muy largos, resolviendo uno de los mayores cuellos de botella de la tecnología actual.

Resumen Técnico: ADASPLASH-2

1. El Problema

El mecanismo de auto-atención en los Transformers es un cuello de botella computacional debido a su complejidad cuadrática ( $O(n^2)$ ) en tiempo y memoria, lo que limita la escalabilidad en secuencias de contexto largo.

Atención Softmax: Aunque implementaciones como FlashAttention-2 han optimizado la atención densa reduciendo la complejidad de memoria mediante fusión de operaciones y uso de memoria SRAM, la atención softmax es inherentemente densa (todos los tokens reciben probabilidad no nula). Esto dispersa la masa de atención en tokens irrelevantes y reduce la distinguibilidad de las representaciones en contextos largos.
Atención $\alpha$ -entmax: Es una alternativa diferenciable y dispersa (sparse) que genera ceros exactos de forma dependiente del input, mejorando la generalización en contextos largos. Sin embargo, su principal desventaja ha sido el alto costo computacional para calcular el normalizador $\tau$ (una constante de normalización similar al log-sum-exp), que requiere métodos iterativos de búsqueda de raíces. Esto ha impedido que iguale el rendimiento de las kernels de softmax optimizadas.
Limitación del trabajo previo (ADASPLASH): La implementación anterior (ADASPLASH) mejoró la situación pero aún requería múltiples pasadas costosas sobre las puntuaciones de atención para refinar $\tau$ , limitando la eficiencia, especialmente en regímenes de dispersión moderada.

2. Metodología: ADASPLASH-2

El artículo introduce ADASPLASH-2, un método consciente del hardware (hardware-aware) que reduce drásticamente el costo de la normalización en la atención $\alpha$ -entmax mediante tres pilares principales:

A. Inicialización mediante Histograma en SRAM

Idea Central: En lugar de calcular $\tau$ directamente sobre todos los puntajes, el método construye un histograma compacto de las puntuaciones de atención "al vuelo" mientras se transmiten los bloques de claves (keys) a través de la memoria SRAM on-chip.
Proceso:
1. Se normalizan y escalan las puntuaciones al intervalo $[0, 1]$ .
2. Se discretizan en $B$ bins (cajas) y se cuentan en un histograma $H$ almacenado enteramente en SRAM.
3. Este histograma permite resolver un problema de normalización reducido (de tamaño $B \ll n$ ) en lugar de sobre los $n$ puntajes originales.
Garantía Teórica: Se demuestra que la solución del histograma ( $\tau_h$ ) proporciona una cota inferior del verdadero $\tau^*$ . Esto asegura que el soporte activo (los tokens no nulos) inducido por la aproximación contenga al soporte real, evitando perder información crítica.
Resultado: La aproximación es tan precisa que el solucionador de raíces posterior converge a la solución exacta en 1 o 2 iteraciones (frente a las múltiples necesarias anteriormente).

B. Solucionador Híbrido con Refinamiento Seguro

Se utiliza un solucionador híbrido que combina métodos de orden superior (Halley, Newton) con biseción, adaptado según el valor de $\alpha$ $α$ :
- $\alpha \le 1.5$ : Pasos de Halley (estables).
- $1.5 < \alpha \le 2$ : Pasos de Newton.
- $\alpha \ge 2$ : Pasos de la secante.
Si un paso propuesto sale de los límites seguros, se retrocede a la biseción. Gracias a la excelente inicialización del histograma, este refinamiento es casi instantáneo.

C. Implementación Eficiente en GPU (Triton)

Máscara de Bloques Bit-packed: Se utiliza una codificación compacta (bit-packing) para almacenar una máscara binaria que indica qué bloques de atención son no nulos. Esto permite saltar bloques enteros de ceros con sobrecarga mínima.
Recorrido Optimizado: En la fase de salida (forward) y en la retropropagación (backward), se utilizan instrucciones nativas de GPU (find-next-set y popc) para iterar solo sobre los bloques activos, logrando una complejidad proporcional a la dispersión real ( $O(|M|)$ ) en lugar de la cuadrática.

3. Contribuciones Clave

Normalización vía Histograma On-Chip: Derivación de cotas teóricas para $\tau$ usando un histograma en SRAM, eliminando la necesidad de materializar intermediarios densos y reduciendo la precisión de error a $1/B$ .
Convergencia en Una Pasada: La combinación de la inicialización del histograma y el solucionador híbrido permite recuperar el normalizador exacto en una sola iteración adicional en la práctica.
Explotación de Dispersión Dinámica: Diseño de kernels en Triton que aprovechan la dispersión dependiente del input con sobrecarga insignificante mediante máscaras bit-packed.
Rendimiento Empírico: Resultados que muestran que ADASPLASH-2 no solo es más rápido que los kernels anteriores de $\alpha$ -entmax, sino que iguala o supera a FlashAttention-2 en regímenes de dispersión moderada a alta (>60%).

4. Resultados Experimentales

Eficiencia de Entrenamiento:
- En regímenes de dispersión moderada y alta (típicos en contextos largos), ADASPLASH-2 supera a las implementaciones de FlashAttention-2 (tanto CUDA como Triton).
- En contextos de 128K tokens, se observa una aceleración de más de 2x en comparación con FlashAttention-2, impulsada principalmente por un paso backward significativamente más rápido debido al uso de la dispersión.
- La inicialización del histograma reduce el error absoluto medio en órdenes de magnitud tras la primera iteración (ver Figura 3 del paper).
Rendimiento en Modelado de Lenguaje (LLMs):
- Se entrenaron modelos de 350M y 1B parámetros (arquitectura LLaMA-3) con $\alpha=1.5$ .
- Contexto Largo: Los modelos con $\alpha$ -entmax (especialmente combinados con la codificación posicional NAPE) superan consistentemente a las baselines de softmax en el benchmark RULER (hasta 32K tokens) y en tareas de aprendizaje en contexto (HELMET).
- Contexto Corto: Los modelos mantienen un rendimiento competitivo o superior en tareas de corto contexto (OLMES), demostrando que la dispersión no sacrifica la capacidad de modelado general.
- Conversión: Se muestra que es posible convertir un modelo preentrenado con softmax a $\alpha$ -entmax mediante Continual Pre-Training (CPT) con una pérdida mínima de rendimiento, recuperando el estado del arte en pocas iteraciones.

5. Significado e Impacto

Viabilidad de la Atención Dispersa Diferenciable: ADASPLASH-2 elimina la barrera computacional que impedía el uso práctico de la atención $\alpha$ -entmax en el entrenamiento a gran escala.
Escalabilidad: Permite entrenar modelos de lenguaje con contextos extremadamente largos (100K+ tokens) de manera eficiente, aprovechando la dispersión natural que emerge en estos escenarios, algo que la atención densa (softmax) no puede hacer sin costos prohibitivos.
Calidad del Modelo: Más allá de la velocidad, la atención dispersa mejora la capacidad de los modelos para identificar y retener información relevante en secuencias largas, mitigando problemas como la "dispersión de atención" y el colapso representacional.
Futuro: Abre la puerta a arquitecturas que combinan la eficiencia de la dispersión dinámica con la flexibilidad del aprendizaje profundo, siendo un paso crucial hacia modelos de IA más eficientes y capaces de manejar documentos completos o conversaciones largas.

En resumen, ADASPLASH-2 es un avance técnico que transforma la atención $\alpha$ -entmax de una curiosidad teórica lenta en una herramienta práctica y de alto rendimiento, capaz de competir e incluso superar a los estándares de la industria (FlashAttention-2) en escenarios de contexto largo.