PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2607.22639v1cs.AI

TL;DR

Enterprise LLMs struggle to retrieve the right tools because embedding retrieval misses business-rule distinctions and parametric retrieval can erase tool knowledge. TRACE trains reasoning traces grounded in enterprise rules before tool-token generation, achieving strong recall while preserving and improving tool understanding under single-beam greedy decoding.

  • Problem

    Existing retrieval methods have limited recall for overlapping enterprise tools, while parametric retrieval can destroy the tool knowledge it is intended to use.

  • Method

    TRACE uses a two-stage curriculum that memorizes the catalog, then trains business-rule-grounded reasoning traces before emitting virtual tool tokens.

  • Results

    TRACE reaches ~86% recall on Domain A and ~60% on Domain B, while improving Stage-1 MCQ and QA probing by +7.6 pp and +4.5 pp under single-beam greedy decoding.

  • Takeaways & Limitations

    TRACE delivers knowledge-preserving parametric tool retrieval that is deployable at production latency across two enterprise domains.

  • Takeaways & Limitations

    Evaluation covers two proprietary enterprise domains at one model size, and Stage 2 depends on nontrivial domain-expert rule curation.

Abstract

from arXiv · show

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.

1 Introduction

TRACE addresses the failure of embedding and parametric retrieval in enterprise copilots by combining business-rule-grounded reasoning with a two-stage curriculum and single-beam greedy decoding. It preserves tool knowledge while improving retrieval for semantically overlapping tools at production scale.

  • Motivation: ∼60% of incorrect tool retrievals come from embedding-based retrieval, making wrong-tool selection the dominant system failure because downstream components cannot recover.This analysis covered ∼39k monthly active users and 609 stratified sampled sessions.
  • Motivation: ∼81.7% of sampled retrieval failures involve more than 10 semantically overlapping tools requiring domain business rules for disambiguation.The finding comes from 609 stratified sampled sessions in a production enterprise AI copilot.
  • TRACE curriculum: TRACE uses a two-stage curriculum that trains models to reason and retrieve from business-grounded and ungrounded query-reasoning pairs.The curriculum is designed to resolve the knowledge-retrieval dissociation while supporting production-scale retrieval.
  • Production deployment: Single-beam greedy decoding reduces inference latency to meet production requirements while TRACE preserves tool knowledge and delivers competitive retrieval performance.The contribution explicitly frames single-beam greedy decoding as a production-latency capability.
  • Business rule grounding: +23 pp is the maximum retrieval lift from injecting enterprise business rules into reasoning traces for semantically overlapping tools.TRACE was validated on a combined catalog of 8,300+ proprietary enterprise tools across two business domains.

2 Related Work

Prior tool-use methods assume tools are already supplied or chosen from small candidate sets, leaving large-catalog retrieval under-studied. ToolGen addresses retrieval with virtual tokens, while TRACE additionally prioritizes tool-knowledge preservation and removes external retrieval at inference.

  • Tool retrieval and function calling: Function-calling and API-generation methods assume the correct tool is provided in context or selected from a small candidate set.They do not address identifying the right tool from a large catalog.
  • Tool retrieval and function calling: Large catalogs make in-context enumeration infeasible, while semantic overlap makes embedding-based tool matching unreliable.The passage characterizes retrieval as a distinct and under-studied problem.
  • Tool retrieval and function calling: ToolGen maps each tool to a unique virtual token and trains the model to generate it, addressing large-catalog retrieval.The supplied passage introduces ToolGen as closing the retrieval gap.
  • Tool retrieval and function calling: TRACE treats tool knowledge preservation as a first-class training objective while eliminating the need for an external retriever at inference time.This distinguishes TRACE from the retrieval approach described immediately before it.
  • Reasoning traces in supervised fine-tuning: Chain-of-thought prompting and subsequent supervised or reinforcement-learning methods establish reasoning traces as a route to eliciting or distilling intermediate reasoning.The cited examples include STaR, Orca, and DeepSeek-R1.
  • Catastrophic forgetting and knowledge retention: LoRA is reported to forget less than full fine-tuning, while probing evaluation measures retained parametric factual knowledge.These findings motivate TRACE’s Stage 1 design and tool-understanding evaluation.

