SALAD: Achieve High-Sparsity Attention via Efficient Linear Attention Tuning for Video Diffusion Transformer

Resumen Técnico: SALAD para Transformadores de Difusión de Video

1. El Problema

Los Transformadores de Difusión (DiT) han demostrado un rendimiento excepcional en la generación de video. Sin embargo, enfrentan un desafío crítico de eficiencia: la complejidad cuadrática ($O(N^2)$) del mecanismo de atención completa (Full Attention) en secuencias largas de video (alta resolución y duración). Esto genera una latencia sustancial y un alto consumo de memoria.

Para mitigar esto, se han propuesto mecanismos de atención dispersa (sparse attention) que limitan las interacciones entre tokens. No obstante, existen dos limitaciones principales en las soluciones actuales:

• Enfoques sin entrenamiento (Training-free): Logran una aceleración moderada pero solo alcanzan una dispersión limitada (40-60%), lo que resulta en una aceleración insuficiente.
• Enfoques basados en entrenamiento: Pueden lograr una alta dispersión (80-95%), pero requieren grandes cantidades de datos y recursos computacionales para el entrenamiento. Además, cuando se intenta recuperar el rendimiento de modelos ultra-dispersos utilizando LoRA (Low-Rank Adaptation) estándar, se observa una degradación significativa en la calidad de generación (inconsistencias temporales, duplicación de objetos, pérdida de coherencia con el texto), ya que LoRA no logra compensar completamente la pérdida de interacciones globales entre tokens.

2. Metodología: SALAD

Los autores proponen SALAD (High-Sparsity Attention paralleling with Linear Attention for Diffusion Transformer), una arquitectura eficiente que combina una rama de atención dispersa con una rama de atención lineal ligera en paralelo.

Componentes Clave:

• Arquitectura Híbrida (Dispersa + Lineal):

  • Se introduce una rama de atención lineal paralela a la atención dispersa.
  • Ambas ramas comparten las matrices de proyección de Query, Key y Value ($W_Q, W_K, W_V$) inicializadas desde el modelo preentrenado.
  • La atención lineal tiene una complejidad computacional $O(N)$, permitiendo la mezcla global de tokens (interacciones a larga distancia) con un costo de inferencia mínimo.
  • La rama dispersa maneja la mayor parte del modelado de la secuencia para mantener la eficiencia.

• Estrategia de Escalado Estático-Dinámico Multi-nivel (Multi-level Static-Dynamic Scaling):

  • Problema detectado: Los autores observaron un desequilibrio de rango (rank imbalance). La salida de la atención lineal suele tener un rango efectivo mucho menor que la dispersa. Sumarlas directamente sin regulación puede suprimir las estructuras de alto rango informativas de la rama dispersa o degradar la calidad.
  • Solución: Se diseña un mecanismo de control en dos niveles para regular la contribución de la rama lineal:
    1. Escalado Estático por Capa (Layer-wise Static Scaling): Una capa de proyección inicializada en cero (zero-initialized) que se ajusta durante el fine-tuning. Esto permite que el modelo comience desde la atención dispersa pura e incorpore gradualmente la rama lineal.
    2. Escalado Dinámico Consciente de Entrada y Tiempo (Input-Timestep-Aware Dynamic Scaling): Un módulo ligero que genera un factor de escala global ($s \in (0, 1)$) basado en el estado oculto actual y el paso de tiempo de denoising ($t$). Esto permite adaptar la influencia de la rama lineal dinámicamente según la necesidad de información global en cada paso de generación.

• Eliminación de Ramas Post-Entrenamiento (Post-tuning Branch Dropping):

  • Dado que el factor de escala dinámico indica la importancia de la rama lineal en cada paso, se puede eliminar (dropear) la ejecución de la rama lineal en los pasos donde su contribución es mínima (por debajo de un umbral $\tau_t$), logrando una mayor eficiencia sin pérdida de calidad.

3. Contribuciones Principales

1. Arquitectura SALAD: Un módulo de atención eficiente que mejora el rendimiento de la atención ultra-dispersa mediante una rama lineal paralela, añadiendo solo un 4.99% de parámetros adicionales al modelo preentrenado.
2. Estrategia de Escalado: La propuesta de una estrategia de escalado estático-dinámico que regula la influencia de la rama lineal, asegurando que actúe como un complemento controlado en lugar de dominar la salida, lo cual es crucial para mantener la calidad.
3. Eficiencia en el Entrenamiento: Demostración de que es posible lograr una alta dispersión (hasta 90%) con un costo de entrenamiento extremadamente bajo:
   • Solo 2,000 muestras de video (OpenSora/Mixkit).
   • Menos de 1,600 pasos de entrenamiento.
   • Máximo 30 horas de GPU (con batch size de 8).

4. Resultados Experimentales

Los experimentos se realizaron en modelos Wan2.1 (1.3B y 14B) con secuencias de 30k y 100k tokens.

• Rendimiento y Calidad:

  • SALAD logra una dispersión del 90% manteniendo una calidad de generación comparable o incluso superior a la del modelo de atención completa (Full Attention) en métricas de VBench (Consistencia de Sujeto, Fondo, Calidad de Imagen y Consistencia de Texto).
  • En comparación con métodos de fine-tuning con LoRA estándar, SALAD supera significativamente el rendimiento, eliminando artefactos como la duplicación de objetos o la pérdida de coherencia temporal.

• Aceleración:

  • Se logra una aceleración de inferencia de 1.52x a 2.03x dependiendo del modelo y la longitud de la secuencia.
  • En el modelo Wan2.1-14B, se alcanzó una aceleración de 1.59x con una dispersión del 88%.

• Eficiencia de Parámetros:

  • SALAD supera a las variantes de LoRA con rangos más altos (hasta 512) utilizando menos parámetros entrenables (165M vs 189M+ de LoRA r=256), demostrando que las mejoras provienen del diseño arquitectónico y no solo de la capacidad de parámetros.

5. Significado e Impacto

El trabajo de SALAD es significativo porque resuelve el dilema entre eficiencia computacional y calidad de generación en video.

• Viabilidad de Alta Dispersión: Demuestra que es posible alcanzar niveles de dispersión extremos (90%) sin sacrificar la fidelidad del video, algo que los métodos anteriores no lograban sin un coste de entrenamiento prohibitivo.
• Accesibilidad: Al requerir recursos de entrenamiento mínimos (2k videos, pocas horas de GPU), hace que la optimización de modelos de difusión de video de gran escala sea accesible para comunidades más amplias, no solo para grandes laboratorios con acceso a miles de GPUs.
• Mecanismo de Control: La introducción de la estrategia de escalado estático-dinámico ofrece una nueva perspectiva sobre cómo integrar mecanismos de atención de diferentes complejidades (lineal vs. dispersa) de manera complementaria, evitando la degradación de la información.

En conclusión, SALAD establece un nuevo estado del arte para la generación de video eficiente, permitiendo la creación de videos de alta resolución y larga duración con una fracción del costo computacional tradicional.