Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection

1. Planteamiento del Problema

Los Modelos de Lenguaje Grande (LLMs) enfrentan un cuello de botella fundamental en la inferencia de contexto largo debido a la complejidad cuadrática ($O(L^2)$) del mecanismo de atención durante la etapa de prellenado. Aunque las optimizaciones de hardware como FlashAttention reducen la sobrecarga de E/S de memoria, no alteran la complejidad algorítmica.

Los métodos de aceleración existentes sufren dos limitaciones principales:

• Dispersión Estructurada: Métodos como Minference o FlexPrefill aplican dispersión a nivel de bloques. Esto a menudo retiene tokens irrelevantes si están agrupados con tokens importantes en el mismo bloque, lo que conduce a una dispersión subóptima.
• Expulsión de Tokens: Métodos como FastKV o GemFilter descartan permanentemente tokens en las primeras capas. Esto asume que la importancia del token es estática, ignorando la naturaleza dinámica de la relevancia del token a través de diferentes capas y cabezas de atención. Una vez expulsado, un token no puede ser reconsiderado, lo que conduce a una degradación de la precisión.

2. Metodología: Token Sparse Attention

Los autores proponen Token Sparse Attention, un mecanismo de dispersión a nivel de token ligero, dinámico y reversible. A diferencia de los métodos basados en expulsión, no elimina tokens permanentemente, sino que comprime el espacio de cálculo y luego restaura la secuencia.

Arquitectura Central

El método opera en dos etapas por cabeza de atención:

1. Compresión (Selección QKV):
   • Para cada cabeza de atención $h$, se selecciona un subconjunto de índices de tokens $S_h$ basado en puntuaciones de importancia.
   • Los tensores de Consulta ($Q$), Clave ($K$) y Valor ($V$) se recopilan para formar tensores comprimidos $\hat{Q}, \hat{K}, \hat{V}$ con una longitud de secuencia reducida $L' \ll L$.
   • Característica Clave: La selección se realiza independientemente por cabeza, permitiendo que diferentes cabezas se enfoquen en diferentes subconjuntos de tokens, abordando la heterogeneidad por cabeza.
   • Los tensores comprimidos permanecen densos y contiguos, permitiendo el uso de núcleos optimizados como FlashAttention.
2. Descompresión (Restauración de Salida):
   • La salida de atención $\hat{O}$ (forma $L' \times d$) se dispersa de nuevo en un tensor inicializado con ceros de la longitud de secuencia original $L \times d$ utilizando los índices específicos de la cabeza.
   • Característica Clave: Las posiciones no seleccionadas se ponen a cero, pero la conexión residual de la capa anterior preserva la información de los tokens omitidos. Esto permite que el modelo "revisite" y reevalúe la importancia del token en capas posteriores.

Cobertura Dinámica de Tokens

Para determinar qué tokens mantener y cuántos, el artículo introduce Cobertura Dinámica de Tokens:

• Puntuación: Se calcula un mapa de atención proxy ligero utilizando consultas recientes contra todas las claves. Las puntuaciones de importancia de los tokens se derivan sumando los pesos de atención verticalmente.
• Asignación de Presupuesto: En lugar de mantener los tokens top-$k$, el método ordena los tokens por importancia ascendente y poda la "cola" (tokens menos importantes) hasta que la masa acumulada de los tokens podados excede un umbral de cobertura $\tau$. Esto trata la poda como regularización estructural contra el ruido de atención.
• Selección por Cabeza: Después de determinar el presupuesto global, cada cabeza selecciona independientemente sus tokens top-$k$ basándose en sus puntuaciones específicas.

Selección de Capas Dispersas

Los autores observan que aplicar dispersión a todas las capas degrada el rendimiento. Introducen la Deriva de Representación Inter-Capa ($R_\ell$) para medir la estabilidad de las representaciones de tokens entre capas.

• Las capas con baja deriva (representaciones estables) se seleccionan para la dispersión.
• Las capas con alta deriva (donde las representaciones cambian significativamente) se mantienen densas para preservar el comportamiento del modelo.
• Esta selección es un paso de preprocesamiento único por modelo.

3. Contribuciones Clave

1. Dispersión a Nivel de Token Reversible: Un diseño novedoso de "Comprimir y Descomprimir" que reduce el costo computacional sin expulsar tokens permanentemente, permitiendo la reevaluación dinámica de la importancia de los tokens a través de las capas.
2. Heterogeneidad por Cabeza: El método soporta la selección independiente de tokens para cada cabeza de atención, capturando el hecho de que diferentes cabezas priorizan diferentes características semánticas.
3. Compatibilidad Transparente: El enfoque es totalmente compatible con núcleos densos existentes (FlashAttention) y dispersos (Minference, FlexPrefill), actuando como un acelerador complementario en lugar de un reemplazo.
4. Presupuestación Dinámica: Un enfoque basado en datos para determinar los niveles de dispersión basados en las distribuciones de importancia de los tokens y la estabilidad por capa, en lugar de ratios fijos.

4. Resultados Experimentales

Los experimentos se realizaron en LLaMA-3.1-8B-Instruct y Mistral-Nemo-12B-Instruct utilizando puntos de referencia como RULER, InfiniteBench y Needle-in-a-Haystack.

• Compensación Precisión-Latencia:
  • Token Sparse Attention mejora consistentemente la compensación velocidad-precisión.
  • En contexto de 128K, logra hasta $\times 3.23$ de aceleración en atención comparado con FlashAttention con menos del 1% de degradación de precisión.
  • Cuando se combina con FlexPrefill, mejora la aceleración de $\times 2.44$ a $\times 2.76$ manteniendo una precisión promedio idéntica (87.27%).
• Comparación con Métodos de Expulsión:
  • Supera significativamente a los métodos de expulsión de tokens (PyramidInfer, FastKV, GemFilter). Por ejemplo, a 128K, Token Sparse Attention logra una precisión del 86.84% frente al 78.49% de PyramidInfer con aceleraciones similares.
• Escalabilidad:
  • Las ganancias de aceleración aumentan con la longitud del contexto (por ejemplo, ganancia mínima a 4K, sustancial a 128K) porque la dispersión de atención aumenta con la longitud de la secuencia.
• Sobrecarga:
  • La sobrecarga para la puntuación y la gestión de índices es mínima, representando <11% de la latencia total de atención incluso en niveles altos de dispersión.

5. Significado

Este artículo aborda una brecha crítica en la inferencia de LLMs de contexto largo alejándose de la expulsión irreversible de tokens hacia una dispersión dinámica e intercalada.

• Insight Teórico: Demuestra que la importancia del token es altamente dinámica a través de las capas y cabezas, necesitando mecanismos de selección flexibles en lugar de poda estática.
• Impacto Práctico: Proporciona una optimización "plug-and-play" que puede superponerse a núcleos de atención existentes de última generación (FlashAttention, Minference, etc.) para lograr aceleraciones masivas en escenarios de contexto largo sin reentrenar modelos ni sacrificar precisión significativa.
• Dirección Futura: Establece un nuevo punto de diseño para la inferencia escalable, sugiriendo que la compresión reversible es una estrategia superior para manejar el cuello de botella $O(L^2)$ en comparación con la eliminación permanente de tokens.