3 TRACE: Two-Stage Curriculum with Reasoning Traces

TRACE uses a two-stage curriculum that first internalizes virtual-tool knowledge, then trains the model to reason over candidate tools and business rules before emitting an ordered JSON tool-token list. Its enterprise data combines rule-grounded tool ambiguity with trace supervision and supports single-beam greedy inference.

  • Stage 1: Stage 1 applies ToolSense’s multi-format memorization SFT with LoRA to bind every virtual token to its tool’s semantics, name, and surface form.Its objective is internalization rather than retrieval, preparing tool knowledge for Stage 2.
  • Setting: TRACE covers 8,283 tools across HR Domain A and Finance Domain B, assigning each tool a unique virtual token appended to the model vocabulary.Domain A contains 918 tools and Domain B 7,365 tools.
  • Setting: The enterprise setting adds 123 expert-curated business rules pairing confusable tool sets with rule text for domain-specific disambiguation.Rules capture API deprecations, versioning changes, and routing constraints.
  • Stage 2: Stage 2 trains the model to emit a free-form thinking trace followed by an ordered list of virtual tool tokens for each query.The trace–answer pair consists of reasoning z and tool-token output y, later compared with ground truth.
  • Stage 2: Training queries come from RRB generation and rule-targeted generation over confusable tools and nearest neighbours, spanning multiple ambiguity levels and phrasing styles.Reference traces scope permissions, compare APIs, narrow endpoint candidates, apply the governing rule by name, and commit to tool tokens.
  • Inference: At inference, TRACE greedily decodes in a single beam, producing the trace and JSON token list without constrained beam search or prefix-trie expansion.This design is intended for direct deployment at production latency.

4 Experiments and Results

TRACE is evaluated across an 8,283-tool enterprise catalog using per-domain retrieval and parametric-knowledge probes. Reasoning retrieval preserves or improves tool knowledge, delivers strong rule-sensitive recall, generalizes across token formats, and substantially reduces inference latency.

  • Evaluation setup: Evaluation uses Gemma4-E4B-it on a combined enterprise catalog, with per-domain metrics and PRB golden production query-tool pairs.PRB contains nA=131 and nB=123 queries; retrieval reports R@10 for constrained beam search and R@gen for greedy reasoning decoding.
  • Parametric knowledge: Non-reasoning retrieval collapses all knowledge probes near random, whereas reasoning retrieval reverses the collapse and improves MCQts and QAts beyond the Stage-1 ceiling.This pattern also holds on the domain-expert MCQexpert probe, where TRACE substantially outperforms the non-reasoning baseline.
  • Rule grounding: 86.3% Domain-A R@gen under replacement rule data represents a +30.6 pp gain over 0 q/rule, while Domain B stays within its baseline confidence interval.At 12 queries per rule, Domain-A R@gen reaches 73.3% (+17.6 pp); gains appear where rules cover confusable clusters.
  • Format generalization: Reasoning retrieval recovers MCQexpert across formats to 59.5, 50.9, and 61.5, while non-reasoning retrieval reaches near-random 34.8%, 30.0%, and 33.7%.F=a and F=c each reach ∼86% Domain-A R@gen under rule grounding, whereas F=b trails by ∼10 pp on every metric.
  • Inference latency: ∼1.9 s single-beam free-form decoding versus ∼19 s constrained beam-10 yields 11.2 qps versus 0.05 qps at concurrency 32, a ∼200× throughput gap.Both modes use vLLM on a single H200 with 100 Domain-A PRB queries.

5 Conclusion

