Source-linked AI summary
Inference-Time Budget Control for LLM Search Agents
Zhengru Fang, Senkang Forest Hu, Zhonghao Chang, Yu Guo, Yihang Tao, Hongyao Liu, Mengzhe Ruan, Jun Huang, Yuguang Fang
TL;DR
Under coupled tool-call and token budgets, LLM search agents must control both search actions and final-answer commitment. The paper proposes task-level VOI action selection plus conservative evidence-grounded finalization, finding useful but non-uniform gains across benchmarks, budgets, and backbones.
Problem
Budgeted search QA lacks explicit joint control over action allocation during search and answer resolution after search under the same hard dual-budget setting.
Method
A training-free two-stage controller uses task-level VOI to choose retrieval, decomposition, or answer commitment, then selectively rewrites only low-risk evidence-supported answer-form errors.
Results
Across four multi-hop QA benchmarks, four budget levels, and three backbones, budget control provides positive aggregate gains, with search-time control contributing the main gains.
Takeaways & Limitations
Inference-time control should govern both how search budget is spent and how the final answer is committed, rather than treating scaling as token-only.
Takeaways & Limitations
Benefits are not uniform: higher budgets, stronger backbones, and some datasets narrow or diminish gains, while finalization cannot recover wrong retrieval paths or unresolved bridge relations.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper frames tool-augmented QA as two-stage inference-time control under coupled tool-call and token budgets. It introduces task-level VOI action selection and conservative answer finalization, with gains that vary by budget, dataset, and backbone.
- Motivation: Search agents must decide both which action receives the next budget unit and when to commit a final answer under explicit dual budgets.The relevant actions are retrieval, decomposition, and answer commitment.
- Challenges: Final-answer errors can remain after successful search, especially in yes/no polarity, binary choices, typed slots, and alias completion.Unconditional rewriting can also damage bridge structure or comparative semantics.
- Method: Task-level VOI ranks retrieval, decomposition, and answer commitment by estimated marginal task value per unit budget.The scorer combines progress, structural signals, cost normalization, budget penalty, and conservative guards.
- Method: The evidence-grounded finalizer rewrites only low-risk answer-form errors and abstains when bridge or comparative reasoning remains unresolved.This limits intervention when the retrieval path or reasoning structure may still be wrong.
- Empirical analysis: Positive aggregate gains were observed across four benchmarks, three backbones, and four budget levels under a shared hard dual-budget protocol.Ablations identify budget-dependent penalty as dominant, while gains diminish under some backbone- and dataset-dependent conditions.
2. Related Work
Prior work covers tool-using agents, test-time scaling, workflow search, and answer verification. This paper connects those strands through explicit action allocation and selective intervention for budgeted search QA.
- Search agents and tool-augmented inference: Tool-using agents have shifted LLM inference from static generation toward interactive reasoning with retrieval, reflection, tree search, collaboration, and software tools.Search-agent systems treat retrieval as part of the inference trajectory rather than fixed preprocessing.
- Test-time scaling and budget-aware agent scaling: Test-time scaling improves outputs through extra inference computation, while budget-aware inference asks where that computation should be spent.Existing approaches include repeated sampling, self-refinement, adaptive allocation, and model routing.
- Workflow search and tree-based control: Workflow-search methods optimize prompts, module graphs, or workflows before deployment, whereas structured-search methods operate over states, thoughts, or plans during inference.The distinction separates deployment-time workflow optimization from online trajectory control.
- Answer verification and selective intervention: Answer-verification research shows that revision quality depends on trigger timing and intervention risk.The paper applies this lesson by retaining the trajectory answer unless an evidence-supported refinement passes a gain–risk check.
3. Problem Formulation
The paper formulates search-agent inference as constrained two-stage decision-making under tool-call, output-token, and answer-replacement-harm limits. Search allocates budget online, then finalization selects between the trajectory answer and a refined candidate.
- Problem setup: The agent operates under dual budgets B = (Btool, Btok), limiting tool calls and budgeted output tokens.At each step, future evidence and the exact next token cost are unknown.
- Search-time decision: During search, the policy allocates remaining budget across retrieval, decomposition, and answer commitment using observed trajectory, evidence, candidate states, and remaining budgets.The decision rule acts online before the next operation’s realized costs are known.
- Answer-time decision: After search, a refinement function constructs a candidate answer and a second policy chooses between the base and refined answers.The formulation treats finalization as selective intervention rather than unrestricted rewriting.
- Constrained objective: The objective maximizes expected bounded QA reward subject to tool-call, token, and expected answer-replacement-harm constraints.The harm constraint limits damage caused by replacing the trajectory answer with a worse final answer.
- Constrained objective: The harm term is zero when finalization preserves or improves answer quality and positive only when replacement makes the answer worse.The resulting two-stage formulation pairs online budget allocation with risk-controlled finalization.
4. Method
The method adds two lightweight layers to a generic tree-search backbone: a task-level VOI controller for online action allocation and a conservative selector for evidence-supported answer refinement.
- Two-stage pipeline: Stage 1 scores feasible operations under the remaining dual budget, executes the selected operation, and debits realized tool-call and token costs.The search layer chooses among evidence acquisition, decomposition, and answer commitment.
- Search-time budget allocation: The task-level VOI score estimates marginal task value per unit budget for action selection under the current state and remaining dual budget.It is an operational control score, not Shannon information gain, Bayesian posterior value, or a learned value model.
- Search-time budget allocation: The controller represents search status, remaining budgets, feasible actions, pre-execution charges, and features for unresolved evidence, structure, readiness, stagnation, and premature-answer risk.The score is formed by budget-shaped utility, value-per-budget normalization, and guards.
- Search-time budget allocation: The budget-dependent penalty is positive for SEARCH and DECOMPOSE under tight budgets and acts as a proxy for budget shadow cost.The controller then clips and normalizes the utility into the task-level VOI score.
- Answer-time finalization: After search, the finalizer compares a base answer with a refined candidate using support, slot type, contradiction, bridge, and comparative-reasoning features.A deterministic gain–risk rule intervenes only for safe, evidence-supported local corrections.
- Answer-time finalization: The safe set excludes unresolved bridge structure, comparative semantics, and missing direct support, so finalization acts as a conservative selector rather than a generic editor.It adds no tool calls or additional LLM call during finalization.
- Inference procedure: The controller repeatedly selects feasible operations, debits realized costs, stops at exhaustion or termination, and then applies final answer selection.The full inference procedure maintains a frontier of trajectories throughout search.
- Theoretical support: Theoretical analysis supports local approximation of oracle budget-charged lookahead and derives the answer-time gain–risk threshold under a harm constraint.The cited support concerns both controller layers.
5. Experiments
Across three backbones, four multi-hop QA benchmarks, and four hard dual-budget levels, the controller improves most audited settings, with gains concentrated at lower budgets and search-time allocation driving most benefit. Answer-time finalization adds targeted exactness corrections, while the controller also reduces measured inference time in the reported probe.
- Main results: VOI improves over AFlow, BATS, Search-o1, and BAVT in 14/16, 10/16, 16/16, and 15/16 Figure 3 cells, respectively.The strongest gains appear at low and lower-mid budgets, but BATS and AFlow lead in some upper- or lower-mid-budget cells.
- Main results: Qwen3.5-122B shows a more mixed regime, with BATS or BAVT leading several upper-budget cells while VOI retains an advantage at lower budgets.The detailed curves indicate substantial dependence on backbone and budget level.
- Final-answer control: Evidence-grounded finalization raises Bamboogle F1 from 0.4382 to 0.4628 at low budget, from 0.5786 to 0.6047 at upper-mid, and from 0.5056 to 0.5576 at high budget.It targets typed-slot, binary-choice, yes/no polarity, and supported factoid completion errors without adding evidence or another search call.
- Search-time ablation: Removing budget-dependent penalty causes the largest ablation drop, while removing normalization, structural signals, or guards also reduces macro average.The ablation keeps the search procedure, retrieval, prompts, sample set, and hard-budget audit unchanged.
- Inference-time cost: The full VOI controller reduces mean inference time from 20.91s to 15.23s per example relative to BAVT, a 27.2% reduction, in the Qwen3-32B upper-mid probe.The reported 2WikiMultihopQA row also shows dataset-dependent runtime behavior among ablated variants.
- Two-stage ablation: Stage 1 alone improves F1 on all four benchmarks, while the full method gains +5.7%, +11.8%, +14.7%, and +18.4% over BAVT on HotpotQA, 2WikiMultihopQA, MuSiQue, and Bamboogle.Stage 2 contributes no additional gain on HotpotQA and accounts for 13.4%, 41.9%, and 27.8% of total gain on the other three benchmarks.
6. Discussion and Limitations
Budget-aware search does not scale monotonically: middle budgets can outperform high budgets, and controller benefits vary by dataset, budget, and backbone. The finalizer remains limited to local answer-form repairs and cannot recover from wrong retrieval paths or unresolved bridge relations.
- Scaling boundary: Middle budgets sometimes outperform high budgets on 2WikiMultihopQA and Bamboogle because larger budgets can add redundant search or noisier finalization.The controller mitigates but does not remove the tension between exploration and commitment.
- Scaling boundary: BATS remains competitive in several upper-budget cells, while Qwen3.5-122B shows a more mixed regime as stronger backbones narrow the relative gap.These results mark the main boundary across high-budget regimes and backbones.
- Finalizer scope: The finalizer can repair local answer-form errors but cannot recover from wrong retrieval paths or unresolved bridge relations.Extending the controller to stronger retrieval, richer query rewriting, and longer-horizon browsing remains future work.
7. Conclusion
The paper presents two-stage budget control for tool-augmented search agents and reports that search-time allocation supplies the main gains, while finalization offers sparse benefits when evidence is adequate. It also identifies deployment benefits and risks requiring explicit auditing and conservative finalization.
- Conclusion: Two-stage budget control combines search-time action allocation with answer-time selection under explicit tool-call and output-token budgets.The approach uses a task-level VOI controller and an evidence-grounded finalizer with abstention.
- Conclusion: Search-time control provides the main gains by improving action allocation under strict budgets.The finalizer adds a sparse exactness benefit mainly when the trajectory already contains adequate evidence.
- Conclusion: The approach is bounded by diminishing returns from additional budget, backbone-dependent behavior, and incorrect retrieval paths that answer-time control cannot repair.
- Broader impacts: Budget-aware search can reduce unnecessary tool use, make inference cost more predictable, and expose whether an agent searches, decomposes, or answers.These properties support applications where latency, token usage, or external tool calls must be controlled.
- Broader impacts: Because efficient search agents could support harmful information-seeking workflows, the paper emphasizes hard audits, explicit accounting, conservative finalization, and failure analysis.
A. Controller Behavior Analysis
The controller behavior analysis visualizes how the implemented rule-based scorer evaluates SEARCH, DECOMPOSE, and ANSWER across budget and compositionality states. Its descriptive results show increasing decomposition scores for more compositional questions and conservative guard activation in high-compositionality states.
- Controller Behavior Analysis: The analysis reports descriptive behavior of the implemented search-time controller rather than introducing additional algorithmic components.It visualizes feasible-action scores and guard activation across budget and compositionality states.
- Controller Behavior Analysis: The mean DECOMPOSE score increases with question compositionality, consistent with decomposition’s intended role in resolving bridge structure.
- Controller Behavior Analysis: The below-floor guard activates mainly in high-compositionality states, reflecting conservative minimum-search behavior that avoids premature answer commitment.
- Controller Behavior Analysis: The controller scores SEARCH, DECOMPOSE, and ANSWER as a fixed rule-based scorer at each decision step.
- Controller Behavior Analysis: It evaluates unresolved evidence, likely retrieval value, question compositionality, and answer readiness when scoring feasible actions.Unresolved bridge structure increases the value of decomposition relative to ordinary search.
B. Inference Algorithm and Prompt Modules
The inference algorithm uses a controller to select budget-aware search actions and a conservative finalizer to intervene only on low-risk answer-form errors. Prompt modules execute the selected action deterministically while preserving abstention when structural reasoning remains risky.
- Inference algorithm: The controller scores feasible actions through a three-stage pipeline and executes the highest-scoring choice while updating the trace, frontier, and remaining budget.The algorithm computes utility, normalized raw scores, and guarded executable scores before selecting an action.
- B.1. System Prompt: The system prompt enforces one action per turn, requiring either exactly one tool call or a final answer under the controller-selected instruction.Dynamic instructions determine whether retrieval or answering is permitted, while explicit tool and token budgets are exposed to the generator.
- B.2. Planning Module: Planning tracks missing evidence and bridge facts, while remaining budget changes which search step the agent may take next.The planner is reused from BAVT; explicit budget-dependent action selection is the new contribution.
- B.3. Stage 1: Search-Time Budget Allocation: The controller, not the prompt, chooses among SEARCH, DECOMPOSE, and ANSWER using task-level VOI derived from trajectory features, costs, fixed coefficients, and guards.The prompt serves as the execution layer that materializes the externally selected action.
- B.4. Stage 2: Answer-Time Finalization: Stage 2 compares the trajectory answer with a refined candidate after search and abstains when unresolved bridge reasoning or comparative semantics make rewriting risky.It adds no new tool calls or LLM calls during finalization and targets local exactness errors.
- B.5. Representative Cases: Representative cases show finalization preserving semantics while correcting supported binary-choice, typed-slot, and other answer-form errors.The binary-choice example keeps the selected option unchanged while replacing it with a better-supported answer span.
- Theoretical statements: The theoretical analysis states that action ranking is reliable when feasible scores are clearly separated, while approximation terms can produce failures.The safe-replacement policy is framed as a guarded threshold decision balancing expected gain against intervention risk.
C.3. Proof of Theorem C.1 and Corollary C.9
This section establishes conditions under which approximate search-time scores preserve oracle action rankings and yield a conditional one-step value-gap guarantee. It then characterizes answer-time replacement as a gain–harm threshold problem, while distinguishing the oracle rule from the released deterministic finalizer.
- Search-time ranking: Normalized score rankings remain stable when the oracle score margin dominates approximation errors and denominators are bounded below.The result applies to the normalized ranking used for action selection, not only the unnormalized utility.
- Search-time ranking: The one-step value-gap guarantee is conditional on guards preserving the utility ordering relevant to the oracle action.It does not establish global optimality for arbitrary guard settings.
- Answer-time selection: Answer replacement is expressed as a gain–harm trade-off between a trajectory answer and a refined candidate.The decomposition separates expected positive reward gain from expected harm relative to the base answer.
- Answer-time selection: Under strong duality and multiplier attainment, the oracle finalizer accepts a refined answer when its penalized expected gain is nonnegative within the safe set.The threshold is represented by G⋆(z) − (1 + γ⋆)H⋆(z) ≥ 0.
- Released finalizer: The released finalizer is a conservative deterministic feature rule rather than a runtime implementation of the oracle threshold rule.Its empirical gap from the oracle performance envelope remains open, and the rule uses explicit gates, priority-ordered branches, and abstention.
F. Qwen3.5-122B Backbone-Sensitivity Results
Figure 6 presents Qwen3.5-122B budget-scaling curves across four benchmarks under a shared dual-budget protocol. Its competitive pattern is mixed: BATS and BAVT lead at several higher budgets, while VOI is advantageous at lower budgets.
- Competitive pattern: At upper budgets, BATS and BAVT lead in several Qwen3.5-122B cells.The backbone therefore exhibits a more mixed competitive regime than the comparison backbones.
- Competitive pattern: At lower budgets, VOI maintains an advantage where explicit budget penalty provides the largest marginal benefit.This identifies a budget-dependent sensitivity in the relative method ranking.
- Figure 6: Figure 6 compares budget-scaling curves for Qwen3.5-122B across four datasets under the shared dual-budget protocol.Each column represents one benchmark; corresponding Qwen3-32B and GPT-5.4-Mini results appear in Figure 3.
G. Stage-1 Component Ablation Protocol
The Stage-1 ablation removes one controller component at a time while holding the search procedure and hard-budget audit fixed. The protocol compares normalization, budget penalty, structural scoring, and guards under strict dual-budget evaluation.
- Ablation protocol: Each ablation variant removes exactly one component while keeping search, prompts, retrieval, sampling, and the hard-budget audit unchanged.This isolates the contribution of individual Stage-1 scoring components.
- Ablation variants: The w/o penalty variant removes the budget-dependent penalty term Πt(k; bt) from Eq. (2).The other variants remove normalization, structural scoring, or guards.
- Evaluation budget: Main strict-budget results report EM/F1 under symmetric tool-call and output-token caps of (1, 100), (2, 200), (2, 300), and (3, 500).VOI denotes the full two-stage budget-control method in the Qwen3-32B table.
- Evaluation summary: Cross-backbone macro deltas average VOI-versus-baseline differences over 48 audited cells with cell-level bootstrap 95% confidence intervals.The cells span three backbones, four benchmarks, and four budgets.
H. Additional Audit Results and Usage Diagnostics
The additional audits clarify token accounting, comparator coverage, and feasibility-conditioned resource usage under a common hard dual-budget protocol. Usage diagnostics are reported separately from scoring because feasible subsets can differ by method.
- Token accounting: BAVT-family average API tokens include prompt and completion tokens across planner, generator, and critic calls.The budget-output-token measure counts only the output-token component debited by the controller-side budget.
- Hard-budget audit: AFlow, Search-o1, BATS, BAVT, and VOI use the same retrieval backend, question-only queries, top_k=5, and tool/output-token caps in the hard audit.Examples exceeding either target cap fail for that cell; AFlow is audited over replayed workflows.
- Usage diagnostics: Feasible-only usage averages tools and output tokens after retaining only examples satisfying both caps for each method and dataset-budget cell.Different methods can therefore have different feasible subsets, making this a secondary diagnostic rather than a scoring rule.