Resumen Técnico: TMPO para Alineación Diversa y Eficiente de Difusión

1. Declaración del Problema

El Aprendizaje por Refuerzo (RL) ha demostrado un potencial significativo en la alineación de modelos de difusión y de ajuste de flujo con tareas posteriores. Sin embargo, los métodos existentes dependen predominantemente de la maximización de recompensas escalares, lo que los autores identifican como inherentemente búsqueda de modos. Este objetivo de optimización hace que los modelos concentren la masa de probabilidad en pocas trayectorias de alta recompensa, lo que conduce a dos modos de fallo críticos:

1. Degradación de la Diversidad Generativa: El modelo colapsa en un conjunto estrecho de composiciones y estilos (colapso de modo visual), ignorando la distribución anisotrópica de salidas plausibles aprendida durante el preentrenamiento.
2. Amplificación de Recompensas No Confiables: Al reforzar en exceso características específicas que favorecen recompensas proxy (por ejemplo, precisión de OCR o puntuaciones estéticas específicas), el modelo amplifica atributos espurios, a menudo a expensas de la calidad general de la generación.

Las estrategias de mitigación actuales, como la regularización KL, no logran alterar la naturaleza fundamental de búsqueda de modos del objetivo, mientras que los enfoques existentes de Red de Flujo Generativo (GFlowNet) suelen depender de estimaciones de flujo de estado intratables o restricciones locales que acumulan errores a lo largo de horizontes de eliminación de ruido prolongados.

2. Metodología: Optimización de Política de Ajuste de Trayectoria (TMPO)

TMPO cambia fundamentalmente el paradigma de optimización de la maximización de recompensas escalares al ajuste de la distribución de recompensas a nivel de trayectoria.

2.1 Balance de Trayectoria con Softmax (Softmax-TB)

En lugar de maximizar la recompensa esperada, TMPO busca igualar las probabilidades de trayectoria de la política a una distribución objetivo inducida por las recompensas.

• Objetivo: Para un grupo de $K$ trayectorias generadas a partir del mismo prompt, el método iguala las probabilidades de la política $P_\theta(\tau)$ a una distribución de Boltzmann $q(\tau) \propto \exp(\beta R(\tau))$.
• Formulación sin Partición: Al normalizar las probabilidades y las recompensas dentro del grupo de $K$ trayectorias, se elimina la función de partición global intratable. El método define una ventaja Softmax-TB:
  $$A_i = \log q_i - \log p_i = \log \frac{\exp(\beta R_i)}{\sum_j \exp(\beta R_j)} - \log \frac{P_\theta(\tau_i)}{\sum_j P_\theta(\tau_j)}$$
• Garantía Teórica: Los autores demuestran que minimizar este objetivo es equivalente a minimizar la divergencia KL hacia adelante ($D_{KL}(q \parallel p)$) dentro del grupo. A diferencia de la KL inversa (utilizada en los gradientes de política estándar), la KL hacia adelante es cobertura de modos: impone una penalización infinita si la probabilidad de la política $p_i$ se acerca a cero para cualquier trayectoria $i$ que tenga un peso objetivo positivo $q_i$. Esto garantiza teóricamente la preservación de la diversidad sin requerir regularización KL auxiliar.

2.2 Muestreo Estocástico de Árbol Dinámico

Para abordar el costo computacional de muestrear múltiples trayectorias independientes para modelos de ajuste de flujo a gran escala, TMPO introduce una estrategia de muestreo con estructura de árbol:

• Compartición de Prefijos: Las trayectorias comparten un prefijo ODE determinista (pasos de eliminación de ruido comunes) y se ramifican en puntos estocásticos programados dinámicamente.
• Ramificación Dinámica: Los puntos de ramificación no son fijos; siguen un currículo que se desplaza desde regiones de alto ruido tempranas (diversidad de grano grueso) hacia regiones de bajo ruido posteriores (variaciones de grano fino) a medida que avanza el entrenamiento. Las posiciones de las ramas se perturban estocásticamente utilizando una distribución Beta para evitar el sobreajuste a una geometría de árbol fija.
• Eficiencia: Este enfoque reduce el cálculo redundante al compartir el paso hacia adelante para el prefijo compartido, reduciendo significativamente el tiempo de entrenamiento mientras mantiene una exploración efectiva.

3. Contribuciones Clave

1. Identificación de la Causa Raíz: El artículo identifica la naturaleza de búsqueda de modos de la maximización de recompensas como la causa principal del hacking de recompensas y la degradación de la diversidad en el RL de difusión, proponiendo un cambio hacia el ajuste de distribuciones.
2. Objetivo Softmax-TB: Introducción de un objetivo de ajuste de distribución a nivel de trayectoria sin partición que hereda las propiedades de cobertura de modos de la divergencia KL hacia adelante, preservando naturalmente la diversidad generativa.
3. Muestreo Estocástico de Árbol Dinámico: Un mecanismo de muestreo escalable que reduce el tiempo de entrenamiento para modelos a gran escala al compartir prefijos de eliminación de ruido y ajustar dinámicamente los calendarios de ramificación.

4. Resultados Experimentales

Los autores evaluaron TMPO en la base FLUX.1-dev a través de tres tareas de alineación: generación de imágenes composicionales (GenEval), renderizado de texto visual (OCR) y alineación de preferencias humanas (PickScore).

• Diversidad: TMPO mejoró la diversidad generativa en un 9.1% en promedio en comparación con los métodos más avanzados (Flow-GRPO, MixGRPO, TreeGRPO, GARDO). Fue el único método que mantuvo consistentemente puntuaciones positivas de Distancia Media Geométrica Logarítmica (LGMD), indicando una diversidad saludable, mientras que las líneas base a menudo exhibieron LGMD negativa (colapso de modos).
• Rendimiento: TMPO logró un rendimiento competitivo o superior en métricas específicas de la tarea (precisión de GenEval, precisión de OCR) y puntuaciones de preferencia humana (HPS-v2.1, ImageReward, PickScore).
• Eficiencia: Al utilizar la compartición de prefijos y el muestreo de árboles, TMPO redujo el tiempo de entrenamiento por iteración hasta en un 27% en comparación con los métodos existentes.
• Compensaciones: El análisis de Pareto demostró que TMPO logra el equilibrio óptimo entre la maximización de recompensas, la preservación de la diversidad y la eficiencia del entrenamiento.

5. Significado y Afirmaciones

El artículo afirma que TMPO ofrece una solución fundamental al dilema del hacking de recompensas en la alineación de IA generativa. Al reformular el objetivo para igualar distribuciones inducidas por recompensas en lugar de maximizar recompensas escalares, TMPO impone estructuralmente la preservación de la diversidad sin depender de regularización externa o reponderación heurística. El método demuestra que es posible lograr puntuaciones de alineación altas mientras se mantiene la rica distribución de salida multimodal aprendida durante el preentrenamiento, todo ello reduciendo la sobrecarga computacional típicamente asociada con los métodos de RL de múltiples trayectorias. Los autores posicionan a TMPO como un marco escalable adecuado para modelos de ajuste de flujo a gran escala, con extensiones potenciales a la generación de video y la generación de acciones robóticas.