TRACE decouples parametric tool knowledge from retrieval through a two-stage, business-rule-grounded reasoning curriculum. Across 8,283 enterprise tools, it achieves strong single-beam recall while improving Stage-1 probing accuracy and enabling production-latency deployment.

  • 5 Conclusion: TRACE decouples tool knowledge from retrieval by memorizing the catalog in Stage 1 and reasoning over business rules before emitting tool tokens in Stage 2.This two-stage curriculum uses a business-rule-grounded reasoning trace.
  • 5 Conclusion: ∼86% recall on Domain A and ∼60% on Domain B are achieved across 8,283 enterprise tools with single-beam greedy decoding.The evaluation covers two domains.
  • 5 Conclusion: +7.6 pp and +4.5 pp are the respective improvements in Stage-1 MCQ and QA probing accuracy.TRACE preserves tool knowledge while improving both probing measures.
  • 5 Conclusion: TRACE delivers knowledge-preserving parametric tool retrieval that is deployable at production latency.Its retrieval uses single-beam greedy decoding.

Limitations

TRACE’s evaluation is limited to two enterprise domains and a single model size, while its proprietary benchmarks cannot be released, leaving broader-domain and scaling behavior open.

  • Evaluation scope: TRACE is evaluated only on HR and Finance at a single model size, leaving higher-parameter models and adjacent verticals such as procurement and healthcare untested.The passage identifies broader scaling and vertical coverage as open questions.
  • Benchmark availability: The benchmarks use proprietary tool catalogs and therefore cannot be released, although the appendix documents the protocol’s full details.This limits direct benchmark release while preserving methodological documentation.

Ethical Considerations · A Multi-Format Tool Memorization

The paper states that its experiments followed ACL Ethics Policy and relevant research-integrity guidelines, with no known unresolved ethical risks. Its flat-token memorization training jointly uses forward, reverse, and discriminative tool-selection formats.

  • Ethical Considerations: The experiments were conducted within the provisions of the ACL Ethics Policy.
  • Ethical Considerations: The authors report no remaining ethical risks that have not been addressed, to the best of their knowledge.
  • A Multi-Format Tool Memorization: Flat-token training jointly uses three formats: forward mapping, reverse mapping, and discriminative Multi-Choice Tool Selection.
  • A Multi-Format Tool Memorization: Forward mapping trains the model to map a tool description desc to its virtual token vt.
  • A Multi-Format Tool Memorization: Reverse mapping trains the model to recover the tool description from its virtual token.
  • A Multi-Format Tool Memorization: In discriminative Multi-Choice Tool Selection, the model selects the gold token vt from K+1 candidates given desct.

B Stage 2 Data Synthesis: Full Pipeline

TRACE’s Stage 2 pipeline synthesizes challenging retrieval queries across ambiguity tiers and business-rule phrasing styles, then generates deliberative traces and filters them for validity. Rule-violation queries specifically target cases where rules collapse multiple candidate tools to one, beyond embedding retrievers’ resolution.

  • Hard negatives: Hard-negative pools retrieve top-K cosine neighbours to force finer-grained learning among tools with similar surface forms.The neighbour pool includes tools whose surface forms are close enough to retrieved tools to make retrieval more challenging.
  • RRB synthesis: RRB generation stratifies anchors into easy, medium, and hard tiers spanning fully specified intent, mild ambiguity, and multi-intent or under-specified queries.The tiers use |A|=1, |A| ∈{2, 3}, and |A| ≥4, respectively, with dynamic few-shot real-user queries grounding generation.
  • Rule-targeted synthesis: Rule-targeted generation uses equal mixes of explicit, implicit, and violation queries, each preserving a pointer to its source rule.Violation queries are deliberately under-specified so applying the rule collapses multiple candidate tools to a single tool, targeting a case embedding retrievers cannot resolve.
  • Trace generation: The trace teacher restates intent, distinguishes plausible candidates, applies the source rule when applicable, and commits to the answer set.RRB samples use r∗= ∅, while rule-targeted samples retain their source rule for trace generation.
  • Filtering and judging: Programmatic validation rejects samples with grounding, leakage, consistency, or rule-attribution failures, while a temperature-0 judge scores four quality axes and permits up to 5 critique-driven retries.The judge evaluates query naturalness, tier compliance, label correctness, and trace faithfulness before rejected samples are regenerated or dropped.

