DyWeight: Dynamic Gradient Weighting for Few-Step Diffusion Sampling

Resumen Técnico: DyWeight

1. El Problema

Los Modelos de Difusión (DMs) han establecido un nuevo estado del arte en la generación de datos, pero su proceso de muestreo es inherentemente lento, requiriendo a menudo cientos de evaluaciones de funciones (NFEs) para generar una muestra de alta calidad.

• Limitación de los Solvers Actuales: Los solvers multi-paso existentes (como DPM-Solver++ o iPNDM) mejoran la eficiencia reutilizando gradientes históricos, pero dependen de coeficientes predefinidos y manuales (handcrafted). Estos coeficientes asumen que los gradientes de desruido en pasos de tiempo cercanos tienen estadísticas similares (estacionariedad).
• La Realidad No Estacionaria: En la práctica, la dinámica de difusión es altamente no estacionaria. Los coeficientes fijos no se adaptan a estos cambios, lo que provoca que los gradientes ruidosos de etapas tempranas dominen o entren en conflicto con los más fiables de etapas posteriores, generando sesgos acumulados y degradando la calidad de la imagen, especialmente en escenarios de pocos pasos (few-step sampling).
• Complejidad de Optimización: Los métodos basados en aprendizaje recientes intentan optimizar estos coeficientes, pero suelen utilizar un paradigma de optimización desacoplado, donde los pesos del solver y la secuencia de tiempo se aprenden por separado. Esto crea un paisaje de optimización complejo que requiere estrategias alternas costosas (búcles de optimización bi-nivel).

2. Metodología: DyWeight

Los autores proponen DyWeight, un solver multi-paso basado en aprendizaje que introduce un paradigma de acoplamiento implícito simplificado.

• Relajación de Restricciones Numéricas:

  • En lugar de forzar a que los pesos de los gradientes históricos sumen 1 (una restricción clásica en métodos como Adams-Bashforth), DyWeight permite pesos no restringidos (unconstrained).
  • Esto permite que el solver aprenda dos cosas simultáneamente:
    1. La parte normalizada: Aprende la combinación óptima de gradientes históricos.
    2. La suma no normalizada: Actúa como un mecanismo de escalado implícito del paso de tiempo, ajustando dinámicamente el tamaño del paso efectivo.

• Calibración del Tiempo (Time Calibration):
  Para corregir la desalineación entre el paso numérico y las expectativas internas de ruido del modelo (especialmente con grandes saltos de tiempo), DyWeight emplea dos mecanismos:

  1. Desplazamiento de Tiempo Implícito (Time Shifting): Utiliza la suma no normalizada de los pesos aprendidos para calcular explícitamente el siguiente paso de tiempo ($\tilde{t}_{n-1}$), ajustando dinámicamente la frontera de integración macro.
  2. Escalado de Tiempo (Time Scaling): Introduce un escalar aprendible por paso ($s_n$) que se aplica al tiempo $t$ antes de ingresarlo a la red neuronal ($\epsilon_\theta(x_{t_n}, s_n t_n)$). Esto corrige el sesgo de exposición y asegura que la red reciba el nivel de ruido correcto.

• Entrenamiento sin Datos (Distillation):

  • Se utiliza un marco de destilación maestro-alumno (teacher-student) sin necesidad de datos de entrenamiento externos.
  • Maestro: Un solver de alta fidelidad (ej. 35 pasos) genera trayectorias de referencia.
  • Alumno (DyWeight): Genera la solución en pocos pasos (NFEs bajos).
  • Supervisión de Punto Final: A diferencia de otros métodos que usan supervisión de trayectoria (forzando al alumno a seguir cada paso del maestro), DyWeight utiliza supervisión de punto final. Solo se minimiza la distancia entre la imagen final generada por el alumno y la del maestro. Esto otorga flexibilidad al solver para encontrar trayectorias más eficientes que se desvían del maestro en el medio pero convergen mejor en el resultado final.

3. Contribuciones Clave

1. Formulación de Pesos Relajada: Reemplaza la parametrización compleja y desacoplada por una formulación de pesos relajada que acopla implícitamente la agregación de gradientes con los ajustes de paso de tiempo.
2. Estrategia de Calibración de Tiempo: Combina el desplazamiento y el escalado de tiempo para mantener la alineación temporal con la dinámica interna del modelo de difusión en regímenes de pocos pasos.
3. Eficiencia y Rendimiento: Logra un nuevo estado del arte (SOTA) en fidelidad visual con significativamente menos NFEs, con una sobrecarga de optimización mínima (entrenamiento rápido y sin modificar el modelo base).

4. Resultados Experimentales

Los autores evaluaron DyWeight en una amplia gama de benchmarks, incluyendo modelos de espacio de píxeles (CIFAR-10, ImageNet, FFHQ) y modelos de espacio latente (Stable Diffusion v1.5, FLUX.1-dev).

• Rendimiento Cuantitativo:

  • En CIFAR-10 con 3 NFEs, DyWeight alcanza un FID de 8.16, superando significativamente a iPNDM (24.55) y a otros solvers aprendibles como EPD-Solver (10.40).
  • En FLUX.1-dev (un modelo de flujo rectificado de última generación), DyWeight supera a DPM-Solver++ e iPNDM en métricas de FID, CLIP Score, HPSv2.1 e ImageReward tanto en 5 como en 7 NFEs.
  • Mantiene un rendimiento superior en LSUN-Bedroom y ImageNet en configuraciones de pocos pasos.

• Estudios de Ablación:

  • La eliminación de la calibración de tiempo (desplazamiento o escalado) degrada drásticamente el rendimiento, confirmando su necesidad.
  • Un orden de paso (K) de 3 o 4 gradientes históricos es óptimo; usar todos los históricos disponibles (K=N) introduce ruido y degrada la fidelidad.
  • La supervisión de punto final supera consistentemente a la supervisión de trayectoria, permitiendo al solver descubrir atajos más eficientes.

• Eficiencia de Entrenamiento:

  • El método converge rápidamente con muy pocos pares maestro-alumno (10k muestras son suficientes).
  • El tiempo de entrenamiento es mínimo (ej. ~3 minutos en CIFAR-10, ~9 minutos en LSUN-Bedroom con GPUs H100).

5. Significado e Impacto

DyWeight representa un avance fundamental en la aceleración de la inferencia de modelos de difusión. Al eliminar la dependencia de coeficientes manuales estáticos y resolver el problema de optimización desacoplado mediante un acoplamiento implícito inteligente, permite generar imágenes de alta calidad en pocos pasos (3-7 NFEs) con una fidelidad superior a los métodos actuales.

Su capacidad para generalizar a modelos complejos como FLUX.1-dev y su bajo costo computacional para el entrenamiento lo convierten en una solución práctica y escalable para aplicaciones en tiempo real e interactivas que requieren generación rápida de imágenes sin sacrificar la calidad visual.