TreeGRPO: Tree-Advantage GRPO per l'Addestramento Online RL Post-Training di Modelli di Diffusione

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

December 9, 2025

Autori: Zheng Ding, Weirui Ye

cs.AI

Abstract

L'addestramento post-rinforzo (RL) è cruciale per allineare i modelli generativi con le preferenze umane, ma il suo costo computazionale proibitivo rimane una barriera principale alla diffusione su larga scala. Introduciamo TreeGRPO, un nuovo framework RL che migliora drasticamente l'efficienza dell'addestramento riconfigurando il processo di rimozione del rumore come un albero di ricerca. Partendo da campioni di rumore iniziali condivisi, TreeGRPO si ramifica strategicamente per generare multiple traiettorie candidate, riutilizzando efficientemente i loro prefissi comuni. Questo approccio ad albero offre tre vantaggi chiave: (1) Alta efficienza campionaria, ottenendo prestazioni migliori a parità di campioni di addestramento; (2) Assegnazione del credito granulare tramite backpropagazione della ricompensa che calcola vantaggi specifici per ogni passo, superando la limitazione dell'assegnazione uniforme del credito dei metodi basati su traiettoria; e (3) Computazione ammortizzata in cui la ramificazione multi-figlio abilita aggiornamenti multipli della politica per ogni passaggio in avanti. Esperimenti estesi su modelli basati sia su diffusione che su flussi dimostrano che TreeGRPO raggiunge un addestramento 2,4 volte più veloce stabilendo al contempo una frontiera di Pareto superiore nello spazio di compromesso efficienza-ricompensa. Il nostro metodo supera costantemente i baseline GRPO su molteplici benchmark e modelli di ricompensa, fornendo un percorso scalabile ed efficace per l'allineamento dei modelli generativi visivi basato su RL. Il sito web del progetto è disponibile all'indirizzo 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.