Inference-Time Budget Control for LLM Search Agents

2026-05-07

AI Summary

This work addresses the challenge of efficiently allocating reasoning resources for large language model (LLM) search agents under dual budget constraints on both tool invocation counts and generated tokens. The authors propose a two-stage budget control mechanism: during the search phase, budgets are dynamically allocated based on task-level value of information (VOI), while in the answer phase, evidence-driven residual error estimation enables selective correction of low-risk errors. The study formulates action selection under dual budgets as a VOI optimization problem and introduces a novel budget-aware penalty strategy. Extensive experiments across four multi-hop question answering benchmarks, three LLM backbones, and diverse budget settings demonstrate that the proposed method consistently outperforms existing baselines, highlighting the synergistic benefits of integrating budget-aware search with answer refinement.

Abstract

LLM search agents increasingly rely on tools at inference time, but their trajectories are often constrained by hard limits on both tool calls and generated tokens. Under such dual budgets, better answers require not only stronger models, but also explicit control over which search action should receive the next budget unit and when the accumulated evidence is sufficient to commit a final answer. We study this problem in multi-hop question answering (QA) and formulate it as two-stage inference-time budget control. At search time, our controller assigns each feasible action a task-level Value-of-Information (VOI) score, defined as an operational estimate of marginal task value per unit budget under the current search state and remaining dual budget, and uses this score to choose among retrieval, decomposition, and answer commitment. After search, a selective evidence-grounded finalizer compares the trajectory answer with a refined candidate and rewrites only when the residual error appears to be a low-risk answer-form error. Across four multi-hop QA benchmarks, three LLM backbones, and four budget levels, the method yields positive aggregate gains over four audited baselines under the same hard dual-budget protocol. Ablations show that search-time budget control, especially budget-dependent penalty, provides the main performance gain, while answer-time control helps mainly when the retrieval path is already adequate. These results suggest that inference-time budget control for LLM search agents should govern both how budget is spent during search and how the final answer is committed.

Problem

• inference-time budget control
• LLM search agents
• multi-hop question answering
• tool calls
• generated tokens

Innovation

• inference-time budget control
• Value-of-Information
• multi-hop QA
• LLM search agents
• dual-budget constraint