Resumen Técnico: Ajuste Fino de Difusión con Destilación de Coincidencia de Momentos con Recompensa (RMMD)

Declaración del Problema

Los modelos de difusión se han convertido en el estándar para la síntesis de imágenes de alta fidelidad y el pronóstico científico debido a sus objetivos de entrenamiento estables. Sin embargo, su despliegue práctico se ve obstaculizado por el alto costo computacional de la inferencia, que requiere docenas o cientos de pasos de eliminación de ruido. Si bien las técnicas de destilación comprimen estos modelos en generadores de pocos pasos, y el ajuste fino mediante Aprendizaje por Refuerzo (RL) los alinea con funciones de recompensa específicas (por ejemplo, preferencias humanas o métricas científicas), la combinación de estas dos fases sigue siendo un desafío.

Los enfoques existentes para fusionar la destilación y la optimización de la recompensa enfrentan limitaciones significativas:

1. Fragilidad en el Entrenamiento Conjunto: Fusionar la destilación y el ajuste fino de la recompensa en una sola fase a menudo provoca que las muestras generadas se desvíen de la distribución de entrada del modelo profesor, invalidando la señal de destilación.
2. Problemas de Memoria y Sesgo: El ajuste fino de un modelo ya destilado mediante retropropagación de múltiples pasos (por ejemplo, DRaFT) es intensivo en memoria. Truncar la retropropagación (por ejemplo, ReFL) introduce sesgo al evaluar las recompensas sobre latentes ruidosos.
3. Restricciones Estructurales: Métodos como HyperNoise evitan la retropropagación en cadena perturbando el ruido inicial, pero están estructuralmente limitados a cambios de baja frecuencia en la imagen, lo que limita su capacidad para optimizar recompensas estructurales complejas de alta frecuencia.

Metodología: Destilación de Coincidencia de Momentos con Recompensa (RMMD)

Los autores proponen RMMD, un marco de dos fases que conecta de manera fundamentada la destilación y el ajuste fino de la recompensa, preservando la "naturalidad" de alta fidelidad de la destilación avanzada.

Fase 1: Destilación de Coincidencia de Momentos (MMD)

El proceso comienza destilando un modelo profesor base utilizando la Destilación de Coincidencia de Momentos (MMD). Este método iguala las distribuciones de eliminación de ruido intermedias a lo largo de la trayectoria de muestreo, produciendo un modelo estudiante de múltiples pasos que sigue de cerca las marginales del profesor. Este estudiante queda congelado para servir como una referencia de distribución estable.

Fase 2: Ajuste Fino de Recompensa On-Policy

En la segunda fase, el modelo estudiante se ajusta para maximizar una función de recompensa mientras mantiene la fidelidad de la distribución. Las innovaciones clave incluyen:

1. Muestreo On-Policy: En lugar de utilizar puntos de datos off-policy, el método muestrea estados ruidosos $x^{pol}_t$ de la política actual del estudiante (ejecutando el estudiante durante $K_{student}$ pasos, re-añadiendo ruido a la salida y deteniendo el gradiente). Esto asegura que el objetivo de recompensa aproxime directamente la recompensa esperada del estudiante y evita que la deriva de la distribución invalide la señal de entrenamiento.
2. Regularización Reciclada: La contribución central es reutilizar la pérdida de Coincidencia de Momentos como un regularizador durante el ajuste fino. En lugar de una penalización heurística, esta pérdida penaliza explícitamente la discrepancia entre las distribuciones de eliminación de ruido intermedias del estudiante y las del modelo de referencia congelado. Esto proporciona un control interpretable (el peso de regularización) para equilibrar la maximización de la recompensa frente a la fidelidad de la distribución.
3. Gradiente de Paso Único: El método captura todo el rango de frecuencias sin retropropagación de múltiples pasos al realizar un único paso de eliminación de ruido sobre muestras on-policy corruptas y evaluar la recompensa.
4. Regularización L2: Un término L2 adicional penaliza desviaciones significativas de las predicciones del modelo inicial destilado por MMD para estabilizar aún más el entrenamiento.

