Collaborative Few-Step Distillation and Low-Bit Quantization for Wan2.2 Dual-Expert Video Diffusion Models

Resumen Técnico: Destilación Colaborativa de Pocos Pasos y Cuantización de Bajo Bit para Modelos de Difusión de Video de Doble Experto Wan2.2

1. Planteamiento del Problema

Los grandes modelos de difusión de video, como el Wan2.2-T2V-A14B, logran una alta calidad visual pero enfrentan barreras de despliegue significativas debido a dos cuellos de botella primarios:

1. Costo Computacional: El proceso iterativo de eliminación de ruido requiere un gran número de pasos de muestreo, lo que hace que la inferencia sea costosa en términos de latencia.
2. Huella de Memoria: Los modelos poseen un recuento de parámetros residentes elevado (aprox. 27B en total, con ~14B activos por paso en la arquitectura de dual-experto), lo que excede los presupuestos de memoria de muchos escenarios de despliegue.

Las estrategias de compresión existentes suelen fallar cuando se aplican directamente a estos modelos. Simplemente reducir los pasos de muestreo sin reentrenamiento conduce a una mala calidad, mientras que cuantizar la trayectoria original de pasos largos suele introducir desajustes en la distribución de las activaciones durante la inferencia, causando la propagación de errores temporales y la degradación de la fidelidad del video. Además, Wan2.2 utiliza una arquitectura de Mezcla de Expertos (MoE) de doble experto (experto de alto ruido y experto de bajo ruido), lo que complica los enfoques de cuantización estándar que tratan al modelo como un tronco denso.

2. Metodología

El artículo propone un pipeline de compresión orientado al despliegue que integra la destilación de pocos pasos y la cuantización de bajo bit de manera codiseñada. El pipeline sigue una secuencia específica: adaptación del modelo $\rightarrow$ destilación de pocos pasos $\rightarrow$ cuantización de bajo bit $\rightarrow$ construcción de calibración $\rightarrow$ generación $\rightarrow$ evaluación.

Componentes Técnicos Clave:

• Gestión de la Arquitectura Consciente de Dual-Experto:

  • El pipeline respeta el mecanismo de enrutamiento específico de Wan2.2, donde un experto de alto ruido maneja el diseño global/movimiento y un experto de bajo ruido refina la textura/detalles.
  • Calibración por Ramas: En lugar de agrupar las activaciones, el método captura y calibra las estadísticas por separado para las ramas de alto y bajo ruido. Esto evita que el cuantizador mezcle distribuciones de activación incompatibles.
  • Protección de Capas de Entrada: Las capas iniciales y los bloques sensibles al enrutamiento se mantienen en una mayor precisión o se calibran de forma más agresiva. Esto es crítico porque los errores en la fase inicial de alto ruido (inicialización de diseño/movimiento) se propagan a través de toda la secuencia, mientras que los errores posteriores suelen ser locales.

• Destilación de Pocos Pasos mediante Coincidencia de Distribución:

  • El método emplea la Destilación de Coincidencia de Distribución (DMD) para aprender una trayectoria de eliminación de ruido más corta en lugar de simplemente truncar el muestreador.
  • El modelo estudiante se optimiza para coincidir con la distribución del profesor a través del número reducido de pasos, con un objetivo consciente de la etapa ($L_{distill} = \lambda_1 L_{high} + \lambda_2 L_{low}$) que respeta los distintos roles de las fases de alto y bajo ruido.

• Cuantización Codiseñada:

  • Calibración Alineada con la Trayectoria: Una innovación central es realizar la calibración de la cuantización en el estudiante destilado de pocos pasos en lugar del modelo original de pasos largos. Esto asegura que el cuantizador observe las estadísticas de activación que encontrará durante la inferencia real, reduciendo el desajuste de la distribución.
  • Representación HiF4 (HiFloat4): El pipeline utiliza un formato de punto flotante de bloque de 4 bits (HiF4) diseñado para la inferencia de bajo bit. Este formato cuenta con una jerarquía de escalado de tres niveles (Global E6M2, Nivel-2 E1, Nivel-3 E1) y una amplitud con signo S1P2. Este diseño proporciona un amplio rango dinámico global (aprox. $2^{-50}$ a $2^{18}$) para manejar rangos de activación heterogéneos a través de las fases de eliminación de ruido, preservando al mismo tiempo el detalle local mediante micro-exponentes.
  • Granularidad: Los pesos se cuantizan por canal, y las activaciones se cuantizan por token para rastrear las variaciones a nivel de prompt y de fotograma.

