Resumen Técnico: PISA - Atención Dispersa por Tramos para Transformers de Difusión Eficientes

1. Planteamiento del Problema

Los Transformers de Difusión (DiTs) se han convertido en el estándar para la generación de imágenes y vídeos de alta fidelidad. Sin embargo, su escalabilidad se ve limitada por la complejidad cuadrática ($O(L^2)$) del mecanismo de auto-atención a medida que las longitudes de la secuencia de entrada ($L$) aumentan para resoluciones más altas o vídeos más largos.

Las soluciones existentes enfrentan limitaciones significativas:

• Atención Dispersa ("Mantener-o-Descartar"): Los métodos que descartan bloques de clave-valor (KV) no críticos aceleran la computación pero sufren una degradación del rendimiento en ratios de dispersión altos debido a la pérdida de contexto. También presentan dificultades con secuencias cortas donde la sobrecarga de selección supera los beneficios.
• Atención Lineal/Híbrida: Los enfoques que alteran la distribución de la atención (por ejemplo, mediante mapeos de características o combinaciones aditivas de ramas dispersas y lineales) suelen romper la normalización intrínseca de los modelos pre-entrenados. Esto requiere un re-entrenamiento o ajuste fino (fine-tuning) costoso, impidiendo el reúso directo de los pesos pre-entrenados.
• Truncamiento Duro: El descarte directo de bloques provoca desviaciones bruscas de la distribución de atención original, causando una pérdida de calidad en el contenido generado.

El desafío central es diseñar un mecanismo que acelere la inferencia con una complejidad sub-cuadrática manteniendo la precisión de la atención completa y siendo compatible con los pesos pre-entrenados sin necesidad de re-entrenamiento.

2. Metodología: Atención Dispersa por Tramos (PISA)

PISA introduce un nuevo paradigma de "exacto-o-aproximado", yendo más allá de la estrategia tradicional de "mantener-o-descartar". Trata la atención como un proceso por tramos que cubre todo el espectro de la atención con una complejidad sub-cuadrática.

Ideas Clave

Los autores analizaron DiTs pre-entrenados y descubrieron dos propiedades críticas:

1. Estabilidad Distribucional: Las puntuaciones de atención pre-softmax de los bloques no críticos exhiben una distribución simétrica en forma de campana centrada en cero o en valores negativos. Esto las hace altamente aptas para una aproximación precisa mediante la expansión de Taylor, en lugar de requerir un cómputo exacto o su descarte.
2. Consistencia de Normalización: Para mantener la compatibilidad sin re-entrenamiento, las aproximaciones deben integrarse internamente en el numerador y denominador del softmax, en lugar de simplemente sumar salidas de diferentes ramas (lo cual viola las reglas de suma ponderada).

Algoritmo Central

PISA particiona los bloques KV en dos conjuntos:

1. Bloques Críticos ($S_i$): Computados de forma exacta para preservar información de alta relevancia.
2. Bloques No Críticos ($U_i$): Aproximados eficientemente mediante expansión de Taylor por bloques.

La salida $o_t$ se deriva normalizando la agregación de valores ponderados:
$$o_t = \frac{N_t}{D_t}$$

Donde el denominador ($D_t$) y el numerador ($N_t$) están compuestos por:

• Términos Dispersos Exactos: Computados para los bloques seleccionados $j \in S_i$.
• Aproximación de Orden Cero por Bloque: Para los bloques no seleccionados, la contribución se aproxima utilizando el centroide del bloque $\bar{k}_j$ y la suma de los valores.
• Corrección Global de Primer Orden: En lugar de computar un término de primer orden distinto para cada bloque no seleccionado (lo cual está limitado por la memoria), PISA emplea una aproximación de orden híbrido. Utiliza una estadística global $\bar{H}$ (la media de las matrices de primer orden por bloque a través de todos los bloques) para aproximar el gradiente de primer orden para todo el conjunto de la cola.

La formulación para el numerador $N_t$ es:
$$N_t = \sum_{j \in S_i} \sum_{n} e^{q_t k_{j,n}^\top} v_{j,n} + \sum_{j \in U_i} \alpha_{t,j} \left(\sum_{n} v_{j,n}\right) + q_t \bar{H} \sum_{j \in U_i} \alpha_{t,j}$$
Donde $\alpha_{t,j} = \exp(q_t \bar{k}_j^\top)$.

Enrutamiento Consciente de la Covarianza

Para minimizar el error de aproximación, PISA utiliza una Estrategia Top-k Consciente de la Covarianza. La puntuación de selección de un bloque incorpora no solo la puntuación de atención, sino también la norma de la matriz de covarianza del bloque ($M_j = \|H_j - \bar{H}\|_2$). Los bloques con alta relevancia semántica o alta heterogeneidad son seleccionados para el cómputo exacto, mientras que los bloques más suaves se delegan a la expansión de Taylor.

