Adaptive Instructed-Retriever：2倍の低レイテンシでフロンティア品質の検索を実現

効果的なエンタープライズデータエージェントには、正確かつ高速な検索が必要です。今年、DatabricksはInstructed-Retriever-1をリリースしました。これは、エンタープライズデータのスキーマやカスタム指示を組み込みながら、並列テスト時スケーリング（parallel test-time scaling）を使用して低レイテンシで検索精度を向上できる検索モデルです。このシングルステップ検索アプローチは、ユーザーリクエストの大部分に対して効果的に機能します。しかし、より複雑なマルチホップの質問では、モデルが段階的に証拠を収集し、複数ステップにわたってクエリを洗練させるシーケンシャル検索（順次検索）が依然として有効です。ただし、これには追加のレイテンシというコストが伴います。

そこで今回、厳格なコストとレイテンシの保証を維持しながら、並列検索のスピードとシーケンシャル検索のパフォーマンスを組み合わせたAdaptive Instructed-Retrieverを導入します。目標はシンプルです。追加の検索ステップは、それが本当に役立つ場合にのみ実行するということです。これは、DatabricksのデータエージェントであるGenie Codeが直面しているものと同じ、検索効率の課題です。つまり、大規模で変化し続けるワークスペースの中から、総当たり的な探索に無駄なステップを費やすことなく、適切なテーブル、ノートブック、ダッシュボード、ドキュメントを見つけ出す必要があります。Adaptive Instructed-Retrieverはこの検索レイヤー向けに設計されており、証拠が明確な場合は迅速に結果を返し、より難しいリクエストで必要な場合にのみシーケンシャル検索を使用します。本ブログ記事で示すように、トレーニングされたモデルは、主要なサードパーティ製モデルと同等の品質を2倍低いレイテンシで実現し、Databricksの検索ベンチマークにおいてシングルステップ検索を大幅に上回る成果を上げています。

Adaptive Instructed-Retrieverを構築するために、シーケンシャルステップの数に固定の上限を設け、各ユーザーリクエストに必要な計算量を適応的に判断するようにエージェントをトレーニングします。十分な証拠がすでに見つかっている場合、エージェントは早期に終了して関連する証拠を返します。追加の検索によって検索品質が向上する可能性が高い場合は、ステップ制限まで検索を継続できます。

Adaptive Instructed-Retrieverのトレーニング

検索品質とシーケンシャルスケーリングによるレイテンシコストのバランスを改善するため、DatabricksはAdaptive Instructed-Retrieverをトレーニングしました。これは、並列シングルステップ検索とシーケンシャル検索の両方をサポートし、主要なサードパーティ製モデルよりも大幅に低いレイテンシで動作する小規模なカスタムモデルです。このモデルを、独自のエンタープライズ検索ベンチマークと公開検索ベンチマーク（マルチホップ推論が有効なタスクを含む）を組み合わせた環境で評価します。これらのベンチマーク全体において、Adaptive Instructed-Retrieverは、主要なサードパーティ製モデルやオープンソースモデルに匹敵するパフォーマンスを達成しつつ、2倍低いレイテンシを実現しています。

トレーニングデータを準備するために、合成エンタープライズ検索環境と、Instructed-Retriever-1に類似し、KARLレポートで公開されているエージェント型データ合成プロセスを利用しています。Instructed-Retriever-1の既存のトレーニングデータを再利用して、高速な並列シングルステップ検索を実行するモデルの能力を維持し、さらに、複数のエージェント型検索ステップからより多くの恩恵を受ける合成マルチホップ質問を導入しています。

ベースモデルから始めて、オンライン強化学習（ORL）を使用し、最終的なパフォーマンスが向上する可能性が高い場合にのみ追加の検索ステップを実行するようにモデルに学習させます。具体的には、CISPO（Clipped Importance Sampling Policy Optimization）を使用してモデルをエンドツーエンドで最適化します。その際、トラジェクトリ（軌跡）の品質と検索コストのバランスをとる報酬設計を採用しています。つまり、パフォーマンスの高いトラジェクトリには報酬を与え、対応するパフォーマンスの向上が得られない追加の検索ステップにはペナルティを科します。

モデルのトレーニングには、AI Runtime（AIR）を使用しています。AIRはDatabricksのお客様も利用できるため、このアプローチは独自のドメインやワークロードに特化したモデルを開発する上で実用的です。トレーニングのレシピは意図的に軽量に設計されています。事前学習済みのベースモデルから開始し、適度な量の合成データを使用してその検索動作を特化させます。図2に示すように、この軽量なアプローチは、新しい検索タスクやドメインに対しても優れた汎化性能を発揮します。

図2は、Adaptive Instructed-Retrieverを、2つの主要なサードパーティ製モデル（Claude Sonnet 5およびGPT-5.6 Luna）と1つのオープンソースモデル（DeepSeek-V4-Flash）と比較したものです。すべてのモデルについて、検索品質と平均エンドツーエンドレイテンシをプロットしています。薄い色のバーはシングル検索ステップによる検索スコアを示し、濃い色のバーはマルチステップ検索による結果を示しており、これらは左軸で測定されます（高いほど良い）。斜線のバーは右軸に対するエンドツーエンドレイテンシを示しています（低いほど良い）。Adaptive Instructed-Retrieverは、主要なサードパーティ製およびオープンソースモデルのパフォーマンスに匹敵しながら、わずか5.8秒で回答を返します。これは、Claude Sonnet 5、DeepSeek-V4-Flash、またはGPT-5.6 Lunaのいずれよりも2倍以上高速です。

