TreeGRPO: Tree-Advantage GRPO für das Online-Reinforcement-Learning-Post-Training von Diffusionsmodellen

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

December 9, 2025

papers.authors: Zheng Ding, Weirui Ye

cs.AI

papers.abstract

Reinforcement Learning (RL) Nachtraining ist entscheidend, um generative Modelle an menschliche Präferenzen anzupassen, doch die prohibitiv hohen Rechenkosten bleiben eine Hauptbarriere für die breite Anwendung. Wir stellen TreeGRPO vor, ein neuartiges RL-Framework, das die Trainingseffizienz dramatisch verbessert, indem es den Denoising-Prozess als Suchbaum umformuliert. Aus gemeinsamen initialen Rauschstichproben verzweigt TreeGRPO strategisch, um mehrere Kandidatentrajektorien zu generieren und dabei deren gemeinsame Präfixe effizient wiederzuverwenden. Dieser baumstrukturierte Ansatz bietet drei Schlüsselvorteile: (1) Hohe Probeneffizienz, die bei gleichen Trainingsstichproben eine bessere Leistung erzielt, (2) Feingranulare Kreditzuweisung durch Reward-Backpropagation, die schrittspezifische Advantages berechnet und so die Beschränkung gleichmäßiger Kreditzuweisung trajektorienbasierter Methoden überwindet, und (3) Amortisierte Berechnung, bei der Verzweigungen mit mehreren Kindern mehrere Policy-Updates pro Vorwärtsdurchlauf ermöglichen. Umfangreiche Experimente mit Diffusions- und Flow-basierten Modellen zeigen, dass TreeGRPO eine 2,4-fach schnellere Trainingszeit erreicht und dabei eine überlegene Pareto-Grenze im Effizienz-Reward-Abwägungsraum etabliert. Unsere Methode übertrifft durchgängig GRPO-Baselines über mehrere Benchmarks und Reward-Modelle hinweg und bietet einen skalierbaren und effektiven Weg für RL-basierte Ausrichtung visueller generativer Modelle. Die Projektwebsite ist unter treegrpo.github.io verfügbar.

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.