Resumen Técnico: AEGPO – Optimización de Políticas Guiada por Entropía Adaptativa para Modelos de Difusión

1. Planteamiento del Problema

El Aprendizaje por Refuerzo a partir de la Retroalimentación Humana (RLHF), particularmente la Optimización de Política Relativa al Grupo (GRPO), se ha convertido en un paradigma estándar para alinear modelos de difusión y de ajuste de flujo (flow-matching) con las preferencias humanas. Sin embargo, las variantes existentes de GRPO sufren de dos ineficiencias fundamentales debido a su dependencia de heurísticas estáticas:

1. Asignación Ineficiente de Muestras: Los métodos estándar asignan un número uniforme de rollouts a cada prompt, ignorando la variación significativa en el "valor de aprendizaje" entre diferentes muestras. Algunos prompts inducen cambios mínimos en la política, mientras que otros impulsan mejoras sustanciales.
2. Exploración de Pasos Temporales Estática: Las estrategias de exploración (por ejemplo, el ramificado en rollouts basados en árboles) suelen seguir esquemas fijos y predefinidos. Estos fallan al no adaptarse a la naturaleza dinámica del proceso de eliminación de ruido (denoising), donde los momentos críticos para la exploración varían significamente entre prompts y etapas de entrenamiento.

Estas limitaciones resultan en un desperdicio de recursos computacionales en muestras de bajo valor y en una exploración subóptima en pasos temporales no críticos, lo que ralentiza la convergencia y limita el rendimiento final de alineación.

2. Metodología: AEGPO

Los autores proponen la Optimización de Políticas Guiada por Entropía Adaptativa (AEGPO), un marco de trabajo plug-and-play que aprovecha las señales de atención intrínsecas para guiar dinámicamente tanto la selección de muestras como la exploración. La idea central es que la Entropía de Atención sirve como un poderoso proxy de señal dual.

2.1 Señales Centrales

El método define dos métricas intrínsecas derivadas de los mapas de atención del modelo:

• $\Delta$Entropía (Señal a Nivel de Muestra): El cambio relativo en la dispersión de la atención entre la política actual y la política base.
  • Definición: $\Delta\text{Entropía} = \frac{1}{T} \sum_{t=1}^{T} |\text{Entropía}_{\theta}(t) - \text{Entropía}_{\theta_{base}}(t)|$.
  • Función: Actúa como un proxy robusto del valor de aprendizaje de la muestra. Una $\Delta$Entropía alta indica que un prompt obliga al modelo a reorganizar significativamente su estrategia de atención, lo que implica un alto potencial de aprendizaje.
• Entropía($t$) (Señal a Nivel de Paso Temporal): La dispersión absoluta de la atención en un paso de eliminación de ruido específico $t$.
  • Definición: La media de la entropía de Shannon de los pesos de atención a través de los tokens de texto para las características de la imagen en el paso $t$.
  • Función: Los picos en Entropía($t$) identifican pasos temporales de exploración críticos. Estos picos corresponden a momentos de alta incertidumbre o de integración semántica amplia (que suelen ocurrir en las etapas tempranas de estructura y en las etapas tardías de refinamiento), donde la exploración genera la mayor varianza y diversidad de recompensa.

2.2 Estrategia Adaptativa de Doble Nivel

AEGPO utiliza estas señales para reasignar la computación sin aumentar el presupuesto total de rollouts:

1. Asignación Adaptativa Global:

   • El lote de prompts se divide en niveles de "alto valor" y "bajo valor" basándose en la mediana de $\Delta$Entropía.
   • Los prompts de alto valor (alta $\Delta$Entropía) reciben un presupuesto de rollout mayor ($r_{high}$), mientras que los de bajo valor reciben menos rollouts ($r_{low}$).
   • Restricción: $r_{high} + r_{low} = 2r_{avg}$, asegurando que la computación total permanezca constante.

