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
Аннотация
Доработка моделей методом обучения с подкреплением (ОП) играет ключевую роль в согласовании генеративных моделей с человеческими предпочтениями, однако запредельные вычислительные затраты остаются главным препятствием для их широкого внедрения. Мы представляем TreeGRPO — новую структуру ОП, которая кардинально повышает эффективность обучения, переосмысливая процесс удаления шума как поисковое дерево. Исходя из общих начальных зашумленных сэмплов, TreeGRPO стратегически ветвится, генерируя множество кандидатных траекторий, при этом эффективно повторно используя их общие префиксы. Этот древовидный подход дает три ключевых преимущества: (1) Высокая эффективность использования сэмплов, достигающая лучших результатов при том же объеме обучающих данных; (2) Точное распределение заслуг через обратное распространение вознаграждения, которое вычисляет преимущества для конкретных шагов, преодолевая ограничение равномерного распределения заслуг в методах, основанных на траекториях; и (3) Усреднение вычислений, при котором ветвление с множеством потомков позволяет выполнять несколько обновлений политики за один прямой проход. Многочисленные эксперименты на диффузионных и flow-моделях демонстрируют, что TreeGRPO обеспечивает ускорение обучения в 2,4 раза, одновременно устанавливая превосходящий парето-фронт в пространстве компромисса между эффективностью и вознаграждением. Наш метод стабильно превосходит базовые алгоритмы GRPO в различных тестах и моделях вознаграждения, предлагая масштабируемый и эффективный путь для согласования визуальных генеративных моделей на основе ОП. Сайт проекта доступен по адресу 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.