Implementación de Hardware

PISA se implementa mediante un kernel fusionado que entrelaza el cómputo exacto y el escaneo aproximado dentro del proceso de softmax en línea (estilo FlashAttention). El pipeline involucra:

1. Pre-cómputo: Cálculo de los centroides de los bloques, sumas de valores y la estadística global $\bar{H}$.
2. Fase 1 (Exacta): Carga de los bloques seleccionados en la SRAM para la atención exacta.
3. Fase 2 (Aproximación): Escaneo de vectores de bloque agregados para los bloques no seleccionados usando la expansión de orden cero.
4. Fase 3 (Corrección): Inyección del término de corrección de primer orden global.

Este diseño asegura que el método sea eficiente en memoria y evite los cuellos de botella asociados con la carga de matrices distintas para cada bloque.

3. Contribuciones Clave

1. Nueva Atención Dispersa por Tramos: Un mecanismo libre de entrenamiento que permite la cobertura de todo el espectro de atención con complejidad sub-cuadrática al unificar el cómputo disperso exacto y la aproximación por bloques.
2. Esquema de Aproximación de Orden Híbrido: Una estrategia que combina la expansión de orden cero por bloque con una corrección global de primer orden. Esto aumenta significativamente la precisión con un costo computacional insignificante en comparación con la atención dispersa estándar.
3. Enrutamiento Consciente de la Covarianza: Una derivación del análisis de error que minimiza la divergencia de la aproximación seleccionando bloques basados tanto en las puntuaciones de atención como en las normas de covarianza de los bloques.
4. Rendimiento de Vanguardia (State-of-the-Art): Demostró una eficiencia y calidad superiores en diversas tareas (generación de vídeo e imagen) sin requerir re-entrenamiento.

4. Resultados Experimentales

Los autores evaluaron PISA en modelos de generación de vídeo (Wan2.1-1.3B/14B, Hunyuan-Video) y modelos de generación de imagen (SD 3.5, FLUX.1).

• Generación de Vídeo:

  • Wan2.1-14B: Logró una aceleración de 1.91× respecto a la atención densa (1468s $\to$ 818s) con un 87.5% de dispersión, manteniendo o mejorando ligeramente las métricas de calidad (VBench, PSNR, SSIM) comparado con la atención densa.
  • Hunyuan-Video: Logró una aceleración de 2.57× (1651s $\to$ 641s) con una preservación de calidad consistente.
  • Comparación: Superó a métodos dispersos existentes como SpargeAttn y SVG2 tanto en velocidad como en calidad, particularmente en la preservación de detalles finos y consistencia temporal.

• Generación de Imagen:

  • FLUX.1-dev: Logró una aceleración de 1.2× (8.32s $\to$ 6.87s) con un 85% de dispersión. Notablemente, PISA alcanzó puntuaciones más altas en métricas de preferencia humana (ImageReward, HPSv2) y métricas de similitud (PSNR, SSIM) en comparación tanto con la atención densa como con SpargeAttn.
  • SD 3.5: Demostró aceleraciones significativas (ej. 1.4× en SD 3.5-M) manteniendo puntuaciones FID comparables o mejores que la atención densa.

• Eficiencia del Kernel:

  • PISA superó consistentemente a FlashAttention-2/3 y SpargeAttn en latencia a través de varias longitudes de secuencia (4K a 32K) y densidades.
  • A diferencia de SpargeAttn, que se degrada en secuencias cortas, PISA mantuvo sus ventajas de velocidad incluso a 4K tokens.

5. Significación y Reivindicaciones

El artículo afirma que PISA representa un cambio de paradigma del "mantener-o-descartar" al "exacto-o-aproximado" en la atención dispersa. Su principal significación radica en:

• Cerrar la Brecha Velocidad-Calidad: Cierra eficazmente la brecha entre la velocidad de inferencia y la calidad de generación, logrando una aproximación casi sin pérdidas con un costo computacional insignificante.
• Compatibilidad sin Entrenamiento: Al integrar las aproximaciones directamente en el marco de normalización de softmax, PISA permite el reúso directo de los pesos pre-entrenados, eliminando la necesidad del costoso re-entrenamiento o ajuste fino requerido por los métodos de atención lineal o híbrida.
• Escalabilidad: Proporciona una solución unificada tanto para la generación de imágenes de alta resolución como para la generación de vídeos de larga duración, abordando el cuello de botella de la complejidad cuadrática sin sacrificar la fidelidad del contenido generado.

Los autores enfatizan que su método es el primero en integrar nativamente la expansión de Taylor por bloques y la atención dispersa en un marco unificado que preserva la distribución intrínseca de los modelos pre-entrenados.