Resumen Técnico: SpotAttention

Planteamiento del Problema

Las capacidades de contexto largo se han convertido en un estándar para los modelos de lenguaje extensos preentrenados (LLMs), permitiendo tareas como el procesamiento de bases de código completas o la fundamentación de respuestas en evidencia extensa. Sin embargo, el costo computacional de la autoatención escala pobremente con la longitud de la secuencia. Durante la fase de prellenado (prefill), el cómputo crece cuadráticamente con la longitud de la secuencia, y durante la fase de decodificación, cada nuevo token requiere volver a leer un caché de Clave-Valor (KV) que crece linealmente con el tamaño del contexto. Aunque FlashAttention mejoró la eficiencia de I/O al evitar la materialización de la matriz de atención completa, no alteró el costo fundamental cuadrático del prellenado ni el crecimiento lineal de las lecturas de memoria por paso.

Los métodos de atención dispersa (sparse attention) existentes intentan mitigar esto atendiendo solo a un subconjunto relevante de tokens pasados. Sin embargo, la selección de este subconjunto suele ser costosa. Métodos como Quest califican bloques de tokens, y Twilight añade la poda top-$p$, pero el selector en sí mismo se convierte en un nuevo cuello de botella. Además, los selectores sin entrenamiento dependen de heurísticas diseñadas a mano que pueden no generalizar bien o requieren múltiples etapas de selección, incurriendo en una sobrecarga acumulativa.

Metodología: SpotAttention

SpotAttention es un mecanismo de enrutamiento de dispersión por bloques (block-sparse) tipo "plug-in", diseñado para integrarse en transformadores ya entrenados y congelados. Consiste en un "selector" (o indexador) ligero acoplado a cada capa de atención completa.

1. Arquitectura del Selector

El selector es un calificador Q-K multicabeza diminuto con una anchura reducida ($d_{idx}$) y $H_{idx}$ cabezas. En lugar de calificar tokens individuales, opera sobre bloques de consultas (queries) y claves (keys) (por ejemplo, 16 tokens por bloque).

• Mecanismo de Calificación: Para una consulta en la posición $q$ y un bloque de claves en la posición $t$, el selector calcula un puntaje ReLU ponderado por cabezas:
  $$s_s[q, t] = \sum_{h=1}^{H_{idx}} w_h \text{ReLU}(\mathbf{q}^I_h \cdot \mathbf{k}^I_{h,t})$$
  donde $\mathbf{q}^I_h$ y $\mathbf{k}^I_{h,t}$ son proyecciones por cabeza, y $w_h$ son pesos de mezcla derivados del estado oculto.
• Eficiencia: Al calificar por bloques, el proceso de selección se ejecuta como una única multiplicación de matrices en núcleos tensoriales (tensor-core), evitando la sobrecarga de las operaciones a nivel de token.

2. Entrenamiento mediante Destilación KL

El selector se entrena de forma independiente contra la distribución de atención densa del backbone congelado (el "maestro").

• Objetivo: El entrenamiento minimiza la divergencia de Kullback-Leibler (KL) entre la distribución de atención del maestro ($p_\tau$) y la distribución estimada del selector ($p_\sigma$).
• Variantes: El artículo explora dos variantes de KL:
  • DenseKL: Calcula la pérdida sobre todas las $T$ claves.
  • SparseKL: Calcula la pérdida solo sobre el propio conjunto seleccionado top-$K$ del selector, con distribuciones renormalizadas.
• Convergencia: El selector converge rápidamente (en aproximadamente 200 pasos) a través de varios tamaños de backbone (4B a 32B) y arquitecturas, aprendiendo a igualar la forma de la distribución del maestro.

3. Regla de Selección Dual Top-$p$

En la inferencia, el selector califica todos los bloques pasados. Para determinar a qué bloques atender, SpotAttention introduce una regla dual top-$p$ que lee el presupuesto directamente de la distribución calibrada del selector:

1. Bloques Reservados: Un prefijo fijo de bloques "sink" (ej. los primeros 128 tokens) y un sufijo fijo de bloques de "recencia" (ej. los últimos 256 tokens) se mantienen siempre para evitar que sean dominados por la selección de núcleo (nucleus selection).
2. Selección de Núcleo (Nucleus Selection): Se aplica el muestreo top-$p$ (núcleo) a la distribución residual (bloques excluyendo sinks y recencia). La selección se detiene cuando la masa de probabilidad acumulada alcanza un umbral $p$.
3. Presupuesto Dinámico: Este enfoque produce un presupuesto variable por consulta y por capa, adaptándose al contenido sin una etapa de poda separada. Un piso mínimo ($K_{min}$) evita casos degenerados donde el núcleo sea demasiado pequeño.