El objetivo total combina el gradiente de recompensa on-policy con el regularizador de coincidencia de momentos y el término L2:
$$\nabla_\theta \mathcal{L}_{RMMD} = \mathbb{E} \left[ \left(\frac{\partial \hat{x}_0}{\partial \theta}\right)^\top \left( m_{student} - m_{teacher} - \lambda \nabla_{\hat{x}_0} R(\hat{x}_0) \right) \right] + \lambda_{reg} \mathcal{L}_{L2}$$

Contribuciones Clave

1. Procedimiento RMMD: Un nuevo método para la destilación de difusión que maximiza simultáneamente una función de recompensa, utilizando una versión on-policy de la pérdida de coincidencia de momentos para la regularización.
2. Superioridad Empírica: Demostró que RMMD supera consistentemente a los baselines de un solo paso (DI++) y a los competidores de múltiples pasos (DRaFT, HyperNoise) en las fronteras de Pareto FID-Recompensa a través de diversas funciones de recompensa, que van desde métricas de nivel de píxel hasta alineación semántica (CLIP).
3. Aplicación Científica: Aplicó con éxito RMMD a GenCast, un modelo de pronóstico meteorológico basado en difusión de vanguardia, optimizando la Puntuación de Probabilidad Clasificada Continuamente (CRPS).

Resultados

Experimentos en ImageNet

• Ventaja de Múltiples Pasos: Utilizar un estudiante destilado por MMD de 8 pasos como punto de partida superó significativamente a los baselines de 1 paso (DI++). El RMMD de 8 pasos logró mejores compromisos en las fronteras de Pareto FID-Recompensa.
• Comparación: RMMD generó mejores fronteras de Pareto que DRaFT y HyperNoise en diversas resoluciones (64x64 y 512x512) y tipos de recompensa (alineación CLIP, Inception Score, suavidad).
• Análisis Cualitativo: Las visualizaciones mostraron que, mientras DRaFT-1 preservaba el contenido pero introducía artefactos adversarios, y DRaFT-2 sufría de "reward hacking" (cambios drásticos en la distribución), RMMD integró con éxito modificaciones sutiles para aumentar las recompensas sin desviarse significativamente de la distribución de datos original.

Pronóstico Meteorológico GenCast

• Aceleración: El modelo destilado logró una aceleración de 7.5× (8 pasos frente a 59 evaluaciones de función) en comparación con el profesor.
• Precisión: El modelo RMMD optimizado con CRPS superó al modelo profesor en el 93% de las variables meteorológicas objetivo (comparado con el 75% del mejor baseline de solo MMD).
• Calibración: RMMD abordó el problema de la sub-dispersión común en la destilación MMD. Aunque la variante on-policy mostró métricas de dispersión ligeramente peores que la variante off-policy, proporcionó las mejores mejoras generales de CRPS.
• Generalización: La optimización del CRPS a un tiempo de adelanto de 12 horas resultó en mejoras consistentes para pronósticos de hasta 7 días, lo que sugiere que el método mejora el modelado de la distribución subyacente en lugar de solo sobreajustarse a la recompensa inmediata.

Significado y Reivindicaciones

El artículo afirma que RMMD proporciona una solución fundamentada al problema de la "deriva" inherente a la destilación y el ajuste fino conjuntos. Al reciclar la pérdida de coincidencia de momentos como un regularizador y utilizar el muestreo on-policy, el método permite la optimización de objetivos complejos y diferenciables (como el CRPS en el pronóstico meteorológico) sin sacrificar la calidad generativa o la estabilidad del modelo destilado.

Los autores enfatizan que este enfoque escala a dominios científicos complejos y de alta dimensión. La aplicación exitosa a GenCast demuestra que los modelos de difusión pueden ser tanto significativamente acelerados como mejorados en sus predicciones probabilísticas, allanando el camino para modelos de difusión más eficientes y precisos en la computación científica donde existen objetivos diferenciables. El trabajo sugiere que mantener una referencia de distribución estable mediante la coincidencia de momentos es crucial para un ajuste fino de recompensa efectivo en los modelos de difusión.