カスタマイズされた品質とレイテンシのトレードオフを実現するオンライン強化学習

Adaptive Instructed-Retrieverをトレーニングすることで、ORL中に使用されるステップペナルティの大きさを調整し、品質とレイテンシのトレードオフを選択できるようになります。したがって、Adaptive Instructed-Retrieverを強化する一連のチェックポイントをトレーニングし、本番環境のワークロードに最適なものを選択できます。

図3は、ORLトレーニング中にペナルティの大きさをスイープ（変化）させることで、品質とレイテンシの平面上の異なるポイントに位置する一連のチェックポイントをどのように取得できるかを示しています。y軸はスコア（高いほど良い）、x軸はエンドツーエンドレイテンシ（右に行くほど高速になるようにプロット）です。赤色の曲線は、これらの動作点をつないでフロンティア（境界線）を形成しています。ステップペナルティを軽くすると、モデルはより多くのステップを実行して高いスコアを達成でき、ペナルティを重くするとレイテンシが短縮されます。

トレーニングされたAdaptive Instructed-Retrieverモデルの完全なフロンティアは、他の代替モデルを凌駕しています。トレーニングされていないInstructed-Retrieverベースモデルと比較すると、すべてのチェックポイントが同等またはそれ以下のレイテンシでより高い品質を提供しており、この向上が「いつ検索し、いつ検索しないか」を学習したことによるものであることを示唆しています。フロンティアの最上部では、Databricksのモデルは他の主要モデルに匹敵するスコアを達成しながら、2倍以上高速です。トレーニングされたフロンティアには異なるトレードオフを持つチェックポイントが含まれているため、インタラクティブな用途には速度を優先し、より困難なオフライン検索には品質を優先するなど、各ワークロードや予算に最適なものを選択できます。

Adaptive Instructed-Retrieverによる、より効率的で的を絞った検索の策定

代替モデルの検索ポリシーと、トレーニングされたAdaptive Instructed-Retrieverモデルの検索ポリシーを比較するいくつかの例を示します。これらの例は、Adaptive Instructed-Retrieverが単純な質問に対してどのように効率的に検索を行い、より難しい質問に対して残りの検索バジェットをどのように効果的に使用できるかを示しています。

網羅的な検索を行わずに否定的な回答を検証する

Adaptive Instructed-Retrieverは、Sonnetよりも1ステップ早く、Lunaよりも2ステップ早く同じ報酬に到達します。

質問：Company Xは、2022年度の損益計算書の項目として事業再編費用を明示的に報告していますか？

次のラウンドを使用して戦略を変更する

Adaptive Instructed-Retrieverは、再現率を向上させるために検索戦略を調整することを学習します。広範な探索から、ターゲットを絞ったアカウント検索へと移行します。最少ステップ数でSonnetと並びつつ、最高の報酬を獲得します。

質問: どの顧客がLiteLLM Proxyを使用しているか、または使用を検討したことがありますか？

結論

Adaptive Instructed-Retrieverは、以前に発表されたInstructed-Retriever-1を、並列シングルステップ検索から適応型マルチステップ検索へと拡張します。これにより、反復的な推論と証拠収集によって検索品質が大幅に向上する複雑なクエリに対しては、モデルが追加の検索ステップを費やすことができるようになり、一方でシンプルなクエリに対しては早期に終了してレイテンシを最小限に抑えることができます。

この適応型アプローチは、シングルステップ検索よりも大幅に強力な検索パフォーマンスを提供する一方で、主要なサードパーティ製モデルと同等のパフォーマンスを2倍低いレイテンシで実現します。

オンライン強化学習（Online Reinforcement Learning）中にステップペナルティを変化させることで、検索品質と推論コストのトレードオフを明示的に最適化できます。これにより、品質とレイテンシのフロンティアに沿ったさまざまなポイントでチェックポイントを生成し、本番環境のワークロードに最も適した動作点を選択できます。

私たちの貢献は、大規模で絶えず変化するワークスペースで動作するDatabricksのデータエージェント向けに、実用的な検索ビルディングブロックを提供することです。

Genie Codeや、Genie OneおよびGenie Agentsなどのエクスペリエンスでは、深い探索に過度な時間と労力を費やすことなく、適切なテーブル、ノートブック、ダッシュボード、ドキュメントを見つける必要があります。

Adaptive Instructed-Retrieverは、この設定に対して境界のあるポリシーを提供します。一般的なルックアップでは高速であり、困難な探索タスクでは徹底的であり、プロダクトのレイテンシ予算内で制御可能です。

より広範には、これらの結果は、マルチステップの推論を必要とするタスクであっても、小規模な特化型モデル\ の競争力を示しています。Databricks AIRを使用することで、お客様は独自のドメインやパフォーマンス要件に合わせて特化型モデルをカスタマイズできます。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事