TreeGRPO: Tree-Advantage GRPO voor Online RL Post-Training van Diffusiemodellen

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

December 9, 2025

Auteurs: Zheng Ding, Weirui Ye

cs.AI

Samenvatting

Versterkend leren (RL) na de training is cruciaal voor het afstemmen van generatieve modellen op menselijke voorkeuren, maar de buitensporige rekenkosten vormen een grote belemmering voor brede adoptie. Wij introduceren TreeGRPO, een nieuw RL-raamwerk dat de trainings efficiëntie aanzienlijk verbetert door het ontruisingsproces te herformuleren als een zoekboom. Uitgaande van gedeelde initiële ruisvoorbeelden vertakt TreeGRPO strategisch om meerdere kandidaat-trajecten te genereren en hergebruikt tegelijkertijd efficiënt hun gemeenschappelijke voorvoegsels. Deze boomstructuuraanpak biedt drie belangrijke voordelen: (1) Hoge steekproef efficiëntie, waarbij betere prestaties worden behaald met hetzelfde aantal trainingsvoorbeelden; (2) Fijnmazige toekenning van verdienste via beloningsbackpropagatie die stap-specifieke voordelen berekent, waardoor de beperking van uniforme toekenning bij trajectgebaseerde methoden wordt overwonnen; en (3) Geamortiseerde berekening waarbij vertakking naar meerdere kinderen meerdere beleidsupdates per voorwaartse pass mogelijk maakt. Uitgebreide experimenten met zowel op diffusie als op stromen gebaseerde modellen tonen aan dat TreeGRPO een 2,4 keer snellere training bereikt en tegelijkertijd een superieur Pareto-frontier vestigt in de efficiëntie-beloning trade-off ruimte. Onze methode presteert consistent beter dan GRPO-basislijnen across meerdere benchmarks en beloningsmodellen, en biedt zo een schaalbare en effectieve route voor RL-gebaseerde afstemming van visuele generatieve modellen. De projectwebsite is beschikbaar op 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.