Resumen Técnico: Aprender a Muestrear de Modelos de Difusión mediante Aprendizaje por Refuerzo Inverso

Planteamiento del Problema

Los modelos de difusión se han convertido en un estándar para la generación de datos de alta dimensión, pero el procedimiento de muestreo sigue siendo un cuello de botella crítico. Aunque la red neuronal de eliminación de ruido suele ser fija tras el entrenamiento, el algoritmo de muestreo requiere un ajuste cuidadoso de los hiperparámetros (por ejemplo, esquemas de ruido, escalas de guía, perfiles de estocasticidad). Actualmente, este ajuste se realiza mediante costosas búsquedas empíricas de cuadrícula (grid search). Para modelos complejos como los entrenados en ImageNet-64, el coste computacional de la búsqueda exhaustiva de hiperparámetros puede aproximarse o incluso exceder el coste de entrenar el eliminador de ruido. Además, los enfoques existentes de aprendizaje por refuerzo (RL) o de control para modelos de difusión suelen requerir funciones de recompensa explícitas o el reentrenamiento del eliminador de ruido, lo que contradice el objetivo de optimizar el muestreador sin modificar el generador preentrenado.

Metodología

Los autores proponen enmarcar el muestreo de difusión como un problema de toma de decisiones secuenciales resoluble mediante Aprendizaje por Refuerzo Inverso (IRL). La innovación central es el aprendizaje de una política de muestreo que iguale las distribuciones de estado objetivo sin acceso a trayectorias de expertos, recompensas explícitas o el reentrenamiento del eliminador de ruido.

1. Formulación del MDP

El proceso de muestreo de difusión se formaliza como un Proceso de Decisión de Markov (MDP) de horizonte finito:

• Estado ($s_t$): Definido como $(x_t, \sigma_t)$, donde $x_t$ es la muestra actual y $\sigma_t$ es el nivel de ruido.
• Acción ($a_t$): Decisiones de control discretas aplicadas a la dinámica de muestreo, tales como la selección de escalas de guía, niveles de inyección de ruido o la decisión de volver a añadir ruido (reiniciar) a un nivel de ruido anterior.
• Dinámica: El núcleo de transición está inducido por la dinámica de difusión discretizada y la acción de control seleccionada. El proceso termina cuando el nivel de ruido alcanza $\Sigma_0$ (un estado absorbente).

2. Aprendizaje por Refuerzo Inverso Basado Solo en Estados

A diferencia del IRL tradicional que a menudo asume acciones o trayectorias de expertos, este método opera bajo un régimen de solo estados. El "experto" se define únicamente por las distribuciones marginales de las muestras de datos en niveles de ruido seleccionados.

• Objetivo: El objetivo es encontrar una política $\pi_\theta$ cuya medida de ocupación de estados $\mu_\theta$ (la distribución marginal de los estados visitados) coincida con la medida de ocupación del experto $\mu_E$.
• Optimización: En lugar de aprender una función de recompensa, los autores minimizan directamente la divergencia-$f$ entre las medidas de ocupación del experto y de la política:
  $$\min_\theta D_f(\mu_E \parallel \mu_\theta)$$
  El artículo explora específicamente las divergencias de Kullback-Leibler (KL) y de Kullback-Leibler inversa (rKL), señalando que la elección de la divergencia-$f$ dicta el compromiso entre la cobertura de la moda (diversidad) y la búsqueda de la moda (fidelidad).

3. Gradiente de Política y Optimización

