TreeGRPO : GRPO à avantage arborescent pour l’apprentissage par renforcement en ligne en post-formation des modèles de diffusion

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

December 9, 2025

papers.authors: Zheng Ding, Weirui Ye

cs.AI

papers.abstract

L'apprentissage par renforcement (RL) post-entraînement est crucial pour aligner les modèles génératifs sur les préférences humaines, mais son coût computationnel prohibitif reste un obstacle majeur à son adoption généralisée. Nous présentons TreeGRPO, un nouveau cadre de RL qui améliore considérablement l'efficacité de l'entraînement en reformulant le processus de débruitage comme un arbre de recherche. À partir d'échantillons de bruit initiaux partagés, TreeGRPO se ramifie stratégiquement pour générer de multiples trajectoires candidates tout en réutilisant efficacement leurs préfixes communs. Cette approche arborescente offre trois avantages clés : (1) Une grande efficacité des échantillons, obtenant de meilleures performances avec le même nombre d'échantillons d'entraînement ; (2) Une attribution de crédit fine via une rétropropagation des récompenses qui calcule des avantages spécifiques à chaque étape, surmontant la limitation de l'attribution uniforme du crédit des méthodes basées sur les trajectoires ; et (3) Un calcul amorti où le branchement multi-enfants permet plusieurs mises à jour de la politique par passe avant. Des expériences approfondies sur des modèles à base de diffusion et de flux démontrent que TreeGRPO permet un entraînement 2,4 fois plus rapide tout en établissant une frontière de Pareto supérieure dans l'espace de compromis efficacité-récompense. Notre méthode surpasse constamment les lignes de base GRPO sur plusieurs benchmarks et modèles de récompense, offrant une voie évolutive et efficace pour l'alignement des modèles génératifs visuels par RL. Le site web du projet est disponible à l'adresse 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.