Large Scale Diffusion Distillation via Score-Regularized Continuous-Time Consistency

Este trabajo presenta el primer marco escalable para la destilación de modelos de difusión de tiempo continuo en tareas de imagen y video a gran escala, introduciendo el modelo de consistencia regularizado por puntuación (rCM) que, mediante un kernel JVP compatible con FlashAttention-2 y un regularizador de destilación de puntuación, logra generar muestras de alta fidelidad en 1-4 pasos con una aceleración de 15x a 50x, superando las limitaciones de calidad y diversidad de métodos anteriores como sCM sin necesidad de ajuste GAN.

Lo esencial

Imagina que tienes un chef de cocina maestro (el modelo de IA original) que puede crear platos espectaculares, pero le toma horas cocinar cada uno porque sigue una receta muy lenta y detallada. Ahora, quieres tener un chef aprendiz que pueda crear platos casi idénticos en cuestión de segundos.

Este paper presenta una nueva técnica llamada rCM (Modelo de Consistencia Regularizado por Puntuación) para entrenar a ese "chef aprendiz" de manera mucho más rápida y eficiente, incluso cuando la cocina es gigante (modelos de miles de millones de parámetros) y los platos son videos complejos.

Aquí te explico cómo funciona, usando analogías sencillas:

1. El Problema: El Aprendiz que se pierde en los detalles

Antes de este trabajo, existía un método llamado sCM (Modelo de Consistencia de Tiempo Continuo).

• La analogía: Imagina que le pides al aprendiz que aprenda a cocinar mirando al maestro y tratando de predecir el resultado final de un solo salto.
• El fallo: Funcionaba bien para recetas simples, pero cuando intentaba hacer algo complejo (como escribir texto en una imagen o mantener la coherencia en un video de 5 segundos), el aprendiz se confundía. Los errores se acumulaban como una bola de nieve: al final, el plato salía borroso, con formas extrañas o el texto ilegible.
• La causa: El método original solo miraba hacia "adelante" (hacia el maestro), pero no tenía un mecanismo para corregir sus propios errores en tiempo real cuando la tarea era muy difícil. Además, calcular las matemáticas necesarias para esto en computadoras gigantes era como intentar adivinar el movimiento de cada gota de agua en un tsunami: muy lento y propenso a errores.

2. La Solución: El "Sistema de Doble Chequeo" (rCM)

Los autores crearon rCM, que combina lo mejor de dos mundos para entrenar al aprendiz.

• El Primer Chequeo (Consistencia): Sigue mirando al maestro para aprender la estructura general y la velocidad. Esto asegura que el aprendiz sea rápido y mantenga la diversidad (que no todos los platos se vean iguales).
• El Segundo Chequeo (Regularización por Puntuación): Aquí está la magia. Se le añade un "inspector de calidad" que le dice al aprendiz: "Oye, mira lo que acabas de cocinar tú mismo. ¿Se parece a lo que debería ser?".
  • Si el aprendiz intenta copiar ciegamente al maestro, el inspector lo corrige para que no pierda detalles finos (como el texto o la textura de la piel).
  • Esto evita que el aprendiz se vuelva "aburrido" o repita siempre lo mismo (un problema llamado "colapso de modos", donde todas las imágenes salen iguales).

En resumen: rCM es como tener un entrenador que te dice "hazlo rápido como el maestro" (velocidad) y un crítico que te dice "asegúrate de que los detalles estén perfectos" (calidad).

3. Los Retos Técnicos: La "Cocina Gigante"

El paper no solo inventó la teoría, sino que resolvió problemas de ingeniería masiva:

• El problema: Los modelos modernos son tan grandes (14 mil millones de "ingredientes" o parámetros) que las matemáticas necesarias para el entrenamiento se rompen o son demasiado lentas.
• La solución: Crearon un nuevo "cuchillo de chef" (un kernel de software llamado FlashAttention-2 JVP) que permite hacer estos cálculos complejos en paralelo, como si tuvieras 100 cocineros trabajando a la vez en lugar de uno. Esto permitió entrenar modelos que antes eran imposibles de acelerar.

