Resumen Técnico: Alineación de Modelos Generativos de Pocos Pasos mediante Amortización de Inferencia Variacional Basada en Muestras

Enunciado del Problema

Alinear modelos generativos de pocos pasos (por ejemplo, GANs, modelos de consistencia, mapas de flujo, modelos de deriva) con objetivos específicos, como funciones de recompensa o preferencias humanas, presenta un desafío significativo. Los marcos de alineación existentes suelen depender de supuestos estructurales restrictivos que los generadores de pocos pasos a menudo violan:

1. Verosimilitudes Tratables: Muchos métodos requieren una evaluación explícita de la verosimilitud sobre trayectorias de eliminación de ruido, la cual no está disponible para modelos implícitos como las GANs o mapeos de un solo paso.
2. Solucionadores Específicos: Los enfoques basados en Control Óptimo Estocástico (SOC) o Aprendizaje por Refuerzo (RL) a menudo asumen solucionadores específicos de Ecuaciones Diferenciales Ordinarias (ODE) o Ecuaciones Diferenciales Estocásticas (SDE) y procesos secuenciales de eliminación de ruido.
3. Especificidad del Modelo: Las formulaciones se adaptan frecuentemente a familias de modelos específicas (por ejemplo, modelos de difusión), fallando en generalizar a arquitecturas diversas de pocos pasos como VAEs o mapas de flujo.

En consecuencia, existe una falta de un marco de alineación general que pueda ajustar finamente diversos generadores de pocos pasos preservando su ventaja central: la inferencia rápida y de un solo paso.

Metodología: FAV

Los autores proponen FAV (Alineación de Modelos Generativos de Pocos Pasos mediante Inferencia Variacional Basada en Muestras), un marco general que desacopla la alineación de la familia del generador y las dinámicas de muestreo. FAV formula la alineación como el muestreo desde una distribución inclinada por recompensa anclada a una distribución de referencia:
$$q^*(x) \propto p_{\text{ref}}(x) \exp(\beta \cdot r(x))$$
donde $p_{\text{ref}}$ es una distribución de referencia (por ejemplo, un generador preentrenado o datos empíricos) y $r(x)$ es la función objetivo.

El método procede en tres etapas clave:

1. Aproximación Basada en Muestras mediante SVGD

