TreePO : Combler l'écart entre l'optimisation de politique, l'efficacité et l'efficience de l'inférence grâce à une modélisation heuristique basée sur les arbres

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

August 24, 2025

papers.authors: 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

papers.abstract

Les récents progrès dans l'alignement des grands modèles de langage via l'apprentissage par renforcement ont permis des gains remarquables dans la résolution de problèmes de raisonnement complexes, mais au prix de déploiements on-policy coûteux et d'une exploration limitée des chemins de raisonnement diversifiés. Dans ce travail, nous introduisons TreePO, impliquant un algorithme de déploiement auto-guidé qui considère la génération de séquences comme un processus de recherche structuré en arbre. Composé d'une politique d'échantillonnage dynamique d'arbre et d'un décodage de segments de longueur fixe, TreePO exploite l'incertitude locale pour garantir des branches supplémentaires. En amortissant le calcul sur les préfixes communs et en élaguant précocement les chemins de faible valeur, TreePO réduit essentiellement la charge de calcul par mise à jour tout en préservant ou en améliorant la diversité de l'exploration. Les contributions clés incluent : (1) un algorithme d'échantillonnage segmenté qui allège la charge du cache KV grâce à des segments contigus et génère de nouvelles branches avec un mécanisme d'arrêt précoce ; (2) une estimation d'avantage au niveau des segments basée sur l'arbre qui prend en compte à la fois l'optimisation de politique proximale globale et locale ; et (3) une analyse de l'efficacité de la divergence dynamique pilotée par la probabilité et la qualité ainsi que de la stratégie de repli. Nous validons empiriquement le gain de performance de TreePO sur un ensemble de benchmarks de raisonnement et l'économie d'efficacité en heures GPU de 22\% à 43\% pour la conception d'échantillonnage des modèles entraînés, tout en montrant une réduction allant jusqu'à 40\% au niveau des trajectoires et 35\% au niveau des tokens pour les modèles existants. Tout en offrant un gain d'efficacité d'inférence sans coût supplémentaire, TreePO révèle une voie pratique pour mettre à l'échelle le post-entraînement basé sur le RL avec moins d'échantillons et moins de calcul. La page d'accueil se trouve à l'adresse 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.