2. Exploración Adaptativa Local:

   • En lugar de esquemas de ramificación fijos, el método identifica dinámicamente los $K$ pasos temporales con los picos de Entropía($t$) más altos para cada rollout.
   • El ramificado (exploración) se realiza exclusivamente en estos pasos temporales de alta dispersión identificados.
   • Esto asegura que la exploración se concentre en los momentos más informativos de la trayectoria de eliminación de ruido.

3. Contribuciones Clave

• Descubrimiento de Señal Dual: El artículo establece la Entropía de Atención como un proxy intrínseco de señal dual, donde el cambio relativo ($\Delta$Entropía) indica el valor de aprendizaje de la muestra y los picos absolutos (Entropía($t$)) identifican los pasos temporales de exploración críticos.
• Marco de Trabajo AEGPO: Introducción de un marco adaptativo de doble nivel que mejora la selección de muestras y la exploración de pasos temporales utilizando mecanismos ligeros guiados por la entropía.
• Generalizabilidad: Demostración de que AEGPO es agnóstico al modelo y al marco de trabajo, mejorando con éxito diversos métodos de optimización de política (DanceGRPO, BranchGRPO, FlowGRPO, DiffusionNFT) y modelos base (FLUX.1-dev, SD3.5-M).

4. Resultados Experimentales

Los experimentos se realizaron en los conjuntos de datos HPDv2.1 y Pick-a-Pic utilizando FLUX.1-dev y SD3.5-M como modelos base.

• Velocidad de Convergencia: AEGPO acelera significativamente la optimización de la política. Logra una convergencia 2$\times$ más rápida en DanceGRPO y una convergencia 5$\times$ más rápida en DiffusionNFT en comparación con las variantes estándar, alcanzando recompensas finales superiores en menos pasos.
• Métricas de Rendimiento:
  • En FLUX.1-dev, BranchGRPO + AEGPO alcanzó un nuevo estado del arte en la puntuación HPS-v2.1 de 0.374 (frente a 0.363) y un Pick Score de 0.232 (frente a 0.229).
  • En SD3.5-M, DiffusionNFT + AEGPO mejoró el HPS-v2.1 a 0.343 e ImageReward a 1.57.
• Compromiso Recompensa-KL (Reward-KL Trade-off): Los modelos AEGPO dominan consistentemente a la línea base en la frontera de Pareto Recompensa-KL, logrando recompensas de alineación significativamente más altas para cualquier nivel dado de desviación de la política (divergencia KL).
• Estudios de Ablación: Tanto la Asignación Global (G) como la Exploración Local (L) proporcionan beneficios independientes, siendo el enfoque combinado (G+L) el que ofrece el mayor rendimiento.
• Sobrecarga Computacional: El método introduce una sobrecarga modesta (un aumento de aproximadamente el 11.1% en el tiempo por paso y un aumento insignificante de VRAM) debido a la extracción de mapas de atención intermedios, lo cual se justifica por la drástica reducción en los pasos totales de entrenamiento requeridos para alcanzar el rendimiento máximo.

5. Significado y Reivindicaciones

El artículo afirma que AEGPO aborda las ineficiencias centrales del muestreo estático en la alineación de modelos de difusión. Al pasar de una asignación basada en heurísticas a una adaptación basada en señales intrínsecas, el método asegura que los recursos computacionales se concentren en las muestras más informativas y en los momentos más críticos del proceso de generación.

Los autores enfatizan que AEGPO es un módulo generalizable y plug-and-play que no requiere nuevas funciones de pérdida ni modelos de recompensa complejos. En su lugar, aprovecha la propia dinámica de atención interna del modelo para guiar la optimización, ofreciendo un camino más eficiente hacia una alineación de alta calidad a través de diversas arquitecturas de difusión y marcos de trabajo de RLHF. El trabajo sugiere que las señales intrínsecas como la entropía de atención son recursos infrautilizados para mejorar la eficiencia de la optimización de políticas en modelos generativos.