OP-GRPO: Efficient Off-Policy GRPO for Flow-Matching Models

1. El Problema

Los modelos de Flow-Matching (como Stable Diffusion 3.5 y Wan) han demostrado ser altamente efectivos para la generación de imágenes y videos. Sin embargo, cuando se aplican métodos de alineación basados en Optimización de Política Relativa por Grupos (GRPO) para mejorar la calidad de la generación (alineación con preferencias humanas, consistencia semántica, etc.), surgen dos limitaciones críticas:

• Ineficiencia de Muestreo (On-Policy): GRPO sigue un paradigma estrictamente on-policy. Esto significa que en cada iteración de entrenamiento, el modelo genera nuevas trayectorias (muestras) bajo la política actual y las descarta al final del paso. Incluso las muestras de alta calidad no se reutilizan, lo que resulta en un costo computacional enorme y una convergencia lenta (requiriendo miles de horas de GPU).
• Inestabilidad en el Aprendizaje Off-Policy: Intentar aplicar técnicas off-policy (reutilizar trayectorias antiguas) directamente en Flow-Matching es problemático debido a:
  • Desplazamiento de Distribución: Reutilizar datos de políticas anteriores distorsiona el objetivo de GRPO, comprometiendo la estabilidad.
  • Problemas Numéricos en Pasos Tardíos: A medida que el proceso de denoising avanza hacia regiones de bajo ruido (pasos finales), la distribución condicional se vuelve casi determinista. Esto provoca que los pesos de importancia (importance weights) se vuelvan mal condicionados (extremadamente grandes o pequeños), causando inestabilidad en las actualizaciones de la política.
  • Colapso de Recompensas: En tareas difíciles, si el modelo no genera trayectorias exitosas, las recompensas tienden a cero, colapsando las ventajas y deteniendo el aprendizaje.

2. Metodología: OP-GRPO

Los autores proponen OP-GRPO, el primer marco de trabajo off-policy diseñado específicamente para modelos de Flow-Matching. La metodología se basa en tres pilares fundamentales:

A. Buffer de Replay de Alta Calidad

Se introduce un buffer de memoria ($B_{off}$) para almacenar y reutilizar trayectorias.

• Selección Activa: No se guardan todas las muestras. Se seleccionan activamente las trayectorias con las recompensas más altas dentro de cada grupo.
• Estrategia de Reemplazo: Se compara la nueva trayectoria seleccionada con la de menor recompensa en el buffer. Si es mejor, se reemplaza.
• Decaimiento por Pasos: Se aplica un decaimiento en las recompensas de las trayectorias almacenadas con el tiempo. Esto asegura que el buffer priorice muestras recientes y cercanas a la política actual, mitigando el riesgo de un desplazamiento de distribución excesivo.
• Rollout Híbrido: Durante la fase de generación (rollout), se mezcla un grupo de muestras: la mayoría proviene de la política actual (on-policy), y una parte se reemplaza por muestras recuperadas del buffer (off-policy).

B. Corrección de Muestreo por Importancia a Nivel de Secuencia

Para corregir el desplazamiento de distribución sin romper el mecanismo de clipping (recorte) de GRPO, se propone una corrección a nivel de secuencia:

• El Problema del Enfoque Naive: Si se intenta corregir el desplazamiento dividiendo por la política antigua en cada paso individual, el término de recorte deja de medir la actualización respecto a la política de referencia correcta ($\pi_{old}$), sino respecto a la política de generación antigua ($\pi_{off}$), lo que causa que más del 40% de las muestras sean recortadas innecesariamente.
• La Solución: Se mantiene el ratio de importancia estándar por paso entre la política actual ($\pi_\theta$) y la de referencia ($\pi_{old}$) para preservar la función de recorte. Simultáneamente, se introduce un término de corrección a nivel de secuencia (producto de los ratios entre $\pi_{old}$ y $\pi_{off}$ a lo largo de toda la trayectoria) que se multiplica a la función de pérdida.
• Resultado: Esto reduce drásticamente la tasa de muestras recortadas (de >40% a ~11.8%), permitiendo un uso eficiente de las muestras off-policy manteniendo la estabilidad.

C. Truncamiento de Trayectorias (Trajectory Truncation)

Para abordar la inestabilidad numérica en los pasos finales de denoising:

• Análisis Teórico y Empírico: Se demuestra que en los últimos pasos (bajo ruido), la varianza de la transición tiende a cero, haciendo que los log-probabilidades sean extremadamente sensibles y mal condicionados.
• Estrategia: Se truncan las trayectorias off-policy en un paso $t_{off}$. Los pasos restantes (desde $t_{off}$ hasta 0) se regeneran utilizando la política actual.
• Beneficio: Elimina los términos mal condicionados de la corrección de importancia, estabilizando el entrenamiento sin degradar significativamente la calidad perceptual de la imagen o video final.

3. Contribuciones Clave

1. Primer Marco Off-Policy para Flow-Matching: OP-GRPO es la primera implementación exitosa de GRPO off-policy adaptada a la dinámica de transporte continuo de los modelos de Flow-Matching.
2. Mecanismo de Corrección Híbrido: Introduce una corrección de muestreo por importancia a nivel de secuencia que equilibra la corrección de sesgo con la preservación de la estabilidad del recorte de GRPO.
3. Estrategia de Truncamiento: Identifica y mitiga teóricamente la inestabilidad numérica en los pasos de bajo ruido mediante el truncamiento de trayectorias.
4. Eficiencia Sin Sacrificar Calidad: Demuestra que se puede acelerar el entrenamiento significativamente manteniendo o mejorando la calidad de generación.

4. Resultados Experimentales

Los experimentos se realizaron en modelos de imagen (SD3.5-M) y video (Wan2.1-1.4B) en tareas de generación composicional, renderizado de texto visual y alineación con preferencias humanas.

• Eficiencia de Entrenamiento: OP-GRPO alcanza el rendimiento final de Flow-GRPO (el método on-policy base) utilizando solo el 34.2% de los pasos de entrenamiento en promedio (y hasta un 30.1% en generación de video).
• Rendimiento: En la mayoría de las métricas (GenEval, OCR, PickScore, ImageReward), OP-GRPO iguala o supera a Flow-GRPO, mostrando una mejora promedio del 1.17% en el rendimiento final.
• Estabilidad: El método evita el colapso de gradientes y la inestabilidad observada en intentos previos de off-policy en este dominio.
• Generalización: Los modelos entrenados con OP-GRPO muestran una mejor capacidad de generalización en métricas no vistas durante el entrenamiento y en benchmarks más complejos (T2I-CompBench++).

5. Significado e Impacto

El trabajo de OP-GRPO es significativo porque resuelve el cuello de botella principal en la optimización de modelos generativos modernos basados en Flow-Matching: el costo computacional prohibitivo del entrenamiento on-policy.

• Escalabilidad: Al permitir la reutilización de datos de alta calidad, reduce drásticamente el tiempo y los recursos necesarios para alinear modelos grandes con preferencias humanas, haciendo viable la optimización de RL en modelos de video y de alta resolución.
• Nueva Dirección: Establece un nuevo estándar para cómo manejar el aprendizaje por refuerzo en modelos de difusión/flow-matching, demostrando que el aprendizaje off-policy es posible y superior si se manejan correctamente los desplazamientos de distribución y las inestabilidades numéricas específicas de estos modelos.
• Aplicabilidad: Los hallazgos sugieren que esta metodología podría extenderse a otros dominios de generación (audio, 3D) en el futuro.