Resumen Técnico: Asignación de Pasos Heterogénea (HSA) para la Generación Eficiente de Video

1. Declaración del Problema

Los Transformadores de Difusión (DiT) han establecido un rendimiento de vanguardia en la generación de video, pero sufren costos computacionales inmensos. Los protocolos de inferencia estándar aplican un número uniforme de pasos de eliminación de ruido ($T$) a cada token espacio-temporal en la secuencia. Este enfoque monolítico ignora la asimetría inherente de los datos visuales: las regiones de fondo, los parches estáticos y las estructuras gruesas son perceptualmente más simples y requieren menos pasos de eliminación de ruido que los objetos detallados del primer plano o las regiones de alto movimiento. Además, la visión humana es menos sensible a la degradación de la calidad en componentes altamente dinámicos en comparación con los estáticos. En consecuencia, aplicar un presupuesto completo de eliminación de ruido uniformemente a todos los tokens es computacionalmente derrochador y no se alinea con las prioridades perceptuales.

Los métodos de aceleración existentes, como la caché de características a nivel de paso (por ejemplo, TeaCache), suelen tomar decisiones globales: en cualquier iteración dada, o bien todos los tokens se calculan o bien todos los tokens reutilizan salidas en caché. Estos métodos no explotan la heterogeneidad espacial o a nivel de token, tratando la secuencia como un bloque monolítico en lugar de una colección de tokens con requisitos dinámicos variables.

2. Metodología: Asignación de Pasos Heterogénea (HSA)

Los autores proponen la Asignación de Pasos Heterogénea (HSA), un algoritmo de inferencia libre de entrenamiento y agnóstico al modelo que asigna trayectorias de eliminación de ruido específicas por token dentro de un DiT preentrenado. La idea central es particionar la secuencia de tokens en grupos basados en sus dinámicas de velocidad, asignando menos pasos a los tokens de evolución más lenta mientras se mantiene un programa completo para un grupo de referencia.

2.1 Partición y Programación de Tokens

El conjunto de $N$ tokens se particiona en $K$ grupos disjuntos ($G_1, \dots, G_K$). A cada grupo $G_k$ se le asigna un presupuesto de pasos $s_k$, que debe ser un divisor del total de pasos $T$.

• Grupo de Referencia ($G_1$): Se le asigna el presupuesto completo $s_1 = T$. Estos tokens están activos en cada iteración.
• Grupos Reducidos: Se les asignan presupuestos más pequeños (por ejemplo, $T/2, T/4$). Estos tokens están activos solo cuando la iteración global $i$ satisface $i \pmod{T/s_k} = 0$.
• Selección Dinámica: El artículo evalúa cuatro estrategias para asignar tokens a grupos. La estrategia de Selección Dinámica de Tokens, que asigna tokens basándose en el cambio relativo L1 por token de las predicciones de velocidad (identificando regiones de evolución lenta como los fondos), produce consistentemente el mejor rendimiento.

2.2 Sincronización de Caché KV

Un desafío principal en la programación heterogénea es que los tokens con diferentes conteos de pasos existen en diferentes niveles de ruido ($\sigma$) en cualquier iteración de tiempo real. La autoatención requiere que todos los tokens se atiendan mutuamente, creando un desajuste de sincronización. HSA resuelve esto mediante un mecanismo de sincronización de caché KV:

1. Tokens Activos: Solo el subconjunto de tokens activos ($A_i$) realiza un paso forward a través de las proyecciones QKV.
2. Actualización de Caché: Las proyecciones de Clave ($K$) y Valor ($V$) frescas para los tokens activos se escriben en una caché por capa, sobrescribiendo las entradas obsoletas.
3. Atención de Contexto Completo: Los tokens activos calculan consultas contra las matrices $K, V$ en caché completas (que contienen tanto entradas frescas para tokens activos como entradas obsoletas pero válidas para tokens omitidos).
4. Reducción de Complejidad: Esto reduce la complejidad de la autoatención de $O(N^2)$ a $O(|A_i| \cdot N)$, donde $|A_i|$ es el número de tokens activos en el paso $i$. Los tokens omitidos no incurren en cálculo de consultas, atención cruzada u operaciones de red neuronal durante las iteraciones inactivas.

2.3 Actualización Euler en Caché

Para avanzar los estados latentes de los tokens omitidos sin evaluaciones adicionales del modelo, HSA emplea una actualización Euler en caché.

