TreeGRPO: 확산 모델의 온라인 강화학습 사후 훈련을 위한 트리 어드밴티지 GRPO
TreeGRPO: Tree-Advantage GRPO for Online RL Post-Training of Diffusion Models
December 9, 2025
저자: Zheng Ding, Weirui Ye
cs.AI
초록
강화 학습(RL) 후속 훈련은 생성 모델을 인간의 선호도에 맞추는 데 중요하지만, 과도한 계산 비용으로 인해 널리 채택되는 데 주요 장벽으로 남아 있습니다. 본 논문에서는 디노이징 과정을 검색 트리로 재구성하여 훈련 효율을 극적으로 개선하는 새로운 RL 프레임워크인 TreeGRPO를 소개합니다. TreeGRPO는 공유된 초기 노이즈 샘플에서 출발하여 공통 접두사를 효율적으로 재사용하면서 여러 후보 궤적을 생성하기 위해 전략적으로 분기합니다. 이 트리 구조 접근법은 세 가지 주요 이점을 제공합니다: (1) 동일한 훈련 샘플 하에서 더 나은 성능을 달성하는 높은 샘플 효율성, (2) 단계별 이점을 계산하는 보상 역전파를 통한 세분화된 신용 할당으로 궤적 기반 방법의 균일한 신용 할당 한계를 극복, (3) 다중 자식 분기를 통해 순방향 패스당 여러 정책 업데이트를 가능하게 하는 분할 계산. 확산 모델과 흐름 기반 모델 모두에 대한 광범위한 실험을 통해 TreeGRPO가 효율성-보상 트레이드오프 공간에서 우수한 파레토 최적선을 구축하면서 2.4배 빠른 훈련을 달성함을 입증했습니다. 본 방법론은 여러 벤치마크와 보상 모델에서 GRPO 기준선을 지속적으로 능가하며, RL 기반 시각 생성 모델 정렬을 위한 확장 가능하고 효과적인 경로를 제공합니다. 프로젝트 웹사이트는 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.