Atkia Mahila
Author directory2026
Discover, Then Diversify: Phase-Aware Control of Tree-of-Thought Search
Atkia Mahila | Avinash Maurya | M. Mustafa Rafique | Bogdan Nicolae
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Atkia Mahila | Avinash Maurya | M. Mustafa Rafique | Bogdan Nicolae
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Tree-of-Thought (ToT) reasoning has significantly enhanced the problem-solving capabilities of Language Models by enabling the structured exploration of reasoning trajectories. However, the efficiency of these frameworks is often bottlenecked by search algorithms that struggle to navigate vast reasoning spaces at a reasonable inference cost. In this work, we address the challenge of optimizing LLM performance on mathematical reasoning tasks under limited token budgets. A key limitation of existing search strategies is that they apply the same exploration and selection rule throughout the entire search, which under a finite budget, often leads to token exhaustion with suboptimal or no candidate answers. Effective search instead requires different behavior at different stages. We introduce PACE (Phase-Aware Controller for Exploration), a strategy that adjusts its behavior based on the current state of the search: it prioritizes progress toward a first complete answer when none exists, and broadens exploration toward structurally distinct regions later on. Across GSM8K and Math500, two model scales, and token budgets from 3k to 10k, PACE consistently matches or improves over recent state-of-the-art baselines. On GSM8K with Llama-3.2-3B at a 3k budget, PACE delivers roughly a 56% relative accuracy improvement over the strongest tree-search baseline, and an early-stopping variant achieves comparable accuracy using about two-thirds of the tokens.