TreeGRPO: Tree-Advantage GRPO para el Entrenamiento Posterior en Línea con Aprendizaje por Refuerzo de Modelos de Difusión

TreeGRPO: Tree-Advantage GRPO for Online RL Post-Training of Diffusion Models

December 9, 2025

Autores: Zheng Ding, Weirui Ye

cs.AI

Resumen

El post-entrenamiento por aprendizaje por refuerzo (RL) es crucial para alinear los modelos generativos con las preferencias humanas, pero su costo computacional prohibitivo sigue siendo una barrera importante para su adopción generalizada. Presentamos TreeGRPO, un novedoso marco de RL que mejora drásticamente la eficiencia del entrenamiento al reformular el proceso de eliminación de ruido como un árbol de búsqueda. A partir de muestras de ruido iniciales compartidas, TreeGRPO se ramifica estratégicamente para generar múltiples trayectorias candidatas mientras reutiliza eficientemente sus prefijos comunes. Este enfoque de estructura arbórea ofrece tres ventajas clave: (1) Alta eficiencia muestral, logrando un mejor rendimiento con la misma cantidad de muestras de entrenamiento; (2) Asignación de crédito de grano fino mediante retropropagación de recompensas que calcula ventajas específicas por paso, superando la limitación de asignación uniforme de crédito de los métodos basados en trayectorias; y (3) Cómputo amortizado, donde la ramificación de múltiples hijos permite múltiples actualizaciones de la política por pasada hacia adelante. Experimentos exhaustivos en modelos basados en difusión y flujos demuestran que TreeGRPO logra un entrenamiento 2.4 veces más rápido mientras establece un frente de Pareto superior en el espacio de compensación eficiencia-recompensa. Nuestro método supera consistentemente los baselines de GRPO en múltiples benchmarks y modelos de recompensa, proporcionando una vía escalable y efectiva para la alineación de modelos generativos visuales basada en RL. El sitio web del proyecto está disponible en treegrpo.github.io.

English

Reinforcement learning (RL) post-training is crucial for aligning generative models with human preferences, but its prohibitive computational cost remains a major barrier to widespread adoption. We introduce TreeGRPO, a novel RL framework that dramatically improves training efficiency by recasting the denoising process as a search tree. From shared initial noise samples, TreeGRPO strategically branches to generate multiple candidate trajectories while efficiently reusing their common prefixes. This tree-structured approach delivers three key advantages: (1) High sample efficiency, achieving better performance under same training samples (2) Fine-grained credit assignment via reward backpropagation that computes step-specific advantages, overcoming the uniform credit assignment limitation of trajectory-based methods, and (3) Amortized computation where multi-child branching enables multiple policy updates per forward pass. Extensive experiments on both diffusion and flow-based models demonstrate that TreeGRPO achieves 2.4times faster training while establishing a superior Pareto frontier in the efficiency-reward trade-off space. Our method consistently outperforms GRPO baselines across multiple benchmarks and reward models, providing a scalable and effective pathway for RL-based visual generative model alignment. The project website is available at treegrpo.github.io.