Source-linked AI summary
Lookahead-R: Budget-Aware Tool Retrieval via Execution-Centric Planning
Zongze Wu, Yani Guo, Runnan Li
TL;DR
Tool retrieval must balance semantic relevance with execution feasibility, while existing semantic and execution-based approaches trade reliability against latency. Lookahead-R uses a predictive world model and budget-aware tree search to plan tool selection, reporting improved ranking and reduced inference latency on ToolBench.
Problem
Semantic retrievers offer low latency but limited execution reliability, while validating tools through real API calls incurs prohibitive latency.
Source: §1 · p. 1Method
Lookahead-R uses an execution-aware surrogate world model to predict success, latency, and utility, then guides cost-sensitive, uncertainty-guided MCTS under resource constraints.
Source: §1 · p. 1Results
1.24% higher NDCG@5 on complex intra-collection queries, with significantly reduced inference latency, is reported over advanced generative baselines.
Source: §1 · p. 1Takeaways & Limitations
The framework operates as an anytime algorithm that the paper describes as suitable for real-world deployment.
Source: §5 · p. 8
Abstract
from arXiv · showhide
Tool retrieval is a critical bottleneck for LLM-based agents operating over large, heterogeneous API ecosystems. Existing approaches face an inherent trade-off: semantic retrievers are fast but suffer from the semantic-functional gap, while execution-based validation improves precision at the cost of prohibitive latency. We propose Lookahead-R, a planning-based framework that reformulates tool retrieval as a resource-constrained sequential decision-making problem. At its core, Lookahead-R introduces a lightweight execution-aware surrogate world model that jointly predicts tool execution success, latency cost, and semantic utility---without invoking real APIs. This world model drives a cost-sensitive, uncertainty-guided Monte Carlo Tree Search that navigates the tool space under strict budget constraints. Evaluated on the large-scale ToolBench benchmark, Lookahead-R achieves a superior accuracy-efficiency trade-off across all test scenarios. On the most challenging I3 split, it attains an NDCG@5 of 91.40\%, outperforming the state-of-the-art ToolGen (90.16\%) by 1.24\%. Ablation studies confirm that explicit latency modeling is the key discriminative signal for identifying high-quality tools under resource constraints.
1 Introduction
Lookahead-R reframes tool retrieval as budget-constrained planning, using a learned world model and uncertainty-guided search to account for execution feasibility. On ToolBench, it reports a stronger accuracy-efficiency trade-off than advanced generative baselines.
1 Introduction
- Tool retrieval must find tools that are both semantically relevant and executable, since parameter mismatches or system constraints can defeat seemingly relevant candidates.Semantic retrievers are fast but limited in execution reliability, whereas invoking candidates for validation incurs prohibitive latency.Source: §1 · p. 1
- Lookahead-R formulates retrieval as resource-constrained sequential decision-making and uses a surrogate world model to predict execution outcomes, latency, and utility without real API calls.The model is trained through a Student-Teacher Mining protocol and supports cost-sensitive, uncertainty-guided MCTS.Source: §1 · p. 1
- 1.24% higher NDCG@5 on complex intra-collection queries, with significantly reduced inference latency, is reported over advanced generative baselines.The paper presents this as a new state-of-the-art Pareto frontier on ToolBench.Source: §1 · p. 1
- Its anytime MCTS can return valid solutions immediately and refine decision quality as the available budget permits.The planner balances exploration against budget constraints and prunes high-risk branches early.Source: §1 · pp. 1–2
2 Related Work
Related work moves from language-model retrieval and structured tool selection to model-based planning, cost-aware inference, and synthetic data. Lookahead-R combines those latter directions to model tool-execution dynamics for budget-aware retrieval.
2.1 LLMs for Multi-modal Retrieval
- LLM retrieval has broadened from interpreting textual intent across modalities to retrieving executable capabilities, where functional affordance matters alongside semantic relevance.The cited work frames tools and APIs as a retrieval target distinct from static documents.Source: §2.1 · p. 2
2.2 Structured Tool Retrieval
- ToolLLM [24] expanded tool-selection benchmarks to over 16,000 real-world APIs, while structured retrievers advanced beyond dense semantic matching; these methods still largely overlook stochastic execution outcomes.The section contrasts approaches including hierarchical routing, tool pointers, and iterative retrieval fine-tuning with Lookahead-R’s prediction grounded in execution outcomes.Source: §2.2 · pp. 2–3
2.3 Model-Based Planning and World Models
- Prior planning frameworks simulate textual reasoning rather than physical tool-execution dynamics, and often rely on costly environment queries for feedback.Lookahead-R instead uses a low-cost surrogate to predict execution success, latency, and utility without expensive real-world interactions.Source: §2.3 · p. 3
2.4 Cost-Aware Inference and Data Synthesis
- Lookahead-R combines cost-aware inference and synthetic-data generation to distill a world model of execution failures and latency profiles, targeting budget-aware retrieval.The approach explicitly models failure cost to optimize the accuracy-efficiency Pareto frontier under strict latency budgets.Source: §2.4 · p. 3
3 Method
Lookahead-R reframes tool retrieval as budget-constrained planning, using an execution-aware world model and virtual MCTS rollouts to account for tool outcomes and costs. Student–Teacher mining supplies execution traces for training the model.
3 Method
- Lookahead-R combines Student–Teacher data synthesis, world-model training, and budget-aware MCTS virtual rollouts to identify tools under budget constraints.The planner uses predicted outcomes to prune costly branches before selecting a tool.Source: §3 · p. 3, Figure 2 · p. 4
3.1 Problem Formulation
- The formulation maximizes expected utility while strictly adhering to a latency budget, treating tool choices as actions with execution costs and risks.The world model predicts execution outcome, utility, and cost without invoking real APIs; failures incur the maximum budget cost.Source: §3.1 · p. 3
3.2 World Model with Critical Judging
- The predictive world model returns structured predictions of responses, utility, risk, and latency cost, with latency represented by categorical buckets including TIMEOUT.The model takes query, tool description, and execution history in a unified textual context.Source: §3.2 · pp. 3–4
- The generative model predicts detailed failure information alongside correctness; Critical Judge instructions raise predicted risk when required arguments are missing.The Llama-3-8B backbone is fine-tuned with QLoRA for the inference efficiency needed by MCTS planning.Source: §3.2 · pp. 3–4
3.3 Uncertainty-Guided MCTS Planning
- MCTS performs virtual rollouts and uses a cost-sensitive selection score that subtracts predicted cost, pruning risky branches while respecting the remaining latency budget.A branch whose predicted cost exceeds the remaining budget is penalized immediately.Source: §3.3 · pp. 4–5
- The planner is anytime: it maintains a valid solution and stops when its computational budget is exhausted or visit counts converge.This supports strict SLA adherence while allowing search time to vary with available resources.Source: §3.3 · p. 5
3.4 Student-Teacher Data Synthesis
- Student–Teacher mining probes sandbox APIs with schema-conditioned dummy parameters, records execution outcomes, and assigns irrelevant tools maximum-cost labels.Recorded signals include success, latency, semantic utility, and structured error category.Source: §3.4 · pp. 5–6
4 Experiments
Experiments use ToolBench APIs and query datasets, with a Llama-3-8B-Instruct backbone fine-tuned as the execution world model. Training traces provide execution status, latency, and semantic utility supervision; evaluation follows disjoint official ToolBench splits.
4.1.1 Model Configuration.
- Llama-3-8B-Instruct is the backbone, selected for its efficiency-capability trade-off and throughput for virtual rollouts.The model is fine-tuned for eight epochs with early stopping; LoRA adapters are applied while backbone parameters remain frozen.Source: §4.1.1 · p. 6
Dataset.
- ToolBench contains more than 16k tool collections and 47k APIs, each documented with structured metadata; each API is treated as an individual tool.Source: Dataset · p. 6
- Training combines G1, G2, and G3 query datasets, using 10% of their combined data as seed instructions for synthetic execution traces.The traces supervise execution status, latency, and semantic utility.Source: Dataset · p. 6
- Evaluation follows the official ToolBench split: 200k query–API pairs across I1, I2, and I3, disjoint from training instructions.The split contains 87,413 I1, 84,815 I2, and 25,251 I3 instances.Source: Dataset · p. 6
Baselines.
Lookahead-R is compared with lexical, embedding, long-context, iterative-refinement, contrastive, and generative retrieval baselines, following the ToolGen comparison setup.
Baselines.
- The comparison includes BM25 lexical retrieval, sentence-embedding similarity, and GPT-4o selection from a 2k-tool prompt.Source: Baselines · p. 6
- Iterative and feedback baselines include Re-Invoke [3], IterFeedback [41], and contrastively trained ToolRetriever [24].Source: Baselines · p. 6
- ToolGen is the generative retrieval baseline, and the comparison set follows its setup.Source: Baselines · p. 6
Settings.
Lookahead-R is trained across mixed query types and evaluated separately by ToolBench complexity split. The reported measures cover ranking quality, execution errors, task completion, and comparative wins.
Settings.
- Training mixes queries from G1, G2, and G3 to represent varied tool-retrieval scenarios.Source: Settings · p. 6
- Evaluation reports I1, I2, and I3 separately, covering single-tool and multi-tool retrieval under the standard ToolBench protocol.Source: Settings · p. 6
Metrics.
- NDCG@k measures ranking across I1–I3; Failure Rate tracks selected-tool execution errors, while Pass Rate and Win Rate capture task completion and baseline comparisons.Source: Metrics · p. 6, Table 1 · p. 7
- Table 1 reports NDCG@1, NDCG@3, and NDCG@5 by complexity level, plus Average Pass Rate and Win Rate against ChatGPT-ReAct.Source: Table 1 · p. 7
4.2 Main Results
Lookahead-R leads the reported retrieval rankings, most notably on complex I3 queries, and also reports higher task success and wins against standard baselines. The cost-penalty ablation reports a lower Win Rate when cost modeling is removed.
4.2 Main Results
- 91.4% NDCG@5 on I3 surpasses ToolGen by +1.24%, the strongest reported ranking result on the most challenging split.Source: §4.2 · p. 7
- 77.4% Average Pass Rate and 66.8% Win Rate are reported against standard baselines.Source: §4.2 · p. 7
- Removing the Cost Penalty lowers Win Rate from 78.5% to 75.2%, supporting the reported importance of modeling tool cost.Source: §4.2 · p. 7
4.3 Anytime Planning Characteristics
Lookahead-R is an anytime planner whose tool-selection quality improves as its latency budget grows. Unlike static and generative baselines, it can provide utility at different points in the computation budget.
4.3 Anytime Planning Characteristics
- Static baselines such as ToolRetriever gain no performance from additional inference time because their results are bounded by initial semantic recall.Their performance remains flat as the latency budget increases.Source: §4.3 · p. 7
- Lookahead-R rapidly reaches high base success and then refines decisions through deeper MCTS lookahead as budget permits.Its accuracy improves progressively with the latency budget, unlike static baselines whose performance stays flat and generative baselines with a high latency floor.Source: §4.3 · p. 7
- Generative baselines ToolGen and IterFeedback have a latency floor above 800ms and offer zero utility in windows below 500ms.They need a complete generation cycle before producing a valid tool call.Source: §4.3 · p. 7
4.4 Mechanism-Level Ablation Studies
Ablations on I1 isolate execution prediction, latency-cost alignment, and uncertainty-guided exploration. Their results associate each mechanism with a distinct benefit to tool selection or efficiency.
4.4 Mechanism-Level Ablation Studies
- On I1 (1,000 queries), Lookahead-R ablations compare the full framework with variants disabling execution prediction, shuffling training-set latency labels, or replacing epistemic uncertainty with policy entropy; hyperparameters remain fixed, including rollout budget N=30.Source: §4.4 · pp. 7–8
- −17.0% Success Rate and a Failure Rate rise from 4.1% to 28.3% follow removal of execution prediction, despite retaining the Llama-3 backbone.The analysis attributes the drop to losing the World Model’s ability to prune functionally invalid tools.Source: Table 2 · p. 8, §4.4 · p. 8
- 680ms to 1120ms: ground-truth average latency more than doubles when latency labels are shuffled, despite the variant still identifying valid tools.The result supports the role of accurate cost estimates in steering search away from physically slow tools.Source: Table 2 · p. 8, §4.4 · p. 8
- Entropy-only exploration yields lower diversity and suboptimal performance, while epistemic uncertainty guides exploration toward high-information “unknown unknowns.”This contrasts uncertainty-guided exploration with randomly sampling known suboptimal actions, which can leave the planner in local optima.Source: §4.4 · p. 8
4.5 Case Study
The case study shows the World Model predicting a modality mismatch before execution and redirecting the planner to an image-capable alternative.
4.5 Case Study
- For a poster-image query, the planner abandons text-only Search_Director_ByName after predicting a parameter error, then predicts success for image-capable Visual_Movie_Lookup.The required movie_name text argument cannot be extracted directly from the raw image input.Source: §4.5 · p. 8
5 Conclusion
The conclusion presents Lookahead-R as resource-constrained planning for tool retrieval, using execution and latency predictions to guide MCTS. ToolBench experiments support its accuracy-efficiency performance and anytime operation.
5 Conclusion
- Lookahead-R reframes tool retrieval as a resource-constrained sequential decision process, using a mined World Model and cost-sensitive, uncertainty-guided MCTS to predict execution outcomes and latency.The conclusion reports that ToolBench experiments show improved results over traditional semantic retrievers and an anytime algorithm.Source: §5 · p. 8