Fast Autoregressive Video Diffusion and World Models with Temporal Cache Compression and Sparse Attention

Resumen Técnico: FAST-AR

Planteamiento del Problema

Los modelos de difusión autorregresivos de video permiten la generación por streaming, facilitando la síntesis de larga duración, los modelos de mundo de video y los motores de juegos neuronales interactivos. Sin embargo, estos modelos enfrentan un cuello de botella crítico durante la inferencia: el caché KV (Key-Value) crece linealmente con el número de fotogramas generados. Este crecimiento provoca dos problemas fundamentales:

1. Latencia Creciente: A medida que el caché se expande, el costo de atención por paso escala linealmente (y el trabajo acumulado escala cuadráticamente, $O(T^2)$), lo que causa que la velocidad de generación se degrade progresivamente en rollouts de larga duración.
2. Explosión de Memoria: El caché en expansión consume una cantidad creciente de memoria GPU, lo que a menudo fuerza ventanas de contexto cortas que perjudican la consistencia temporal a largo plazo.

Las soluciones existentes de NLP (compresión de KV para LLMs) no se transfieren directamente a la difusión de video, y los métodos recientes específicos para video ofrecen solo mejoras modestas, a menudo a costa de la calidad visual o al no lograr estabilizar el uso de memoria a lo largo del tiempo.

Metodología: FAST-AR

Los autores proponen FAST-AR (Fast AutoRegressive diffusion), un marco de atención unificado, libre de entrenamiento (training-free), diseñado para explotar la redundancia en la generación de video autorregresiva. El marco consta de tres componentes principales, todos ellos aprovechando el emparejamiento rápido de Vecinos Más Cercanos Aproximados (ANN) para reducir la computación y la memoria sin reentrenar el modelo subyacente.

1. Motivación: Fuentes de Redundancia

Los autores identifican tres fuentes persistentes de redundancia en la difusión de video autorregresiva:

• Redundancia de Claves Temporales (Temporal Key Redundancy): Las claves almacenadas en caché a través de los fotogramas suelen ser casi duplicados, ya que el contenido semántico evoluciona lentamente.
• Evolución Semántica de Consultas/Claves (Semantic Query/Key Evolution): Las consultas (queries) y las claves (keys) son mayoritariamente semánticas y evolucionan lentamente, lo que hace que muchos cálculos de puntuación de atención sean redundantes.
• Dispersión de la Atención Cruzada (Cross-Attention Sparsity): En prompts largos, solo un pequeño subconjunto de tokens es relevante para la síntesis de un fotograma específico; atender a todos los tokens es computacionalmente ineficiente.

2. Componentes Principales

A. TempCache (Compresión de Caché Temporal)

• Objetivo: Comprimir el caché KV identificando y fusionando tokens temporalmente correspondientes.
• Mecanismo: En lugar de calcular matrices de atención completas para encontrar correspondencias, el método utiliza ANN ligero (Hashing de Sensibilidad Local o Cuantización) para encontrar el vecino más cercano (top-1) de la clave para cada consulta del fotograma actual en los fotogramas anteriores.
• Base Teórica: Basándose en el Lema 5.1, los autores demuestran que fusionar claves idénticas no incurre en errores de aproximación si los valores se agregan (media) y los logits se desplazan según el tamaño del grupo. En la práctica, las claves se fusionan basadas en un umbral de similitud.
• Beneficio: Esto limita el crecimiento del caché KV, asegurando que el uso de memoria y la latencia permanezcan constantes independientemente de la longitud de la generación.

B. AnnCA (Atención Cruzada de Vecinos Más Cercanos Aproximados)

• Objetivo: Acelerar la atención cruzada seleccionando solo los tokens del prompt relevantes para el fotograma.
• Mecanismo: El método proyecta tanto las consultas latentes como las claves del prompt en un espacio de incrustación (embedding) compartido (LSH o cuantizado). Identifica los tokens del prompt que comparten un cubo (bucket) con al menos una consulta del fotograma actual. Los tokens que no tienen un cubo coincidente son podados para ese fotograma.
• Beneficio: Reduce drásticamente el número de tokens del prompt procesados por fotograma sin calcular mapas de atención densos.

C. AnnSA (Atención de Auto-Atención de Vecinos Más Cercanos Aproximados)

• Objetivo: Dispersar la auto-atención restringiendo las consultas a claves semánticamente emparejadas.
• Mecanismo: Reutilizando los cubos semánticos descubiertos durante la fase de poda de la atención cruzada, el método restringe la consulta de cada token para que atienda solo a las claves dentro del mismo cubo (o cubos) semántico.
• Beneficio: Fuerza la localidad semántica y reduce el costo cuadrático de la auto-atención a una operación de dispersión de bloques (block-sparse) utilizando kernels eficientes (por ejemplo, FlashInfer).

Resultados Clave

Los experimentos se realizaron en modelos autorregresivos de última generación (Rolling-Forcing y LongVie2) utilizando los conjuntos de datos LongVBench y LongVGenBench.

• Aceleración (Speedup): FAST-AR logra aceleraciones de extremo a extremo de 5× a 10.8× en comparación con FlashAttention-3 (FA3) denso.
  • TempCache por sí solo proporciona una aceleración de ~3×.
  • Combinar TempCache con AnnSA y AnnCA produce las ganancias máximas (hasta ×10.8).
• Preservación de la Calidad: El método mantiene una calidad visual casi idéntica a la del baseline denso.
  • Las puntuaciones PSNR, SSIM y LPIPS son comparables a la atención densa.
  • Las puntuaciones de VBench (calidad perceptual) se preservan (~84.0 frente a 84.1 del baseline).
  • En contraste, los baselines de atención dispersa existentes (SVG, RadialAttn) degradan significativamente la calidad (VBench cae a ~33–61) o no logran mantener el rendimiento (throughput).
• Estabilidad:
  • Rendimiento (Throughput): Mientras que la atención densa y otros baselines se ralentizan a medida que el contexto crece (3000 fotogramas), FAST-AR mantiene un FPS constante.
  • Memoria: El uso de memoria GPU pico permanece casi constante durante rollouts largos, mientras que los baselines exhiben un crecimiento de memoria lineal.

Significación y Reivindicaciones

El artículo afirma que FAST-AR es la primera aplicación libre de entrenamiento basada en ANN específicamente para modelos de difusión autorregresivos. Su principal significación radica en:

1. Escalabilidad: Permite la generación de videos de varios minutos (por ejemplo, 3000 fotogramas) en una sola GPU con latencia y memoria estables, una capacidad previamente restringida por el escalado $O(T^2)$ de la atención.
2. Eficiencia sin Reentrenamiento: A diferencia de los métodos que requieren ajuste fino (fine-tuning) o recetas de entrenamiento complejas, FAST-AR opera como un marco de inferencia tipo "plug-in" compatible con backbones existentes.
3. Robustez: Aborda con éxito los patrones específicos de redundancia de la difusión de video (correspondencia temporal y dispersión semántica) que los métodos genéricos de compresión de NLP o de aceleración de video offline no logran manejar eficazmente.

Los autores enfatizan que su enfoque desbloquea la viabilidad práctica de los modelos de mundo de video autorregresivos y la síntesis de larga duración al eliminar las barreras computacionales y de memoria que actualmente limitan su longitud de rollout.