Difficulty-Adaptive Tree-Structured Policy Optimization for Expanding Reasoning Coverage in RLVR

AI summary

Overview

Research area: Reinforcement learning for large reasoning models (RLVR), specifically the structure of train-time rollouts and their effect on intrinsic reasoning coverage (pass@k), with evaluation on mathematical reasoning benchmarks.

Technical level: Advanced. The paper assumes familiarity with policy-gradient methods (PPO/GRPO), advantage estimation, tree search, entropy-based branching, and pass@k / avg@k metrics.

Scope: The paper empirically analyzes three structural design choices for train-time rollouts (budget allocation across difficulty, tree vs. parallel topology, and where to fork), then packages the resulting principles into a single framework called DATPO.

What This Paper Is About

Reinforcement Learning with Verifiable Rewards (RLVR) reliably improves how often a reasoning model gets a single answer right (avg@k), but recent work argues it often fails to broaden the set of distinct problems the model can solve at all (pass@k, the "intrinsic reasoning coverage"). The authors argue this is because the rollout structure used during training — how many samples are drawn, whether they branch into trees, and where branching occurs — is treated as a fixed, uniform parallel sampling procedure. The goal is to redesign that train-time rollout structure so training explicitly expands reasoning coverage, which in turn should improve test-time scaling.

Key Contributions

1. An empirical and theoretical analysis of difficulty-adaptive rollout showing it is not merely a compute-efficiency heuristic: increasing the rollout budget (GRPO group size G) raised pass@256 for models trained on the Hard subset, degraded it on the Easy subset, and was non-monotonic (peaking at G=8) on the Medium subset. A theorem (Theorem A.1) formalizes the avg@k behavior.

2. A comparison of rollout topologies showing tree-structured rollout achieves a higher PassRate per generated token than parallel sampling, and that the choice of forking strategy matters within trees (fixed-seg beat random).

3. Identification of a "localization" phenomenon in token-level entropy forking — high-entropy tokens cluster in narrow segments and monopolize the search budget — and a proposed sentence-level entropy forking strategy (sent-entropy) that raised PassRate from 12.0 (tok-entropy) to 17.3 while losing little sibling diversity (0.1065 vs. 0.1049).

4. DATPO (Difficulty-Adaptive Sentence-entropy-guided Tree-structured Policy Optimization), which combines difficulty-adaptive tree search, sent-entropy forking, block-level Monte Carlo advantage estimation, and an annealed sibling-diversity bonus applied only to blocks with positive base advantage.

Main Findings

• Difficulty-adaptive rollout changes pass@k, not just cost. Across 9 GRPO training runs with G ∈ {4, 8, 16} on Easy, Medium, and Hard subsets of the training data, avg@256 improved consistently but marginally with larger budgets, while pass@256 degraded with larger budgets on Easy data, improved with larger budgets on Hard data, and peaked at G=8 on Medium data. Evaluation used avg@256 and pass@256 on 3 benchmarks.

• Trees beat parallel sampling per token. Under a fixed token budget, both tree variants (fixed-seg and random) showed a steeper growth of PassRate versus generation tokens than parallel sampling, attributed to prefix sharing that avoids regenerating identical early segments. Within tree methods, fixed-seg consistently achieved a higher PassRate than random.

• Token-level entropy forking is dominated by localization. tok-entropy achieved the highest sibling diversity (SibDiv 0.1065) but a low PassRate (12.0), because its diversity is confined to a narrow segment of the trajectory. sent-entropy achieved the highest PassRate (17.3) with SibDiv 0.1049. Baselines: random (10.0 / 0.0796), fixed-seg (13.3 / 0.0853), ATB (14.7 / 0.0874).

• DATPO leads on pass@k with marginal avg@k gains. On Qwen2.5-3B-Base, DATPO scored average avg@k 22.4 and pass@k 54.9 versus AttnRL's 21.3 and 53.0 (+1.1 avg@k, +1.9 pass@k). On Qwen3-4B-Base, DATPO scored 31.3 and 60.4 versus AttnRL's 30.7 and 57.4 (+0.6 avg@k, +3.0 pass@k).

• Per-benchmark highlights. DATPO on Qwen2.5-3B-Base reached MATH500 avg@8 63.5 / pass@8 81.7, AIME26 2.4 / 33.3, AIME25 1.6 / 33.3, AIME24 4.8 / 35.6, and AMC23 39.8 / 90.8. On Qwen3-4B-Base: MATH500 76.4 / 87.9, AIME26 6.5 / 33.3, AIME25 10.9 / 42.2, AIME24 14.0 / 46.7, AMC23 48.4 / 91.7.

• Coverage gains translate to test-time scaling. Applying majority voting on Qwen2.5-3B-Base (k=8 for MATH500, k=64 for the other four benchmarks), DATPO showed the largest maj@k gain, improving by +7.2 over its own avg@k.

• Training is more stable. Learning curves on MATH500 show TreeRL and AttnRL plateauing or degrading while DATPO continued to trend upward, which the authors attribute to the annealed sibling-diversity term preventing premature convergence.

• Resource allocation differs from AttnRL. DATPO and AttnRL both adapt rollouts to difficulty, but DATPO allocates fewer rollouts to easy problems and more to hard problems; AttnRL discards problems with below-average attention scores, which the authors describe as a rigid cut-off. DATPO achieved a higher PassRate on harder problems (Levels 4 and 5 of MATH).

