TreePO: Preenchendo a Lacuna entre Otimização de Políticas e Eficácia e Eficiência de Inferência com Modelagem Heurística Baseada em Árvores

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

Resumo

Os avanços recentes no alinhamento de grandes modelos de linguagem por meio de aprendizado por reforço têm alcançado ganhos notáveis na resolução de problemas complexos de raciocínio, mas ao custo de rollouts on-policy dispendiosos e uma exploração limitada de caminhos de raciocínio diversos. Neste trabalho, apresentamos o TreePO, que envolve um algoritmo de rollout autoguiado que enxerga a geração de sequências como um processo de busca estruturado em árvore. Composto por uma política de amostragem dinâmica em árvore e decodificação de segmentos de comprimento fixo, o TreePO aproveita a incerteza local para garantir ramificações adicionais. Ao amortizar a computação em prefixos comuns e podar caminhos de baixo valor precocemente, o TreePO essencialmente reduz o custo computacional por atualização, preservando ou até mesmo aumentando a diversidade de exploração. As principais contribuições incluem: (1) um algoritmo de amostragem segmentada que alivia a carga do cache KV por meio de segmentos contíguos e gera novos ramos juntamente com um mecanismo de parada antecipada; (2) uma estimativa de vantagem em nível de segmento baseada em árvore que considera tanto a otimização de política proximal global quanto local; e (3) uma análise sobre a eficácia da divergência dinâmica orientada por probabilidade e qualidade, além de uma estratégia de fallback. Validamos empiricamente o ganho de desempenho do TreePO em um conjunto de benchmarks de raciocínio e a economia de eficiência em horas de GPU, que variou de 22% a 43% no design de amostragem para os modelos treinados, ao mesmo tempo em que mostramos uma redução de até 40% no nível de trajetória e 35% no nível de token na computação de amostragem para os modelos existentes. Ao oferecer uma melhoria gratuita na eficiência de inferência, o TreePO revela um caminho prático para escalar o pós-treinamento baseado em RL com menos amostras e menos computação. A página inicial está localizada em 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.