TreePO: Преодоление разрыва между оптимизацией политик, эффективностью и производительностью вывода с помощью эвристического древовидного моделирования

TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling

August 24, 2025

Авторы: Yizhi Li, Qingshui Gu, Zhoufutu Wen, Ziniu Li, Tianshun Xing, Shuyue Guo, Tianyu Zheng, Xin Zhou, Xingwei Qu, Wangchunshu Zhou, Zheng Zhang, Wei Shen, Qian Liu, Chenghua Lin, Jian Yang, Ge Zhang, Wenhao Huang

cs.AI

Аннотация

Последние достижения в области выравнивания больших языковых моделей с использованием обучения с подкреплением позволили добиться значительных успехов в решении сложных задач логического рассуждения, однако ценой дорогостоящих on-policy прогонов и ограниченного исследования разнообразных путей рассуждения. В данной работе мы представляем TreePO, включающий самоуправляемый алгоритм прогона, который рассматривает генерацию последовательностей как процесс поиска с древовидной структурой. Состоящий из динамической политики выборки деревьев и декодирования сегментов фиксированной длины, TreePO использует локальную неопределенность для обеспечения дополнительных ветвей. Путем амортизации вычислений на общих префиксах и раннего отсечения низкоценных путей, TreePO существенно снижает вычислительную нагрузку на каждое обновление, сохраняя или даже повышая разнообразие исследования. Ключевые вклады включают: (1) алгоритм пошаговой выборки, который снижает нагрузку на кэш KV через непрерывные сегменты и создает новые ветви вместе с механизмом ранней остановки; (2) оценку преимуществ на уровне сегментов на основе дерева, учитывающую как глобальную, так и локальную оптимизацию проксимальной политики; и (3) анализ эффективности динамического расхождения и стратегии отката, основанных на вероятности и качестве. Мы эмпирически подтверждаем прирост производительности TreePO на наборе тестов для логического рассуждения и экономию вычислительных ресурсов GPU от 22\% до 43\% для обученных моделей, одновременно демонстрируя снижение вычислительной нагрузки на уровне траекторий до 40\% и на уровне токенов до 35\% для существующих моделей. Предлагая бесплатное повышение эффективности вывода, TreePO открывает практический путь к масштабированию пост-обучения на основе RL с меньшим количеством образцов и меньшими вычислительными затратами. Домашняя страница находится по адресу https://m-a-p.ai/TreePO.

English

Recent advancements in aligning large language models via reinforcement learning have achieved remarkable gains in solving complex reasoning problems, but at the cost of expensive on-policy rollouts and limited exploration of diverse reasoning paths. In this work, we introduce TreePO, involving a self-guided rollout algorithm that views sequence generation as a tree-structured searching process. Composed of dynamic tree sampling policy and fixed-length segment decoding, TreePO leverages local uncertainty to warrant additional branches. By amortizing computation across common prefixes and pruning low-value paths early, TreePO essentially reduces the per-update compute burden while preserving or enhancing exploration diversity. Key contributions include: (1) a segment-wise sampling algorithm that alleviates the KV cache burden through contiguous segments and spawns new branches along with an early-stop mechanism; (2) a tree-based segment-level advantage estimation that considers both global and local proximal policy optimization. and (3) analysis on the effectiveness of probability and quality-driven dynamic divergence and fallback strategy. We empirically validate the performance gain of TreePO on a set reasoning benchmarks and the efficiency saving of GPU hours from 22\% up to 43\% of the sampling design for the trained models, meanwhile showing up to 40\% reduction at trajectory-level and 35\% at token-level sampling compute for the existing models. While offering a free lunch of inference efficiency, TreePO reveals a practical path toward scaling RL-based post-training with fewer samples and less compute. Home page locates at https://m-a-p.ai/TreePO.