Difficulty adaptive tree search improves reasoning coverage in reinforcement learning
Difficulty-Adaptive Tree-Structured Policy Optimization for Expanding Reasoning Coverage in RLVR
Machine LearningArtificial IntelligenceComputation and Language
Summary
Finding the best answers with AI models can be hard because they might not explore many different possibilities during training. The authors show that changing how these models explore ideas—by focusing on harder challenges and using a tree-like method—helps discover more correct answers. They also found that looking at whole sentences instead of single words during this search helps the model be more creative and cover more reasoning options. Their new approach, called DATPO, leads to better performance especially when considering multiple attempts to solve problems. This method was tested on math puzzles and did better than older methods.
What this means in practice
- •For ai model trainers: Train reasoning models with adaptive tree search to improve multi-answer accuracy on complex tasks.
- •For automated tutoring systems: Enhance problem-solving diversity in educational AI by increasing reasoning coverage through difficulty-adaptive rollout policies.
Authors
Youngjun Yu, Sanghwan Jang, Hwanjo Yu
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.