LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models

1. El Problema: La Inflación de KV (KV Inflation) en Modelos de Difusión

Los Modelos de Lenguaje de Difusión (DLMs) basados en bloques representan una alternativa prometedora a la decodificación autoregresiva tradicional, ya que pueden generar múltiples tokens en cualquier orden dentro de un bloque. Sin embargo, enfrentan un cuello de botella crítico en escenarios de contexto largo: la atención basada en memoria.

• Fallo de la Atención Esparsa Ingenua: Al adaptar métodos de atención esparsa (diseñados para LLMs autoregresivos) a DLMs, surge un problema llamado "Inflación de KV".
  • En un bloque de tamaño $B$, cada token de consulta (query) selecciona un subconjunto diferente de posiciones de claves/valores (KV) en el caché previo.
  • Dado que la atención está limitada por el ancho de banda de memoria, el costo real no depende del presupuesto por consulta, sino del tamaño de la unión de todos los índices KV seleccionados por las $B$ consultas del bloque.
  • Esta unión puede crecer hasta ser $B$ veces más grande que el presupuesto individual, anulando las ganancias de velocidad esperadas por la esparsidad y aumentando significativamente el tráfico de memoria.

2. Metodología: LoSA (Atención Esparsa Consciente de la Localidad)

Los autores proponen LoSA, un método que explota una observación clave sobre la estructura de los DLMs: la localidad de los cambios de representación entre pasos de denoising consecutivos.

Observación Fundamental

Entre dos pasos de denoising ($t-1$ y $t$):

1. Tokens Activos: Solo una pequeña fracción de tokens (aquellos que cambian de [MASK] a un token real o sufren actualizaciones significativas) experimentan cambios grandes en sus estados ocultos (queries).
2. Tokens Estables: La mayoría de los tokens mantienen estados ocultos casi constantes. Sus patrones de atención y salidas de atención no cambian significativamente.

El Algoritmo LoSA

Basado en esta observación, LoSA reduce el tráfico de memoria mediante los siguientes pasos:

1. Poda por Localidad (Locality Pruning):

   • Se calcula una puntuación de localidad para cada token en el bloque, midiendo el error cuadrático medio (MSE) entre las consultas del paso actual y el anterior.
   • Se selecciona un subconjunto pequeño de tokens con los cambios más grandes como el conjunto $A$ (tokens activos). El resto se considera tokens estables.

2. Reutilización de Caché para Tokens Estables:

   • Para los tokens estables, no se realiza ningún cálculo de atención. En su lugar, se reutilizan directamente las estadísticas de atención del prefijo (salidas $o_p$ y normalizadores logarítmicos $L_p$) guardadas en el paso anterior.
   • Esto preserva la información de atención completa (densa) para estos tokens, evitando la pérdida de precisión típica de la atención esparsa.

3. Atención Esparsa Solo para Activos:

   • Se aplica un selector de atención esparsa (como QUEST) solo a los tokens activos del conjunto $A$.
   • Se calcula la unión de los índices KV seleccionados por este subconjunto reducido ($I = \bigcup_{i \in A} I_i$).
   • Al reducir el número de consultas que participan en la selección ($|A| \ll B$), el tamaño de la unión de índices KV se reduce drásticamente, disminuyendo el tráfico de memoria.

4. Fusión y Actualización:

   • Se fusionan las contribuciones del prefijo (reutilizadas o calculadas) y del bloque actual mediante online softmax.
   • Se actualiza la caché con las nuevas estadísticas para el siguiente paso.

3. Contribuciones Clave

• Identificación del Problema de Inflación de KV: Demostraron que la aplicación directa de atención esparsa en DLMs de bloques falla debido a la unión de índices KV, lo que genera un costo de memoria mayor que el esperado.
• Descubrimiento de la Localidad de Representación: Identificaron que los cambios de representación en los DLMs son altamente localizados en el tiempo, permitiendo distinguir entre tokens activos y estables.
• Propuesta de LoSA: Un método que combina la reutilización de caché (para tokens estables) con la atención esparsa dinámica (para tokens activos), logrando reducir la densidad de atención y mejorar la precisión.
• Evaluación Exhaustiva: Validación en múltiples modelos (Trado-8B, Trado-4B, SDAR-8B) y benchmarks de contexto largo y razonamiento.

4. Resultados Experimentales

Los resultados muestran que LoSA supera consistentemente a las líneas base (QUEST, SparseD, Sparse-dLLM):

• Precisión (LongBench):
  • En el modelo Trado-8B, LoSA logra una mejora de +10.43 puntos en precisión promedio sobre QUEST con un presupuesto de recuperación bajo (128).
  • Mantiene una precisión cercana a la atención densa incluso en niveles de esparsidad agresivos, superando a los métodos que pierden información al aplicar esparsidad a todos los tokens.
  • En SDAR-8B, se observa una mejora de +8.38 puntos en el presupuesto 128.
• Eficiencia y Densidad:
  • LoSA logra una densidad de atención 1.54× menor que las líneas base en promedio, lo que significa que se cargan menos datos de KV.
  • Aceleración: En GPUs NVIDIA RTX A6000, LoSA logra una aceleración de 4.14× en la latencia de atención frente a la atención densa, superando el 3.26× de QUEST.
  • La tendencia de aceleración se mantiene en hardware más nuevo (RTX 5090), con un speedup de 3.67×.
• Generalización: Los resultados son consistentes en diferentes tamaños de modelo (4B, 8B) y tamaños de bloque (16, 32).

5. Significado e Impacto

El trabajo de LoSA es significativo por varias razones:

1. Habilitación de Contextos Largos en DLMs: Resuelve el principal obstáculo de memoria que impedía que los modelos de difusión de lenguaje fueran competitivos en tareas de contexto largo, donde la carga de KV es dominante.
2. Mejora de Precisión y Velocidad Simultánea: A diferencia de muchas técnicas de optimización que sacrifican precisión por velocidad, LoSA mejora la precisión al evitar la aplicación de esparsidad a tokens estables (manteniendo su información densa) mientras acelera el proceso al ignorar cálculos innecesarios.
3. Cambio de Paradigma en Optimización: Introduce el concepto de reutilizar salidas de atención completas basándose en la estabilidad de la representación, en lugar de simplemente seleccionar subconjuntos de KV para todos los tokens. Esto es fundamental para la arquitectura de bloques de difusión.
4. Escalabilidad: Demuestra que la eficiencia en inferencia de IA no solo depende de hardware más rápido, sino de algoritmos que comprenden la dinámica temporal de los modelos generativos.

En resumen, LoSA transforma la inferencia de modelos de difusión de bloques al mitigar la inflación de memoria mediante la inteligencia de la localidad de los cambios, ofreciendo una solución práctica y eficiente para la próxima generación de modelos de lenguaje de larga duración.