TRACE: Business Rule-Grounded Reasoning Curriculum for Knowledge-Preserving Parametric Tool Retrieval in Enterprise LLMs
Abstract
Parametric retrieval enables LLMs to retrieve tools implicitly by assigning each API a unique virtual token and training the model to generate it via constrained beam search. Toolsense shows that this regime has two critical drawbacks: it destroys parametric tool knowledge during training, and its beam-search decoding is too slow for real-time deployment. We introduce TRACE (Tool Retrieval via Augmented Chain-of-thought and Enterprise rules), a two-stage curriculum that resolves this dissociation. Stage 1 reuses the multi-format memorization SFT from ToolSense to seed tool knowledge with LoRA. Stage 2 is our core contribution: the model is trained to emit a thinking trace before producing a JSON list of tool tokens, using two data sources -- RRB pairs from ToolSense and queries synthesized to target business rules curated by domain experts -- both augmented with reasoning traces. This training objective preserves Stage 1 MCQ and QA probing accuracy while enabling single-beam greedy decoding at production latency. Evaluated on a combined enterprise catalog of 8,300+ tools across two enterprise product lines, TRACE training for Stage 2 not only preserves but improves tool understanding: MCQ accuracy gains +3.2 pp and QA probing gains +9 pp over Stage 1. On retrieval, TRACE achieves ~86% recall on Domain A and ~60% on Domain B -- compared to embedding baseline performance of ~27% & ~52% -- both with single-beam greedy decoding, making it directly deployable at production latency.
Community
Parametric tool retrieval lets LLMs retrieve APIs by generating virtual tokens — but existing training objectives (ToolGen) catastrophically destroy the model's tool knowledge in the process, and their constrained beam-search decoding is too slow for production. TRACE introduces a two-stage curriculum that resolves this dissociation: Stage 1 seeds tool knowledge via multi-format memorization (LoRA), and Stage 2 trains the model to emit a reasoning trace before committing to tool tokens, grounded in domain-expert-curated business rules. On a combined enterprise catalog of 8,283 tools (HR + Finance), TRACE achieves 86% recall under single-beam greedy decoding at production latency (1.9s vs ~19s for beam search, a 200× throughput gain), while improving MCQ and QA probing accuracy by +7.6 pp and +4.5 pp over Stage 1 — the first parametric retrieval system that is both knowledge-preserving and deployable at scale.
Get this paper in your agent:
hf papers read 2607.22639 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper