ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule

Resumen Técnico: ART para Muestreo de Difusión

Enunciado del Problema
Los modelos de difusión basados en puntuación generan muestras integrando una dinámica de tiempo inverso aprendida. Una elección numérica crítica en este proceso es la discretización del horizonte temporal en una cuadrícula de pasos de tiempo. Los enfoques estándar suelen emplear cuadrículas uniformes o horarios diseñados a mano (por ejemplo, EDM, DPM-Solver). Los autores argumentan que estos horarios fijos a menudo son subóptimos dado un presupuesto fijo de evaluaciones de función (NFE). Las cuadrículas uniformes asignan el esfuerzo computacional de manera uniforme, incluso cuando la dinámica inversa varía significativamente a lo largo de la trayectoria, mientras que las cuadrículas diseñadas a mano pueden no adaptarse a distribuciones de datos específicas o a complejidades de la función de puntuación. El objetivo es redistribuir sistemáticamente el esfuerzo computacional a lo largo de la trayectoria de muestreo para minimizar el error agregado de discretización de Euler.

Metodología
El artículo introduce Tiempo Reparametrizado Adaptativo (ART), un marco de teoría de control para la programación óptima de pasos de tiempo.

1. Formulación de Control (ART):

   • En lugar de seleccionar puntos de cuadrícula directamente, ART trata la "velocidad del reloj" del proceso de difusión como una variable de control.
   • Introduce una variable de tiempo reparametrizada $t$ mapeada al tiempo de difusión original $\tau$ mediante una aplicación continua $\psi(t)$.
   • El control $\theta(t) = \dot{\psi}(t)$ representa la tasa de cambio instantánea del tiempo de difusión.
   • El objetivo es elegir $\theta(t)$ para minimizar el error local agregado de discretización de Euler, sujeto a la restricción de que el presupuesto total de tiempo se preserve ($\int_0^T \theta(t) dt = T$).
   • El proxy de error local se deriva como proporcional a $|Q(x, \psi)| \theta(t)^2$, donde $Q$ depende de la función de puntuación y sus derivadas.

2. Solución de Aprendizaje por Refuerzo (ART-RL):

   • El problema determinista de ART es intratable de resolver directamente debido al espacio de estados de alta dimensión y la falta de una solución de forma cerrada para la ecuación de Hamilton-Jacobi-Bellman (HJB).
   • Los autores proponen ART-RL, un enfoque de aprendizaje por refuerzo de tiempo continuo (CTRL).
   • Introducen un problema auxiliar aleatorizado donde el control $\theta$ se muestrea de una política gaussiana $\pi(\cdot | t, x, \psi) = \mathcal{N}(\mu, \lambda/|Q|)$.
   • Esta formulación permite el uso de algoritmos actor-crítico (adaptando específicamente el marco de Jia y Zhou [14]) para aprender la media de la política $\mu$ y la función de valor mediante aproximación estocástica.

3. Puente Teórico:

   • Una contribución teórica central es el establecimiento de un puente bidireccional entre el óptimo determinista de ART y la política gaussiana óptima en el problema auxiliar de ART-RL.
   • Suficiencia: El control óptimo determinista $\mu^*$ se eleva a una política gaussiana óptima cuya media es exactamente $\mu^*$.
   • Necesidad: Cualquier política gaussiana óptima para el problema auxiliar debe tener una media que recupere el control determinista de ART $\mu^*$.
   • Este resultado justifica parametrizar al actor como una media gaussiana y desplegarlo de manera determinista después del entrenamiento, proporcionando garantías rigurosas de que el horario aprendido es el óptimo verdadero para el problema de control original.

4. Destilación:

   • Después del entrenamiento fuera de línea, la política estocástica se destila en una cuadrícula determinista basada solo en el tiempo. Esto elimina la necesidad de evaluaciones de la red del actor o cálculos de varianza por paso durante la inferencia, asegurando un costo computacional adicional cero en comparación con los programadores estándar.

Contribuciones Clave

• Formulación: Propone ART, el primer marco de teoría de control para la programación de pasos de tiempo de difusión mediante reparametrización de tiempo adaptativa.
• Teoría: Establece la equivalencia bidireccional entre un problema de control óptimo determinista y un procedimiento de aprendizaje actor-crítico de tiempo continuo con políticas gaussianas. Esto transforma aplicaciones heurísticas de RL en una ruta principista para encontrar óptimos deterministas.
• Algoritmo: Desarrolla ART-RL, un algoritmo actor-crítico que aprende horarios de tiempo adaptativos basados en datos.
• Generalización: Demuestra que el horario aprendido puede ser destilado y transferido entre diferentes conjuntos de datos y conteos de pasos de tiempo sin reentrenamiento.

Resultados Experimentales
Los autores evalúan ART-RL dentro de la tubería oficial de EDM en CIFAR-10, AFHQv2, FFHQ e ImageNet.

• CIFAR-10: ART-RL supera consistentemente a los horarios Uniforme, log-SNR de DPM-Solver y EDM estándar en una amplia gama de presupuestos de NFE.
  • Con presupuestos bajos (por ejemplo, NFE=3), ART-RL produce objetos reconocibles donde las líneas base permanecen cerca del ruido.
  • Con NFE=35, ART-RL alcanza un FID de 1.82, mejorando a EDM (1.85) y DPM-logSNR (1.89).
  • La significancia estadística se confirma mediante múltiples semillas, donde la peor ejecución de ART-RL supera a la mejor de EDM.
• Generalización:
  • Interpolación/Extrapolación: Un horario aprendido en $K=18$ pasos en CIFAR-10 mantiene su ventaja de rendimiento cuando se interpola o extrapola a otros conteos de pasos ($K \in \{4, 9, 12, 15, 20\}$).
  • Transferencia entre Conjuntos de Datos: El horario destilado de CIFAR-10 se aplica a AFHQv2, FFHQ e ImageNet sin reentrenamiento. Mejora consistentemente el FID sobre EDM y DPM-logSNR, particularmente en NFEs bajos y moderados.
• Análisis Unidimensional: En un escenario con una función de puntuación analítica, ART-RL reduce significativamente la distancia de Wasserstein ($W_2$) en comparación con horarios diseñados a mano, confirmando que la mejora proviene de una mejor asignación de tiempo en lugar de errores de estimación de puntuación.

Significado y Afirmaciones
El artículo afirma ser el primero en formular la programación de pasos de tiempo de difusión como un problema de control mediante reparametrización de tiempo adaptativa y el primero en proporcionar un puente bidireccional riguroso entre un óptimo determinista y un procedimiento de aprendizaje actor-crítico.

Los autores enfatizan que su método es basado en datos y amortizado:

• Basado en datos: El horario se aprende a partir de la función objetivo y la dinámica de los datos, no se diseña a mano.
• Transferencia Amortizada: El costo de entrenamiento fuera de línea (aproximadamente 1–2 horas en una Colab T4 para CIFAR-10) es un gasto único. La cuadrícula determinista destilada resultante es reutilizable en diferentes conjuntos de datos y conteos de pasos con ningún costo de inferencia adicional, igualando el tiempo de ejecución de programadores estándar como Uniforme o EDM.

El trabajo posiciona la reparametrización de tiempo adaptativa como una dirección incipiente pero prometedora para el diseño sistemático de horarios en modelos de difusión generativos, ofreciendo un punto de partida de teoría de control para futuras investigaciones.