TreeAdv: Tree-Structured Advantage Redistribution for Group-Based RL

TL;DR

TreeAdv tackles the inefficiencies of GRPO-style group-based RL by making rollouts tree-structured and redistributing rewards at the token level along shared prefixes. By entropy-guided branching and prefix-aware credit assignment, it yields structure-aware advantages that emphasize informative reasoning steps while downweighting redundant ones. Across 10 math-focused benchmarks and diverse Qwen baselines, TreeAdv achieves higher final accuracy with substantially fewer generated tokens and more stable training than GRPO/GSPO. The method offers a practical, drop-in improvement for long-context reasoning in LLMs, improving both efficiency and reliability under existing training budgets.

Abstract

Reinforcement learning with group-based objectives, such as Group Relative Policy Optimization (GRPO), is a common framework for aligning large language models on complex reasoning tasks. However, standard GRPO treats each rollout trajectory as an independent flat sequence and assigns a single sequence-level advantage to all tokens, which leads to sample inefficiency and a length bias toward verbose, redundant chains of thought without improving logical depth. We introduce TreeAdv (Tree-Structured Advantage Redistribution for Group-Based RL), which makes the tree structure of group rollouts explicit for both exploration and advantage assignment. Specifically, TreeAdv builds a group of trees (a forest) based on an entropy-driven sampling method where each tree branches at high-uncertainty decisions while sharing low-uncertainty tokens across rollouts. Then, TreeAdv aggregates token-level advantages for internal tree segments by redistributing the advantages of complete rollouts (all leaf nodes), and TreeAdv can easily apply to group-based objectives such as GRPO or GSPO. Across 10 math reasoning benchmarks, TreeAdv consistently outperforms GRPO and GSPO, while using substantially fewer generated tokens under identical supervision, data, and decoding budgets.

Figures (9)

• Figure 1: Comparison of trajectory generation. Left: Standard GRPO samples independent rollouts that redundantly repeat prefixes. Right: TreeAdv constructs a tree by sharing prefixes and branching only at high-entropy states ($H > \tau$).
• Figure 2: Illustration of TreeAdv's tree-aware advantage construction. Scalar rewards are computed only on complete rollouts (leaves) and converted to sequence-level advantages. These leaf advantages are aggregated into segment-level advantages $A^{\text{tok}}(n)$. Each token then receives the advantage of the unique segment node that covers its time step.
• Figure 3: Test-set accuracy versus training steps on AIME25, BRUMO, and GPQA. All runs are initialized from Qwen3-8B-Base and compare GRPO against TreeAdv-GRPO, illustrating how accuracy evolves during training under the two objectives.
• Figure 4: Training and exploration dynamics across model scales. We compare TreeAdv (Ours, red curves) with standard GRPO/GSPO baselines (blue curves) on three model sizes: Qwen3-4B-Instruct-2507, Qwen3-8B-Base, and Qwen3-30B-Instruct-2507. The top row shows reward trajectories, where TreeAdv consistently achieves higher rewards and faster convergence across scales. The bottom row shows entropy trajectories, characterizing how each method maintains exploration while optimizing the policy throughout training.
• Figure 5: Performance and cost analysis across training steps. The top row, (a) and (b), illustrates the Average Accuracy trends, while the bottom row, (c) and (d), shows the corresponding Average Token consumption. The left column depicts the Qwen3-8B-Inst Non-think model, and the right column depicts the Qwen3-8B-Base model. Curves represent comparisons between the Baseline and varying tree search strategies (1 Tree, 2 Trees, and 4 Trees)