4. Los Resultados: ¿Qué logramos?

Gracias a esta mezcla de "velocidad del maestro" + "ojo crítico":

• Velocidad extrema: El chef aprendiz puede crear imágenes o videos de alta calidad en 1 a 4 pasos (antes necesitaba 20 o 50). Es como pasar de cocinar una cena en 2 horas a hacerlo en 5 minutos.
• Calidad: Los videos y textos generados son nítidos, sin las distorsiones extrañas que tenía el método anterior.
• Diversidad: A diferencia de otros métodos rápidos que hacían que todos los videos se vieran iguales, este mantiene la variedad creativa.
• Escalabilidad: Funciona en modelos gigantes (como Wan2.1 o Cosmos-Predict2) que generan videos de 5 segundos, algo que nadie había logrado antes con esta tecnología.

Conclusión

Este trabajo es un gran paso para que la IA generativa de video e imágenes sea rápida, barata y de alta calidad. Logró lo que muchos pensaban imposible: tener la velocidad de un rayo sin sacrificar la belleza y el detalle de la obra final, todo gracias a enseñarle al modelo a mirarse al espejo (corregirse a sí mismo) mientras aprende del maestro.

Resumen técnico

1. El Problema

Aunque los Modelos de Consistencia de Tiempo Continuo (como sCM y MeanFlow) son teóricamente sólidos y eficaces para la aceleración de difusión a escala académica, su aplicación a modelos de difusión a gran escala (imágenes y video de alta resolución) ha sido limitada debido a dos desafíos principales:

1. Desafíos de Infraestructura: El cálculo del Producto Jacobiano-Vector (JVP), esencial para el entrenamiento de sCM, es incompatible con las infraestructuras modernas de entrenamiento masivo (como FlashAttention-2, precisión BF16 y paralelismo de contexto/FP). Esto impide escalar estos modelos a más de 10 mil millones de parámetros o a tareas de video de alta dimensionalidad.
2. Limitaciones de Calidad y Diversidad:
   • Acumulación de Errores: sCM sufre de una degradación de calidad en la generación de detalles finos (como texto o texturas complejas) debido a la acumulación de errores numéricos a lo largo de la trayectoria de la EDO del profesor.
   • Naturaleza "Mode-Covering": La divergencia hacia adelante utilizada en sCM tiende a cubrir todos los modos de la distribución, lo que genera alta diversidad pero baja calidad (imágenes borrosas o distorsionadas).
   • Fallas de los Métodos Actuales: Los métodos actuales de estado del arte (como DMD2) utilizan divergencia inversa ("mode-seeking"), lo que mejora la calidad pero a menudo provoca colapso de modos (baja diversidad), haciendo que las generaciones sean repetitivas.

2. Metodología: Score-Regularized Continuous-Time Consistency Model (rCM)

Los autores proponen rCM, un marco que integra la distilación de consistencia de tiempo continuo con la distilación de puntuación (score distillation) para combinar lo mejor de ambos mundos.

A. Infraestructura Escalable (FlashAttention-2 JVP)

Para permitir el entrenamiento de sCM en modelos masivos (>10B parámetros) y video:

• Desarrollaron un kernel personalizado de JVP basado en FlashAttention-2 (usando Triton).
• Este kernel integra el cálculo del JVP directamente en el paso de atención, soportando tanto atención auto-atentiva como cruzada.
• Garantizan compatibilidad con FSDP (Fully Sharded Data Parallel) y CP (Context Parallelism/Ulysses), permitiendo el entrenamiento en modelos de hasta 14 mil millones de parámetros y videos de 5 segundos.

B. Regularización de Puntuación (Score Regularization)

Para abordar la baja calidad de sCM, introducen un término de regularización basado en la distilación de puntuación (DMD):

• Objetivo Híbrido: El modelo optimiza una combinación de dos pérdidas:
  1. Pérdida sCM (Divergencia hacia adelante): Asegura la consistencia en la trayectoria del profesor y mantiene la diversidad.
  2. Pérdida DMD (Divergencia inversa): Actúa como un regularizador de "salto largo" (long-skip) que guía al modelo hacia modos de alta calidad, corrigiendo las distorsiones y el borrosidad.
• Mecanismo: Utilizan una red de puntuación falsa (fake score network) entrenada adversarialmente para estimar la puntuación del estudiante, similar a los métodos GAN, pero sin necesidad de un discriminador complejo o ajuste fino de GAN.

C. Estabilidad Numérica

• Implementan técnicas para estabilizar el cálculo de la derivada temporal en el JVP, que es propenso a inestabilidades en precisión BF16.
• Utilizan derivadas de tiempo semi-continuas (diferencias finitas) para modelos más pequeños y precisión FP32 en las capas de incrustación temporal para modelos masivos y video, evitando el colapso del entrenamiento.

3. Contribuciones Clave

1. Primera Escalabilidad Masiva: Es el primer trabajo que escala la consistencia de tiempo continuo y el JVP a modelos de difusión generales de imágenes y video (hasta 14B parámetros y 5 segundos de video).
2. Resolución del Compromiso Calidad-Diversidad: rCM logra superar la limitación fundamental de los métodos de consistencia (baja calidad) y los métodos de puntuación/adversarios (colapso de modos), ofreciendo simultáneamente alta fidelidad visual y alta diversidad.
3. Infraestructura Eficiente: La implementación del kernel JVP en FlashAttention-2 permite entrenar estos modelos sin sacrificar la eficiencia de memoria o el rendimiento, algo que los métodos anteriores no lograban en escalas industriales.
4. Simplicidad: A diferencia de otros métodos que requieren entrenamiento en múltiples etapas o búsquedas exhaustivas de hiperparámetros, rCM es un enfoque de un solo paso con una configuración robusta.

4. Resultados Experimentales

El modelo fue validado en modelos de vanguardia como Cosmos-Predict2 (T2I) y Wan2.1 (T2V).

• Calidad y Métricas:
  • En tareas de Texto-a-Imagen (T2I), rCM iguala o supera a DMD2 en métricas como GenEval, logrando un puntaje general de 0.83 en el modelo de 14B en solo 4 pasos.
  • En tareas de Texto-a-Video (T2V), rCM supera al modelo profesor Wan2.1 en las métricas de VBench, alcanzando un puntaje total de 85 (vs 83.58 del profesor) con 4 pasos.
• Velocidad de Inferencia:
  • Los modelos destilados generan muestras de alta fidelidad en 1 a 4 pasos.
  • Esto representa una aceleración de 15x a 50x en comparación con los modelos de difusión originales (que requieren ~50 pasos).
  • Se logran resultados competitivos en T2I con 1 paso y en T2V con 2 pasos.
• Diversidad:
  • A diferencia de DMD2, que tiende a colapsar modos (objetos en posiciones/orientaciones similares), rCM mantiene la diversidad de sCM, generando variaciones únicas en cada muestra sin sacrificar la calidad.
  • En videos, rCM resuelve problemas de distorsión temporal y geometría inestable presentes en sCM puro.

5. Significado e Impacto

Este trabajo establece un nuevo paradigma para la distilación de difusión a gran escala:

• Unificación Teórica: Demuestra que combinar divergencias hacia adelante (consistencia) e inversas (puntuación) es una estrategia superior para la generación generativa, superando las limitaciones de los enfoques puros.
• Viabilidad Industrial: Al resolver los problemas de infraestructura del JVP, rCM hace viable la implementación de modelos de generación de video e imágenes ultra-rápidos en entornos de producción reales.
• Futuro: Abre la puerta a la exploración de modelos autoregresivos de video y mundos interactivos, donde la velocidad de inferencia y la diversidad son críticas.

En resumen, rCM es un marco práctico y teóricamente fundamentado que permite la transición de modelos de difusión lentos a generadores de un solo paso de alta calidad y alta diversidad, escalando exitosamente a los modelos más grandes y complejos disponibles actualmente.