TRACE retrieves the right API for enterprise LLMs, recall from 27% to 86% with about 200x the throughput of beam search

TRACE: Business Rule-Grounded Reasoning Curriculum for Knowledge-Preserving Parametric Tool Retrieval in Enterprise LLMs

Sai Shruthi Sistla, Ashutosh Hathidara, Christopher Toukmaji, Mayank Shrivastava, Karthikeyan Asokkumar

cs.AI

2026-06-23

TRACE adds a business-rule-grounded reasoning trace to parametric tool retrieval, lifting recall from the embedding baseline's 27% to 86% on one enterprise domain while replacing beam search with single-beam greedy decoding (about 200x throughput).

What problem this solves

An enterprise LLM copilot must route a user query to the right API out of thousands. Embedding-based retrieval (vector similarity over tool descriptions) dominates, but the authors' SAP production data (about 39k monthly active users) shows it causes about 60% of wrong-tool errors. Worse, roughly 81.7% of those failures involve more than 10 semantically overlapping tools that need business rules (API deprecations, versioning changes, domain routing) to disambiguate, and embeddings cannot capture rules.

Parametric retrieval (ToolGen / ToolSense) was meant to fix this: assign each tool a unique virtual token such as «WeatherAPI/GetForecast» and train the model to emit it. But ToolSense exposed two flaws: the retrieval training objective catastrophically destroys the model's parametric tool knowledge, and constrained beam-search decoding is too slow for real-time serving. You either retrieve well and forget the tools, or understand the tools and cannot deploy. TRACE fixes both.

Method

A two-stage curriculum on Gemma4-E4B-it with LoRA (r=64, α=128), over a catalog of 8,283 tools (Domain A, HR, 918 tools; Domain B, Finance, 7,365 tools).

Stage 1: multi-format memorization SFT (reused from ToolSense), seeding parametric tool knowledge back in with LoRA. It trains desc-to-virtual-token and reverse mappings plus an MCTS objective. LoRA is chosen because it forgets less than full fine-tuning.

Stage 2 (the core contribution): reasoning-augmented retrieval. The model is trained to emit a thinking trace before producing a JSON list of tool tokens. Two data sources: RRB (Realistic Retrieval Benchmark) pairs from ToolSense, and queries synthesized to target 123 business rules curated by domain experts (confusable tool clusters plus natural-language disambiguation rules, with hard-negative pools and easy/medium/hard tiers, validated by a programmatic filter and an LLM judge).

Why emit a trace first: it lets the model reason about which rule applies before committing to a token, and because the retrieval objective is not retrained directly, the Stage-1 knowledge survives. More importantly, emitting a trace plus JSON lets you drop constrained beam search for plain greedy single-beam decoding, fast enough to serve.

The Stage-2 mode is the R axis: ∅ none, n non-reasoning, r reasoning, R rule-grounded data. The authors ablate these.

Results

Knowledge preservation (MCQexpert probe, random baseline 29.3%): ∅ scores 69.2, n drops to 33.7, r recovers to 61.5, r+R is 56.4. Non-reasoning retrieval (n) wrecks tool knowledge down to near random, and adding the reasoning trace (r) pulls it back to 61.5, so the trace is what preserves knowledge. Stage 2 also improves on Stage 1: MCQ +3.2 points, QA probing +9 points.

Retrieval (single-beam greedy R@gen):

Domain ADomain B
Embedding (text-embedding-3-large)27.5%52.7%
TRACE (c, m, r+R)85.5%60.2%

Rule grounding pays off: Domain A climbs from 55.7% at 0 queries per rule to 73.3% at 12 queries per rule. Traces that cite a rule hit 94.6% recall versus 63.2% for those that do not; when a trace actually invokes a rule, it is almost always right.

Latency: single-beam free-form decoding about 1.9 s versus about 19 s for constrained beam-10 (single user); at concurrency 32, 11.2 qps versus 0.05 qps, roughly 200x throughput. That is the deployability win.

Why it matters

For enterprise LLM builders, this is a concrete recipe to make parametric tool retrieval both accurate and real-time, which the prior regime (ToolSense) could not. Two ideas worth stealing even outside this exact setup: a reasoning trace before the tool token preserves knowledge that raw retrieval training destroys, and that trace buys greedy decoding instead of beam search, which is where latency goes from unusable to production-grade. The rule-grounded data synthesis (confusable clusters plus expert rules plus hard negatives) transfers to any retrieval-over-overlapping-tools problem.

Caveat: this is an industry paper on proprietary SAP catalogs you cannot access, built on a ToolSense that is itself anonymized and apparently under review. So it is a recipe and a result on internal data, not a benchmark you can reproduce.

Limitations

Terms

Source

Related papers

All paper explainers