Resumen Técnico: MobileWan

Planteamiento del Problema

Los avances recientes en la difusión de video han sido impulsados por el escalado de arquitecturas basadas en transformadores (DiT) a miles de millones de parámetros, mejorando significativamente la fidelidad visual y la coherencia del movimiento. Sin embargo, la generación de video móvil está considerablemente rezagada respecto a los sistemas a escala de servidor debido a las estrictas restricciones en el pico de memoria, la complejidad cuadrática de la atención propia (self-attention), la sobrecarga del decodificador y la latencia de extremo a extremo. Los modelos móviles existentes suelen estar restringidos a backbones pequeños (0.4–1.8 mil millones de parámetros), lo que inevitablemente limita la calidad de la generación y crea una brecha pronunciada entre el rendimiento móvil y el de escala de servidor. El desafío central es que los cuellos de botella no se limitan a un solo módulo; simplemente reemplazar la atención softmax con atención lineal a menudo requiere reentrenar desde cero o incurre en una brecha de expresividad, mientras que los enfoques híbridos aún no han convertido completamente el backbone en una arquitectura recurrente de baja memoria adecuada para el despliegue móvil.

Metodología

Los autores proponen MobileWan, un sistema que demuestra que un transformer de difusión de video de 5 mil millones de parámetros (basado en Wan2.2-5B) puede desplegarse eficientemente en hardware móvil con restricciones de memoria mediante una combinación de reformulación recurrente y compresión estructurada. La metodología consta de cuatro etapas principales de optimización:

1. Poda de Cabezales Aprendible (Learnable Head Pruning)

En lugar de la poda de bloques heurística, a la cual los autores encontraron que Wan2.2-5B es altamente sensible, introducen un método de poda de cabezales aprendible.

• Mecanismo: Se introduce una compuerta escalar aprendible para cada cabezal de autoatención, modulando su contribución. Estas compuertas se optimizan de extremo a extremo utilizando un regularizador de dispersión y una estrategia de aprendizaje con sesgo de ruido.
• Estrategia de Entrenamiento: El modelo se entrena con una distribución de alto sesgo de ruido (extrayendo ruido $\sigma \sim \text{LogitNormal}(1.5, 1)$) durante la fase de aprendizaje de la compuerta. Esto enfatiza la estructura global y el movimiento, haciendo que las decisiones de poda sean más robustas. El proceso involucra un esquema de tres fases: exploración de compuerta suave, el de la temperatura hacia valores binarios y ajuste fino final con una arquitectura podada fija.
• Resultado: Este enfoque supera a la poda heurística, particularmente en ratios de poda agresivos, produciendo reducciones medibles de FLOPs con una pérdida de calidad limitada.

2. Destilación Recurrente

Para abordar la complejidad de memoria cuadrática de la autoatención, los autores convierten el transformer preentrenado en un modelo autorregresivo por fragmentos (chunk-wise).

• Atención Híbrida: Emplean un mecanismo de atención híbrida que combina atención lineal causal con atención softmax local. El modelo atiende a un pequeño conjunto de tokens locales mediante softmax para el contexto inmediato y utiliza atención lineal para el contexto histórico.
• Reformulación RNN: En la inferencia, la atención lineal causal permite que el modelo opere como una Red Neuronal Recurrente (RNN). Mantiene agregados de ejecución compactos (vectores de estado $s_t$ y $z_t$) para los fragmentos pasados en lugar de almacenar la matriz de atención completa.
• Beneficio: Esto reduce la huella de memoria para que sea constante independientemente de la longitud de la secuencia, permitiendo la generación de videos largos en dispositivos con RAM limitada.

3. Destilación de Pasos (Step Distillation)

Para reducir la latencia, los autores aplican Destilación de Coincidencia de Distribución (DMD) y su variante D-DMD.

