Resumen Técnico: Veda - Difusión de Video Escalable mediante Atención Escasa Destilada

1. Declaración del Problema

Escalar los Transformadores de Difusión (DiT) para la generación de video de alta resolución y larga duración está actualmente limitado por la complejidad computacional y de memoria cuadrática ($O(N^2)$) de los mecanismos de autoatención en relación con la longitud de la secuencia de tokens espacio-temporal. Aunque la atención escasa ofrece una vía teórica para aliviar esto, los métodos existentes enfrentan una compensación crítica:

• Patrones Estáticos: Los métodos que utilizan patrones escasos predefinidos (por ejemplo, ventanas deslizantes) carecen de adaptabilidad a las estructuras de atención dinámicas y específicas de cada cabeza aprendidas por los DiT.
• Selección Dinámica: Los métodos que aprenden máscaras dinámicas (por ejemplo, VSA, VMOBA) a menudo sufren de estimaciones inexactas. Dependen de una supervisión implícita a través del objetivo de difusión y estadísticas gruesas (por ejemplo, agrupación promedio), lo que lleva a una desalineación con la geometría real de la atención completa.
• El Problema Central: La evidencia empírica muestra que, en niveles altos de escasez (≥90%), estos enfoques existentes degradan la calidad de generación, produciendo artefactos estructurados como distorsiones espaciales, patrones de "onda de agua" y parpadeo temporal. El artículo postula que la degradación no es causada por la escasez en sí misma, sino por la falla de las máscaras escasas para alinearse con la geometría por baldosas de la atención completa.

2. Metodología: Marco Veda

Veda es un marco de atención escasa destilada diseñado para preservar la estructura y el ranking por baldosas de la atención completa incluso bajo escasez extrema. Trata la selección de baldosas como un problema de reconstrucción explícito en lugar de un subproducto implícito del objetivo de difusión.

2.1. Puntuación de Baldosas Destilada

En lugar de depender de la pérdida de difusión para dar forma a la escasez, Veda emplea un estimador ligero para aprender explícitamente puntuaciones de atención a nivel de baldosa a partir de una columna vertebral de atención completa.

• Construcción del Objetivo: Se construye una máscara de referencia "Oráculo" aplicando agrupación máxima (max-pooling) a la matriz de atención completa sobre las regiones de baldosas de consulta-clave. Se elige la agrupación máxima sobre el promedio para preservar los picos de señal salientes y de alta frecuencia que a menudo se diluyen por el ruido de fondo.
• Estimador Consciente de Estadísticas: Para reconstruir estas puntuaciones, Veda utiliza un descriptor TripPool para cada baldosa, concatenando estadísticas de Promedio, Máximo y Mínimo. Esto enriquece la representación de la baldosa más allá de la simple agrupación promedio.
• Proyecciones Específicas por Cabeza: El estimador utiliza proyecciones MLP distintas para cada cabeza de atención para capturar patrones de dependencia heterogéneos.
• Optimización: El modelo se entrena utilizando un objetivo de destilación de divergencia KL por fila ($L_{distill}$) para alinear la distribución de puntuaciones de baldosas predicha con la referencia de atención completa. Crucialmente, se aplica una operación de stop-gradient a las características de la columna vertebral que alimentan el estimador. Esto desacopla el aprendizaje de la máscara del aprendizaje de características, evitando la interrupción del manifold generativo preentrenado.

2.2. Baldosas Conscientes de Cabeza

Reconociendo que las cabezas de atención exhiben una heterogeneidad significativa en las dependencias espacio-temporales (algunas capturan interacciones espaciales locales, otras dependencias temporales de largo alcance), Veda abandona las estrategias de baldosas uniformes.

• Búsqueda de Configuración: Para cada capa y cabeza, Veda busca una configuración de baldosas óptima $\pi_{l,h} = (p_t, p_h, p_w)$ que factorice el tamaño de la baldosa de hardware.
• Selección Offline: La configuración óptima se selecciona offline en un conjunto de calibración minimizando el error de aproximación entre la salida de atención completa y la salida de atención escasa, asegurando que la estrategia de baldosas coincida con las necesidades estructurales específicas de cada cabeza.

2.3. Implementación Eficiente para Hardware

Para traducir las reducciones teóricas de FLOP en aceleraciones prácticas de tiempo de reloj, Veda implementa un kernel de atención escasa con salto de baldosas personalizado utilizando el DSL ThunderKittens.

