TreePO: Überbrückung der Lücke zwischen Policy-Optimierung, Effizienz und Inferenzeffizienz durch heuristische baumbasierte Modellierung

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

Jüngste Fortschritte bei der Ausrichtung großer Sprachmodelle durch Reinforcement Learning haben bemerkenswerte Erfolge bei der Lösung komplexer Denkprobleme erzielt, jedoch auf Kosten teurer On-Policy-Rollouts und einer begrenzten Exploration verschiedener Denkpfade. In dieser Arbeit stellen wir TreePO vor, einen selbstgesteuerten Rollout-Algorithmus, der die Sequenzerzeugung als einen baumstrukturierten Suchprozess betrachtet. TreePO besteht aus einer dynamischen Baum-Sampling-Politik und einer festen Segment-Decodierung und nutzt lokale Unsicherheit, um zusätzliche Verzweigungen zu gewährleisten. Durch die Amortisierung von Berechnungen über gemeinsame Präfixe und das frühzeitige Beschneiden von Pfaden mit geringem Wert reduziert TreePO die Rechenlast pro Update erheblich, während es die Explorationsvielfalt bewahrt oder sogar verbessert. Zu den wichtigsten Beiträgen gehören: (1) ein segmentweiser Sampling-Algorithmus, der die KV-Cache-Last durch zusammenhängende Segmente verringert und neue Verzweigungen mit einem Frühstopp-Mechanismus erzeugt; (2) eine baumbasierte Segment-Level-Vorteilsschätzung, die sowohl globale als auch lokale Proximale Policy-Optimierung berücksichtigt; und (3) eine Analyse der Wirksamkeit von Wahrscheinlichkeits- und Qualitäts-getriebener dynamischer Divergenz und Fallback-Strategie. Wir validieren empirisch die Leistungssteigerung von TreePO anhand einer Reihe von Denk-Benchmarks und die Effizienzsteigerung bei GPU-Stunden von 22\% bis zu 43\% des Sampling-Designs für die trainierten Modelle, während gleichzeitig eine Reduzierung von bis zu 40\% auf Trajektorienebene und 35\% auf Token-Ebene bei der Sampling-Berechnung für bestehende Modelle gezeigt wird. Während TreePO eine kostenlose Verbesserung der Inferenzeffizienz bietet, zeigt es einen praktischen Weg zur Skalierung von RL-basiertem Post-Training mit weniger Proben und geringerem Rechenaufwand. Die Homepage befindet sich unter 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.