• Ablations confirm the design choices. On MATH500, sent-entropy forking gave the best training result (avg@8 63.5 / pass@8 81.7) versus random (61.5 / 81.6), fixed-seg (61.1 / 80.6), ATB (62.1 / 80.7), and tok-entropy (61.3 / 80.3). For the diversity coefficient, the 0.2 → 0 annealing schedule was best (63.5 / 81.7), beating no diversity (0 → 0: 62.1 / 81.4) and a higher coefficient (0.4 → 0: 61.9 / 81.3); applying the bonus to all blocks hurt (0.2 → 0, all: 61.6 / 81.3); a constant α hurt avg@8 (0.2 → 0.2: 61.5 / 82.0); and annealing to negative reduced pass@8 (0.2 → −0.2: 63.4 / 81.0).

Methodology in Plain English

The work proceeds in two stages. First, a diagnostic stage. The authors split the training data by how often the base model solved each problem (Easy, Medium, Hard) and trained with GRPO at different rollout budgets to see how avg@k and pass@k respond. They then compared rollout topologies: standard parallel sampling versus a two-phase tree rollout that first generates N independent base trajectories and then picks K forking points in each, generating B extra branch rollouts from each forking point. They measured the probability of finding at least one correct answer (PassRate) as a function of generated tokens, and separately measured the semantic diversity of sibling branches using an embedding-based metric (SibDiv), defined as the average pairwise cosine distance between sibling blocks of text.

Second, a construction stage. DATPO derives the empirical difficulty of each prompt from the average verifiable reward of N base rollouts, V(root) ∈ [0, 1], and scales the tree budget down as difficulty falls: K̂ = ⌈K_max(1 − V(root))⌉ forking points and B̂ = ⌈B_max(1 − V(root))⌉ branches each, giving up to N(1 + K̂B̂) leaves, with expansion skipped entirely only when V(root) = 1. Forking points are the starts of the top-K̂ highest-average-entropy sentences rather than the highest-entropy tokens. Trajectories are divided into contiguous blocks bounded by forking points or the end of the trajectory, and each block receives one shared advantage computed from Monte Carlo state values (the average verifiable reward of all descending terminal blocks; a terminal state with no further rollouts has value 0). The advantage for a block is its reward plus the endpoint value minus the start-point value, augmented by α times the block's average cosine distance to its sibling blocks — but only when the base advantage is positive, so incorrect paths are never rewarded for being different. α is linearly annealed to zero over training, and the policy is optimized with a clipped ratio objective summed over all blocks and tokens of the tree.

Why This Matters

Impact on research: The paper reframes rollout structure as a first-class design variable in RLVR rather than a sampling detail, and provides a theoretical complement (Theorem A.1, with Ψ_G(p) = 1 − p^G − (1−p)^G) to prior empirical claims that RLVR does not expand pass@k beyond the base model. It also names and diagnoses the localization failure mode of entropy-based branching, which affects existing tree-search RL methods.

Real-world applications:

• Mathematical reasoning assistants that must solve competition-style problems (AIME, AMC, MATH) where any single correct solution suffices.
• Test-time scaling pipelines that use majority voting or best-of-N, since those pipelines can only aggregate correct answers if the model's coverage includes them.
• Automated tutoring or grading systems that need breadth of solution paths rather than one memorized template.
• Any deployment where generating many candidate answers is cheap but recurring errors cannot be corrected by sampling more.

Industry relevance: Training compute is a direct cost, and the paper's finding that larger rollout budgets can actively reduce coverage on easy data suggests current uniform-budget training pipelines may be misallocating compute. Difficulty-adaptive allocation concentrates budget on unsolved problems, and the reported +7.2 maj@k improvement over avg@k on Qwen2.5-3B-Base quantifies the downstream payoff for inference-time ensembling.

Future Directions

• Reducing the overhead of the sibling-diversity term, which currently requires extra forward passes through an external embedding model.
• Reducing estimation noise from the small rollout samples used to compute empirical difficulty and Monte Carlo value estimates (e.g., N=4, B=4), which the authors note can cause high-variance policy updates.
• Verifying scalability beyond the 3B–4B parameter regime tested here, since models at or above 7B parameters were not evaluated.
• Extending the framework beyond mathematical reasoning to other rigorous domains such as complex logical reasoning and code generation, which the authors flag as unverified.

Target Audience

Researchers and engineers working on RL post-training for language models, particularly those building RLVR pipelines with GRPO-style objectives, tree-search or MCTS-based rollout methods, or test-time scaling systems. It is also relevant to practitioners deciding how to allocate rollout budgets across heterogeneous training data, and to readers interested in pass@k as a training target rather than purely an evaluation metric. The paper is not written for beginners: it assumes fluency with policy-gradient advantage estimation and the standard RLVR literature.

Authors’ abstract

Reinforcement Learning with Verifiable Rewards (RLVR) has been central to the recent success of Large Reasoning Models. However, while RLVR significantly improves single-sample accuracy, it often fails to expand the model's intrinsic reasoning coverage (pass@k) due to limited exploration during training. To address this, we optimize the structural design of train-time rollouts to enhance pass@k. Our analysis identifies three key design principles: (1) difficulty-adaptive rollout can play an important role in expanding pass@k, beyond serving as an efficiency heuristic; (2) tree-based rollout outperforms parallel sampling in discovering correct answers; and (3) sentence-entropy-guided forking overcomes the localization phenomenon of token-level branching to maximize semantic diversity. Building on these insights, we propose DATPO (Difficulty-Adaptive Sentence-entropy-guided Tree-structured Policy Optimization). DATPO integrates difficulty-adaptive tree search with a sibling-diversity advantage term, explicitly promoting semantic diversity to expand reasoning coverage during training. Experiments on mathematical reasoning benchmarks demonstrate that DATPO outperforms baselines especially in pass@k, which directly translates to superior test-time scaling performance.