• Ingeniería y Evaluación:

  • Gestión de Memoria: Para evitar errores de falta de memoria (OOM) en GPUs individuales (ej. H100 de 80GB), la calibración y el reemplazo se realizan capa por capa, con búferes auxiliares cargados bajo demanda.
  • Evaluación de Ciclo Cerrado: El sistema utiliza VBench para evaluar cinco dimensiones específicas: Consistencia del Sujeto (SC), Calidad Estética (AQ), Calidad de Imagen (IQ), Consistencia General (OC) y Suavidad de Movimiento (MS). El protocolo de evaluación asegura una comparación justa utilizando prompts, resoluciones y esquemas de muestreo idénticos para las líneas base.

3. Contribuciones Clave

1. Codiseño de Destilación y Cuantización: El artículo demuestra que calibrar la cuantización en la trayectoria destilada de pocos pasos (en lugar del modelo original de pasos largos) reduce significativamente el desajuste de la distribución de activación, permitiendo que el modelo comprimido mantenga una calidad cercana a la del estudiante de precisión completa.
2. Estrategia de Compresión Consciente de MoE: El método introduce un esquema de calibración por ramas y protección de capas de entrada diseñado específicamente para la estructura de dual-experto de Wan2.2, evitando la mezcla de estadísticas de activación de alto y bajo ruido.
3. Integración de HiF4 para Difusión de Video: Se demuestra que la aplicación del escalado jerárquico de HiF4 es efectiva para los modelos de difusión de video, equilibrando la necesidad de un amplio rango global (para evitar el desbordamiento en las fases de alto ruido) con la precisión local (para preservar la textura en las fases de bajo ruido).
4. Validación Empírica de Recuentos de Pasos: El estudio evalúa sistemáticamente las configuraciones de 4, 8, 20 y 40 pasos, identificando el punto de equilibrio óptimo entre calidad y eficiencia.

4. Resultados Experimentales

Los experimentos se realizaron en el modelo Wan2.2-T2V-A14B utilizando el conjunto de datos OpenS2V para la calibración y VBench para la evaluación.

• Rendimiento vs. Líneas Base:

  • A 4 pasos, el modelo cuantizado mantiene un error pequeño respecto al modelo de precisión completa de los mismos pasos, pero queda por detrás de la línea base original de pasos largos.
  • A 8 pasos, el modelo cuantizado (Puntuación Promedio: 0.7002) comienza a superar a la línea base original de precisión completa (0.6976).
  • A 20 pasos, el modelo cuantizado logra el mejor rendimiento (Puntuación Promedio: 0.7074), superando tanto a la línea base original de precisión completa como al modelo destilado de precisión completa de los mismos pasos (0.7051).
  • A 40 pasos, el rendimiento se estanca, lo que indica rendimientos decrecientes al aumentar la longitud de la trayectoria.

• Análisis de Métricas (Configuración de 20 Pasos):

  • El modelo cuantizado mejoró respecto a la línea base original en Consistencia del Sujeto (0.9629 vs. 0.9495), Calidad Estética, Calidad de Imagen y Suavidad de Movimiento.
  • La Consistencia General se mantuvo comparable a la del modelo de los mismos pasos.
  • Los resultados indican que la calibración alineada con la distribución recupera con éxito las pérdidas de bajo bit sin sacrificar la estabilidad temporal.

• Sensibilidad al Error:

  • El análisis revela que las perturbaciones de un solo fotograma afectan principalmente a la Calidad de Imagen y la Calidad Estética, mientras que los errores acumulados a través de los pasos se manifiestan como Deriva del Sujeto y Vibración de Movimiento (Motion Jitter). La protección de las capas de entrada propuesta mitiga eficazmente esto último.

5. Significado y Reivindicaciones

El artículo sostiene que el pipeline propuesto ofrece un camino práctico para desplegar modelos de difusión de video de alta calidad bajo presupuestos limitados de memoria GPU y latencia.

• Compromiso Calidad-Eficiencia: Los autores identifican la configuración de 20 pasos como la configuración óptima, proporcionando suficiente capacidad de eliminación de ruido para recuperar la calidad temporal y visual, requiriendo significativamente menos pasos que los muestreadores convencionales.
• Cambio Estratégico: El trabajo argumenta que la cuantización no debe ser un "añadido con pérdida" (lossy afterthought), sino una parte integral de la estrategia de aceleración. Al codiseñar la destilación y la cuantización, el método convierte la conversión de bajo bit en un componente estable del pipeline de inferencia.
• Preparación para el Despliegue: El uso de la protección de capas de entrada y la calibración por ramas aborda los modos de fallo específicos de los modelos de video de dual-experto (ej. deriva del sujeto e inestabilidad del movimiento), haciendo que el modelo comprimido sea adecuado para la generación interactiva y el procesamiento por lotes con recursos restringidos.

El artículo concluye que, si bien el pipeline actual se centra en la cuantización post-entrenamiento, el marco establece un bucle de verificación donde los fallos específicos de métricas (ej. caídas en la Consistencia del Sujeto) proporcionan señales de depuración directas para refinar la calibración y los recuentos de pasos.