TreePO: Colmare il Divario tra Ottimizzazione delle Politiche, Efficacia e Efficienza nell'Inferenza con Modellazione Euristica Basata su Alberi

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

August 24, 2025

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

Abstract

I recenti progressi nell'allineamento di modelli linguistici di grandi dimensioni tramite apprendimento per rinforzo hanno ottenuto risultati notevoli nella risoluzione di problemi complessi di ragionamento, ma al costo di rollout on-policy dispendiosi e di un'esplorazione limitata di percorsi di ragionamento diversi. In questo lavoro, introduciamo TreePO, che coinvolge un algoritmo di rollout autoguidato che considera la generazione di sequenze come un processo di ricerca strutturato ad albero. Composto da una politica di campionamento dinamica ad albero e da un decoding a segmenti di lunghezza fissa, TreePO sfrutta l'incertezza locale per garantire ramificazioni aggiuntive. Ammortizzando il calcolo attraverso prefissi comuni e potando precocemente i percorsi a basso valore, TreePO riduce sostanzialmente il carico computazionale per aggiornamento, preservando o migliorando la diversità dell'esplorazione. I contributi chiave includono: (1) un algoritmo di campionamento segmentale che allevia il carico della cache KV attraverso segmenti contigui e genera nuovi rami insieme a un meccanismo di arresto anticipato; (2) una stima del vantaggio a livello di segmento basata su albero che considera sia l'ottimizzazione della politica prossimale globale che locale; e (3) un'analisi sull'efficacia della divergenza dinamica guidata da probabilità e qualità e della strategia di fallback. Validiamo empiricamente il guadagno di prestazioni di TreePO su una serie di benchmark di ragionamento e il risparmio di efficienza in ore GPU dal 22% fino al 43% del design di campionamento per i modelli addestrati, mostrando nel frattempo una riduzione fino al 40% a livello di traiettoria e del 35% a livello di token nel calcolo di campionamento per i modelli esistenti. Offrendo un "pranzo gratis" in termini di efficienza di inferenza, TreePO rivela un percorso pratico verso il ridimensionamento del post-training basato su RL con meno campioni e meno calcolo. La pagina principale si trova all'indirizzo 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.