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

ArXiv: 2508.17445

🎯 Pitch

Reinforcement learning for LLM reasoning wastes massive compute by generating independent rollouts that re-compute shared prefixes from scratch. TreePO replaces this with tree-structured sampling that shares KV-cache computation across common branches and prunes low-value paths early, cutting GPU hours by up to 43% while matching or exceeding accuracy—and it reveals that forcing models toward unlikely reasoning paths actually harms performance.

---

1. Executive Summary

This paper introduces TreePO, a reinforcement learning framework that replaces standard i.i.d. sequential rollouts with a heuristic tree-based sampling mechanism paired with a hierarchical tree-based advantage estimator to improve both training stability and computational efficiency when aligning LLMs for mathematical reasoning. The tree-based sampling mechanism leverages shared KV-cache computation across common prefixes and implements dynamic divergence, probability-based fallback, and early stopping to amortize generation costs, while the tree-based advantage estimator computes per-trajectory credit assignment by aggregating subgroup-specific advantages across multiple depths of the shared tree hierarchy rather than using a single group-wide normalization. Experiments on Qwen2.5-7B base models trained with MATH and DeepScaler data and evaluated on AIME 2024, AMC 2023, MATH500, MINERVA, and Olympiad Bench show that TreePO reduces GPU hours by 12% to 43% compared to sequential sampling baselines while achieving comparable or improved accuracy—for instance, TreePO with more initial divergence reaches 58.21% overall versus 46.63% for vanilla GRPO—and provides up to 40% trajectory-level and 35% token-level sampling compute reduction during inference. The paper establishes that tree-structured rollouts with subgroup-aware advantage estimation can stabilize RL training and improve efficiency, but that the benefits depend critically on matching segment length and tree depth to the model type, and that forced exploration toward low-probability reasoning paths degrades performance rather than improving it.

2. Context and Motivation

The Core Problem: RL for Reasoning Is Computationally Inefficient and Exploration-Limited

The fundamental problem this paper addresses emerges from the intersection of two trends in LLM post-training. First, reinforcement learning—particularly outcome-supervised RL with verifiable rewards (RLVR)—has proven remarkably effective at eliciting complex reasoning capabilities from base models, as demonstrated by DeepSeek-R1, DeepSeekMath, and related systems. Second, the computational cost of running these RL pipelines at scale is enormous because each training update requires generating fresh on-policy rollouts from the current model checkpoint. The paper's central observation is that the standard rollout mechanism is structurally wasteful: it generates multiple complete, independent trajectories per prompt, each starting from scratch, even when those trajectories share extensive common prefixes in their reasoning chains.

This inefficiency manifests in two forms:

• Computational redundancy: When a model receives the same prompt and samples 16 independent completions (a typical group size for GRPO), it recomputes identical or near-identical tokens across all 16 forward passes. The KV-cache for the prompt prefix is shared across samples in modern inference engines, but the KV-cache for generated tokens—the actual reasoning steps—is not. If the model tends to begin its reasoning with the same initial deductions (problem interpretation, variable assignment, logical setup), each independent rollout redundantly computes and stores these tokens.

• Exploration myopia: Independent sampling treats every token position identically regardless of whether the model is at a high-uncertainty decision point or confidently progressing along a well-established reasoning path. There is no mechanism to allocate more sampling budget to the parts of the reasoning space that need exploration and less to the parts that are already reliably correct. The model may waste compute exploring variations of an already-correct prefix while failing to adequately explore alternative high-level strategies.

The paper frames this as a dual challenge in the opening paragraph of Section 1: "How can we enable LLMs to explore potentially correct reasoning paths while maintaining or reducing computational costs?" and "How can we accurately attribute sparse outcome rewards to the specific tokens that contributed to correct answers?" The first question is about the sampling mechanism itself; the second is about the advantage estimation that uses the sampled data. TreePO proposes a unified solution to both through its tree-structured rollout architecture.

Why This Problem Matters

The paper's motivation is grounded in practical scaling concerns rather than purely theoretical interest. As the authors note, RL-based post-training for reasoning has become a dominant paradigm—but its cost scales poorly. Each training step requires sampling $G$ complete trajectories per query (typically $G = 16$ for GRPO), running those trajectories through a reward function, computing advantages, and updating the policy. The sampling phase dominates wall-clock time because it is inherently sequential: the model must generate, the engine must return log probabilities, and only then can training proceed.

If the RL pipeline is to scale to longer reasoning horizons (multi-turn dialogue, tool use, multi-agent systems, as the conclusion envisions), the sampling cost becomes the primary bottleneck. The paper's efficiency results—22% to 43% reduction in GPU hours during training, 40% trajectory-level and 35% token-level reduction during inference—translate directly to either lower dollar costs for training runs or the ability to train for more steps within a fixed budget. This is not a marginal optimization; it potentially changes what scale of RL training is economically feasible for a given team or organization.

Beyond cost, there is a training stability argument. The validation curves in Figure 1 (Left, Mid) show that vanilla GRPO (blue line) exhibits substantial volatility in performance, while TreePO variants produce smoother, more monotonic improvement. Unstable training is not just a cosmetic issue—it means that checkpoint selection becomes unreliable, hyperparameter tuning is harder, and the risk of catastrophic forgetting or reward hacking increases. If tree-structured rollouts inherently stabilize training by providing more diverse and better-structured advantage signals, that is a significant practical benefit independent of computational savings.

Prior Approaches and Their Limitations

The paper identifies three categories of prior work and explains where each falls short:

1. Independent i.i.d. Sampling (the GRPO/DAPO status quo). The dominant paradigm, as used in DeepSeekMath, DAPO, SimpleRL, and related systems, samples $G$ complete trajectories per query as independent sequences. From a computational perspective, DAPO and related methods do share the prompt's KV-cache across samples within a group, but once the model begins generating, each trajectory's KV-cache is stored separately. The paper's key insight (validated by the case study in Section 2.1, Figure 2) is that reasoning trajectories from an aligned model share extensive common prefixes in their generated content—not just the prompt. For instance, when solving a math problem, the model might consistently start with "Let x be..." and follow the same variable assignment and initial equation setup before diverging at a critical decision point. Standard sequential sampling recomputes this shared prefix $G$ times. The paper states this bluntly in Section 1:

"From a computational perspective, this approach creates paths with separate Key-Value (KV) caches, failing to utilize shared KV caching mechanisms that could significantly accelerate inference. Conceptually, continuing to explore paths already known to be impossible or incorrect, without early termination, represents a critical limitation in adaptability."

DAPO improved GRPO's stability through dynamic sampling, clip-higher, and token-level loss, but did not address the fundamental structure of how rollouts are generated. TreePO adopts DAPO's policy objective as its starting point (Section 2.3), indicating that the contribution is orthogonal to and compatible with those improvements.

2. Monte Carlo Tree Search (MCTS) for LLMs. MCTS and its variants provide the tree structure that TreePO aims for, but the paper identifies a fundamental mismatch with LLM inference hardware:

"Despite its promise, MCTS is often inefficient for LLM inference, requiring numerous sequential rollouts that are poorly suited for parallelized engines."

MCTS operates one node at a time—select a leaf, expand it, roll out a completion, backpropagate the value, repeat. This sequential dependency means that the inference engine cannot batch multiple forward passes efficiently. Each MCTS step incurs the overhead of a separate inference call, which is catastrophic for GPU utilization. The paper references recent systems like TreeRL (Hou et al., 2025) and SPO (Guo et al., 2025) that adapt tree search for LLM training, but identifies specific limitations:

• TreeRL couples on-policy tree expansion with process-level rewards, but "its trees stay shallow and the algorithm rolls one full answer to compute log-probabilities before it can branch again, which doubles the running time." The shallow trees limit exploration depth, and the one-at-a-time branching eliminates batching efficiency.

• SPO (Segment Policy Optimization) uses segment-level credit assignment but adopts MCTS-like advantage calculations that focus on "the value difference between a parent and its child node." The paper argues this is too local—it doesn't capture the collective signal from an entire subtree of descendants.

• Broader applicability gap: The paper makes a specific and important distinction that both TreeRL and SPO are "demonstrated on models that have already undergone SFT" (supervised fine-tuning), whereas TreePO is designed to work "directly from a base model, aligning with the 'RL-zero' paradigm where reasoning capabilities are elicited without prior supervised fine-tuning." This matters because a base model's outputs are noisier and less structured, placing higher demands on the robustness of the sampling and advantage estimation mechanisms.

3. Efficient batching and scheduling approaches. Recent work on efficient RL sampling includes Infinite Sampling (Wang et al., 2025), which breaks groups into micro-batches with continuous interleaving and length-aware scheduling, and Truncated PPO (Fan et al., 2025), which cuts sampled chains after a short window and backpropagates early. The paper acknowledges these improvements but identifies gaps:

• Infinite Sampling's approach "does not look at the partial trajectories while they are generated, introduces extra scheduling logic, and leaves the advantage estimator untouched." It improves hardware utilization but doesn't leverage the content of the generated sequences to make smarter branching or pruning decisions.

• Truncated PPO improves wall-clock time by shortening rollouts, but "the price is that long-range information is lost and credit assignment becomes harder." Truncating before the model reaches a conclusion means the reward signal may be incomplete or misleading.

• Neither approach modifies the advantage estimation to account for structural relationships among trajectories. They optimize the scheduling of computation but not the information structure of the data being generated.

4. Process reward models and dense supervision. The paper implicitly distinguishes its approach from methods that require intermediate supervision. MCTS-based methods often rely on process reward models (PRMs) trained to score intermediate reasoning steps, providing dense signals to guide tree expansion. Training a reliable PRM requires either expensive human annotation or sophisticated Monte Carlo rollout procedures (as in the companion paper example from the prompt). TreePO avoids this entirely—its heuristic branching decisions ("dynamic divergence" based on local uncertainty, early stopping for flawed paths) require no external verifier, no PRM, and no additional training. The tree structure emerges from the model's own generation behavior combined with lightweight rule-based heuristics (repetitive substring detection, answer format checking) that incur negligible overhead.

How TreePO Positions Itself

The paper positions TreePO not as a completely new RL algorithm, but as a data generation architecture that improves the rollout mechanism underlying existing policy optimization objectives. This is a crucial distinction. TreePO takes the optimization objective from DAPO (which itself refines GRPO) and proposes a new way to generate the data that feeds into that objective. The advantage estimation function is similarly a refinement of the group-normalized advantage from GRPO/DAPO, leveraging the structural information provided by the tree hierarchy.

This positioning is visible in the architecture: the core contribution (Sections 2.2 and 2.3) describes how rollouts are generated (Algorithm 1, the branching/fallback protocol, segment-level tree sampling) and how advantages are computed from tree-structured rollouts (Equations 3–5, the subgroup aggregation). The policy gradient itself remains standard PPO clipping. The claim is that better-structured data—with shared prefixes exploited, redundant computation avoided, and subgroup-aware advantages—improves both training efficiency and stability without requiring changes to the core optimization algorithm.

The paper's relationship to prior work is best understood as synthesizing orthogonal ideas: it takes the segment-level decomposition from works like SPO, combines it with the batching efficiency insights from Infinite Sampling, adds lightweight heuristics inspired by MCTS (branching, early stopping), and implements everything within the established GRPO/DAPO training framework. The novel combination—rather than any single component—is what distinguishes TreePO.

The paper also positions itself within the broader trend toward "RL-zero" training (eliciting reasoning from base models through RL without SFT). This is significant because it means TreePO must handle less-structured outputs, more diverse reward signals, and less stable training dynamics than methods applied to already-instruction-tuned models. The heuristic pruning (repetitive substring detection) and careful fallback design (only fallback from properly formatted answers) are explicitly motivated by the challenges of base model training, where the model in early steps may produce gibberish, get stuck in loops, or fail to produce parseable answers.

3. Technical Approach

3.1 Reader Orientation

TreePO is a data generation and credit assignment architecture that sits underneath standard policy-gradient RL for LLMs: rather than generating multiple independent complete reasoning traces for each prompt, it generates a tree of partial reasoning segments that share computation across common prefixes. The system solves the dual problem of training-time sampling inefficiency (redundant computation of shared reasoning prefixes across independent rollouts) and noisy credit assignment (attributing a single outcome reward to individual tokens) by reshaping the rollout process into a controllable tree search whose hierarchical structure provides natural granularity for advantage estimation.

3.2 Big-Picture Architecture (Diagram in Words)

The system has four major components:

1. Segment-Level Tree Sampler — a rollout engine that generates reasoning traces not as complete sequences but as a tree of fixed-length segments, managing branching at decision points, pruning flawed paths early, and maintaining a queue of active prompts. This is where the computational savings come from (shared KV-caching across common prefixes).

2. Heuristic Branching and Fallback Controller — a lightweight policy layer that decides, at each tree depth, how to allocate a fixed branching budget across active search paths (using log-probability signals or uniform assignment) and when to trigger fallback generation to ensure each query produces the required number of complete trajectories.

3. Tree-Based Advantage Estimator — a credit assignment function that decomposes each trajectory into segments organized by shared ancestors, computes subgroup-specific advantages at each tree depth, and aggregates them into a per-trajectory advantage signal used in the policy gradient update.

4. Policy Optimizer (DAPO objective) — a standard PPO-clipped policy gradient with DAPO's improvements (clip-higher, dynamic sampling, token-level loss) that consumes the tree-structured rollouts and subgroup-aware advantages to update the base model.

Information flows as follows: a batch of queries enters the system → each query is forked into an initial set of branching prompts → the inference engine generates one segment (up to l tokens) for each active prompt → segments are checked for early stopping (EOS, answer boxed, or repetitive patterns) and partial responses are concatenated to form new prompts or finalize as leaf trajectories → when all active paths terminate, the fallback mechanism ensures sufficient trajectories → completed trajectories are scored with outcome rewards → the tree hierarchy (which segments share which ancestors) is used to compute subgroup advantages → the DAPO loss backpropagates through the policy model.

3.3 Roadmap for the Deep Dive

• First, the DAPO policy objective (Equation 1–2), because TreePO adopts this as its optimization backbone and all subsequent advantage estimation modifications are relative to this baseline. Understanding what advantage \hat{A}_{i,t} means in GRPO/DAPO is prerequisite to understanding what the tree-based variant changes.