• Enfoque: Destilan el modelo maestro de múltiples pasos en un modelo estudiante de pocos pasos.
• Optimización: Los experimentos compararon DMD estándar, D-DMD (dividiendo la pérdida en términos de Coincidencia de Distribución y de aumentación de Guía Libre de Clasificador) y DMD2 (incorporando pérdida adversarial).
• Selección: Aunque DMD2 con 2 pasos logró el mayor grado dinámico y puntuaciones de VBench, los autores señalaron que D-DMD produjo muestras visualmente atractivas con grados dinámicos ligeramente menores. El sistema final utiliza una configuración optimizada para 3 pasos de eliminación de ruido (NFEs) para equilibrar calidad y velocidad.

4. Optimización del Decodificador

Los VAE de video estándar utilizan convoluciones 3D causales que son computacionalmente costosas.

• Modificación: Los autores construyen sobre el decodificador eficiente LightX2V pero modifican su arquitectura para extender la ventana de mirada temporal hacia atrás (look-back window).
• Mecanismo: En lugar de condicionarse en los cuadros pasados inmediatos, la predicción del cuadro $F_t$ se condiciona en un contexto temporal más largo (hasta $F_{t-4}$).
• Resultado: Esto reduce los artefactos temporales y el parpadeo (flickering) mientras mantiene la eficiencia de las convoluciones 2D, mejorando significativamente la calidad de reconstrucción (PSNR y VBench) en comparación con otros decodificadores eficientes.

5. Pipeline de Cuantización

Para permitir el despliegue en el Snapdragon® 8 Gen 5 NPU, el modelo se somete a un proceso de cuantización de dos etapas:

1. Cuantización Post-Entrenamiento (PTQ): Comprime los pesos a 8 bits.
2. Entrenamiento Consciente de la Cuantización (QAT): Un proceso de destilación de dos etapas (primero por bloques y luego de extremo a extremo) alinea el modelo cuantizado con el maestro de precisión completa, recuperando la precisión perdida durante el PTQ.

Resultados Clave

MobileWan logra las siguientes métricas de rendimiento en un dispositivo móvil (Snapdragon® 8 Gen 5 NPU):

• Especificaciones de Generación: Genera videos de 5 segundos a una resolución de 480×832, 16 FPS.
• Latencia: Latencia de extremo a extremo de 20 segundos (incluyendo 6.65 segundos para el paso de eliminación de ruido del DiT).
• Calidad: Logra una puntuación VBench de 83.79, que es comparable al modelo de escala de servidor Wan2.2 (83.12) y supera significativamente al estado del arte móvil anterior (Neodragon: 81.61).
• Memoria: La reformulación recurrente reduce el requerimiento de pico de memoria de "Fuera de Memoria" (OOM) en móvil a 9.9 GB para el DiT y 1.9 GB para el decodificador.
• Estudio de Usuario: En un estudio de usuario ciego, los participantes prefirieron MobileWan sobre Neodragon el 80% de las veces. En comparaciones contra el Wan2.2 de servidor ajustado, MobileWan no fue preferido en el 53% de los casos, pero sí fue preferido en el 22%, indicando que es competitivo con el modelo de servidor en la mitad de los casos.

Significancia y Reivindicaciones

El artículo afirma cerrar la brecha de calidad entre la generación de video móvil y la de escala de servidor desafiando la suposición de que la generación de video móvil requiere modelos pequeños.

• Escalabilidad: Demuestra que un video DiT de 5 mil millones de parámetros puede desplegarse nativamente en hardware móvil comercial mediante compresión fina y reformulación recurrente.
• Eficiencia vs. Calidad: El trabajo prueba que se pueden aplicar optimizaciones agresivas (poda, recurrencia, destilación) sin una pérdida extrema de calidad, logrando una puntuación VBench (83.79) que rivaliza con los modelos de servidor.
• Integración de Sistema: MobileWan se presenta como el primer modelo de difusión de video de escala de 5B desplegable en un dispositivo móvil comercial, estableciendo un nuevo estado del arte en la generación de video móvil.

Los autores reconocen limitaciones, incluyendo una calidad subóptima para ciertos tipos de escenas (ej. rostos humanos a media distancia), sobresaturación ocasional o reducción de movimiento dependiendo del método de destilación, y posibles discontinuidades temporales o parpadeo en algunos videos generados. Pretenden abordar esto en trabajos futuros.