Los autores derivan un teorema de gradiente de política para minimizar las divergencias-$f$ directamente. El gradiente se expresa como:
$$\nabla_\theta L_f(\theta) = \mathbb{E}_{s_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \sum_{t' \ge t} h_f\left(\frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})}\right) \right]$$
donde $h_f$ es una función derivada del generador de la divergencia-$f$.

• Estimación de la Relación de Densidades: Dado que $\mu_E$ y $\mu_\theta$ son desconocidos, la relación de densidades se estima mediante un discriminador entrenado para distinguir entre los estados del experto (datos con ruido) y los estados generados por la política.
• Estabilización: Para manejar la alta varianza en las estimaciones de gradiente, el método emplea un objetivo recortado tipo PPO, correcciones de muestreo de importancia y una línea base de media de lote global.

4. Estrategias de Muestreo Evaluadas

El marco se aplica a tres intervenciones específicas en el tiempo de inferencia:

1. Guía de Clasificación Libre Adaptativa (Adaptive Classifier-Free Guidance): Aprendizaje de una escala de guía dependiente del estado $\omega(x, \sigma)$.
2. Inyección de Estocasticidad: Aprendizaje de un nivel de inyección de ruido dependiente del estado $\gamma_{EDM}(x, \sigma)$.
3. Renoising (Muestreo de Reinicio): Aprendizaje de cuándo revertir el proceso a un nivel de ruido anterior para escapar de malas generaciones locales.

Contribuciones Clave

1. Formulación de MDP para el Muestreo: El artículo plantea el muestreo de difusión como un MDP para adaptar los muestreadores en el tiempo de inferencia sin reentrenar el eliminador de ruido.
2. Igualación Directa de la Marginal de Estado: Deriva un objetivo de gradiente de política que minimiza las divergencias-$f$ entre las medidas de ocupación de los estados, evitando la necesidad de modelos de recompensa explícitos o datos de acciones de expertos.
3. Compromisos Interpretables: El marco expone un compromiso controlable entre la diversidad de la muestra (Recall) y la fidelidad (Precision) a través de la elección de la divergencia-$f$ (KL vs. rKL).
4. Control Dependiente del Estado: Las políticas aprendidas actúan como controles en el tiempo de inferencia que intervienen selectivamente (por ejemplo, aplicando una guía alta solo en niveles de ruido específicos o para tipos de muestras específicos) en lugar de aplicar esquemas globales fijos.

Resultados Experimentales

El método fue evaluado en CIFAR-10, FFHQ e ImageNet-64 utilizando modelos EDM preentrenados.

• Rendimiento: Las políticas aprendidas igualan o superan el rendimiento de las mejores líneas de base ajustadas manualmente (encontradas mediante búsqueda de cuadrícula) en las métricas de FID, Precisión y Recall.
  • Guía Adaptativa: Mejoró el compromiso calidad-diversidad en comparación con la guía fija. La divergencia KL mejoró el Recall (diversidad), mientras que la rKL mejoró la Precisión (fidelidad).
  • Inyección de Estocasticidad: Los perfiles dependientes del estado aprendidos recuperaron los patrones cualitativos de los esquemas óptimos de EDM (baja estocasticidad al final de la trayectoria) y superaron a las líneas de base fijas en conjuntos de datos complejos como ImageNet.
  • Renoising: Redujo significativamente el Número de Evaluaciones de Función (NFE) en ImageNet manteniendo o mejorando el FID en comparación con las estrategias de reinicio estándar.
• Eficiencia Computacional:
  • Entrenamiento: Una única ejecución de entrenamiento reemplaza la búsqueda de cuadrícula exhaustiva. En ImageNet-64, esto reduce el presupuesto de optimización hasta en un 9$\times$ (de ~1437 EFLOPs para una búsqueda de cuadrícula de 3 hiperparámetros a ~161 EFLOPs para el entrenamiento de IRL).
  • Inferencia: La política añade una sobrecarga modesta de entre el 16% y el 40% en FLOPs respecto al eliminador de ruido por sí solo, lo cual se considera insignificante dado que elimina el ajuste manual.
• Hallazgos de Ablación:
  • Condicionamiento: Las políticas condicionadas al estado completo $(x, \sigma)$ superan significativamente a aquellas condicionadas solo al nivel de ruido $\sigma$, y funcionan casi tan bien como las que tienen etiquetas de clase explícitas $(x, \sigma, c)$, lo que sugiere que la política aprende implícitamente características discriminativas de clase a partir del contenido de la imagen.
  • Elección de la Divergencia: Se demostró que KL y rKL son las más efectivas para controlar el compromiso precisión-recall, mientras que las divergencias simétricas (TV, Hellinger) produjeron resultados menos interpretables.

Significado y Reivindicaciones

El artículo afirma que este marco proporciona una alternativa sistemática y automatizada al ajuste manual de hiperparámetros para el muestreo de difusión. Al tratar el muestreo como un problema de IRL, permite la optimización de estrategias de muestreo complejas y dependientes del estado que son difíciles de diseñar manualmente.

Los autores enfatizan que el método es particularmente valioso para modelos a gran escala donde el coste de la evaluación por cada prueba en una búsqueda de cuadrícula es prohibitivo. El enfoque no modifica el modelo generativo en sí, sino que aprende una "política de control" óptima para la trayectoria de inferencia, ofreciendo un camino práctico para mejorar la calidad y diversidad de las muestras sin la carga computacional de un reentrenamiento o una búsqueda exhaustiva. El artículo mantiene la modestia, señalando que los beneficios dependen del conjunto de datos y que el marco actual depende de espacios de acción discretos, lo que sugiere trabajos futuros en espacios de acción continuos y una inicialización más robusta.