• Cada token almacena la velocidad ($\hat{v}_n$) predicha en su paso activo más reciente.
• En cada iteración global, todos los tokens (activos y omitidos) se avanzan mediante el mismo paso incremental Euleriano $(\sigma_{i+1} - \sigma_i)$.
• Los tokens activos utilizan su velocidad recién calculada; los tokens omitidos reutilizan su velocidad en caché.
• Esta actualización se formula como una única operación tensorial sobre todos los $N$ tokens, evitando bifurcaciones y manteniendo la eficiencia de la GPU.

2.4 Ventana de Caché Consciente de la Fase

Reconociendo que las etapas tempranas y tardías de eliminación de ruido son críticas para la estructura global y los detalles finos, respectivamente, HSA restringe la caché a la porción media de la trayectoria. Los primeros y últimos $m$ pasos ejecutan pasadas forward completas y sin caché para todos los tokens, asegurando que la integridad estructural y el refinamiento de detalles no se vean comprometidos por aproximaciones obsoletas.

3. Contribuciones Clave

1. Algoritmo HSA: Un paradigma de inferencia completamente libre de entrenamiento, plug-and-play, que asigna trayectorias de eliminación de ruido heterogéneas a grupos de tokens espacio-temporales basándose en dinámicas de velocidad.
2. Sincronización de Caché KV: Un mecanismo que preserva el contexto completo de atención de $N$ tokens para los tokens activos mientras elude por completo los tokens inactivos, reduciendo el costo de atención por iteración a $O(|A_i| \cdot N)$.
3. Actualización Euler en Caché: Una operación tensorial simple que avanza los tokens omitidos utilizando velocidades en caché, manteniendo la calidad de generación sin evaluaciones adicionales del modelo.
4. Evaluación Exhaustiva: Un análisis de cuatro estrategias de agrupación de tokens y cuatro configuraciones de presupuesto, demostrando robustez a través de diferentes presupuestos de inferencia.

4. Resultados Experimentales

Los autores evaluaron HSA en Wan-2.1-1.3B y Wan-2.2-A14B (Texto-a-Video e Imagen-a-Video) y LTX-2 (Audio-Video).

• Compensación Calidad-Tiempo de Ejecución: HSA logra una frontera de Pareto superior en comparación con la Coincidencia de Flujo (FM) uniforme y métodos de caché de vanguardia como TeaCache.
• Aceleración Agresiva: Las ventajas son más pronunciadas en regímenes de aceleración agresiva (50% y 25% de tiempo de ejecución). Al 25% de tiempo de ejecución, HSA supera significativamente a las líneas base en la métrica VBench (Puntuación Total: 79.87% vs. 75.68% para FM y 78.33% para TeaCache).
• Robustez Dimensional: Mientras que los métodos base sufren de "colapso dimensional catastrófico" (degradación severa en dimensiones específicas de VBench como suavidad del movimiento o consistencia temporal) bajo presupuestos ajustados, HSA rastrea robustamente la referencia de presupuesto completo a través de las 16 dimensiones de evaluación.
• Imagen-a-Video (I2V): HSA preserva con éxito la alineación de la condición de imagen y la fidelidad visual en tareas I2V, mientras que las líneas base FM uniformes con pasos reducidos ($T=10$) exhiben un colapso estructural severo.
• Sin Perfilado Offline: A diferencia de TeaCache, que requiere un perfilado offline costoso para estimar las fluctuaciones de salida del modelo, HSA opera sin ningún ajuste offline ni ajuste fino.

5. Significado y Afirmaciones

El artículo afirma que HSA aborda una ineficiencia fundamental en la inferencia de DiT al alinear los recursos computacionales con la complejidad perceptual y dinámica de los datos visuales. Al tratar los tokens individualmente en lugar de globalmente, HSA logra aceleraciones significativas sin la degradación de calidad típicamente asociada con la aceleración agresiva.

Los autores enfatizan que HSA es un marco flexible que puede instanciarse con diversas estrategias de agrupación de tokens. Si bien la selección dinámica actualmente produce los resultados más sólidos, el marco anticipa mejoras adicionales con estrategias de agrupación más sofisticadas. El trabajo sienta las bases para una nueva clase de algoritmos de inferencia que aprovechan la heterogeneidad inherente de los datos espacio-temporales para democratizar la generación de video de alta fidelidad.