今回の論文
今回取り上げるのは、Yuanhang Zheng、Peng Li、Wei Liu、Yang Liu、Jian Luan、Bin Wang による論文「ToolRerank: Adaptive and Hierarchy-Aware Reranking for Tool Retrieval」です。2024年3月11日に arXiv で公開され、その後 LREC-COLING 2024 でも発表されています。公開元は arXiv、研究分野は AI エージェント、ツール検索、情報検索ベースの API 選択です。URL は https://arxiv.org/abs/2403.06551 です。
この論文を選んだ理由は、エージェントが使えるツール数が増えたときの、かなり現実的なボトルネックを正面から扱っているからです。今の AI エージェントは「推論」以前に「どの API を候補に渡すか」で失敗することが多く、この論文はそこを検索と再ランキングの設計で改善しています。
どんな技術か
ToolRerank は、AI エージェントが大量の API やツール群から適切なものを選ぶときに使う再ランキング手法です。ベースになる検索器で候補を広めに取ったあと、その候補をもう一段精密に並び替えます。
ただし、普通の reranker をそのまま使うだけでは足りない、というのがこの論文の出発点です。ツール選択には、通常の文書検索と違う癖があります。たとえば、学習時に見たことがあるツールと、未学習の新規ツールでは最適な候補数が違います。また、ユーザーの依頼によっては「1つのツールの複数 API をまとめて使うべき」場合と、「複数ツールをまたいで組み合わせるべき」場合があります。
ToolRerank はこの違いを明示的に扱います。具体的には、見たことのあるツールと未知のツールで reranker に渡す候補数を変える Adaptive Truncation と、単一ツール型クエリと複数ツール型クエリで並べ替え方を変える Hierarchy-Aware Reranking の 2 段構えになっています。
課題
この技術が解決しようとしているのは、ツール数が多いエージェントで、正しい API を上位候補に残し続けることの難しさです。
何が難しいのかというと、ツール選択は単なるキーワード一致では済まないからです。ユーザーの依頼文と API ドキュメントの意味的な対応を取る必要がありますし、同じような機能を持つ API が複数あると、少しの表現差で誤ったツールに寄りやすくなります。
既存の方法ではどこに限界があるのかというと、まず BM25 のような語句ベース検索は意味の近さを十分に拾えません。逆に dense retrieval は効率は良いものの、上位候補の中で細かな API 差分を見分けるのが弱いことがあります。そこで cross-encoder reranker を重ねるのが自然ですが、論文では「候補数を固定する」「ツール階層を無視する」設計だと tool retrieval では最適にならないと指摘しています。
なぜこの課題を解く必要があるのかというと、エージェントの失敗はしばしば最初のツール候補選定で決まるからです。後段の LLM が優秀でも、必要な API が top-k に入っていなければ正しい実行計画は立てにくくなります。特に MCP、社内 API、SaaS 連携のようにツールカタログが大きい環境では、この問題が先に効きます。
実際の AI システムでは、同じツール内の別 API がまとまって必要な問い合わせもあれば、検索 API と変換 API をまたいで使うような問い合わせもあります。それなのに結果上位が似た API で埋まったり、逆に必要な 1 つのツールの関連 API がばらけて落ちたりすると、エージェントは不安定になります。ToolRerank はこの「候補集合の作り方の粗さ」を改善しようとしています。
用語解説
- Dual-Encoder Retriever
- クエリと API ドキュメントを別々にベクトル化し、類似度で候補を引く検索器です。大量 API に対して高速ですが、細かな相互作用を見にくいので、ToolRerank では一次候補を集める役割として使われます。
- Cross-Encoder Reranker
- クエリと候補文書を連結して一緒に読ませ、関連度を精密に点数化するモデルです。ToolRerank の土台ですが、この論文では単純に top-k を並べ替えるだけでは tool retrieval に最適化しきれない点が重要です。
- Seen Tool / Unseen Tool
- 学習データ中で登場したツールか、未登場の新規ツールかという区別です。ToolRerank は、この違いによって「候補をどこまで広げるべきか」が変わると見ています。
- Single-tool Query / Multi-tool Query
- 1つのツール群だけで解決しやすい依頼か、複数ツールを組み合わせる依頼かを表します。ToolRerank はこの判定結果で再ランキング方針を切り替えます。
- Recall@5
- 正解 API が上位 5 件にどれだけ含まれるかを見る指標です。エージェントに top-5 API を渡すような設定では、そのまま実用上の成功率に近い意味を持ちます。
技術の仕組み
ToolRerank の核は、「候補数」と「候補の多様性・集中度」をクエリとツールの性質に合わせて動的に制御することです。
基本アイデア
全体の流れは 4 段階です。まず dual-encoder retriever で API 候補を粗く集めます。次に Adaptive Truncation で、reranker に渡す候補数をツールの seen / unseen に応じて変えます。その後 cross-encoder reranker で関連度順に並べます。最後に Hierarchy-Aware Reranking で、単一ツール型か複数ツール型かに応じて並び順をもう一度調整します。
この設計は、検索精度を上げるだけでなく、「LLM にどんな API セットを見せるか」を制御していると見ると理解しやすいです。
モデル構造
この論文は新しい基盤モデルそのものを提案しているわけではありません。使っている主な構成要素は、dual-encoder retriever、cross-encoder reranker、そして単一ツール型か複数ツール型かを判定する classifier です。いずれも BERT ベースで初期化されています。
つまり新規性は Transformer の内部構造ではなく、検索パイプラインのオーケストレーション側にあります。既存の retriever / reranker を持つチームでも取り入れやすいのが利点です。
Adaptive Truncation
論文の観察では、seen tool の問い合わせでは reranker に多くの候補を渡すとかえってノイズが増え、未学習の unseen tool では候補を広めに渡したほうが正解 API に届きやすくなります。
そのため ToolRerank は、seen 用の打ち切り位置 k_s と unseen 用の打ち切り位置 k_u を別々に持ちます。学習済みツールに属する API なら浅い位置で打ち切り、新規ツールに属する API ならより深い位置まで reranker に回します。要するに、「既知ツールは絞り、未知ツールは拾いにいく」という方針です。
これは実務でもかなり自然です。社内で頻出 API は coarse retrieval の時点で十分上位に来やすい一方、新規連携ツールやニッチなカテゴリは一次検索で埋もれやすいからです。
Hierarchy-Aware Reranking
ToolRerank のもう 1 つの柱が、ツールライブラリの階層構造を使った再配置です。ここではクエリを single-tool と multi-tool に分類して、別アルゴリズムを使います。
単一ツール型クエリの処理
単一ツール型クエリでは、最終的な上位 API が同じツールに集中していたほうがよい場面があります。たとえば映画 API から検索 API と詳細取得 API を続けて使うようなケースです。
しかし通常の reranker だと、似た機能を持つ別ツールの API が混ざり、上位候補が散ってしまいます。そこで ToolRerank は、高い関連度を持つ API を起点に「このツールは有望だ」と判断し、そのツールに属する API を優先的に前へ寄せます。
さらに unseen tool については、top-k の外側にも正解 API がある可能性を考慮して、同じツールに属する API をライブラリ全体から追加探索する extended API list も使います。これは、新規ツールに対して coarse retrieval の取りこぼしを補う工夫です。
複数ツール型クエリの処理
複数ツール型クエリでは、逆に同じような API ばかり並ぶのが問題です。たとえば QR コード関連 API が上位を埋めると、本当は別カテゴリの分析 API も必要なのに top-5 に入らなくなります。
そこで ToolRerank は、候補 API をノードにしたグラフを作り、同じツールに属するか、意味的に近い API 同士を辺でつなぎます。そのうえで連結成分ごとに上位候補を少数ずつ選び、結果集合の多様性を確保します。
言い換えると、「似た API の束からは取りすぎず、異なる役割のツール群をバランスよく残す」仕組みです。エージェントの実行計画に必要な操作の幅を保つための設計と言えます。
学習方法とデータ
実験では ToolBench を使い、約 16,000 API を含むツールライブラリ上で評価しています。訓練用クエリは 169,287 件、開発用は 600 件、テストは 6 分割で各 100 件です。未学習ツール性能を見るために、いくつかのツールやカテゴリを訓練データから外した unseen 設定も作られています。
重要なのは、この論文が「ツールが増え続ける環境」を前提にしていることです。単一の固定ベンチマークだけでなく、新しいツールが後から入る現実に合わせた評価になっています。
実験と結果
論文では、ToolRerank が本当に API 検索とエージェント実行結果を改善するかを、検索品質と実行品質の両方で検証しています。
何を検証したのか
主な検証点は 3 つです。1 つ目は、ToolRerank が BM25、DPR、固定 top-k reranking を上回るかです。2 つ目は、Adaptive Truncation と Hierarchy-Aware Reranking がそれぞれ効いているかです。3 つ目は、検索精度の改善が本当に LLM のツール実行結果まで波及するかです。
データセットと評価指標
検索品質の評価には ToolBench の 6 テスト分割を使い、指標は NDCG@5 と Recall@5 です。実行品質の評価では ToolLLaMA-7B をバックボーンにし、Pass Rate と Win Rate を測っています。Pass Rate は有効な解決策を返せた割合、Win Rate はベースライン解より良い結果を返せた割合です。
検索品質は seen / unseen の両方で改善
全テスト平均では、ToolRerank の NDCG@5 は 82.1、Recall@5 は 84.2 でした。これは DPR の 71.0 / 75.2、固定 reranking の最良だった Rerank-30 の 78.8 / 80.3 を上回っています。
特に unseen 側の改善が分かりやすく、unseen 平均 Recall@5 は Rerank-30 の 74.7 に対して ToolRerank は 78.7 でした。未知ツールで候補を広く見る設計が、きちんと効いていると読めます。
seen 側でも Recall@5 は 89.6 で、Rerank-10 の 88.9 を上回っています。つまり ToolRerank は unseen に寄せすぎて既知ツールを落とすのではなく、両立させています。
Adaptive Truncation が固定 top-k より有効
アブレーションでは、seen に 10 件、unseen に 50 件を渡す組み合わせが平均 Recall@5 84.2 と最良でした。固定で 10 件だと 82.0、30 件だと 81.9、50 件だと 81.0 まで落ちています。
この結果から言えるのは、「候補を増やすほど良い」でも「絞るほど良い」でもなく、既知か未知かで最適点が違うということです。これは大規模ツールカタログの設計で、そのまま使える示唆です。
階層を使った再配置も効いている
Hierarchy-Aware Reranking を外すと、平均 Recall@5 は 84.2 から 82.8 に下がっています。単一ツール型用アルゴリズムだけを全件に使っても 83.5、複数ツール型用だけを全件に使っても 82.6 で、クエリ種別ごとの切り替えが効いています。
さらに、単一ツール型か複数ツール型かを判定する classifier の精度は平均 93.3% で、oracle に置き換えても平均改善は 0.3 ポイントだけでした。つまり分類器が多少外れても、全体設計として十分堅いです。
実行結果も改善
ToolLLaMA-7B を使った実行評価では、ToolRerank の Pass Rate は 61.5、Win Rate は 57.0 でした。固定 reranking の Rerank-50 は 59.8 / 55.2、DPR は 57.2 / 49.8 なので、検索改善がそのままエージェント実行の改善につながっています。
論文では、より良い retriever が必ずしもより良い execution result になるわけではないとも指摘しています。ここが重要で、ToolRerank は単に検索指標を上げたのではなく、LLM に渡す API 集合の質を改善したから実行も良くなった、と解釈できます。
何に使える？
ToolRerank は、ツール数が増えてきた AI エージェント全般に応用しやすいです。
社内業務エージェントの API 選択
社内の CRM、請求、在庫、承認、検索 API などが増えると、単純な embedding 検索だけでは誤選択が出やすくなります。ToolRerank の発想を入れると、単一システム内で完結する依頼では API を集中させ、横断業務では多様な API を残しやすくなります。
MCP やプラグインの大規模カタログ
MCP サーバーやプラグインが多い環境では、同系統のツールが大量に並ぶ問題が起きます。複数ツール型クエリ向けの多様化ロジックは、似たツールばかり候補に出る状態を抑えるのに向いています。
エージェントの planner 前段
ツール選択を planner の前段に置く構成では、最初に渡す top-k が計画の質をほぼ決めます。ToolRerank は planner を変えずに候補セットだけ改善できるので、既存エージェントの前処理強化として入れやすいです。
新規 SaaS 連携が多いプロダクト
新しい連携先が頻繁に追加されるプロダクトでは、unseen tool への強さが重要です。Adaptive Truncation や extended API list の考え方は、「新しい連携が増えるたびに fine-tune し直したくない」現場に合っています。
API 推薦付き開発支援
開発者向けの API 推薦や関数呼び出し支援でも使えます。単一ライブラリ内の複数メソッドをまとめて使うケースと、複数ライブラリをまたぐケースを分けると、IDE 補完やコードエージェントの候補提示も改善しやすいはずです。これは論文の直接検証ではありませんが、構造的にはかなり近い応用です。
開発や事業へのヒント
この論文から得られるヒントは、エージェント品質を上げるには、モデル本体より前段の候補形成を見直す価値が大きいということです。
top-k を固定しない設計が効く
多くの実装では「とりあえず top-5」「とりあえず top-10」で固定しがちです。しかしこの論文は、既知ツールと未知ツールで最適な深さが違うことを示しています。プロダクトでも、利用頻度、学習履歴、カテゴリ新規性に応じて候補数を変えるだけで改善余地があります。
クエリ種別ごとのオーケストレーションが重要
エージェント開発では、すべての依頼を同じ planner に流しがちです。ですが実際には、単一システム完結の依頼と横断ワークフローの依頼では、必要な候補集合の性質が違います。ToolRerank の single-tool / multi-tool 分岐は、その差を軽量な分類器で吸収している点が参考になります。
新しいツール追加時の UX 改善に使える
新規ツールは学習履歴が薄く、既存人気ツールに埋もれやすいです。事業的には、せっかく追加した連携が使われない原因になります。unknown / unseen への探索を少し広げる設計は、ツールカタログの活性化にも効きそうです。
評価指標は実行結果まで見るべき
論文では retrieval 指標だけでなく Pass Rate と Win Rate まで見ています。エージェント開発でも、検索ベンチだけで満足せず、最終タスク成功率に接続して評価すべきだという良い例です。検索改善が execution 改善に変わる設計を選ぶことが重要です。
限界
ToolRerank にも限界はあります。まず、dual-encoder、cross-encoder、classifier の複数段を回すため、単純な dense retrieval より推論コストは増えます。ツール数がさらに大きくなると、reranking コストの管理が課題になります。
また、seen / unseen の区別をどう持つかは実運用では簡単ではありません。論文では学習データに基づいて明確に区切れますが、現場ではツールの更新頻度や API 追加単位がまちまちで、どこまでを unseen とみなすか設計が必要です。
single-tool / multi-tool 分類にもデータ依存性があります。論文では 93.3% の分類精度が出ていますが、社内業務や複雑な SaaS ワークフローでは、依頼文だけで完全に見分けるのは難しいかもしれません。
さらに、評価は ToolBench と ToolLLaMA-7B が中心です。MCP、コードエージェント、ブラウザ操作、長期メモリ付きエージェントでも同じ改善幅が出るかは追加検証が必要です。この応用可能性は高そうですが、一部は推測を含みます。
最後に、この論文は「どのツール候補を渡すか」を改善する技術であって、後段の reasoning や execution policy を直接改善するものではありません。したがって、ツール候補が正しくても計画器が弱ければ失敗は残ります。
よくある質問
Q. ToolRerank は LLM 自体を再学習する手法ですか？
A. いいえ、主役はツール検索パイプラインです。dual-encoder retriever、cross-encoder reranker、query classifier を組み合わせて、LLM に渡す API 候補集合を改善します。既存のエージェントに前段として追加しやすいのが利点です。
Q. なぜ unseen tool では深く候補を見る必要があるのですか？
A. 未学習ツールは coarse retrieval の上位に出にくいからです。論文でも unseen テストでは正解 API が top-10 に入る割合が seen より低く、固定 10 件では取りこぼしやすいことが示されています。そのため unseen だけ候補幅を広げる設計が効きます。
Q. single-tool と multi-tool を分ける意味は何ですか？
A. 必要な候補集合の形が違うからです。単一ツール型では同じツール内の関連 API をまとめて残したい一方、複数ツール型では似た API の重複を減らして役割の違うツールを残したいです。ToolRerank はこの違いを利用しています。
Q. 小規模なツール数でも導入する価値はありますか？
A. 数十個程度なら単純検索で十分なこともあります。ただし、同機能の API が増えてきた段階や、新規連携が多い段階では価値が出やすいです。まずは固定 top-k と query-type 分岐だけ試す簡略版でも効果検証しやすいです。
Q. RAG の文書検索にも同じ考え方は使えますか？
A. 一部は使えます。たとえば、単一ソースを深掘りしたい質問と複数ソースを横断したい質問で reranking 方針を変える発想は転用できます。ただし ToolRerank は API 階層を前提にした設計なので、文書検索にそのまま当てはめるには調整が必要です。
今日の学び
この論文は、大量のツールを持つ AI エージェントで、正しい API を top-k に残し続ける難しさを扱っています。そこで、既知ツールと未知ツールで候補深さを変える Adaptive Truncation と、単一ツール型か複数ツール型かで並び替えを変える Hierarchy-Aware Reranking を組み合わせました。
ここから得られるヒントは、エージェント品質の改善はモデル本体の強化だけではなく、「何を候補として見せるか」の設計でもかなり伸ばせるということです。特にツールや API が増えてきたプロダクトでは、固定 top-k を疑うこと自体が有効な改善ポイントになりそうです。