TreeGRPO: Tree-Advantage GRPO para Pós-Treinamento Online de RL em Modelos de Difusão

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

December 9, 2025

Autores: Zheng Ding, Weirui Ye

cs.AI

Resumo

O pós-treinamento por aprendizagem por reforço (RL) é crucial para alinhar modelos generativos com as preferências humanas, mas o seu custo computacional proibitivo permanece uma grande barreira para a adoção generalizada. Apresentamos o TreeGRPO, uma nova estrutura de RL que melhora drasticamente a eficiência do treinamento ao reformular o processo de remoção de ruído como uma árvore de busca. A partir de amostras de ruído inicial compartilhadas, o TreeGRPO ramifica-se estrategicamente para gerar múltiplos percursos candidatos, reutilizando eficientemente os seus prefixos comuns. Esta abordagem estruturada em árvore oferece três vantagens principais: (1) Alta eficiência de amostragem, alcançando melhor desempenho com a mesma quantidade de amostras de treinamento; (2) Atribuição de crédito refinada através da retropropagação de recompensas que calcula vantagens específicas por passo, superando a limitação de atribuição uniforme de crédito dos métodos baseados em trajetória; e (3) Computação amortizada, onde a ramificação com múltiplos filhos permite várias atualizações da política por passagem forward. Experimentos extensivos em modelos baseados em difusão e fluxo demonstram que o TreeGRPO alcança um treinamento 2,4 vezes mais rápido, estabelecendo simultaneamente uma fronteira de Pareto superior no espaço de compromisso eficiência-recompensa. O nosso método supera consistentemente as linhas de base do GRPO em múltiplos benchmarks e modelos de recompensa, fornecendo um caminho escalável e eficaz para o alinhamento de modelos generativos visuais baseado em RL. O site do projeto está disponível em 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.