TreeGRPO: 拡散モデルのオンライン強化学習事後学習のためのTree-Advantage 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は、共有された初期ノイズサンプルから出発し、戦略的に分岐して複数の候補軌道を生成すると同時に、それらの共通接頭辞を効率的に再利用する。この木構造アプローチは3つの主要な利点をもたらす：（1）同じ学習サンプル数条件下でより優れた性能を達成する高いサンプル効率、（2）ステップごとのアドバンテージを計算する報酬の逆伝播による細粒度の信用割り当て（軌道ベース手法の均一な信用割り当ての限界を克服）、（3）マルチチャイルド分岐により1回の順伝播で複数のポリシー更新を可能とする償却計算。拡散モデルとフローベースモデル双方における大規模な実験により、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.