SLA2: Sparse-Linear Attention with Learnable Routing and QAT

Resumen Técnico: SLA2

1. El Problema

Los modelos de difusión, especialmente para la generación de video, requieren una atención masiva que es computacionalmente costosa ($O(N^2)$). Para mitigar esto, se han desarrollado métodos de atención dispersa (sparse) y lineal. Sin embargo, el enfoque previo conocido como Atención Lineal Dispersa (SLA) presenta dos limitaciones fundamentales que degradan su eficiencia y precisión:

1. Desajuste en la Descomposición Teórica: SLA intenta descomponer el mapa de atención completo ($P$) en una parte dispersa ($P_1$) y una parte de bajo rango ($P_2$). Sin embargo, la implementación actual de SLA introduce un error de escalado constante en la rama dispersa debido a la re-normalización de softmax. Para compensar esto, SLA depende de una proyección lineal aprendible adicional que intenta corregir el error, lo cual es una solución subóptima y difícil de aprender.
2. Enrutamiento Heurístico Subóptimo: SLA utiliza una regla heurística fija para decidir qué cálculos van a la rama dispersa y cuáles a la lineal (basándose en la magnitud de los pesos de atención). Esta división no es óptima; mover ciertos pesos de la parte dispersa a la lineal podría mejorar la dispersión sin aumentar el rango de la parte lineal, pero la heurística actual no lo detecta.
3. Falta de Optimización para Cuantización: No se aprovecha plenamente el potencial de la atención de baja precisión (low-bit) para acelerar aún más el proceso sin sacrificar calidad.

2. Metodología: SLA2

Los autores proponen SLA2, una reformulación de la atención lineal dispersa que aborda las limitaciones anteriores mediante tres componentes clave:

• A. Formulación de Atención Directa y Consistente:
  En lugar de la formulación de SLA que requiere una proyección de corrección, SLA2 introduce una combinación directa basada en una razón aprendible ($\alpha$). La salida se formula como:
  $$O = \alpha \odot O_s + (1 - \alpha) \odot O_l$$
  Donde $O_s$ es la salida de la atención dispersa y $O_l$ es la de la atención lineal. El vector $\alpha$ (aprendible) ajusta dinámicamente la contribución de cada rama, eliminando el error de escalado de fila y alineándose matemáticamente con la descomposición original de $P = P_1 + P_2$.

• B. Enrutador Aprendible (Learnable Router):
  Se introduce un módulo $R$ que decide dinámicamente qué posiciones de la matriz de atención deben procesarse por la rama dispersa y cuáles por la lineal.

  • Entrada: Utiliza las matrices de consulta ($Q$) y clave ($K$), ignorando $V$ (ya que la formación de la atención depende de $Q$ y $K$).
  • Eficiencia: Para evitar el costo $O(N^2)$, aplica pooling (agrupamiento) a tokens adyacentes en $Q$ y $K$ y utiliza proyecciones lineales aprendibles ($proj_q, proj_k$).
  • Entrenamiento: Utiliza un operador SoftTop-k diferenciable durante el entrenamiento para permitir la retropropagación de gradientes, cambiando a un Top-k duro (hard) durante la inferencia para mantener la eficiencia.

• C. Atención Dispersa + Cuantización Consciente del Entrenamiento (QAT):
  Para lograr una aceleración adicional, SLA2 integra la cuantización de baja precisión (ej. INT8/FP8) en la rama de atención dispersa.

  • Se utiliza Entrenamiento Consciente de la Cuantización (QAT): Durante el entrenamiento, la parte forward usa cuantización de baja precisión, mientras que el backward se mantiene en FP16. Esto permite que el modelo se adapte a los errores de cuantización, minimizando la pérdida de calidad en la inferencia final.

3. Contribuciones Clave

1. Análisis y Reformulación: Identificaron y corrigieron el desajuste matemático en SLA, proponiendo una formulación que combina ramas dispersas y lineales mediante un factor de mezcla aprendible ($\alpha$), eliminando la necesidad de proyecciones de corrección complejas.
2. Enrutamiento Dinámico: Diseñaron un enrutador aprendible que optimiza la división de tareas entre las ramas dispersa y lineal basándose en un objetivo de optimización claro, superando a las heurísticas estáticas.
3. Aceleración Extrema con QAT: Integraron exitosamente la cuantización de baja precisión en el flujo de trabajo de atención dispersa mediante QAT, logrando aceleraciones significativas sin degradar la calidad de generación.
4. Rendimiento en Video: Demostraron que SLA2 mantiene la calidad de generación de video incluso con niveles de dispersión extremadamente altos.

4. Resultados Experimentales

Los experimentos se realizaron en modelos de difusión de video Wan2.1 (versiones de 1.3B y 14B parámetros).

• Eficiencia y Velocidad:
  • Logró un 97% de dispersión en la atención.
  • Aceleración de tiempo de ejecución de atención de 18.6x (y hasta 18.7x en kernels) en comparación con la atención completa (Full Attention).
  • Reducción de la latencia de extremo a extremo: En el modelo 1.3B, la latencia bajó de 97s a 7s (2.3x más rápido en total); en el modelo 14B, la reducción fue de 4.35x.
• Calidad de Generación:
  • SLA2 superó consistentemente a los métodos de base (SLA, VSA, VMoBA) en todas las métricas de calidad de video (IQ, OC, AQ, MS, SC, VR) incluso al 97% de dispersión.
  • Sorprendentemente, en muchos casos, SLA2 con alta dispersión superó a la atención completa (0% de dispersión) tras el ajuste fino, atribuido a la mayor calidad del conjunto de datos de ajuste fino.
• Ablación:
  • La eliminación de QAT resultó en una caída notable en la calidad, confirmando su importancia.
  • El enrutador aprendible superó significativamente al enrutador Top-k estático.

5. Significado e Impacto

SLA2 representa un avance significativo en la eficiencia de los modelos de difusión para video. Al resolver los defectos teóricos de la atención dispersa-lineal anterior y combinarlo con técnicas de cuantización avanzadas, permite:

• Generación de video de alta calidad en hardware limitado (ej. una sola GPU RTX 5090) que antes requería recursos masivos.
• Escalabilidad: Hace viable el entrenamiento y la inferencia de modelos de video grandes (14B+ parámetros) con una fracción del costo computacional.
• Paradigma de Diseño: Establece un nuevo estándar para el diseño de mecanismos de atención en transformadores, demostrando que la combinación de enrutamiento aprendible, formulaciones matemáticas correctas y cuantización consciente del entrenamiento es crucial para el rendimiento óptimo.

En resumen, SLA2 ofrece una solución práctica y altamente eficiente para el cuello de botella computacional en la generación de video, logrando un equilibrio superior entre velocidad y calidad.