TreePO: Cerrando la brecha entre la optimización de políticas y la eficacia, y la eficiencia en la inferencia mediante modelado heurístico basado en árboles

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

August 24, 2025

Autores: 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

Resumen

Los avances recientes en la alineación de modelos de lenguaje grandes mediante aprendizaje por refuerzo han logrado mejoras notables en la resolución de problemas de razonamiento complejo, pero a costa de costosos despliegues en política y una exploración limitada de diversos caminos de razonamiento. En este trabajo, presentamos TreePO, que incluye un algoritmo de despliegue autoguiado que considera la generación de secuencias como un proceso de búsqueda estructurado en árbol. Compuesto por una política de muestreo dinámico de árboles y decodificación de segmentos de longitud fija, TreePO aprovecha la incertidumbre local para garantizar ramificaciones adicionales. Al amortizar el cálculo a través de prefijos comunes y podar tempranamente las rutas de bajo valor, TreePO esencialmente reduce la carga computacional por actualización mientras preserva o mejora la diversidad de exploración. Las contribuciones clave incluyen: (1) un algoritmo de muestreo por segmentos que alivia la carga de la caché KV mediante segmentos contiguos y genera nuevas ramificaciones junto con un mecanismo de parada temprana; (2) una estimación de ventaja a nivel de segmento basada en árboles que considera tanto la optimización de políticas proximales global como local; y (3) un análisis sobre la efectividad de la divergencia dinámica impulsada por probabilidad y calidad, junto con una estrategia de retroceso. Validamos empíricamente la mejora en el rendimiento de TreePO en un conjunto de benchmarks de razonamiento y el ahorro de eficiencia en horas de GPU desde un 22% hasta un 43% en el diseño de muestreo para los modelos entrenados, mostrando además una reducción de hasta un 40% a nivel de trayectoria y un 35% a nivel de token en el cálculo de muestreo para los modelos existentes. Al ofrecer una mejora gratuita en la eficiencia de inferencia, TreePO revela un camino práctico hacia la escalabilidad del post-entrenamiento basado en RL con menos muestras y menor cálculo. La página principal se encuentra en 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.