C Training Setup · D Inference · E Detailed Evaluation Results

TRACE uses masked autoregressive SFT with LoRA across two stages, then performs single-pass greedy inference without constrained beam search. Evaluation reports retrieval and probing metrics with confidence intervals, alongside a point estimate for MCQexpert.

  • C Training Setup: Both stages use chat-style autoregressive SFT with cross-entropy loss masked on prompt tokens, supervising only assistant completions.Stage 2 trains trace–answer pairs, while Stage 1 uses format-specific completions such as desc→v_t.
  • C Training Setup: Gemma4-E4B-it is trained with LoRA, fully finetuning embedding layers using r = 64 and α = 128 in both stages.Both stages use AdamW with cosine scheduling; learning rates are 5e-5 for Stage 1 and 1e-4 for Stage 2.
  • C Training Setup: 16 is the Stage 2 batch size, compared with 8 for Stage 1, using bf16 precision on a single H200 GPU.The batch-size settings differ between stages while hardware and numerical precision remain fixed.
  • C Training Setup: 95% confidence intervals are reported for retrieval and probing metrics using paired bootstrap over per-example predictions.The bootstrap uses B = 1,000; MCQexpert is instead reported as a point estimate over n = 454.
  • D Inference: TRACE decodes greedily in a single forward pass, emitting a reasoning trace followed by a JSON list of virtual tool tokens.Generation terminates at the closing JSON bracket.
  • D Inference: TRACE does not use constrained beam search or prefix-trie expansion over the virtual-token vocabulary during inference.The passage attributes valid token-list generation to Stage 1 parametric grounding combined with the Stage 2 deliberative scaffold.
  • E Detailed Evaluation Results: The detailed evaluation section provides full results for the main-body experiments, including confidence intervals and all metrics.This section supplements the summarized experimental results with complete statistical reporting.

E.1 Probing Evaluations

Probing evaluations show that reasoning retrieval restores tool knowledge lost under non-reasoning retrieval across token formats, while reasoning during Stage 1 further improves MCQexpert accuracy. Rule-grounded scaling keeps in-distribution probes stable but trades held-out knowledge for rule-data volume.

  • Format sweep: Non-reasoning retrieval collapses every probe across all three token formats, whereas reasoning retrieval recovers in-distribution probes to within a few points of the Stage-1 ceiling.This dissociation is consistent across flat, bare-hierarchical, and wrapped-hierarchical token formats.
  • Format sweep: On flat tokens, n→r partially recovers MCQ accuracy, but it fails entirely for hierarchical formats, indicating irrecoverable damage to description–token binding.The failure occurs for both bare-hierarchical and wrapped-hierarchical formats.
  • Stage-1 ablation: 67.8 MCQexpert accuracy is achieved by direct reasoning retrieval without Stage 1, within 1.4 pp of the strongest Stage-1 ceiling at 69.2, but with the lowest retrieval recall.The no-Stage-1 result is reported for (a, ∅, r), while the ceiling is (c, m, ∅) = 69.2.
  • Stage-1 ablation: Switching Stage 1 from non-reasoning memorization to reasoning memorization yields consistent 2–13 pp gains on MCQexpert across token formats.This comparison holds reasoning retrieval fixed in the Stage-1 ablation.
  • Rule-grounded scaling: Each rule-data increment costs roughly 1–2 pp of held-out general-domain knowledge, with MCQexpert reaching 55.7 for F=a and 56.4 for F=c, while in-distribution probes remain stable.These values remain above the R=n collapse but below the no-rules ceiling.