• Ejecución Asíncrona: El kernel aprovecha el Acceso a Memoria Tensor (TMA) de NVIDIA Hopper y la Especialización de Warps. Desacopla el movimiento de datos del cálculo utilizando un paradigma productor-consumidor, donde los warps productores traen baldosas de clave/valor no contiguas a la memoria compartida mientras los warps consumidores ejecutan operaciones de núcleo tensorial.
• Eficiencia: Este diseño oculta la latencia de memoria detrás de matemáticas de matrices densas, logrando aproximadamente el 80% de la Utilización de FLOP de Memoria (MFU) de FlashAttention-3.
• Eficiencia de Entrenamiento: Un kernel especializado TileLang genera mapas de calor de verdad fundamental en dos pasadas, permitiendo una supervisión escasa (subconjuntos aleatorios de baldosas de consulta) para reducir la sobrecarga de entrenamiento sin comprometer el rendimiento.

3. Contribuciones Clave

1. Insight Empírico: El artículo demuestra que la calidad de generación en la difusión de video escasa está determinada por la alineación de la máscara escasa con la geometría por baldosas de la atención completa, en lugar de la propia tasa de escasez.
2. Atención Escasa Destilada: Veda introduce un marco que formula la selección de baldosas como un problema de reconstrucción explícito, utilizando un estimador consciente de estadísticas y proyecciones específicas por cabeza para minimizar el error de estimación.
3. Baldosas Conscientes de Cabeza: Una estrategia novedosa que asigna factorizaciones espacio-temporales de baldosas distintas a diferentes cabezas de atención, abordando la heterogeneidad de los patrones de dependencia en los Video DiT.
4. Kernel Optimizado para Hardware: Un kernel personalizado con salto de baldosas que logra una escalabilidad casi lineal con la longitud de la secuencia y un alto uso de hardware, convirtiendo la escasez algorítmica en ganancias tangibles de latencia.

4. Resultados Experimentales

Los experimentos se realizaron en modelos de difusión de video a gran escala, incluidos Waver-T2V-12B y Wan2.1-T2V-14B, en resoluciones de hasta 720P y longitudes de secuencia de hasta 241 cuadros.

• Aceleración: En Waver-T2V-12B (720P, 241 cuadros), Veda logra una aceleración de extremo a extremo de 5.1× (reduciendo el tiempo de muestreo de 19.4 a 3.8 minutos) y una aceleración de autoatención de 10.5×. La sobrecarga de atención se reduce del 92% al 50%.
• Escalabilidad: Las aceleraciones aumentan con la longitud de la secuencia. A 245K tokens, Veda escala de forma casi lineal (309.1ms) en comparación con el crecimiento cuadrático de la Atención Completa (1576.5ms), generando una aceleración de 5.1×.
• Preservación de Calidad:
  • Evaluación Humana: En Waver-bench 1.0, Veda al 90% de escasez logra paridad perceptual con la Atención Completa. Al 95% de escasez, supera significativamente al método VSA de última generación (incluso cuando VSA opera con menor escasez).
  • Reducción de Artefactos: A diferencia de las líneas base que sufren patrones de onda de agua y parpadeo temporal en alta escasez, Veda mantiene una alta fidelidad visual y coherencia temporal.
  • Métricas Cuantitativas: Las evaluaciones de VBench muestran que Veda mantiene la consistencia del sujeto y la suavidad del movimiento comparables a las líneas base de Atención Completa, incluso al 95% de escasez.

5. Significado y Afirmaciones

El artículo afirma que Veda aborda con éxito el cuello de botella de escalabilidad de la generación de video de alta resolución sin comprometer la integridad estructural. Al replantear la atención escasa como un problema de destilación en lugar de una búsqueda heurística o una tarea de aprendizaje implícito, Veda permite una escasez agresiva (hasta el 95%) que anteriormente era inalcanzable sin una degradación severa de la calidad.

Los autores enfatizan que las ganancias no son meramente teóricas; el kernel personalizado asegura que la escasez se traduzca en aceleración práctica de tiempo de reloj. El marco se presenta como una solución modular y consciente del hardware que escala favorablemente con la resolución espacio-temporal, haciendo viable la generación de video de alta fidelidad y larga duración bajo las restricciones actuales de hardware. El trabajo sugiere que las mejoras futuras podrían centrarse en una fusión de kernels más estrecha, estrategias de escasez adaptativa a través de pasos de tiempo y almacenamiento en caché temporal de puntuaciones de baldosas.