TRACE: エンタープライズLLMにおける知識保存型パラメトリックツール検索のためのビジネスルールに基づく推論カリキュラム

要旨

パラメトリック検索により、LLMは各APIに固有の仮想トークンを割り当て、モデルに制約付きビーム探索でそれを生成するよう訓練することで、ツールを暗黙的に検索できるようになる。ToolSenseは、この手法に2つの重大な欠点があることを示している。すなわち、訓練中にパラメトリックなツール知識が破壊されることと、ビーム探索によるデコードがリアルタイム展開には遅すぎることである。我々は、この乖離を解決する2段階カリキュラムであるTRACE（拡張思考連鎖とエンタープライズルールによるツール検索）を導入する。第1段階では、ToolSenseのマルチフォーマット記憶SFTを再利用し、LoRAを用いてツール知識を注入する。第2段階が我々の中核的貢献である。モデルは、ツールトークンのJSONリストを生成する前に思考トレースを出力するよう訓練される。その際、ToolSenseのRRBペアと、専門家がキュレーションしたビジネスルールを対象とする合成クエリという2つのデータソースを使用し、両方に推論トレースを付加する。この訓練目的により、第1段階のMCQおよびQAプローブ精度を維持しつつ、生産レイテンシでの単一ビーム貪欲デコードを可能にする。2つのエンタープライズ製品ラインにわたる8,300以上のツールからなる複合エンタープライズカタログで評価した結果、TRACEの第2段階訓練はツール理解を維持するだけでなく向上させ、第1段階と比較してMCQ精度が+3.2ポイント、QAプローブ精度が+9ポイント向上した。検索において、TRACEはドメインAで約86%の再現率、ドメインBで約60%の再現率を達成した。一方、埋め込みベースラインの性能はそれぞれ約27%と約52%であり、両方とも単一ビーム貪欲デコードを用いており、生産レイテンシで直接デプロイ可能である。

English

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.