E.2 Retrieval Evaluations · E.3 Qualitative Trace Example · F MCQexpert Evaluation Details

The retrieval evaluations compare checkpoints under decoding modes matched to their reasoning design, while a qualitative example shows rule-grounded traces correcting ambiguous-tool retrieval. MCQexpert provides expert-crafted multiple-choice evaluation for Domain A and Domain B knowledge.

  • E.2 Retrieval Evaluations: Tables 7 and 8 report retrieval results for every trained (F, M, R) checkpoint.These tables complete the comparisons described in §4.3.
  • E.2 Retrieval Evaluations: Non-reasoning checkpoints use constrained decoding, whereas reasoning checkpoints use single-beam reasoning decoding.The evaluation decoding mode is matched to each checkpoint’s design.
  • E.2 Retrieval Evaluations: Two TRACE-class checkpoints were additionally reevaluated under constrained decoding, marked with the suffix cstr and an asterisk.For these checkpoints, both decoding modes are reported.
  • E.3 Qualitative Trace Example: Figure 5 compares model reasoning traces before and after business-rule grounding for an ambiguous query.The example illustrates the effect of the <think> trace mechanism.
  • E.3 Qualitative Trace Example: The ungrounded model reasons plausibly but selects the wrong tool in the representative comparison.This behavior appears in the left panel of Figure 5.
  • E.3 Qualitative Trace Example: Citing the governing business rule within the reasoning trace corrects retrieval for the ambiguous query.This corrected behavior appears in the right panel of Figure 5.
  • F MCQexpert Evaluation Details: MCQexpert is a multiple-choice question dataset handcrafted by domain experts.It evaluates model performance on knowledge from Domain A and Domain B.
  • F MCQexpert Evaluation Details: MCQexpert evaluates model performance on Domain A and Domain B knowledge.The dataset is used specifically for multiple-choice assessment across both domains.

F.1 Dataset Construction … H System Prompts

The appendix details expert-curated evaluation data, constrained MCQ inference, business-rule catalog structure, and Stage 2 prompt templates for rule-targeted queries and reasoning traces. It specifies 454 filtered questions, 123 business rules, and procedural constraints governing tool-selection reasoning.

  • F.1 Dataset Construction: 454 multiple-choice questions remain after filtering Domain A and Domain B lesson groups by keywords and deduplicating questions across lessons.The questions were written by domain experts for certification-exam domain-understanding assessment.
  • F.2 Model Inference: MCQ inference uses greedy constrained decoding over valid indexed answer letters, with accuracy computed as the fraction of correct predictions across N questions.The model receives question content and answer choices, each prefixed by an indexed English-alphabet letter.
  • F.3 Example Question: The example MCQ asks which step is not typical in indirect-materials financing, with the provided correct answer being C.The displayed choices are requirements determination, source-of-supply determination, production planning and control, and purchase order processing.
  • G Business Rule Catalog: 123 business rules span the two domains, covering API deprecations, versioning, product-line routing, and data-model boundaries beyond tool descriptions.The catalog contains 20 rules for Domain A and 103 for Domain B, authored by domain experts.
  • G.1 Rule Structure: Each rule is stored as JSON with tool_name, confusables, and rule_text fields defining the target, overlapping tools, and routing directive.The confusable set captures tools that overlap semantically with the governed target.
  • G.2 Example Rule: An example routes current-state queries to EndpointA, historical or date-range queries to EndpointB, transactional details to EndpointC, and otherwise to EndpointD.The example is defined for API1/EndpointA and lists API1/EndpointB, API1/EndpointC, and API1/EndpointD as confusables.
  • H System Prompts: Stage 2 prompts generate explicit, implicit, and exception queries, then produce traces that filter permissions and reason hierarchically from APIs to endpoints before selection.Query generation requires realistic, varied, nontechnical, semantically distinct questions whose targets are tools mentioned in the rule; trace generation requires exact provided tool names and business-rule adherence.
Loading 2607.22639v1…