Inference-Time Budget Control for LLM Search Agents

TL;DR

This paper introduces a novel inference-time budget control method for LLM search agents in multi-hop QA, optimizing resource allocation and answer finalization to improve performance under dual constraints.

Contribution

It formulates a two-stage budget control approach using VOI scores for action selection and a selective finalizer, demonstrating significant gains across benchmarks.

Findings

Method outperforms baselines under same dual-budget constraints.

Search-time budget control, especially penalty, is the main performance driver.

Answer-time control is beneficial when retrieval is already adequate.

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.