• Second, the segment-level tree sampling algorithm (Algorithm 1 and the surrounding discussion), because this is the core architectural innovation that generates the tree-structured data which everything else depends on. The branching budget transfer, depth-first fallback, and early stopping mechanisms are defined here. This is the longest subsection.

• Third, the tree-based advantage estimation (Equations 3–5), because this consumes the tree structure produced by the sampler and computes the credit assignment signal. The subgroup definition, aggregation strategy, and variance normalization are detailed here.

• Fourth, the heuristic probability-based branching assignment (Section 4.4 material, Equation for softmax allocation), because this is the mechanism for steering exploration within the tree beyond uniform budget assignment, and its failure modes are informative about what kinds of exploration are productive.

3.4 Detailed, Sentence-based Technical Breakdown

This is primarily a systems and algorithms paper whose core idea is that organizing RL rollouts as a tree of segments—rather than independent complete sequences—enables both computational efficiency (through shared KV-cache prefixes) and finer credit assignment (through subgroup-aware advantage estimation), and that lightweight heuristics for branching and pruning can replace expensive MCTS-style value estimation without sacrificing training quality.

---

3.4.1 The Policy Optimization Objective (DAPO Baseline)

TreePO does not introduce a new policy gradient formula. It inherits the objective from DAPO (Yu et al., 2025), which itself refines the GRPO objective from DeepSeekMath (Shao et al., 2024). The paper states this explicitly: "We take the GRPO optimizing objective and adopt the improved modifications proposed in DAPO as our starting point, which further highlights clip-higher gradient, dynamic sampling, and token-level loss." Understanding this baseline is essential because the tree-based advantage function in Section 2.3 is a drop-in replacement for the standard group-normalized advantage \hat{A}_{i,t} used in this objective.

The DAPO policy objective is:

$$\mathcal{J}_{\text{TreePO}}(\theta) = \mathbb{E}_{(q,a) \sim \mathcal{D}, \{o_i\}_{i=1}^{G} \sim \pi_{\theta_{\text{old}}}(\cdot \mid q)} \left[ \frac{1}{\sum_{i=1}^{G} |o_i|} \sum_{i=1}^{G} \sum_{t=1}^{|o_i|} \min \left( r_{i,t}(\theta) \hat{A}_{i,t}, \text{clip}\left( r_{i,t}(\theta), 1-\varepsilon_{\text{low}}, 1+\varepsilon_{\text{high}} \right) \hat{A}_{i,t} \right) \right]$$

subject to a dynamic sampling constraint:

$$0 < \left| \{ o_i \mid \texttt{is\_equivalent}(a, o_i) \} \right| < G$$

where:

• $\mathcal{D}$ is the training data distribution over query-answer pairs $(q, a)$ (the answer is used by the reward function, not for supervised training);
• $G$ is the group size—the number of trajectories sampled per query (set to 16 in the experiments);
• $\{o_i\}_{i=1}^{G}$ are the $G$ complete output trajectories sampled from the old policy $\pi_{\theta_{\text{old}}}$;
• $|o_i|$ is the length in tokens of trajectory $i$;
• $r_{i,t}(\theta) = \frac{\pi_{\theta}(o_{i,t} \mid q, o_{i,<t})}{\pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,<t})}$ is the per-token importance sampling ratio between the current policy and the old policy;
• $\varepsilon_{\text{low}}$ and $\varepsilon_{\text{high}}$ are asymmetric clipping thresholds (DAPO's "clip-higher" modification that applies different bounds for ratios below and above 1.0);
• $\hat{A}_{i,t}$ is the per-token advantage estimate, which in standard GRPO is the group-normalized outcome reward, and in TreePO is replaced by the tree-based estimator (Equation 5).

The constraint $0 < |\{o_i \mid \texttt{is\_equivalent}(a, o_i)\}| < G$ enforces DAPO's dynamic rejection sampling: a query is only used for training if its $G$ sampled trajectories include at least one correct and at least one incorrect answer, ensuring that the advantage signal has contrast.

What this objective computes: For each query, the model samples $G$ trajectories from the old policy (the checkpoint before the current update). Each trajectory receives a scalar outcome reward $R_i$ (typically 1.0 for a correct final answer and 0.0 otherwise, or a rule-based partial-credit score). These rewards are normalized into advantages $\hat{A}_{i,t}$ (all tokens in trajectory $i$ receive the same advantage value in standard GRPO; TreePO modifies how this value is computed). The PPO objective then increases the probability of tokens from trajectories with positive advantage and decreases the probability of tokens from trajectories with negative advantage, subject to the clipping constraint that prevents the policy from changing too much in a single update.

Why this form: The token-level loss—averaging over all tokens rather than per-trajectory—is a DAPO refinement that gives longer trajectories proportionally more weight in the gradient, which the authors found stabilizes training. The asymmetric clipping ($\varepsilon_{\text{low}} \neq \varepsilon_{\text{high}}$) allows the policy to increase probabilities more aggressively than it decreases them, which DAPO showed helps when the model needs to learn new correct behaviors faster than it unlearns incorrect ones. The dynamic rejection sampling ensures that every training batch contains a mix of successful and failed trajectories, preventing the degenerate case where all samples are correct (advantage = 0 everywhere, no learning signal) or all wrong (no positive examples to reinforce).

Critically, the advantage $\hat{A}_{i,t}$ in standard GRPO/DAPO is defined as:

$$\hat{A}_{i,t} = \frac{R_i - \text{mean}(\{R_i\}_{i=1}^{G})}{\text{std}(\{R_i\}_{i=1}^{G})}$$

All tokens in trajectory $i$ share the same group-normalized advantage. This is a coarse signal: if two trajectories have different reasoning paths that lead to different answers but happen to get the same outcome reward in a particular group, the advantage signal cannot distinguish them. TreePO's contribution is to replace this flat per-trajectory advantage with a hierarchically computed value that leverages the tree structure (Equation 5), which we detail in Section 3.4.3.

---

3.4.2 The Segment-Level Tree Sampling Algorithm

This is the core architectural innovation of TreePO and the source of its computational efficiency gains. The paper reformulates the process of generating $G$ trajectories per query from "generate $G$ independent complete sequences" to "grow a tree of partial segments, branching at decision points, pruning flawed paths early, and reusing computation across shared prefixes."

3.4.2.1 Terminology

The paper establishes a consistent vocabulary mapping between RLVR (reinforcement learning with verifiable rewards) and tree search (Section 2.2, Preliminaries):

• Query $q$ → the root node at depth 0. This is the prompt, potentially including additional context beyond the bare question.
• Number of complete trajectories → the tree width $\boldsymbol{w}$, set to 16 in training (matching the group size $G$ in DAPO).
• Maximum decoding steps of a trajectory → the depth $\boldsymbol{d}$, which combined with segment length determines total response budget.
• Maximum decoding tokens per step → the segment length $\boldsymbol{l}$, the number of tokens generated in one inference call before the system decides whether to continue, branch, or stop.
• Branching budget per segment node → $\boldsymbol{b}$, the number of child branches to fork from each active search path at a given depth. The total branching budget at depth $d$ is $b^d$ before early stopping adjustments.

The total generation budget per query is bounded by $d \times l$ tokens per trajectory, with $w$ trajectories required. The paper experiments with three depth–segment configurations: $28 \times 256$, $14 \times 512$, and $7 \times 1024$, all yielding $28 \times 256 = 14 \times 512 = 7 \times 1024 = 7168 \approx 7 \times 1024$ tokens maximum per trajectory.

3.4.2.2 Algorithm Walkthrough

Algorithm 1 in the paper provides the pseudocode. Here is the line-by-line operational description:

Initialization (Lines 3–4). The prompt queue $P$ is initialized with the batch of queries $Q$. Immediately, $\textsc{Branching}(P)$ forks each query $b$ times, producing $b \cdot |Q|$ active prompts. In the fixed divergence setting, $b = 2^d$ at depth $d$ (binary tree, so $b=2$ at depth 1, $b=4$ at depth 2, etc., with the budget allocated across active paths). In the "More Init Divergence" setting, additional random branches (2–8) are allocated at the root to increase initial diversity. Each forked prompt is an identical copy of the query text but tracked as a separate path in the tree.

Main generation loop (Lines 5–18). While the prompt queue is non-empty:

• Segment inference (Line 6): $\textsc{Inference}(P)$ sends all active prompts to the inference engine in one batch. The engine generates exactly one segment of up to $l$ tokens for each prompt using the current policy model $\pi_{\theta_{\text{old}}}$ with sampling temperature (0.8 in the case study, though training temperature is not explicitly specified). The engine returns $S$, a set of $|P|$ generated text segments along with their token-level log-probabilities. This is the key efficiency step: because all prompts in $P$ that share a common prefix (e.g., prompts forked from the same parent at the previous depth) share their KV-cache up to that fork point, the inference engine avoids recomputing attention for the shared portion.

• Queue reset (Lines 7–8): The current prompt set is saved as $P^{last}$ (needed to reconstruct full trajectories when segments complete), and $P$ is cleared to prepare for the next generation round.

• Segment processing (Lines 9–15): For each generated segment $s_k$:

  • Termination check (Line 10): $\textsc{Finish}(s_k)$ returns true if the segment contains an end-of-sequence token (EOS) or a boxed answer (\boxed{...} in math reasoning). $\textsc{FailedNode}(s_k)$ returns true if the segment contains repetitive substrings (detected by scanning for repeated token patterns within the new segment), indicating the model has entered a degenerate loop. The repetitive substring check is the "simple early stopping trick" described in Section 2.2's "Heuristic Sampling" paragraph and is specifically motivated by base model training, where unaligned models frequently produce mumbling or repeated text.
  • If terminal or failed (Line 11): The full trajectory $p_k^{last} \oplus s_k$ (the concatenation of all accumulated prefix segments plus this final segment) is added to the output set $O$. This path is now a leaf node in the tree.
  • If still active (Line 13): The concatenated partial trajectory $p_k^{last} \oplus s_k$ is added back to the prompt queue $P$ as a new prompt for the next depth level. At this point, the model has generated $l$ more tokens of reasoning but has not yet reached an answer or terminal state.

• Branching (Line 16): Before the next round of inference, $\textsc{Branching}(P)$ forks each active prompt into multiple copies according to the branching budget allocation policy. In the default binary tree setting, the total branching budget at depth $d$ is $2^d$. With the Branching Budget Transfer mechanism (described in Section 2.2's "Branching and Fallback"), this total budget is allocated evenly across all currently active paths. For example, if at depth 2 there are 2 active paths and the total budget is $2^2 = 4$, each path receives 2 branches. If some paths have terminated early, their unused budget is redistributed to the surviving paths, keeping the total inference batch size balanced and preventing low GPU utilization from small batches.

• Fallback (Line 17): $\textsc{Fallback}(P, O)$ triggers only when (a) the prompt queue for a query is empty (all active paths have terminated or been pruned) and (b) the number of completed trajectories for that query $w_q < w$ (the required width, 16). Fallback selects terminated paths that contain either a properly formatted answer (\boxed{...}) or an EOS token, and randomly branches from them at the segment level—meaning it goes back to the segment before termination and generates alternative continuations. The paper specifies that fallback is depth-first: "TreePO launches the fallback mechanism only when there is no active path for $q$ and the tree does not have enough trajectories." This means fallback prioritizes completing existing promising paths before starting new ones from the root, preserving the depth of search.

Output (Line 19). The algorithm returns $O$, a set of $w$ (or as close as possible to $w$) complete trajectories per query, each structured as a sequence of segments connected by shared ancestor nodes in the tree.

3.4.2.3 Design Choices and Their Justifications

Why segment-level rather than token-level? If the system branched at every token, the tree would grow exponentially and inference batching would fragment into many tiny requests. By decoupling into segments of length $l$, TreePO batches $|P|$ forward passes per segment round, where $|P|$ is typically in the tens or hundreds, maintaining high GPU utilization. The segment length is a tunable hyperparameter that trades off batching efficiency (longer segments = fewer rounds, larger batches) against branching granularity (shorter segments = more branching opportunities, finer-grained early stopping).

Why binary tree as default? The paper states: "we define a vanilla $\boldsymbol{N}$-ary tree as a baseline searching strategy, i.e., the branching budget for a the root node $q$ at depth $\boldsymbol{d}$ is $N^d$ until it reaches the maximum width $\boldsymbol{w}$." With $N=2$ and $w=16$, the tree reaches full width by depth 4 ( $2^4 = 16$ ), after which no further branching occurs—all active paths simply continue generating until they terminate or hit the token budget. The binary tree is the simplest non-trivial tree structure that enables shared prefix reuse (at depth 1, two paths share the root segment's KV-cache; at depth 2, four paths share the root and second-level segments in pairs, etc.).

Why Branching Budget Transfer? The paper notes that "early stopped short search paths could derive a small request batch to the inference engine and thus cause low utilization." If some paths terminate early (e.g., reaching an answer in 3 segments while others need 7), the active batch size shrinks and GPU utilization drops. Transferring the unused budget of terminated paths to surviving paths maintains a more constant batch size and prevents the sampling process from being bottlenecked by a few very long trajectories.

Why depth-first fallback? If fallback always starts from the root (breadth-first), the system would generate many short, shallow trajectories at the expense of deep, complex ones. By waiting until all active paths are exhausted before fallback triggers, TreePO ensures that the model first attempts to generate complete reasoning chains at the full depth, and only resorts to fallback when necessary. This preserves the model's ability to learn long-range reasoning, which the paper explicitly identifies as a concern: "To avoid sampling progress overly conducts fallback on the early stopped short paths and lose the capability of long complex reasoning."

Why early stopping on repetitive substrings? Base models trained without SFT frequently produce degenerate outputs—repeating the same token or phrase, getting stuck in loops. Without early stopping, these paths would consume the full token budget while producing useless training data. The substring repetition check is a computationally cheap heuristic (no model inference required) that prunes these paths immediately, saving compute and preventing the policy from being updated on degenerate examples.

Why only fallback from properly formatted or EOS-terminated paths? The fallback candidate selection rule ("only those stopped paths containing formatted answer... or ending with [EOS] can be selected") ensures that fallback branches from reasoning paths that reached some kind of conclusion rather than from paths that were pruned for repetitive output. The intuition is that a path that produced a boxed answer (even if wrong) at least demonstrated the ability to complete the reasoning format, and branching from just before that conclusion may yield alternative correct answers. Paths pruned for gibberish have no useful structure to branch from.

3.4.2.4 Relationship to MCTS

The paper is explicit that TreePO is not MCTS. MCTS selects nodes to expand based on value estimates (typically the UCB formula: upper confidence bound balancing exploitation of high-value paths with exploration of under-visited paths). TreePO's branching is heuristic, not value-guided: the budget is allocated either uniformly or based on log-probability signals (the "Heuristic Sampling" mechanisms), not on learned value functions. There is no backpropagation of value estimates through the tree. The paper states that MCTS is "often inefficient for LLM inference, requiring numerous sequential rollouts that are poorly suited for parallelized engines." TreePO's segment-batched generation avoids this sequential dependency entirely by decoupling tree management from generation scheduling.

---

3.4.3 The Tree-Based Advantage Estimator

The tree structure produced by the sampling algorithm provides hierarchical information about how trajectories relate to one another. TreePO exploits this to compute a more granular advantage signal than the flat group normalization used in GRPO/DAPO.

3.4.3.1 Segment Decomposition

Given a complete trajectory $o_i$ produced by the tree sampler, it can be decomposed into segments $S_j$ corresponding to the inference steps at which each segment was generated:

$$o_i = s_1 \oplus s_2 \oplus \dots \oplus s_{j-1} \oplus s_j$$

where $\{j \in J \mid j \leq \text{max depth}\}$. Each $s_j$ is a block of up to $l$ tokens generated in one inference call. The operators $\oplus$ denote concatenation. The index $j$ corresponds to the tree depth at which that segment was generated.

What this decomposition enables: Because the tree sampler tracks which segments share common ancestors, we can group trajectories based on their shared prefixes at each depth. For example, if two trajectories $o_1$ and $o_2$ were both forked from the same parent at depth 1 (they share $s_1$), they belong to the same subgroup at depth 1 even if they diverge at depth 2. If three trajectories share $s_1$ and $s_2$ but diverge at $s_3$, they form a subgroup at depth 2. This nested structure provides multiple levels of "local group normalization" for advantages.

3.4.3.2 Subgroup Hierarchy

Let $G$ be the full group of $w$ trajectories for a query. The tree structure defines nested subgroups:

$$G_{|J|} \subseteq G_{|J|-1} \subseteq \dots \subseteq G_2 \subseteq G_1 \subseteq G$$

where $\{j \in J \mid j < \text{max depth}\}$. $G_j$ is the set of trajectories that share the same predecessor node at inference depth $j$. $G_1$ contains all trajectories sharing the same root (query) branch at depth 1—this is the largest subgroup below the full group. $G_2$ contains trajectories sharing the same depth-2 node, which is a subset of $G_1$. At the deepest level, $G_{|J|}$ is a singleton or very small group containing only trajectories that followed identical paths through all branching points (this could be multiple trajectories if fallback generated siblings from the same terminal node).

Figure 3 in the paper illustrates this with a concrete example: 8 leaf trajectories organized under a shared root $q$, with intermediate ancestor nodes $c$, $c_1$, $c_2$, $c_{2,1}$, $c_{2,2}$, etc. A trajectory ending at leaf $c_{2,2}$ has subgroups defined by ancestors $c_2$ (the immediate parent), $c$ (the grandparent), and $q$ (the root). Each of these ancestor nodes defines a set of sibling trajectories that form a comparison group for advantage computation.

3.4.3.3 The TreePO Advantage Formula

The final advantage for token $t$ in trajectory $i$ is computed as the averaged and normalized sum of subgroup-specific advantages across all depths:

$$\hat{A}_{i,t} = \frac{\sum^{J}_{j=1} \hat{A}_{i,t,j}}{|J| \cdot std(\{\hat{A}_{i,t,j}\}^{J-1})}$$

where $\hat{A}_{i,t,j}$ is the subgroup advantage at depth $j$:

$$\hat{A}_{i,t,j} = R_i - mean(\{R_{i,j}\}^{G_j})$$

subject to the constraint that $std(\{R_i\}^{G}) \neq 0$ (the full group has non-zero reward variance, i.e., not all trajectories are correct or all incorrect—the DAPO filtering condition).

Symbol definitions:

• $R_i$ is the scalar outcome reward for trajectory $i$ (typically 1.0 for correct final answer, 0.0 otherwise);
• $\{R_{i,j}\}^{G_j}$ is the set of outcome rewards for all trajectories in subgroup $G_j$ (the trajectories that share the same ancestor at depth $j$ as trajectory $i$);
• $mean(\{R_{i,j}\}^{G_j})$ is the average reward within that subgroup—the baseline that trajectory $i$'s reward is compared against;
• $\hat{A}_{i,t,j}$ is the subgroup-specific advantage: how much better or worse trajectory $i$'s outcome is compared to the average outcome of its depth-$j$ siblings;
• $|J|$ is the number of depth levels in the trajectory (the number of segments), so the numerator averages over all depths;
• $std(\{\hat{A}_{i,t,j}\}^{J-1})$ is the standard deviation of the subgroup advantages across depths $j=1$ to $J-1$ (the paper uses $J-1$ in the standard deviation, excluding the root-level $j=J$ which would be the full-group normalization already covered by the DAPO baseline);

What this computes, operationally: For a given trajectory $i$ with outcome reward $R_i$ (say, 1.0 because it was correct):

1. At depth 1, look at all other trajectories that branched from the same root-level fork. Compute their mean reward (say, 0.6 because these trajectories were mixed). The depth-1 advantage is $R_i - 0.6 = +0.4$—this trajectory outperformed its immediate siblings.

2. At depth 2, look at the tighter subgroup of trajectories that share the same depth-2 ancestor. Their mean reward might be 0.8 (these are trajectories that all went down the same "promising" subtree). The depth-2 advantage is $R_i - 0.8 = +0.2$—still positive but less extreme, because the comparison group is more similar.

3. At depth 3, the subgroup might be just this trajectory and one sibling that also happened to be correct (mean = 1.0). The depth-3 advantage is $1.0 - 1.0 = 0.0$—no advantage because the local comparison group was uniformly correct.

4. These per-depth advantages $[+0.4, +0.2, 0.0, \dots]$ are summed and divided by $|J|$ to get a mean per-depth advantage, then divided by the standard deviation of these advantages across depths to normalize the scale.

The result $\hat{A}_{i,t}$ is a scalar assigned to every token in trajectory $i$ that reflects not just "was this trajectory better than the average of all trajectories for this query?" (the GRPO advantage) but "was this trajectory consistently better than its progressively more similar sibling groups?" The division by $std(\{\hat{A}_{i,t,j}\}^{J-1})$ normalizes the variance across depths, similar to how REINFORCE++ (Hu et al., 2025) uses global variance normalization to improve robustness.

Why this form over alternatives:

• Versus flat GRPO advantage: The GRPO advantage $R_i - mean(\{R_i\}^G) / std(\{R_i\}^G)$ compares each trajectory against the full group, which is coarse. If the query is hard and most trajectories fail, all correct trajectories get the same large positive advantage regardless of why they succeeded. If the query is easy and most succeed, all incorrect trajectories get the same negative advantage. The tree-based estimator provides multiple levels of contrast: a trajectory that succeeded by taking an unusual path (high advantage against the root subgroup but low advantage against the immediate siblings, who also succeeded) receives a different signal from one that succeeded along the common path (low advantage at all levels).

• Versus MCTS-style advantage (value difference between parent and child): TreeRL and SPO compute advantage as the difference in value estimates between a parent node and its immediate child. This is a local signal that captures the incremental contribution of the last decision. TreePO's subgroup averaging captures a global signal: it compares the final outcome of the entire trajectory against the average outcome of all trajectories that made the same decision at a particular decision point. This is more robust to noise in individual value estimates and does not require a learned value function.

• Versus subgroup-size weighted averaging (Equation 6): The paper explicitly tests an alternative where subgroup advantages are weighted by $|G_j|$ (the number of trajectories in the subgroup), then sums and normalizes:

  $$\hat{A}_{i,t} = \frac{\sum^{J}_{j=1} |G_j| \cdot \hat{A}_{i,t,j}}{std(\{\hat{A}_{i,t,j}\}^{J-1}) \sum^{J}_{j=1} |G_j|}$$

  The paper finds that simple averaging outperforms size weighting, because weighting by subgroup size "over-emphasizes large/easy subgroups and down-weights informative small/hard ones." A large subgroup near the root (many trajectories share the same first segment) dominates the weighted average, drowning out the more granular signal from deeper subgroups with fewer members. Equal averaging preserves the signal from all depths.

• Why exclude the root-level subgroup from the standard deviation? The root-level subgroup $G_J$ (or the full group $G$) is the same as the DAPO normalization baseline. Including it in the $std$ computation would double-count the full-group variance and potentially inflate the normalization factor. The paper's ablation (Section 4.2, "Removing the root-group advantage does not degrade performance") shows that dropping the root-group term from the aggregation entirely (using only $j=1$ to $J-1$) yields comparable curves, suggesting the subgroup signals already capture the relevant variance and the full-group term is partially redundant.

• Why incorporate global variance normalization: The paper mentions that "probability-based branching could bring potential turbulent rollout rewards across queries." With heuristic branching, some queries may receive more exploration budget than others, leading to higher variance in reward distributions across queries. The $std(\{\hat{A}_{i,t,j}\}^{J-1})$ term provides per-trajectory normalization that partially compensates for this query-level variance, borrowing from REINFORCE++'s finding that global variance normalization improves robustness.

---

3.4.4 Heuristic Probability-Based Branching Assignment

Section 4.4 of the paper explores using the model's own token-level log-probabilities (which are returned from the inference engine by default during sampling) to guide how the branching budget is allocated among active paths. This is not the default TreePO setting (the default is uniform allocation or fixed binary tree), but it represents the paper's investigation into whether more sophisticated heuristic control can improve exploration.

3.4.4.1 Mechanism

At a given tree depth $d$, the total branching budget is $2^d$ (the number of branches the binary tree would naturally have at that depth). This budget must be allocated among the currently active search paths. The log-probability-based allocation works as follows:

1. After generating a segment for each active path, the inference engine returns the aggregated log-probability of that segment under the policy model. This is the sum of log-probabilities of all tokens in the segment given the prefix.

2. These log-probabilities are passed through a softmax function with a specified temperature $T$:

   $p_k = \frac{\exp(\text{logprob}_k / T)}{\sum_{m} \exp(\text{logprob}_m / T)}$

   where $\text{logprob}_k$ is the log-probability of the last generated segment for active path $k$, $T$ is the temperature (set to 2.0 in the static experiments, scheduled from 5.0 to 1.0 in the scheduled variant), and $p_k$ is the fraction of the total branching budget allocated to path $k$.

3. Each active path receives $\max(1, \text{round}(p_k \times 2^d))$ branches, with the constraint that every active path is guaranteed at least one branch (to prevent entire subtrees from being pruned prematurely).

What this computes: Paths with higher log-probability (the model is more confident in the segment just generated) receive either more or fewer branches depending on the policy setting:

• "High Prob Encourage": The softmax is applied directly to log-probabilities, so high-probability paths get more budget. This is exploitation-biased: invest compute in paths the model already thinks are promising.
• "Low Prob Encourage": The negative log-probability is used instead (or the probabilities are inverted), so low-probability paths get more budget. This is exploration-biased: invest compute in paths the model is uncertain about or finds surprising.
• "Scheduled Low Prob Encourage": Starts with high temperature ($T=5.0$) to strongly encourage low-probability paths early in training, then anneals to $T=1.0$ over the course of training to gradually shift toward a more balanced allocation.

3.4.4.2 Results: Why Forced Exploration Fails

The paper finds that both static heuristic controls underperform the uniform baseline, and the "Low Prob Encourage" strategy is particularly harmful:

• "Low Prob Encourage" consistently yields the lowest accuracy on both MATH and AIME benchmarks, accompanied by significantly increased response length and entropy loss. The paper interprets this as evidence that "forcing the model to explore low-probability states leads to less efficient and coherent search trajectories." Low-probability paths in a base model are often low-probability for good reason—they represent nonsensical continuations, arithmetic errors, or logically invalid steps. Allocating more compute to these paths generates longer, rambling, incorrect trajectories that provide poor training signal.

• "High Prob Encourage" performs better but still underperforms the baseline, with the lowest entropy and shortest responses, "indicating a potentially overly greedy search that may prune promising, less obvious paths too early." The model becomes overconfident and fails to discover alternative reasoning strategies.

• The scheduled variant does not rescue the approach, despite starting with similar branching behavior to the uniform baseline (high temperature means near-uniform allocation early in training). The paper speculates that the harm from forced exploration compounds over training because the poor data generated from low-probability paths updates the policy in detrimental directions, creating a feedback loop.

The paper's conclusion from this ablation is important for understanding the limits of heuristic control: "merely forcing the model to explore more diverse paths is not beneficial; the exploration must be meaningful." The uniform branching allocation, despite being simpler, appears to strike an effective balance—it provides enough exploration through randomness in the sampling temperature and the inherent stochasticity of the policy, without pushing the model into degenerate regions of the output space.

4. Key Insights and Innovations

Innovation 1: Test-Time Compute Can Be Reorganized into a Tree, Not Just a Set of Independent Chains, Without Loss of Training Quality

The dominant mental model for on-policy RL rollouts in LLM training is embarrassingly parallel: sample G independent complete trajectories per prompt, score them, compute advantages, and update. This model is simple and maps cleanly onto batch inference, but it treats every token of every trajectory as computationally independent, ignoring what the paper's case study (Section 2.1, Figure 2) makes empirically visible — that reasoning trajectories from an aligned model consistently share extensive common prefixes in their generated content, not just in the prompt. The field has known about prompt-level KV-cache sharing for some time, and systems like DAPO, Infinite Sampling, and Truncated PPO have optimized around that. But the realization that the model's own generated reasoning — the variable assignments, the initial equation setups, the problem restatements — is also largely shared across independent samples changes what is possible.

Prior work on tree search for LLMs (TreeRL, SPO, MCTS variants) recognized the value of tree structure but paid for it with sequential node-by-node expansion that killed batching efficiency. The conceptual move TreePO makes is to decouple tree management from generation scheduling: the tree is a data structure describing how segments relate to one another, but the inference engine still processes all active segments in one batched forward pass per depth level. This means the tree provides structural information for credit assignment and exploration control, while the actual computation remains batch-parallel and GPU-efficient. The paper does not claim this decoupling is theoretically profound — it is an engineering insight — but it is the enabling idea that makes tree-structured rollouts practical at training scale. Without it, trees are too slow; with it, the 22–43% GPU-hour reductions in Table 2 become achievable.

The significance is not just the efficiency numbers (which are downstream of the architecture), but the reframing of on-policy sampling as a controllable search process rather than a black-box generation step. Prior work treated the rollout phase as something to optimize for throughput (better batching, shorter truncation windows) but not as something to inform. TreePO treats the rollout structure as first-class information: where trajectories diverge tells us something about uncertainty; which prefixes are shared tells us something about consensus; which paths terminate early tells us something about dead ends. This transforms sampling from a passive data-generation step into an active exploration mechanism, which is a conceptual shift with implications beyond this specific implementation.

Evidence: Figure 1 (Right) visualizes the architectural contrast, but the core empirical support is Table 2, which shows that tree-based sampling maintains or improves accuracy (58.21% → 58.06% at b=8 with More Init Divergence) while reducing GPU hours by 22%, and at more aggressive efficiency settings (b=2) achieves comparable or slightly lower accuracy with 43% GPU-hour savings. The fact that tree sampling can match sequential accuracy while being substantially cheaper validates the decoupling claim.

---

Innovation 2: A Hierarchical, Subgroup-Based Advantage Estimator That Uses the Tree Structure for Multi-Resolution Credit Assignment Without Learned Value Functions

In standard GRPO (and DAPO), the advantage for every token in trajectory i is a single scalar: the group-normalized outcome reward (R_i - mean({R})) / std({R}). This is computationally trivial but informationally coarse — all trajectories with the same outcome get the same advantage, whether they arrived at that outcome through dramatically different reasoning paths or nearly identical ones. This creates a credit assignment blind spot: if two trajectories both produce correct answers, the policy receives identical positive reinforcement for all their tokens, even if one trajectory contained a critical reasoning step that the other lacked, or if both trajectories are actually identical up to a late-stage branching point where one made the right decision and the other got lucky.

The field has responded to this blind spot in two ways. One is process reward models (PRMs), which require expensive training data (human annotations or Monte Carlo rollouts) and add inference overhead. The other is segment-level methods like SPO that compute MCTS-style parent-child value differences, which capture local credit but require a value function and are sensitive to estimation error. TreePO's innovation is a third path: use the tree structure that already exists from the sampling process to compute multiple advantage signals at different resolutions, then aggregate them.

The idea is that a trajectory belongs to multiple nested comparison groups — all trajectories from this query (coarsest), all trajectories sharing the first reasoning segment, all trajectories sharing the first two segments, etc. — and its advantage at each level is "how much better or worse was my outcome than the average outcome of trajectories that made the same decisions up to this point?" A trajectory that succeeded by taking an unusual path will have high advantage at coarse levels (it outperformed most trajectories, which failed) but moderate advantage at fine levels (its immediate siblings, who made the same unusual choice, also succeeded). A trajectory that succeeded on the common path will have moderate advantage at all levels. A trajectory that failed despite following the common path will have negative advantage at coarse levels but near-zero advantage at fine levels (its immediate siblings also failed). These signals are qualitatively different and provide the policy with richer feedback about which decisions mattered.

The finding that simple averaging across subgroup depths outperforms subgroup-size weighting (Section 4.2) is important here: it shows that the signal from small, tight subgroups (which may contain only 2–3 trajectories that made the same rare decision) is disproportionately informative and should not be drowned out by the signal from large, loose subgroups. The further finding that removing the root-group advantage entirely does not degrade performance suggests that the multi-resolution signal is substitutive for the coarse group normalization, not merely supplementary — the tree structure provides enough granularity that the traditional group normalization becomes partially redundant.

This is a diagnostic contribution as much as a methodological one: it demonstrates that the structure of how rollouts are generated can provide useful credit assignment signal without additional models, annotations, or value functions. It does not claim to be as precise as a well-trained PRM, but it is essentially free — the tree structure is a byproduct of the sampling algorithm.

Evidence: The validation curves in Figure 1 (Left, Mid) show that adding the TreePO advantage estimator to tree sampling ("TreePO w/ Fixed Init Divergence" and "TreePO w/ More Init Divergence") produces smoother, more stable training than tree sampling alone with standard advantages ("GRPO w/ TreePO Sampling"), which already stabilizes over vanilla GRPO. Table 1 shows the accuracy improvement from adding the estimator: 54.61% → 56.88% for Fixed Init Divergence and 54.61% → 58.21% for More Init Divergence. The ablation in Section 4.2 (Figure 6) demonstrates that simple averaging over subgroups outperforms size-weighted averaging and that subgroup-level rejection degrades performance, providing evidence for the specific design choices.

---

Innovation 3: A Negative Result Establishing That Forcing Exploration Toward Low-Probability Reasoning Paths Is Counterproductive

The reinforcement learning literature, particularly in continuous control and game-playing domains, has developed a rich set of exploration bonuses, entropy regularization techniques, and curiosity-driven objectives designed to prevent policies from prematurely converging to suboptimal deterministic strategies. The intuition — that encouraging the agent to visit low-probability states prevents it from missing high-reward regions of the space — is deep and well-supported. It is natural to ask whether a similar principle applies to LLM reasoning: should we allocate more sampling budget to reasoning paths the model considers unlikely, on the theory that these might represent creative or non-obvious solutions?

TreePO's segment-level branching control provides a clean experimental testbed for this question. Because the inference engine returns token-level log-probabilities by default, the system can compute the model's confidence in each generated segment and redirect the branching budget toward high-confidence paths (exploitation), low-confidence paths (exploration), or maintain a uniform allocation (baseline). Section 4.4 reports the results of these experiments, and they are decisively negative for forced exploration.

The "Low Prob Encourage" strategy — allocating more branches to segments with lower model probability — produces the worst accuracy on both MATH and AIME across all tested variants. It also produces substantially longer responses and higher entropy loss, indicating that the model is not just exploring but floundering: the low-probability paths it is forced down are low-probability for valid reasons (they represent arithmetic errors, logical contradictions, or nonsensical continuations), and the additional compute spent exploring them generates poor training data that actively degrades the policy. Even a scheduled variant that starts with strong exploration (temperature 5.0) and anneals to balanced allocation fails to match the uniform baseline.

The "High Prob Encourage" strategy — allocating more branches to high-confidence paths — performs better than forced exploration but still underperforms the uniform baseline, with the lowest entropy and shortest responses, suggesting it is overly conservative and fails to discover alternative strategies.

This is a significant finding because it pushes back against a natural extrapolation from the RL exploration literature. In LLM reasoning, probability and quality are correlated in a way they are not in, say, Atari games: a low-probability action in an Atari game might be a rare but crucial strategic move; a low-probability token sequence in a math reasoning chain is more likely to be a mistake. The paper's result suggests that the model's own uncertainty estimates, at least for base models in mathematical reasoning, are reasonably well-calibrated for distinguishing promising from unpromising paths, and that uniform allocation with stochastic sampling temperature already provides sufficient exploration diversity. More sophisticated exploration mechanisms would need to distinguish between "low-probability because creative" and "low-probability because wrong" — a distinction that raw log-probability cannot make.

This is properly understood as a negative result with diagnostic value rather than a performance improvement. It identifies a boundary condition for heuristic exploration strategies and provides empirical guidance for future work: exploration in LLM reasoning spaces should be guided by something other than (or in addition to) raw model confidence.

Evidence: Figure 8 in Section 4.4 shows the training curves for the different branching strategies. "Low Prob Encourage" is consistently the lowest accuracy on both benchmarks, accompanied by the highest response lengths and entropy. "High Prob Encourage" shows the lowest entropy and shortest responses, underperforming the baseline. The scheduled variant tracks near the baseline early in training (when temperature is high and allocation is near-uniform) but diverges downward as training progresses, supporting the interpretation that forced exploration compounds harm over time.

---

Innovation 4: Reframing the Pretraining–Inference Tradeoff as an Online Training–Efficiency Tradeoff Enabled by Tree-Structured Rollouts

The companion paper summary (the reference example in the prompt) examined the tradeoff between spending compute on pretraining a larger model versus spending it on test-time inference strategies for a fixed model. TreePO addresses a related but distinct question: given a fixed model architecture and training compute budget, how should the training-time sampling be structured to maximize the policy improvement per GPU-hour? This is a training-efficiency question, not a deployment-efficiency question, and it has been relatively underexplored in the RL-for-reasoning literature, which has focused more on algorithmic improvements (better clipping, dynamic sampling, reward design) than on the structure of the data generation process itself.

The paper's efficiency analysis in Section 4.1 provides a nuanced picture that goes beyond "tree sampling is faster." It shows that the optimal tree configuration (depth vs. segment length) is model-specific: Qwen2.5-7B-Instruct peaks at depth 28, Qwen2.5-Math-7B peaks at depth 14, and Qwen2.5-Math-7B-Instruct shows a split where token-level throughput peaks at depth 14 but trajectory-level throughput peaks at depth 28 or 56. This model-specificity is not a weakness of the approach — it is a diagnostic finding that reveals how different model types (base vs. instruct, general vs. math-specialized) generate reasoning traces with different structural properties. Instruction-tuned models produce more structured, aligned prefixes that benefit from shallower trees with longer segments (better KV-cache reuse); math-specialized base models produce more diverse traces that benefit from deeper trees with more branching opportunities.

The rollout scaling analysis (Figure 5) adds further nuance: Qwen2.5-7B-Instruct shows nearly linear throughput gains as rollout count increases under tree sampling, while Qwen2.5-Math-7B peaks at around 16 rollouts and then degrades. The interpretation — that instruction tuning produces more structured, cache-friendly outputs while base models diverge more, fragmenting the KV-cache and increasing management overhead — is a concrete, empirically grounded insight about how model alignment state interacts with inference efficiency.

This reframing matters because it suggests that the choice of sampling architecture should be part of the training design, not an afterthought. Current practice treats the inference engine as a black box that generates completions on demand; TreePO argues that by exposing the tree structure and tuning the depth–segment configuration to the model type, practitioners can extract substantially more training progress from a fixed compute budget. The 22–43% GPU-hour savings are not just a cost reduction — they mean that for the same budget, a team can train for more steps, explore more hyperparameter configurations, or scale to larger models. This has practical implications for how RL training pipelines should be architected.

Evidence: Figure 4 shows the depth–segment tradeoff for three model variants, with peak throughput configurations labeled. Figure 5 shows rollout scaling behavior. Table 2 provides the concrete GPU-hour comparisons, ranging from 3.65 hours (b=2, More Init Divergence) to 6.40 hours (Sequential baseline) for achieving comparable accuracy. The fact that tree sampling at b=8 achieves 58.06% accuracy (within 0.15 percentage points of sequential's 58.21%) while saving 22% GPU hours (5.05 vs. 6.40) is the cleanest demonstration of the training-efficiency tradeoff in action.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The training data draws from two sources: the MATH dataset (Hendrycks et al., 2021), contributing approximately 8,000 queries of difficulty levels 3–5 following the same setting as SimpleRL (Zeng et al., 2025), and the DeepScaler collection (Luo et al., 2025), contributing approximately 40,000 additional samples. For evaluation, the paper uses five benchmarks: AIME 2024, AMC 2023, MATH500, MINERVA (Lewkowycz et al., 2022), and Olympiad Bench (He et al., 2024). The overall metric is a weighted average across these benchmarks based on test set sizes.

• Base model(s). The primary RL training experiments use Qwen2.5-7B base model (Yang et al., 2025) — a 7-billion-parameter model trained without instruction tuning or SFT — chosen because the work targets the "RL-zero" paradigm where reasoning capabilities are elicited from a base model without prior supervised fine-tuning. For offline efficiency analyses, the paper additionally uses Qwen2.5-7B-Instruct and Qwen2.5-Math-7B-Instruct to compare sequential versus tree-based sampling throughput on already-aligned models.

• Metrics. The primary performance metric is Major@16 accuracy: for each query, the model generates 16 complete trajectories and the final answer is selected by majority voting among the 16 candidates; the reported percentage is the fraction of test queries where this majority answer matches ground truth. To produce this metric, the paper samples 1,000 sets of 16 rollouts per query (with majority voting applied to each set) and reports the average accuracy. For efficiency, the paper reports GPU hours (total wall-clock time multiplied by number of GPUs used, as recorded in Table 2), Tokens per second (TokenPS; total model-processed tokens including prefill and decode, measured as wall-clock throughput), and Trajectories per second (TrajPS; completed continuations per second). Validation curves during training track accuracy on held-out evaluation sets along with entropy loss and average response length.

• Baselines. The primary baselines are:

  • GRPO (Shao et al., 2024, "DeepSeekMath"): the original Group Relative Policy Optimization objective with group-normalized outcome advantages, standard i.i.d. sequential sampling, and without DAPO's improvements.
  • GRPO w/ TreePO Sampling: the GRPO/DAPO objective with tree-based sampling replacing sequential sampling, but using the standard (flat) group-normalized advantage rather than the tree-based advantage estimator. This isolates the effect of tree sampling alone.
  • TreePO w/ Fixed Init Divergence: the full TreePO system (tree sampling + tree-based advantage estimator) with fixed initial branching at the root (binary tree structure).
  • TreePO w/ More Init Divergence: the full TreePO system with additional random initial branching (2–8 divergences allocated at the root) to increase diversity.
  • Sequential sampling (Table 2): a baseline that uses the same DAPO objective and training setup as TreePO but generates 16 complete independent trajectories per query via standard autoregressive decoding rather than tree-structured generation. This is the primary efficiency comparison point.

• Generation budget / compute accounting. For training, the per-query generation budget is fixed at w = 16 complete trajectories with a maximum response length of 7 × 1024 = 7,168 tokens per trajectory. The tree depth–segment configurations define how this budget is structured: 28 × 256, 14 × 512, and 7 × 1024 are the three main search configurations, all yielding the same total token budget per trajectory. The branching budget at depth d is 2^d by default (binary tree). For test-time compute scaling analysis (Section 4.5, Figure 9), the x-axis represents the number of rollouts on a log scale, starting from d (the initial divergence factor) rather than from 1, to account for the minimum branching inherent in the tree structure. For offline efficiency experiments (Section 4.1), compute is measured as wall-clock throughput (TokenPS, TrajPS) under controlled conditions: single NVIDIA H100 80GB GPU, no tensor or data parallelism, 60% GPU utilization, batch size of 64 prompts with 64 rollouts per prompt for tree sampling.

• Cross-validation / statistical protocol. The training protocol follows DAPO's dynamic sampling: 3× batch size queries are sent to the sampler to acquire groups of 16 trajectories each, and from these, 512 queries with non-zero reward standard deviation (std({R_i}^G) ≠ 0) are randomly selected to form a training batch. When insufficient queries meet this criterion, up to two additional sampling rounds are conducted, which can reduce the total number of training steps due to the data loader's enumeration logic. Checkpoints are saved every 50 training steps. The paper does not report confidence intervals or statistical significance tests for the major accuracy numbers; results are presented as point estimates. For the efficiency experiments, the paper uses a held-out prompt pool independent of model training, used solely for throughput measurement.

Main Quantitative Results

Training Performance and Stability (Table 1, Figure 1)

Headline finding: TreePO with More Initial Divergence achieves 58.21% overall Major@16 accuracy across all five benchmarks, compared to 46.63% for vanilla GRPO — an 11.58 percentage point absolute improvement. The gain comes primarily from two compounding changes: adopting tree-based sampling over i.i.d. sequential generation, and adopting the tree-based advantage estimator over flat group normalization.

Detailed results from Table 1:

What each comparison reveals:

• Tree sampling alone versus GRPO: Moving from 46.63% (GRPO) to 54.61% (GRPO w/ TreePO Sampling) is a 7.98 percentage point gain attributable entirely to the tree-based sampling mechanism with standard flat advantages. This is the paper's strongest evidence that tree-structured rollouts are not just more efficient but also qualitatively better for training — the tree structure itself produces more informative data even before advantage estimation is modified. The gains are particularly large on MATH (72.89% → 81.85%, +8.96 points) and AMC (44.42% → 51.63%, +7.21 points), suggesting tree sampling particularly benefits benchmarks where reasoning diversity matters.

• Adding the tree-based advantage estimator on top of tree sampling: From 54.61% (GRPO w/ TreePO Sampling) to 56.88% (Fixed Init Divergence) is a 2.27 point gain, and to 58.21% (More Init Divergence) is a 3.60 point gain. These are the marginal contributions of the hierarchical advantage estimation over the flat group normalization, given that tree sampling is already active. The gains are concentrated in AIME (19.66% → 28.89% for Fixed, 27.83% for More) and Olympiad Bench (44.76% → 47.75% / 49.15%), the two hardest benchmarks, suggesting the multi-resolution advantage signal is most valuable when outcome rewards are sparse and trajectories are diverse.

• Fixed versus More Initial Divergence: The More Init Divergence variant (additional 2–8 random branches at the root) achieves slightly higher overall accuracy (58.21% vs. 56.88%) but is not uniformly better — Fixed Init Divergence wins on AIME (28.89% vs. 27.83%) and AMC (56.63% vs. 55.53%). This suggests that additional initial branching helps on some problem distributions but can add noise on others; the overall benefit is modest (1.33 percentage points) and inconsistent across benchmarks.

Training stability evidence from Figure 1 (Left, Mid): The validation performance curves show qualitative differences in training dynamics:

• GRPO (blue line) exhibits high volatility with sharp peaks and troughs, making checkpoint selection unreliable and suggesting the training process is fragile to the specific batch composition.
• GRPO w/ TreePO Sampling (orange line) is substantially more stable, with smoother improvement and fewer sharp drops, even though it uses the same flat advantage estimator as GRPO. This indicates that tree sampling alone provides a stabilizing effect, likely because the shared prefixes and early stopping produce more consistent trajectory qualities across batches.
• TreePO with the full advantage estimator (green line) shows the most stable and consistently high-performing trajectory, sustaining higher validation accuracy throughout training with minimal regression. The paper states this indicates "the estimator component provides a more precise reward signal based on the tree hierarchy, guiding the model, and leading to more reliable convergence."

Important caveat from Figure 1: The paper notes that "although replaced additional tree-based sampling causes a slower convergence, it could stabilize the training." The tree-based variants initially lag behind GRPO in the very early steps of training, suggesting that the structural constraints of tree sampling (fixed segment boundaries, early stopping) may slow the initial exploration burst, but this is compensated by superior long-term stability.

---

Training Efficiency: GPU Hour Savings (Table 2)

Headline finding: Tree-based sampling reduces training GPU hours by 12% to 43% compared to sequential sampling at the same batch size and rollout count, while maintaining comparable or only slightly reduced accuracy.

Detailed results from Table 2 for the Fixed Init Divergence setting:

And for More Init Divergence:

The efficiency–accuracy tradeoff is visible in the branching budget b:

• At b=2 (most aggressive pruning, fewest branches per node), GPU hours drop dramatically (26–43% reduction) but accuracy also drops (56.03% vs. 56.88% for Fixed; 54.67% vs. 58.21% for More). The smaller branching budget means the tree explores fewer alternative paths, missing potentially correct reasoning branches, but the inference batches are larger and more efficient because KV-cache sharing is maximized.

• At b=8 (least aggressive pruning, most branches), accuracy is nearly identical to sequential (56.60% vs. 56.88% for Fixed; 58.06% vs. 58.21% for More) while still saving 12–22% GPU hours. This is the paper's headline efficiency claim: tree sampling can match sequential training quality while being measurably cheaper.

• At b=4, the tradeoff is intermediate: moderate accuracy differences with moderate GPU savings (17–29%).

Key interpretation: The GPU hour savings come primarily from the KV-cache reuse across shared prefixes. When many trajectories share the same initial reasoning segments, the tree sampler computes those segments once and reuses the cached attention keys and values for all descendant branches, avoiding redundant forward passes. The More Init Divergence setting (which has extra random branches at the root) shows larger absolute GPU hour savings at b=2 (3.65 vs. 6.40, 43% reduction) but also larger accuracy degradation (54.67% vs. 58.21%, a 3.54 point drop), suggesting that the additional initial branches create more divergent trajectories that reduce cache-hit rates at aggressive pruning levels.

Cross-reference with validation curves: The accuracy numbers in Table 2 are final checkpoint values evaluated with sequential sampling (Major@16), meaning they reflect the policy's capability after training, not the training-time accuracy. The fact that tree-sampled models perform well when evaluated with sequential decoding (which they were not trained with) indicates that tree sampling does not cause the policy to overfit to the tree structure — the learned reasoning capabilities transfer to standard autoregressive generation.

---

Offline Sampling Efficiency (Section 4.1, Figures 4 and 5)

Headline finding: Across three Qwen2.5 model variants, tree-based sampling yields on average +40% TrajPS (trajectories per second) and +30% TokenPS (tokens per second) compared to conventional sequential sampling at the same batch size, rollout count, and per-trajectory token budget (7,000 tokens). Efficiency peaks at intermediate depth–segment configurations rather than growing monotonically with tree depth, and the optimal configuration is model-specific.

Efficiency by depth–segment configuration (Figure 4):

The paper tests configurations from 7 × 1024 (shallow, long segments) to 112 × 62 (deep, short segments), all with the same total token budget of approximately 7,000 tokens per trajectory. The key patterns:

• Qwen2.5-7B-Instruct (Figure 4a): Both TokenPS and TrajPS peak at depth 28 (28 × 250 tokens per segment). The instruct model benefits from a mid-depth balance where "segments are not too short (better batched prefilling and context retention) while depth still yields sufficient decoding parallelism." Throughput declines at both shallower and deeper configurations — too shallow means less KV-cache reuse from shared prefixes; too deep means very short segments that incur repeated prefilling overhead with each inference call.

• Qwen2.5-Math-7B-Instruct (Figure 4b): TokenPS peaks at depth 14 (14 × 500) while TrajPS peaks at both depth 28 and 56 (28 × 250, 56 × 125). The paper interprets this split as a tradeoff between token-level efficiency (longer segments reduce recomputation overhead, favoring shallower trees) and trajectory-level parallelism (deeper trees enable more branching and parallel rollouts, favoring more depth). The math-instruct model's structured outputs (consistent mathematical formatting) maintain high cache-hit rates even at moderate depths.

• Qwen2.5-Math-7B (Figure 4c): Both metrics peak at depth 14 with a clear dropoff at higher depths. The base math model, lacking instruction tuning, produces less structured reasoning traces that diverge more at early stages, reducing the benefit of deeper trees (fewer shared prefixes to cache) and making longer segments preferable to amortize prefill costs.

Why the peak exists: The paper explains that "prefill prefers longer segments and shallower trees, which reduces repeated KV cache and attention computation; decoding prefers deeper trees with more branches and parallel rollouts, better exploiting speculative execution and batched sampling. If segments are too short, the extra recomputation offsets the gains from depth, and the peak appears where these opposing effects balance." This is a concrete engineering insight: tree sampling is not universally faster; it is faster only when the depth–segment configuration is matched to the model's output structure.

Rollout scaling behavior (Figure 5):

• Qwen2.5-7B-Instruct (Figure 5a): Tree-based TokenPS and TrajPS grow nearly linearly as rollout count increases (from 8 to 128 rollouts per prompt, with query count fixed at 64 and tree depth 28), reaching roughly 2× the baseline throughput. Standard autoregressive decoding shows only modest gains from more rollouts (batched decoding helps marginally). The linear scaling indicates that the instruct model's outputs share large common prefixes consistently, enabling near-constant cache-hit rates regardless of how many trajectories are generated.

• Qwen2.5-Math-7B-Instruct (Figure 5b): Tree sampling maintains a stable ≈2× speedup over sequential across all rollout counts. The paper attributes this to "structured, semantically aligned math trajectories" that "sustain high cache-hit rates and efficient KV reuse, keeping batched decoding effective."

• Qwen2.5-Math-7B (Figure 5c): Throughput is non-monotonic with rollout count: TokenPS and TrajPS peak around 16 rollouts, then decline. The paper explains this as a trajectory divergence effect: as more rollouts are generated, the base model's outputs increasingly diverge from shared prefixes, reducing cache-hit rates. Additionally, "KV-cache fragmentation and management overhead grow, memory pressure rises, and batching efficiency degrades; the lack of instruction tuning further loosens output structure." This is a critical finding: tree sampling's efficiency advantage depends on trajectory convergence, and base models without alignment produce less convergent traces.

Practical implication: The model-specific optimal configuration means that deploying tree sampling in a training pipeline requires configuration tuning — there is no universally optimal depth–segment pair. The instruction-tuned models peak at depth 28, the math-specialized base model peaks at depth 14, and the instruct-math hybrid shows a split optimum. The paper does not provide a method for predicting the optimal configuration without benchmarking, which is a limitation for practitioners.

---

Test-Time Compute Scaling (Section 4.5, Figure 9)

Headline finding: TreePO generates a family of compute-accuracy scaling curves parameterized by the tree divergence factor d (the number of branches at each divergence point), rather than the single curve produced by sequential sampling's N parameter. At low compute budgets, smaller d is more efficient; at high compute budgets, larger d achieves higher peak accuracy. This enables "compute-optimal inference" where the tree structure is selected based on the available budget.

Results from Figure 9: The x-axis is the compute budget (number of rollouts) on a log scale, with tree-based sampling curves starting from d rather than 1 (because the minimum rollout count for a tree with divergence factor d is d). The y-axis is average performance on the aggregated benchmark.

• At the lowest compute budget (roughly 4–8 rollouts), d=2 achieves the highest accuracy — the narrower tree uses its limited budget more efficiently by focusing on a few deeper paths rather than spreading thin across many branches.
• As the budget increases (roughly 32–64 rollouts), d=4 overtakes d=2, reaching higher accuracy at the same budget.
• At the highest budgets tested (256+ rollouts), d=8 achieves the highest peak performance, though it required the largest compute budget to reach that peak.

What this demonstrates: Sequential sampling scaling (increasing N, the number of independent samples) traces a single curve because the generation strategy doesn't change — only the quantity changes. TreePO's tree-structured generation traces different curves depending on the internal search strategy (how many branches, how deep, how aggressively to prune). This means that for any given compute budget, there is a "compute-optimal" tree configuration that maximizes accuracy, and this configuration changes with the budget. The paper does not provide a method for predicting which d is optimal at a given budget without running the experiment, but the existence of the family of curves is the claim being demonstrated.

The paper explicitly states: "Rather than simply scaling the number of samples, one can select the optimal tree structure to maximize performance for a given computational constraint." This frames tree sampling not just as a training efficiency tool but as an inference-time strategy selection mechanism, analogous to the compute-optimal test-time scaling framework from the companion paper reference example, but parameterized by tree structure rather than search algorithm choice.

Limitation: Figure 9 reports only the single aggregated benchmark curve. Per-benchmark scaling curves are not provided, so it is unclear whether the d-optimality pattern holds uniformly across easy, medium, and hard problems, or whether (as in the companion paper's findings) the optimal strategy depends on problem difficulty. Given that Section 4.4 found difficulty-dependent effects for branching strategies, it is plausible that the optimal d for test-time scaling also varies by problem difficulty, but this is not tested.

---

Ablation Studies and Robustness Checks

Subgroup-size weighted vs. simple averaging in advantage estimation (Section 4.2, Figure 6): Simple averaging across subgroup depths (Equation 5) outperforms weighting by subgroup size (Equation 6) on both MATH and AIME accuracy, with lower and more stable entropy and no unnecessary growth in response length. Size-weighting over-emphasizes large, loose subgroups near the root (which contain many trajectories sharing only an early common prefix) and down-weights small, tight subgroups deeper in the tree (which contain few trajectories that made similar decisions and thus provide more diagnostic contrast). This ablation validates a specific design choice in the advantage estimator and demonstrates that subgroup informativeness is inversely related to subgroup size — the smaller, deeper subgroups carry disproportionately useful signal.

Subgroup-level dynamic rejection sampling (Section 4.2, Equation 7, Figure 6): Applying DAPO-style rejection sampling at the subgroup level — discarding subgroups where all trajectories are correct or all are incorrect — degrades performance relative to keeping all subgroups. The paper explains that "these 'extreme' subgroups actually calibrate margins; removing them strips away high-signal cases." A uniformly-correct subgroup provides a meaningful baseline: trajectories in that subgroup that are individually correct don't get an artificial positive advantage, and trajectories that are individually incorrect get a clear negative signal. Removing these subgroups eliminates informative contrast. This finding mirrors DAPO's original dynamic sampling insight but at a finer granularity — the principle that "all-correct" and "all-incorrect" groups are useful for calibration applies not just at the query level but at the subgroup level as well.

Removing the root-group advantage (Section 4.2, Figure 6): Using only the aggregated subgroup advantages j = 1 to J-1 (excluding the full-group comparison j = J) yields comparable training curves to the full estimator. This suggests that the multi-resolution subgroup signals already capture the variance that the traditional group normalization provides, making the root-group term partially redundant. The paper frames this as "a promising direction for further analysis of credit assignment" — it implies that the traditional GRPO advantage normalization can be entirely replaced by tree-structured subgroup comparisons without loss of signal, though this is not tested as a standalone configuration in the main experiments.

Misaligned fallback segments (Section 4.2, Figure 6): When the tree uses 7 × 1024 segments but the fallback mechanism uses 512-token segments (creating a mismatch where trajectories share an abstract tree structure but segments are not token-aligned), AIME accuracy drops and response length rises sharply. The paper interprets this as evidence that "token-aligned segments are important for stable optimization and precise stopping behavior." When segments at the same tree depth contain different numbers of tokens or branch at semantically different points, the subgroup definitions become noisy — trajectories grouped together at depth j may not actually share the same reasoning prefix in a meaningful sense, weakening the advantage signal. This is a robustness check that validates the design choice to use equal-length segments.

Depth–segment configuration sweep (Section 4.3, Figure 7): Under the subgroup-size weighted advantage (the same setting used for Section 4.2), four depth–segment configurations are compared: 56 × 128, 28 × 256, 14 × 512, and 7 × 1024, all with the same total token budget per trajectory. The 14 × 512 configuration attains the highest final MATH and AIME accuracy, while 7 × 1024 (the shallowest tree with longest segments) significantly lags, especially on AIME. The paper interprets this as evidence that "deeper trees with moderate segments provide stronger credit assignment than shallow rollouts with very long segments." The best-performing 14 × 512 also drives the largest growth in response length and highest entropy, suggesting that "online TreePO benefits from more exploratory, longer reasoning traces; shorter traces trade accuracy for brevity." This is a crucial tradeoff: the tree configuration that maximizes accuracy also produces longer outputs, which has implications for deployment latency and inference cost.

Probability-based branching assignment (Section 4.4, Figure 8): As discussed in the Key Insights section (Innovation 3), both "Low Prob Encourage" and "High Prob Encourage" heuristic branching strategies underperform the uniform baseline, and the "Low Prob Encourage" variant is consistently worst. The "High Prob Encourage" variant produces the lowest entropy and shortest responses, suggesting overly conservative search. Even the scheduled variant (temperature annealing from 5.0 to 1.0 for "Low Prob Encourage") fails to match the baseline. The ablation establishes that forced exploration toward low-probability segments is counterproductive, and that uniform branching allocation (the default TreePO setting) provides a better exploration–exploitation balance without requiring heuristic tuning.

Training from base model versus instruction-tuned model (implicit across experiments): All training experiments (Tables 1 and 2, Figures 1, 6–8) use the Qwen2.5-7B base model, while the efficiency experiments (Figures 4 and 5) compare base, instruct, and math-instruct variants. This is not a formal ablation but reveals an important robustness dimension: tree sampling's efficiency gains are consistent across model types (all three variants show positive throughput improvements), but the magnitude and optimal configuration vary. The base math model shows the most constrained scaling (throughput peaks at 16 rollouts and degrades after), while the instruct variants scale better. This suggests that tree sampling is broadly applicable but benefits most from models with structured, convergent reasoning patterns — which instruction tuning or domain-specific fine-tuning tends to produce.

---

Critical Assessment

Claim 1: "TreePO reduces GPU hours by 12% to 43% compared to sequential sampling baselines while achieving comparable or improved accuracy."

This claim is the paper's central practical contribution and is supported with qualifications. Table 2 provides clear evidence that tree-based sampling reduces GPU hours, but the "comparable or improved accuracy" part depends on configuration. At b=8 in the More Init Divergence setting, tree sampling achieves 58.06% versus sequential's 58.21% — a 0.15 percentage point difference that is likely within noise (no confidence intervals are reported), and the 22% GPU hour savings is genuine. At b=4 in the Fixed Init Divergence setting, tree sampling achieves 57.50% versus sequential's 56.88% — slightly better accuracy with 17% savings. However, at b=2 in the More Init Divergence setting, tree sampling achieves 54.67% versus 58.21% — a 3.54 point accuracy loss for 43% savings. The claim of "comparable or improved accuracy" thus holds for b=4 and b=8 but not for the most aggressive pruning setting. The paper does not provide guidance on how to select b a priori to achieve the desired accuracy–efficiency tradeoff, which limits the practical applicability.

A deeper concern: The GPU hour measurements in Table 2 are for the training runs reported. It is unclear whether these are single-run measurements or averages over multiple runs. If single runs, the GPU hour differences could reflect variance in training dynamics (some runs converge faster than others for reasons unrelated to sampling efficiency) rather than structural efficiency gains. The paper does not report error bars, standard deviations, or repeated measurements for the GPU hour figures.

Missing baseline: The paper does not compare against simpler efficiency improvements that could be applied to sequential sampling, such as reducing the maximum response length (which would reduce generation cost for both sequential and tree sampling), using a smaller group size G (e.g., G=8 instead of G=16), or applying the same early stopping heuristics to sequential sampling. The sequential baseline generates full trajectories regardless of content; adding the repetitive substring detection from TreePO to sequential sampling might close some of the efficiency gap without requiring tree structure at all.

Claim 2: "Tree-structured rollouts stabilize RL training compared to vanilla GRPO."

This claim is strongly supported by the validation curves in Figure 1 (Left, Mid). The GRPO baseline shows visible volatility with sharp performance drops, while both TreePO sampling variants (with and without the tree-based advantage estimator) show smoother, more monotonic improvement. The visual evidence is clear. However, the paper does not quantify "stability" — there is no metric such as variance of validation accuracy across training steps, number of performance regressions exceeding a threshold, or checkpoint-to-checkpoint consistency. The claim is qualitative and visually evident but not statistically characterized.

A contributing factor that is not disentangled: DAPO's dynamic sampling (rejecting all-correct and all-incorrect groups) is active in the GRPO w/ TreePO Sampling and TreePO configurations, which itself stabilizes training relative to vanilla GRPO (as demonstrated in the DAPO paper). Some of the stability gain attributed to tree sampling may actually be from DAPO's improvements. The paper compares against vanilla GRPO, not DAPO with sequential sampling (except in Table 2, which compares DAPO+sequential vs. DAPO+tree, but only reports final accuracy, not training stability). The stability claim would be stronger with a DAPO+sequential baseline in Figure 1.

Claim 3: "The tree-based advantage estimator enables more precise credit assignment, evidenced by improved training stability and final accuracy when added to tree sampling."

This claim is supported by the ablation in Table 1 and Figure 1. Adding the tree-based advantage estimator to tree sampling (moving from GRPO w/ TreePO Sampling to TreePO w/ Fixed or More Init Divergence) improves overall accuracy by 2.27–3.60 points and further smooths the validation curves. The gains are concentrated on the hardest benchmarks (AIME, Olympiad Bench), which is consistent with the interpretation that multi-resolution advantage estimation is most valuable when rewards are sparse and trajectories diverse — precisely the conditions under which flat group normalization provides the weakest signal.

However, the evidence is somewhat confounded. The TreePO configurations include both the tree-based advantage estimator and the specific branching/fallback policies (Fixed vs. More Init Divergence). The 2.27–3.60 point gain is the combined effect of the estimator and the divergence policy, not a clean ablation of the estimator alone. A cleaner comparison would be: tree sampling with flat advantage and uniform branching versus tree sampling with tree-based advantage and uniform branching, holding the branching policy constant. This is approximated by comparing GRPO w/ TreePO Sampling (which likely uses uniform branching, though the paper does not explicitly state the branching policy for this configuration) against TreePO w/ Fixed Init Divergence, but the Fixed Init Divergence configuration includes the "fixed" divergence policy at the root, which is a confound.

Claim 4: "TreePO provides up to 40% trajectory-level and 35% token-level sampling compute reduction during inference."

This claim is supported by the offline efficiency experiments in Section 4.1, with important caveats. The +40% TrajPS and +30% TokenPS figures (geometric means across configurations in Figure 4) represent throughput improvements under specific controlled conditions: single GPU, no parallelism, 60% utilization, batch size 64, 64 rollouts per prompt for tree sampling. These are throughput benchmarks, not end-to-end latency measurements for a single query. The 40% trajectory-level reduction means the system generates 40% more complete trajectories per second; it does not necessarily mean a single user query receives its answer 40% faster, because tree sampling may increase latency for the first few trajectories (the tree must grow to sufficient width before many trajectories complete) while providing higher throughput in aggregate.

The efficiency gains are model- and configuration-dependent. Figure 4 shows that the throughput advantage varies substantially across depth–segment configurations and model types. The "up to 40%" figure is a best-case number; the actual gain in a specific deployment depends on whether the configuration is tuned to the model. The paper does not provide a methodology for configuration selection, which means practitioners must run their own efficiency benchmarks to identify the optimal depth–segment pair for their model.

A critical missing analysis: The paper does not report whether the tree-sampled model's inference outputs are different in quality from the sequentially-trained model's outputs when both are evaluated with the same decoding strategy (sequential, Major@16). Table 2 compares the trained models using sequential evaluation, which shows they are comparable in accuracy. However, there is no analysis of whether tree sampling during training changes the distribution of output lengths, the diversity of generated reasoning paths, or the calibration of the model's confidence. If tree-trained models produce longer, more verbose reasoning chains (as Figure 7 suggests for the 14 × 512 configuration), this could increase inference latency at deployment even if accuracy is matched, partially offsetting the training-time efficiency gains.

Missing experiment: combining tree sampling with tree-structured inference at test time. The paper evaluates trained models using sequential sampling (Major@16), not tree-based sampling. Since the models were trained with tree-structured rollouts, they might perform better or worse with tree-based inference — this interaction is not tested. If tree-trained models benefit from tree-structured decoding at test time, the inference efficiency story becomes stronger; if they perform worse, it introduces a training–inference mismatch that practitioners would need to navigate.

Generalizability concern: single model family, single domain. All training experiments use Qwen2.5-7B base on mathematical reasoning. The paper's claims about training stability, accuracy improvements, and efficiency gains are demonstrated on one model scale (7B parameters) in one domain (math). It is unknown whether the benefits scale to larger models (where KV-cache memory pressure is higher and batching dynamics differ), to other architectures (where attention patterns may affect shared-prefix overlap), or to other reasoning domains (code generation, logical reasoning, scientific QA) where the structure of shared prefixes may be qualitatively different. The offline efficiency experiments on three model variants provide some evidence of generalizability for inference throughput, but the training quality results are from a single training setup.

The "RL-zero" claim is partially validated. The paper emphasizes that TreePO works "directly from a base model, aligning with the 'RL-zero' paradigm." This is demonstrated — the Qwen2.5-7B base model is indeed trained without SFT. However, there is no comparison showing that TreePO's advantages relative to sequential sampling are larger for base models than for instruction-tuned models. It is possible that tree sampling provides even larger benefits for already-tuned models (since their outputs are more structured, increasing shared-prefix overlap), or that the advantage estimator is less necessary for tuned models (since their outputs are already higher-quality). This is an untested dimension.

Scale of evaluation: The test benchmarks (AIME 2024, AMC 2023, MATH500, MINERVA, Olympiad Bench) are standard in the mathematical reasoning literature, and the use of weighted averaging based on test set sizes is reasonable. However, the paper does not report per-benchmark test set sizes, making it impossible to verify the weighted average computation or assess whether the overall metric is dominated by a single large benchmark. AIME 2024, for example, contains 30 problems, while MATH500 contains 500 — if the weighted average is proportional to test set size, MATH500 would dominate the overall metric by a factor of ~16× over AIME. This would mean the "Overall" column primarily reflects MATH500 performance, which may be less informative about the model's reasoning capability on competition-level problems.

6. Limitations and Trade-offs

6.1 Difficulty Estimation Cost Is Unmeasured and Potentially Dominant

The assumption or constraint. TreePO's entire sampling framework depends on the existence of shared reasoning prefixes across trajectories to amortize computation via KV-cache reuse. The case study in Section 2.1 argues that "even under stochastic sampling, the model consistently follows a common path for the initial stages of reasoning before diverging at later decision points," and uses this observation to motivate the tree-structured rollout design. However, the paper only demonstrates this phenomenon qualitatively (Figure 2) on an already-aligned or partially-trained model — it does not characterize how the degree of prefix sharing evolves during training, how it varies across problem difficulties, or whether the base model in early RL steps (when outputs are noisy and unstructured) shares sufficient prefixes to realize the advertised efficiency gains.

The consequence. If shared prefixes are sparse or shallow early in training — before the policy has learned to produce consistent reasoning structures — tree sampling may provide little to no KV-cache reuse benefit during the initial training phase, precisely when computational efficiency matters most for rapid exploration. The branching budget transfer mechanism (Section 2.2) is designed to maintain batch sizes when paths terminate early, but it cannot compensate for a fundamental lack of prefix overlap: if trajectories are highly divergent from the first generated token onward, the tree degenerates into independent sequential generation with additional overhead from tree management. The paper's own rollout scaling experiment (Figure 5c) shows that Qwen2.5-Math-7B base model throughput peaks at 16 rollouts and degrades thereafter due to trajectory divergence and KV-cache fragmentation, directly demonstrating that base models without instruction tuning produce less convergent traces. This is precisely the model type used in the main training experiments (Qwen2.5-7B base), yet the training-time efficiency numbers in Table 2 do not report per-step GPU hour breakdowns that would reveal whether the savings are concentrated in later training (when the policy has converged to more structured outputs) rather than distributed evenly.

What evidence exists in the paper. Figure 5c documents the degradation explicitly for Qwen2.5-Math-7B, but this is an offline efficiency measurement using a frozen base model, not an online measurement during RL training where the policy evolves. The paper notes that "the lack of instruction tuning further loosens output structure" and causes KV-cache fragmentation, but does not measure how output structure changes over the course of a TreePO training run. The validation curves in Figure 1 (Left, Mid) show that tree-based methods initially lag behind GRPO in convergence speed — the paper attributes this to tree sampling "causing a slower convergence" — but does not investigate whether this initial slowdown is due to low prefix-sharing rates early in training that reduce the effective compute budget per step.

Mitigation status. Not addressed. The paper does not propose any mechanism for estimating prefix-sharing rates before or during training, does not adjust the tree configuration dynamically based on observed cache-hit rates, and does not report per-step GPU utilization measurements that would reveal when and whether tree sampling actually delivers savings during training. The "More Init Divergence" setting, which adds random branches at the root, likely reduces early-stage prefix sharing by spreading the initial budget across more divergent paths — the paper does not analyze this interaction. A practitioner implementing TreePO would need to instrument their training pipeline to verify that prefix-sharing rates justify the tree overhead at each training phase, but the paper provides no methodology or heuristics for doing so.

---

6.2 Single Model Scale, Single Model Family, Single Domain

The assumption or constraint. All training experiments — the core evidence for TreePO's performance and stability claims — use a single model: Qwen2.5-7B base (7 billion parameters). All training and evaluation data concern mathematical reasoning. The paper acknowledges neither of these as limitations in the main text or conclusion, though the conclusion does gesture toward future applications in "multi-turn dialogue, tool use, and multi-agent systems" without claiming that the current results support those domains.

The consequence. Three separate generalizability questions are conflated:

1. Model scale: KV-cache memory pressure scales with model size. A 7B model's KV-cache for 7,168 tokens is manageable on a single H100; for a 70B or 405B model, the memory overhead of maintaining dozens of partial trajectories with their KV-caches simultaneously may fundamentally change the efficiency calculus. The paper's finding that deeper trees (more branching, more parallel active paths) improve credit assignment (Section 4.3) would, on larger models, compete with memory constraints that force shallower trees. The optimal depth–segment configuration that balances accuracy and memory is unknown at larger scales.

2. Model family: Qwen2.5 uses a specific architecture with specific attention patterns. Models with different attention mechanisms (grouped-query attention ratios, sliding window, sparse attention) will have different KV-cache sharing characteristics and different prefill-vs-decode cost ratios. The paper's offline efficiency experiments (Section 4.1) test three Qwen2.5 variants, all from the same family, so the finding that optimal depth varies by model variant (28 for Instruct, 14 for Math) may not extrapolate to other architectures.

3. Domain: Mathematical reasoning has a distinctive property: solutions follow relatively structured formats (variable definitions, equation setup, algebraic manipulation, final answer) that make shared prefixes more likely. Code generation may share syntax structures; creative writing or dialogue may not. The heuristic early stopping mechanism (repetitive substring detection, Section 2.2) is specifically motivated by base model math training where unaligned models "get stuck in loops" — this failure mode may be domain-specific. The paper provides no evidence about whether shared reasoning prefixes are a general phenomenon or a math-specific one. If non-math domains exhibit substantially less prefix sharing, tree sampling's efficiency advantage would shrink or vanish.

What evidence exists in the paper. The offline efficiency experiments on three Qwen2.5 variants (Figure 4) show that tree sampling provides throughput improvements across all three — this is partial evidence for cross-model generalizability within the Qwen2.5 family. The rollout scaling experiment (Figure 5) shows model-dependent behavior: instruct variants scale nearly linearly while base math peaks early, suggesting that model properties (instruction tuning, domain specialization) interact with tree sampling efficiency in ways the paper documents but does not model. There is zero evidence regarding larger models or non-math domains.

Mitigation status. Not addressed. The paper does not discuss scale limitations, does not analyze memory overhead as a function of model size or tree width, and does not test on any non-math task. The conclusion's mention of "multi-turn dialogue, tool use, and multi-agent systems" is aspirational, not supported by data. A practitioner considering TreePO for a different model family, scale, or domain would need to replicate the case study (Section 2.1) and efficiency benchmarks (Section 4.1) from scratch to determine applicability.

---

6.3 Latency vs. Throughput Tradeoff Is Not Characterized

The assumption or constraint. TreePO's efficiency is measured in GPU hours (total compute over the training run, Table 2), Tokens per second (TokenPS), and Trajectories per second (TrajPS) — all throughput metrics that aggregate over large batches. The paper does not measure latency: the wall-clock time to generate w complete trajectories for a single query. This distinction matters because tree sampling introduces serial dependencies that sequential sampling avoids.

The consequence. Sequential sampling with w = 16 can generate all 16 trajectories in parallel if 16 independent inference requests are issued simultaneously — the latency is bounded by the time to generate the single longest trajectory. Tree sampling, by contrast, must grow the tree depth-first: generation at depth d + 1 cannot begin until all active paths at depth d have completed their segments, because the branching decisions at depth d determine which prompts exist at depth d + 1. This creates a serial bottleneck: for a tree of depth D, the minimum latency is at least D sequential inference rounds, each generating one segment of length l. For the 14 × 512 configuration, this is 14 sequential rounds; for 28 × 256, it is 28 rounds. Each round incurs inference engine overhead (scheduling, memory management, communication), which adds up across many rounds.

Additionally, the depth-first fallback mechanism (Section 2.2, "Depth-First Search Fallback") introduces further latency variance: if some paths terminate early and the tree has not yet reached w = 16 complete trajectories, the system must wait until all active paths terminate before fallback triggers, generate additional segments from fallback paths, and potentially repeat this process multiple times. A query that requires multiple rounds of fallback could have substantially higher latency than a query where the tree naturally reaches width w without fallback. The paper provides no distributional characterization of fallback frequency or its impact on per-query latency.

What evidence exists in the paper. None. All efficiency measurements (Figures 4 and 5, Table 2) use batch processing with 64 queries and 64 rollouts per query. These are throughput-optimized measurements that amortize overhead over many parallel operations. Single-query latency — which is the relevant metric for interactive applications — is not reported, not analyzed, and not discussed as a limitation.

Mitigation status. Not addressed. The paper makes no distinction between throughput and latency, does not measure per-query generation time, and does not analyze how tree depth or fallback frequency affects the time-to-first-answer or time-to-final-trajectory. A practitioner deploying TreePO in a latency-sensitive setting (e.g., an interactive math tutor that needs to generate and verify a response within a few seconds) would need to measure latency separately and could find that tree sampling's throughput advantages are accompanied by latency penalties that make it unsuitable for their use case. The paper's claim of "up to 40% trajectory-level and 35% token-level sampling compute reduction" applies to throughput (more trajectories per second, more tokens per second), not to per-query completion time.

---

6.4 Hard Problems Remain Unsolved, and There Is No Mechanism to Detect Them Before Wasting Compute

The assumption or constraint. TreePO's heuristic branching and pruning depend entirely on surface-level features of the generated text: repetitive substrings trigger early pruning, boxed answers trigger termination, and fallback branches from paths that produced parseable answers. There is no mechanism — learned or heuristic — to estimate whether a partial reasoning path is making progress toward a correct answer versus proceeding confidently in the wrong direction. The system can detect that the model is stuck in a loop (repetitive substrings) or that it has produced an answer (boxed text), but it cannot distinguish a promising partial solution from a flawed one that happens to contain novel tokens.

The consequence. This limitation has a sharp practical consequence that the paper does not quantify: when the base model's probability of producing a correct answer on a given problem is near zero (the hardest problems, or problems outside its training distribution), TreePO will still grow a full tree, explore branches, trigger fallback, and ultimately produce w incorrect trajectories at substantial computational cost. The tree structure may even amplify the waste: if the model consistently produces wrong but non-repetitive reasoning (e.g., a systematic algebraic error propagated through multiple steps), the tree will branch from these wrong prefixes, producing many variations of the same error rather than terminating early. Sequential sampling would also produce incorrect trajectories, but without the tree management overhead — so for hard problems where no amount of exploration helps, tree sampling may be strictly less efficient than sequential sampling because it adds tree construction cost without any accuracy benefit.

The paper's own data suggests this is not a hypothetical concern. The companion paper summary (reference example) demonstrated that on the hardest quintile of MATH problems, no test-time compute strategy improved performance above near-zero levels because the base model's pass@1 was ~1–3%. TreePO's evaluation benchmarks (AIME 2024, AMC 2023) are substantially harder than MATH — AIME problems especially would have low base-model pass@1. The paper does not break down TreePO's performance by problem difficulty quantile, so it is impossible to know whether the reported accuracy gains (e.g., 27.83% on AIME for TreePO w/ More Init Divergence) come primarily from problems where the base model already had non-trivial pass@1 (and tree sampling refined those paths) or whether the tree structure helped on genuinely hard problems. If the gains are concentrated on medium-difficulty problems — as the companion paper found for test-time compute scaling — then tree sampling is adding overhead to hard problems without corresponding benefit.

What evidence exists in the paper. The probability-based branching ablation (Section 4.4, Figure 8) provides indirect evidence: allocating more budget to low-probability paths (which are more common on hard problems where the model is uncertain) degrades performance, indicating that the tree exploration mechanisms are not well-suited to problems where the model's confident paths are wrong and its uncertain paths are also wrong. The paper does not provide a difficulty-stratified accuracy breakdown, does not measure per-difficulty computation cost, and does not analyze whether problems that received more fallback rounds (indicating harder problems where initial paths terminated without reaching width w) showed different accuracy patterns.

Mitigation status. Not addressed. The paper does not propose difficulty estimation (unlike the companion paper, which made difficulty estimation central to its compute-optimal framework), does not suggest adaptive budget allocation based on early reward signals, and does not characterize the cost–benefit ratio as a function of problem difficulty. A practitioner concerned about compute waste on hard problems would need to combine TreePO with an external difficulty estimation mechanism (e.g., running a small number of cheap samples first to estimate pass@1, then deciding whether to invoke tree sampling) — but this combination is neither proposed nor tested.

---

6.5 The Revision Model Analog (Fallback) Is Fragile and Untested for Quality

The assumption or constraint. TreePO's fallback mechanism (Section 2.2, "Depth-First Search Fallback") is designed to ensure that each query produces the required w = 16 trajectories even when many paths terminate early. Fallback selects from "those stopped paths containing formatted answer... or ending with [EOS]" and randomly branches from the segment before termination. The assumption is that a path that produced a properly formatted (boxed) answer, even if wrong, has a useful reasoning prefix from which alternative continuations might yield correct answers.

The consequence. This assumption has a failure mode that the paper acknowledges indirectly but does not quantify: just as the companion paper's revision model had a ~38% correct-to-incorrect reversion rate, TreePO's fallback mechanism can branch from a path that produced a correct answer (which happened to be selected for fallback because the tree hadn't yet reached width w) and generate incorrect alternatives from the same prefix, or branch from an incorrect path and generate more incorrect variations. The fallback mechanism has no quality filter — it does not prefer paths with higher reward or leverage the reward model to guide which terminated paths to branch from. The only selection criterion is format validity: "contains a legal answer surrounded by \boxed{} or ending with EOS."

The risk is that fallback generates low-quality training data — trajectories that branch from flawed reasoning prefixes and produce incorrect answers that look superficially plausible — which then feed into the policy update and reinforce the flawed reasoning patterns. This is analogous to the companion paper's finding that "prompting existing LLMs to self-correct their own mistakes tends to be largely ineffective for obtaining performance improvements on reasoning problems" — forcing the model to generate alternatives from a flawed prefix may not produce useful corrections, and may instead produce variations that share the same underlying error.

What evidence exists in the paper. The paper provides no ablation or analysis of fallback quality. It does not report what fraction of completed trajectories come from fallback versus natural tree completion, whether fallback-generated trajectories have different accuracy rates than non-fallback trajectories, or whether excluding fallback entirely (accepting fewer than w trajectories per query) would affect training quality. The "More Init Divergence" setting, which adds branches at the root, probably reduces reliance on fallback (more initial branches mean the tree is more likely to reach width w naturally), but this interaction is not analyzed. The "Misaligned fallback" ablation (Section 4.2, Figure 6) tests only the effect of segment-length mismatch during fallback, not the effect of fallback itself versus no fallback.

Mitigation status. Not addressed. The paper treats fallback as a necessary mechanism to maintain batch size consistency but does not evaluate its impact on training data quality, does not compare fallback-based tree sampling against tree sampling without fallback (accepting variable trajectory counts per query), and does not propose quality-guided fallback (e.g., biasing fallback toward paths whose partial segments received higher reward under the outcome verifier). This is a significant gap because fallback is not a rare edge case — for hard problems or early in training when the model frequently fails to produce properly formatted answers, fallback could be the dominant source of completed trajectories, making its quality properties central to TreePO's training dynamics.

---

6.6 No Comparison Against Simpler Efficiency Baselines That Would Partially Close the Gap

The assumption or constraint. TreePO's efficiency claim — 12–43% GPU hour reduction compared to sequential sampling — is measured against a sequential baseline that generates complete trajectories without any of TreePO's lightweight heuristics. The sequential baseline does not use early stopping on repetitive substrings, does not prune paths that fail to produce formatted answers, and always generates the full 7,168 tokens per trajectory if the model doesn't stop earlier. Some fraction of TreePO's efficiency gain may come from these heuristics (which are independent of the tree structure) rather than from KV-cache sharing (which is the tree-specific contribution).

The consequence. A practitioner could implement early stopping and length pruning on top of standard sequential sampling — detect repetitive substrings and terminate the generation early, set a shorter maximum response length, or abort generations that haven't produced a boxed answer after a certain token count — and achieve some portion of TreePO's efficiency savings without implementing tree-structured generation at all. The paper does not disentangle how much of the 12–43% GPU hour reduction comes from:

• Early termination of repetitive or unformatted paths (independent of tree structure).
• Reduced maximum effective response length (tree paths that terminate early before hitting the 7,168 token budget).
• KV-cache sharing across common prefixes (the tree-specific mechanism).

Without this decomposition, the marginal benefit of the tree structure over simpler heuristic improvements applied to sequential sampling is unknown. It is possible that a sequentially-sampled baseline with the same early stopping heuristics would close most of the efficiency gap, leaving only a modest residual benefit from KV-cache sharing — which might not justify the implementation complexity of tree-structured generation.

Additionally, the sequential baseline does not use techniques from recent efficient sampling literature that could improve its throughput: Infinite Sampling's micro-batching and length-aware scheduling, Truncated PPO's early chain cutting, or standard prompt KV-cache batching (which the VeRL framework and vLLM engine already support for sequential sampling). The paper references these works in Section 5 but does not include them as baselines. A sequential baseline augmented with these techniques could have substantially higher throughput than the "vanilla sequential" configuration used in Table 2.

What evidence exists in the paper. None. The paper does not ablate the contribution of early stopping heuristics independent of the tree structure, does not compare against a sequential baseline with the same heuristics applied, and does not benchmark against sequential sampling with modern efficiency optimizations (Infinite Sampling, Truncated PPO). Table 2 compares against a single sequential configuration with the same total token budget and batch size, but without any heuristic optimizations. The offline efficiency experiments (Figures 4 and 5) compare tree sampling against "conventional sampling" with the same budget constraints but do not describe the conventional sampling implementation in sufficient detail to know whether it includes standard KV-cache optimizations for sequential batching.

Mitigation status. Not addressed. The paper treats "sequential sampling" as a monolithic baseline without considering that the heuristics it introduces (early stopping, length management) are orthogonal to tree structure and could be applied to any sampling method. The claim that TreePO "replaces standard i.i.d. sequential sampling with a heuristic tree-based rollout mechanism" conflates the tree structure with the heuristics, making it impossible to attribute efficiency gains to the novel contribution (the tree) versus the non-novel components (early stopping heuristics). A proper ablation would compare four conditions: sequential without heuristics (current baseline), sequential with heuristics, tree without heuristics (tree structure only, no early stopping), and tree with heuristics (full TreePO). The difference between sequential+heuristics and tree+heuristics would isolate the marginal benefit of the tree structure.

7. Implications and Future Directions

How This Work Changes the Landscape

This is a systems-level reframing, not an algorithmic paradigm shift. TreePO does not introduce a new policy gradient objective, a new reward model, or a new exploration principle. What it introduces is a re-architecture of the data generation pipeline underlying standard RL-for-reasoning methods: instead of treating rollouts as independent sequences to be batched for throughput, it treats them as a structured tree whose topology carries information that can be exploited for both efficiency (shared KV-caching across common prefixes) and credit assignment (multi-resolution subgroup advantages). The conceptual move is to elevate the rollout structure from an implementation detail to a first-class design dimension — something to be configured, optimized, and analyzed on equal footing with the learning algorithm itself.

The magnitude is best characterized as opening a new axis for optimization within an existing paradigm. Prior work optimized the policy objective (GRPO → DAPO → REINFORCE++), the reward signal (outcome → process → learned verifiers), or the inference scheduling (Infinite Sampling, Truncated PPO). TreePO argues that the structure of how data is generated is an independent axis that interacts with all of these, and that optimizing it can yield gains orthogonal to algorithmic improvements. The evidence that tree sampling alone — without the tree-based advantage estimator — improves GRPO from 46.63% to 54.61% (Table 1) shows that the data structure itself carries signal, independent of how advantages are computed. This is analogous to how data augmentation or curriculum learning improved supervised training without changing the loss function: the arrangement of data matters, not just the quantity.

The paper resolves a tension between tree search and batch efficiency. Prior work treated these as in opposition: MCTS gives you structure but kills throughput (sequential node expansion); batched sampling gives you throughput but abandons structure (independent sequences). TreePO's segment-batched tree sampler shows this tradeoff is not fundamental — you can maintain batched GPU utilization while building a tree, by decoupling tree management (a logical data structure) from generation scheduling (which remains batch-parallel per depth level). This is an engineering insight with conceptual consequences: it means that tree-structured exploration is not limited to low-throughput research prototypes but can scale to training runs with 64 GPUs and millions of rollouts. This makes tree-based methods practical for production RL training, which was not obvious before.

It provides the first empirical characterization of how model alignment state interacts with sampling efficiency. The finding that instruction-tuned models scale nearly linearly with rollout count under tree sampling while base models peak and degrade (Figure 5) is not just a throughput benchmark — it is evidence that the convergence of reasoning traces (how much trajectories share common prefixes) is a property that changes with training and differs across model types. This gives practitioners a diagnostic: if your tree sampling throughput is not improving with more rollouts, your model's outputs are too divergent, and you may need longer segments, better instruction tuning, or an earlier training stage before tree sampling becomes effective. Prior work treated inference throughput as a fixed property of the model and hardware; TreePO shows it is dynamic and trainable.

It establishes a negative result that will redirect exploration research. The finding that forcing exploration toward low-probability paths degrades performance — and that even scheduled annealing from strong exploration to balanced allocation doesn't help — provides a concrete boundary condition for heuristic exploration in LLM reasoning. The natural impulse from the RL literature ("add an exploration bonus," "encourage high-entropy states") does not transfer straightforwardly to token-space reasoning, where low probability correlates with error rather than with undiscovered opportunity. This narrows the space of promising next steps: future exploration mechanisms for LLM reasoning will need to distinguish between epistemic uncertainty (the model doesn't know but could figure it out with more thought) and aleatoric uncertainty (the token is low-probability because it's wrong), which raw log-probability cannot do. Process reward models or value functions trained to predict eventual correctness, rather than raw model confidence, become the natural next step.

It partially resolves the conflicting evidence on self-correction and tree search in RL training. Prior work on tree-based RL for LLMs (TreeRL, SPO) demonstrated benefits but required SFT-initialized models and faced throughput bottlenecks. Prior work on efficient sampling (Infinite Sampling, Truncated PPO) achieved throughput gains but left the advantage estimator unchanged and didn't exploit trajectory structure. TreePO shows that these are not separate conversations — efficient sampling, structured exploration, and fine-grained credit assignment are three facets of the same design problem, and addressing them jointly (tree structure provides all three) yields better outcomes than optimizing each in isolation.

Follow-Up Research This Work Enables

Process reward model guidance within the tree structure. TreePO's branching decisions are currently heuristic (uniform allocation or log-probability-based, both of which have limitations demonstrated in Section 4.4). A natural extension is to replace these heuristics with a lightweight process reward model (PRM) trained on the tree's own outcome data: for each partial segment, use the average reward of all trajectories descended from that segment as a training signal for a value head, then use that value head at deployment to guide branching toward high-value prefixes and prune low-value ones. The tree structure makes this training data essentially free — every node already has an empirical success rate computed from its subtree's leaf rewards. A strong follow-up would train a value head jointly with the policy during TreePO training (requiring only a small additional output head on the policy model, not a separate PRM network), compare PRM-guided branching against the uniform and probability-based baselines on the same Qwen2.5-7B setup, and measure whether value-guided branching can push accuracy beyond the 58.21% ceiling reported in Table 1. The key metric is whether PRM guidance allows the tree to achieve the same accuracy as b=8 (the least pruned configuration) at the compute cost of b=4 — i.e., maintain accuracy while increasing efficiency beyond the 22% GPU-hour reduction reported.

Difficulty-adaptive tree configuration during training. The paper shows (Figure 7) that tree depth–segment configuration affects accuracy and that the optimal depth is configuration-dependent, and the offline efficiency experiments (Figures 4–5) show that optimal depth varies by model type. A natural question is whether the optimal tree configuration varies by problem difficulty — analogous to the companion paper's central finding. A strong follow-up would add difficulty estimation (PRM-based pass@1 estimation on a small number of initial samples, as done in the companion paper) before each TreePO rollout batch, then select the tree depth–segment configuration per-query based on estimated difficulty: shallow trees with long segments for easy problems (where the model is confident and needs less exploration), deep trees with short segments for medium problems (where branching helps discover alternative strategies), and sequential fallback for hard problems (where tree overhead provides no benefit because all paths are wrong). The hypothesis is that difficulty-adaptive tree configuration would close the accuracy gap between b=2 (43% faster but 3.54 points less accurate, Table 2) and b=8 (22% faster, 0.15 points less accurate) by spending tree compute only where it helps, making the aggressive pruning setting viable without accuracy loss.

Fine-grained analysis of when shared prefixes emerge during training. The case study in Section 2.1 demonstrates shared reasoning prefixes on a trained model, and Figure 5c shows that base model throughput degrades with rollouts due to trajectory divergence. What the paper does not characterize is the emergence of shared prefixes during RL training: at what training step do trajectories begin to consistently share initial reasoning segments? Does this emergence correlate with accuracy improvements? A strong follow-up would instrument the TreePO training pipeline to log, at each training step, the average prefix-sharing depth (how many initial segments, on average, are identical across trajectories from the same query) and the KV-cache hit rate, then plot these against validation accuracy over the course of training. The hypothesis is that prefix sharing increases as the policy improves (the model learns stable reasoning patterns), creating a positive feedback loop where tree sampling becomes more efficient over the course of training. If this holds, practitioners could start with sequential sampling or shallow trees early in training (when outputs are divergent and tree overhead isn't justified) and switch to deeper tree sampling later (when outputs have converged and KV-cache reuse is high) — a curriculum over sampling strategies. The paper's observation that tree sampling causes "slower convergence" early in training (Figure 1) would be explained if early-stage prefix sharing is low.

Stress test on a non-math reasoning domain with different structural properties. Mathematical reasoning has a distinctive property that favors tree sampling: solutions follow a relatively linear structure (problem restatement → variable assignment → equation → solution → boxed answer) where initial steps are formulaic and likely to be shared across trajectories. A domain with fundamentally different structure — e.g., code generation (where solutions branch at the algorithm-choice level, not the syntax level), multi-hop question answering (where each hop may require different retrieval), or creative writing (where high-level structure is shared but surface form varies wildly) — may show very different prefix-sharing patterns and thus different tree sampling benefits. A strong follow-up would run the identical TreePO pipeline (same Qwen2.5-7B base model, same GRPO/DAPO objective, same depth–segment sweep, same branching configurations) on a code generation benchmark (HumanEval or MBPP), measure prefix-sharing rates and throughput gains, and determine whether the 22–43% GPU-hour savings replicate or collapse. A negative result (tree sampling provides minimal benefit for code generation because coding trajectories diverge at the first function name or API choice) would be valuable: it would establish that TreePO's efficiency gains are domain-conditional, not universal, and would motivate domain-specific tree heuristics (e.g., branching at syntax boundaries rather than fixed token counts for code).

Combining TreePO with model scaling to test whether the efficiency gap widens or narrows. The paper's experiments are on 7B-parameter models, where KV-cache memory pressure is moderate. As model size increases, the memory cost of storing KV-caches for dozens of active tree paths grows linearly with hidden dimension and number of layers. At 70B or 405B parameters, the maximum tree width (number of simultaneous active paths) may be memory-limited even on H100 clusters with 80GB per GPU. A strong follow-up would measure tree sampling throughput on Qwen2.5 at 0.5B, 1.5B, 3B, 7B, 14B, and 32B (within a single model family to control architecture) at fixed per-GPU batch sizes, determining whether the throughput advantage over sequential sampling scales with model size or saturates. The hypothesis is that memory pressure will force shallower trees at larger scales, reducing the effective depth for credit assignment and potentially requiring model-parallel tree management (different GPUs handle different subtrees with KV-cache transfer). This experiment would determine whether TreePO is primarily a technique for the 7B–13B scale (where it was validated) with decreasing returns at larger scales, or whether the gains compound with model size.

Practical Applications and Downstream Use Cases

Cost-reduced RL post-training for reasoning models at the 7B–13B scale. The most direct application of TreePO is as a drop-in replacement for the sampling backend in existing GRPO/DAPO-based RL training pipelines targeting mathematical reasoning. The paper's numbers in Table 2 provide a concrete economic argument: for the same final model accuracy (58.06% vs. 58.21%), TreePO with b=8 and More Init Divergence reduces GPU hours by 22% (5.05 vs. 6.40 hours). For a team running weekly RL training jobs on 64 GPUs, this translates to recovering roughly 14 GPU-hours per run — enough to run additional ablation experiments, extend training by 20% more steps within the same budget, or reduce cloud costs proportionally. The implementation path is relatively low-risk because TreePO adopts DAPO's objective unchanged and only modifies the data generation layer, meaning existing hyperparameter tuning and reward designs carry over. The VeRL + vLLM stack used in the paper is open-source, and the algorithm (Algorithm 1) is specified clearly enough for reimplementation. The main adoption cost is the depth–segment configuration tuning (Section 4.3), which requires running a sweep on the target model and domain to find the optimal depth–segment pair — but the paper's finding that 14 × 512 is the sweet spot for Qwen2.5-7B on math provides a reasonable default starting point.

Batch inference for large-scale evaluation or data generation. For organizations that need to run inference on large sets of math problems — e.g., evaluating model checkpoints across thousands of test problems, generating synthetic training data for distillation, or scoring candidate solutions — tree sampling's throughput advantage applies during inference as well as training. The offline efficiency results (Section 4.1, Figures 4–5) show that tree-based sampling yields on average +40% trajectory throughput and +30% token throughput across three Qwen2.5 variants at the same batch size and budget. If an evaluation pipeline needs to sample 16 trajectories per problem on a 10,000-problem test set (160,000 total trajectories), tree sampling at the optimal depth configuration would complete the evaluation in roughly 70% of the wall-clock time of sequential sampling, assuming batch sizes are large enough to amortize tree management overhead. This is a direct cost saving for evaluation infrastructure that requires no changes to the model or evaluation metrics — only to the inference serving configuration. The main caveat is the latency caveat from Section 6.3: single-query latency may be higher under tree sampling, so this benefit applies primarily to batch evaluation, not interactive serving.

Diagnostic tool for analyzing model reasoning convergence during training. Beyond its use as a production sampling method, TreePO's tree structure provides a diagnostic window into how a model's reasoning evolves during RL training that sequential sampling obscures. By logging the prefix-sharing depth, branching points, and fallback frequency during TreePO training runs, practitioners can answer questions that are currently invisible: At what training step does the model start producing consistent initial reasoning? Are there specific decision points (e.g., choice of algebraic manipulation, selection of theorem) where the model consistently branches, and do these points shift over training? Do accuracy improvements correlate with deeper shared prefixes (the model getting better at the "setup" phase) or with better decisions at branching points (the model getting better at critical reasoning steps)? The paper already demonstrates that the tree structure carries this information (the subgroup advantage estimator exploits it), but the information itself — visualized as a tree topology over training steps — could be valuable for understanding training dynamics even if TreePO's efficiency gains are not the primary motivation for adoption.