4. Cuantización

El caché K del selector puede cuantizarse a microescala INT4 o FP4. Dado que el caché del selector es distinto del caché KV del backbone, esto reduce la huella de memoria del selector aproximadamente 3.5$\times$ sin afectar la precisión o exactitud del backbone.

Contribuciones Clave

1. Atención Dispersa por Bloques Plug-In: Un selector ligero entrenado mediante destilación KL que se acopla a transformadores congelados. Se ejecuta como una única multiplicación de matrices en tensor-core, evitando la pesada sobrecarga de selectores previos.
2. Regla Dual Top-$p$: Un mecanismo de selección adaptativo al contenido que deriva el presupuesto por consulta y por capa directamente de la distribución aprendida. Reserva bloques de sink y recencia para asegurar la estabilidad, igualando la precisión de la top-$K$ estática con un presupuesto medio menor.
3. Comportamiento de Entrenamiento e Inferencia: El artículo demuestra que el selector aprende la asimetría del presupuesto por capa del backbone (las capas tempranas atienden a más contexto que las posteriores), un patrón que la $K$ estática no puede replicar. El método generaliza de una longitud de entrenamiento de 16K a una de evaluación de 128K sin degradación de la exactitud.

Resultados

El método fue evaluado en cinco backbones (Qwen3: 4B–32B; Qwen3.5: 4B–9B) a través de cuatro benchmarks (RULER, BABILong, InfiniteBench, LongBench-v2) con longitudes de contexto de hasta 128K.

• Exactitud: SpotAttention iguala la exactitud del backbone denso en longitudes de contexto de hasta 128K (8$\times$ la longitud de entrenamiento). Con un presupuesto de $K=0.5L$, se mantiene dentro de un error estándar de bootstrap del baseline denso en todos los modelos.
• Latencia y Rendimiento (Throughput):
  • Con $L=128K$ en Qwen3-8B, SpotAttention es 3.9$\times$ más rápido que FlashAttention y 1.8$\times$ más rápido que Twilight (el baseline sin entrenamiento más fuerte).
  • El rendimiento de decodificación se mantiene en 70 tokens/s a 128K y 36 tokens/s a 512K, mientras que FlashAttention y Twilight sufren errores de memoria (OOM) a 512K y 256K respectivamente.
  • El propio selector es significativamente más rápido que el selector de Twilight (1.55 ms vs 4.59 ms a 128K).
• Cuantización: La cuantización del caché K del selector a INT4 o FP4 reduce su huella de memoria en 3.5$\times$ sin pérdida medible de exactitud. Aunque esto no mejora la latencia de decodificación de flujo único (dominada por el ancho de banda del KV del backbone), reduce significamente el uso de memoria y los costos de prellenado.

Significancia y Reivindicaciones

El artículo posiciona a SpotAttention como una solución al "cuello de botella del selector" en la atención dispersa. Al aprender una distribución calibrada mediante la destilación KL, el selector permite un presupuesto basado en el núcleo y dinámico que se adapta a las necesidades específicas de cada consulta y capa.

Los autores afirman que:

• Eficiencia: Logra el mayor rendimiento entre los métodos comparados en contextos largos (hasta 512K), superando las limitaciones de memoria que detienen a los baselines densos y otros métodos dispersos.
• Generalización: Un selector entrenado con un contexto de 16K se generaliza eficazmente a 128K y más allá sin necesidad de reentrenamiento.
• Modularidad: Es una solución "plug-in" que no requiere reentrenar el backbone, lo que la hace aplicable a modelos preentrenados existentes.
• Limitaciones: Los autores señalan que el backbone más grande evaluado es de 32B, y los resultados están actualmente limitados a la familia Qwen en benchmarks de inglés. También reconocen que la latencia de prellenado disperso no fue evaluada, aunque la arquitectura lo soporta, y que la regla dual top-$p$ depende de que el selector esté bien calibrado (lo cual se cumple bajo el entrenamiento SparseKL, pero colapsa bajo DenseKL).