FAV aproxima la distribución objetivo $q^*$ utilizando Descenso de Gradiente Variacional de Stein (SVGD). En lugar de requerir evaluación de densidad, SVGD representa la distribución con partículas y las transporta iterativamente hacia el objetivo. El campo de velocidad óptimo de Stein $\phi^*$ se deriva para maximizar la disminución de la divergencia KL:
$$\phi^*_{q_\theta, q^*}(x) = \mathbb{E}_{x' \sim q_\theta} \left[ k(x', x) \nabla_{x'} \log p_{\text{ref}}(x') + \beta \cdot k(x', x) \nabla_{x'} r(x') + \nabla_{x'} k(x', x) \right]$$
Este campo comprende tres componentes:

• Alineación Prioritaria: Impulsa las muestras hacia la distribución de referencia $p_{\text{ref}}$.
• Guía de Recompensa: Empuja las muestras hacia regiones de alta recompensa.
• Aplicación de Diversidad: Previene el colapso de modos mediante el término del gradiente del núcleo.

2. Estimación de Puntuación No Paramétrica (KDE)

Un desafío crítico es que la puntuación $\nabla \log p_{\text{ref}}$ a menudo es intratable para datos empíricos o generadores implícitos. FAV aborda esto estimando la puntuación de forma no paramétrica utilizando Estimación de Densidad de Núcleo (KDE). Al convolucionar la distribución de referencia con un núcleo gaussiano, la puntuación se aproxima como un promedio ponderado de diferencias entre muestras de referencia y partículas actuales, evitando efectivamente la necesidad de verosimilitudes explícitas.

3. Amortización mediante Regresión de Punto Fijo

Para preservar la velocidad de inferencia rápida de los modelos de pocos pasos, FAV no realiza actualizaciones iterativas de SVGD en el momento de la inferencia. En su lugar, amortiza el transporte en los parámetros del generador $\theta$ mediante regresión de punto fijo. El generador se entrena para minimizar la distancia entre su salida y el objetivo "transportado":
$$\mathcal{L}(\theta) = \mathbb{E}_{\epsilon} \left[ \| \text{stopgrad}(f_\theta(\epsilon) + \hat{\phi}^*_{q_\theta, q^*_\sigma}(f_\theta(\epsilon))) - f_\theta(\epsilon) \|_2^2 \right]$$
Aquí, la muestra transportada actúa como un objetivo de regresión estable. Esto permite que el generador aprenda a producir muestras alineadas en una sola pasada hacia adelante.

Extensión a Recompensas de Caja Negra: Para recompensas no diferenciables, FAV emplea estimación de gradiente de orden cero (suavizado gaussiano) para aproximar $\nabla r(x)$, permitiendo la alineación incluso cuando la función de recompensa es una caja negra.

Contribuciones Clave

1. Marco de Alineación General: FAV es el primer marco que alinea diversos generadores de pocos pasos (incluyendo VAEs, GANs, modelos de consistencia y mapas de flujo) sin requerir verosimilitudes tratables, solucionadores específicos de ODE/SDE o trayectorias de múltiples pasos.
2. Formulación de IV Basada en Muestras: Al plantear la alineación como muestreo desde una distribución inclinada por recompensa utilizando SVGD y KDE, FAV desacopla el procedimiento de alineación de las dinámicas internas del generador.
3. Preservación de la Eficiencia: Mediante la amortización mediante regresión de punto fijo, FAV mantiene la eficiencia de inferencia de un solo paso del generador base, evitando la sobrecarga computacional de procesos de muestreo iterativo o destilación.
4. Compatibilidad con Caja Negra: La integración de la estimación de gradiente de orden cero permite que FAV optimice recompensas arbitrarias de caja negra.

Resultados Experimentales

Los autores evaluaron FAV en dos dominios principales: manipulación robótica y alineación de generadores de imágenes.

1. Manipulación Robótica (RL Offline y Offline-a-Online)

• Benchmarks: 56 tareas de RL offline (OGBench y D4RL) y 30 tareas de RL offline-a-online.
• Rendimiento: FAV logró un rendimiento de vanguardia en todas las tareas, superando a 10 líneas base que incluyen gradientes de política reparametrizados, Coincidencia Adyunta y políticas generativas destiladas (por ejemplo, FQL, IFQL).
• Hallazgo Clave: FAV demostró capacidades superiores de extracción de políticas y rendimiento de transición desde el aprendizaje offline al online, incluso sin modelos generativos auxiliares o eliminación de ruido de múltiples pasos.

2. Alineación de Generadores de Imágenes

• Escalas: Los experimentos variaron desde ImageNet-256 (256x256) hasta síntesis texto-a-imagen de alta resolución (1024x1024).
• Modelos: Probado en StyleGAN-XL, modelos de deriva, iMeanFlow y SANA-Sprint.
• Métricas: Evaluado frente a puntuaciones estéticas, puntuaciones de preferencia humana (HPS) y detectores de seguridad (NSFW).
• Rendimiento:
  • FAV logró consistentemente la mejor frontera de Pareto entre la optimización de recompensa y la diversidad/calidad de las muestras, mitigando la sobreoptimización de recompensas mejor que líneas base como DRaFT y Coincidencia Adyunta.
  • Eficiencia de Inferencia: FAV generó muestras alineadas aproximadamente 180 a 280 veces más rápido que los métodos de alineación en tiempo de inferencia (por ejemplo, Best-of-N, ReNO) mientras mantenía una calidad de imagen competitiva.
  • Escalabilidad: El método escaló con éxito desde resoluciones de 256x256 hasta 1024x1024 y se generalizó a través de diferentes arquitecturas de modelos.

Significado y Afirmaciones

El artículo afirma que FAV proporciona una solución unificada y agnóstica al modelo para alinear modelos generativos de pocos pasos. Su significado radica en:

• Eliminación de Barreras Estructurales: Elimina la necesidad de verosimilitudes tratables o trayectorias de muestreo específicas, haciendo la alineación accesible a una clase más amplia de modelos generativos implícitos.
• Equilibrio entre Eficiencia y Calidad: Reconcilia con éxito el compromiso entre la generación de alta recompensa y la eficiencia computacional de la inferencia de pocos pasos, un desafío donde los métodos anteriores a menudo degradaban el rendimiento o requerían destilación compleja.
• Robustez: El marco demuestra robustez en diversas tareas, desde control de baja dimensión hasta síntesis de imágenes de alta resolución, y maneja tanto funciones de recompensa diferenciables como de caja negra.

Los autores concluyen que FAV representa un enfoque fuerte y de propósito general para la extracción de políticas en RL y la alineación de recompensas en la generación de imágenes, ofreciendo una alternativa práctica a métodos más complejos y cargados de supuestos.