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

arXiv: https://arxiv.org/abs/2508.17445 · HTML: https://mlsanity.com/paper/2508.17445/

## Key facts

- Cited by 61 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 18 of its 25 references that are in the map.
- First posted on arXiv 24 Aug 2025.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)
- LLM alignment / RLHF / preference optimisation / [Process reward models for reasoning](https://mlsanity.com/alignment/process-reward-models-for-reasoning/)

## Builds on (most cited first, 18 of 18)

- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [DAPO: An Open-Source LLM Reinforcement Learning System at Scale](https://mlsanity.com/paper/2503.14476/) (2025)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [Solving Quantitative Reasoning Problems with Language Models](https://mlsanity.com/paper/2206.14858/) (2022)
- [OlympiadBench: A Challenging Benchmark for Promoting AGI with Olympiad-Level Bilingual Multimodal Scientific Problems](https://mlsanity.com/paper/2402.14008/) (2024)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models](https://mlsanity.com/paper/2310.10505/) (2023)
- [Monte Carlo Tree Search: A Review of Recent Modifications and Applications](https://mlsanity.com/paper/2103.04931/) (2021)
- [TreeRL: LLM Reinforcement Learning with On-Policy Tree Search](https://mlsanity.com/paper/2506.11902/) (2025)
- [Segment Policy Optimization: Effective Segment-Level Credit Assignment in RL for Large Language Models](https://mlsanity.com/paper/2505.23564/) (2025)
- [SPO: Multi-Dimensional Preference Sequential Alignment With Implicit Reward Modeling](https://mlsanity.com/paper/2405.12739/) (2024)
- [Value-Guided Search for Efficient Chain-of-Thought Reasoning](https://mlsanity.com/paper/2505.17373/) (2025)
- [Infinite Sampling: Efficient and Stable Grouped RL Training for Large Language Models](https://mlsanity.com/paper/2506.22950/) (2025)
- [Multi-Bin Batching for Increasing LLM Inference Throughput](https://mlsanity.com/paper/2412.04504/) (2024)

## Built on by (most cited first, 30 of 61)

- [The Landscape of Agentic Reinforcement Learning for LLMs: A Survey](https://mlsanity.com/paper/2509.02547/) (2025)
- [A Survey of Reinforcement Learning for Large Reasoning Models](https://mlsanity.com/paper/2509.08827/) (2025)
- [Tree Search for LLM Agent Reinforcement Learning](https://mlsanity.com/paper/2509.21240/) (2025)
- [BranchGRPO: Stable and Efficient GRPO with Structured Branching in Diffusion Models](https://mlsanity.com/paper/2509.06040/) (2025)
- [Agent-World: Scaling Real-World Environment Synthesis for Evolving General Agent Intelligence](https://mlsanity.com/paper/2604.18292/) (2026)
- [Agentic Entropy-Balanced Policy Optimization](https://mlsanity.com/paper/2510.14545/) (2025)
- [Exploiting Tree Structure for Credit Assignment in RL Training of LLMs](https://mlsanity.com/paper/2509.18314/) (2025)
- [Reinforce-Ada: An Adaptive Sampling Framework under Non-linear RL Objectives](https://mlsanity.com/paper/2510.04996/) (2025)
- [Knapsack RL: Unlocking Exploration of LLMs via Optimizing Budget Allocation](https://mlsanity.com/paper/2509.25849/) (2025)
- [Plan Then Action:High-Level Planning Guidance Reinforcement Learning for LLM Reasoning](https://mlsanity.com/paper/2510.01833/) (2025)
- [Rethinking Sample Polarity in Reinforcement Learning with Verifiable Rewards](https://mlsanity.com/paper/2512.21625/) (2025)
- [From Code Foundation Models to Agents and Applications: A Comprehensive Survey and Practical Guide to Code Intelligence](https://mlsanity.com/paper/2511.18538/) (2025)
- [ParallelMuse: Agentic Parallel Thinking for Deep Information Seeking](https://mlsanity.com/paper/2510.24698/) (2025)
- [Attention as a Compass: Efficient Exploration for Process-Supervised RL in Reasoning Models](https://mlsanity.com/paper/2509.26628/) (2025)
- [Lookahead Tree-Based Rollouts for Enhanced Trajectory-Level Exploration in Reinforcement Learning with Verifiable Rewards](https://mlsanity.com/paper/2510.24302/) (2025)
- [Tree-OPO: Off-policy Monte Carlo Tree-Guided Advantage Optimization for Multistep Reasoning](https://mlsanity.com/paper/2509.09284/) (2025)
- [SpatialDreamer: Incentivizing Spatial Reasoning via Active Mental Imagery](https://mlsanity.com/paper/2512.07733/) (2025)
- [Spark: Strategic Policy-Aware Exploration via Dynamic Branching for Long-Horizon Agentic Learning](https://mlsanity.com/paper/2601.20209/) (2026)
- [Why Tree-Style Branching Matters for Thought Advantage Estimation in GRPO](https://mlsanity.com/paper/2509.24494/) (2025)
- [Let it Calm: Exploratory Annealed Decoding for Verifiable Reinforcement Learning](https://mlsanity.com/paper/2510.05251/) (2025)
- [FastGRPO: Accelerating Policy Optimization via Concurrency-aware Speculative Decoding and Online Draft Learning](https://mlsanity.com/paper/2509.21792/) (2025)
- [Generate, Filter, Control, Replay: A Comprehensive Survey of Rollout Strategies for LLM Reinforcement Learning](https://mlsanity.com/paper/2605.02913/) (2026)
- [Tree-based Dialogue Reinforced Policy Optimization for Red-Teaming Attacks](https://mlsanity.com/paper/2510.02286/) (2025)
- [Reinforced Efficient Reasoning via Semantically Diverse Exploration](https://mlsanity.com/paper/2601.05053/) (2026)
- [Learning Decentralized LLM Collaboration with Multi-Agent Actor Critic](https://mlsanity.com/paper/2601.21972/) (2026)
- [ATPO: Adaptive Tree Policy Optimization for Multi-Turn Medical Dialogue](https://mlsanity.com/paper/2603.02216/) (2026)
- [Distilled Reinforcement Learning for LLM Post-training](https://mlsanity.com/paper/2607.17247/) (2026)
- [DPWriter: Reinforcement Learning with Diverse Planning Branching for Creative Writing](https://mlsanity.com/paper/2601.09609/) (2026)
- [ERPO: Token-Level Entropy-Regulated Policy Optimization for Large Reasoning Models](https://mlsanity.com/paper/2603.28204/) (2026)
- [Teaching Language Models to Reason with Tools](https://mlsanity.com/paper/2510.20342/) (2025)

## Abstract (arXiv)

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.
