AI・機械学習

より良く、より効率的なRAGのための高速な検索

Faster Retrieval for Better, More Efficient RAG (machinelearning.apple.com)

142 pointsby zerop2026年11月12日24 コメント

要約

Appleの研究者たちは、最大内積検索（MIPS）を高速化する新しいアプローチを提案しています。これは、検索対象のキーとクエリの分布の違いや、クエリの特性を考慮することで、従来のクラスタリングや量子化手法よりも効率的です。学習ベースのアプローチであるSupportNetとKeyNetは、オフラインで学習したニューラルネットワークを用いて、推論時の検索速度と精度を向上させます。

全文翻訳

より良く、より効率的なRAGのための高速な検索

最大内積検索（MIPS）は、レコメンデーションシステム、検索、検索拡張生成（RAG）の中核をなしています。MIPSシステムは、クエリベクトルが与えられると、数百万の候補キーベクトルデータベース内でクエリに最も一致するものを検索します。その最も一致するキーに関連付けられた元のデータポイントがクエリの処理に使用されます。RAGでは、質問に最も適したドキュメントが通常クエリの前にプレフィックスとして付けられます。

キーの数が数百万に達し、ベクトルの次元が数千に達すると、クエリとすべてのキーとの網羅的な比較を実行する計算コストは劇的に増加します。クラスタリングと量子化を使用して精度と速度をトレードオフする多数の近似手法が提案されてきました。これらのベースライン手法は、巧妙なインデックス構造を構築するためにキーセットのプロパティに焦点を当てていますが、それらの幾何学的議論は、クエリセット自体がキーセットに似ていると緩やかに仮定するか、クエリがどのようになるかについての考慮事項を回避します。これは2つの点で問題となる可能性があります。重要なRAGの場合、クエリは短いユーザー生成の質問から生じますが、キーは長く、情報量が多く、キュレーションされたテキストの断片を表すため、キーとクエリは大きく異なります。さらに、本番環境で展開されるほとんどのMIPSユースケースは、安定した、よく特徴付けられたクエリ分布を持つ既存のユーザーベースを対象としています。これは、汎用的なクエリに依存しない方法が活用されていない信号です。良いパフォーマンスをすべての可能な入力に対して目指すのではなく、入力の可能性のあるソースに最適な方法を考案することは、計算リソースをより効率的にリダイレクトするのに役立つ現代の機械学習の中心的な教義です。

Amortizing Maximum Inner Product Search with Learned Support Functions（ICML 2026で発表された論文）では、Appleの研究者たちが、小さなニューラルネットワーク（NN）を、固定されたキーデータベースに対して一度オフラインでトレーニングし、クエリをより速く最適な一致にルーティングする方法を学習させます。彼らのアプローチは、典型的なクエリがどのように一致する傾向があるかを理解するための初期の償却コストをトレードオフし、これらの小さなNNを介して推論時に、より高速で正確な検索を実現します。

最大内積検索の基本

MIPSの基本的な原則は、検索を高速化するためにキーセットの幾何学的構造を活用することです。典型的なMIPSシステムは、キーセットをk個のクラスタに分割します（kは通常、合計キー数nよりもはるかに小さい）。各クラスタの代表的なセントロイドを追跡します。新しいクエリを受け取ると、システムはそのクエリをこれらのk個のセントロイドと比較します。次に、クエリはそれに最も近い代表的なセントロイドのクラスタにルーティングされ、その後、そのクラスタに格納されているすべてのキーと網羅的に比較されます。データベースにn個のキーとk個の均等に分散したクラスタがある場合、2番目のステップで考慮されるキーの実効数はn/kのオーダーになります。これは、ナイーブな網羅的アプローチで行われるn回の比較と比較してk倍の高速化です。残念ながら、このルーティングメカニズムはエラーを引き起こす可能性があります。なぜなら、セントロイドルーティングは正確なアプローチではないからです。これは以下の図1に示されています。

［動画］

図1: 検索の幾何学的アプローチは、セントロイドへの距離に基づいて空間を分割することに依存しているため、失敗して最近傍を見逃す可能性があります。アニメーションでは、最近傍は最も近いセントロイドに関連付けられていないため、見逃されています。

MIPSをサポート関数として

機械学習の観点からMIPSを興味深いものにしているのは、この作業以前にはほとんど見過ごされていた数学的構造に基づいていることです。クエリをセット内の任意のキーとの最大内積にマッピングする関数は、そのセットのサポート関数であり、凸解析でよく研究されているオブジェクトです。この関数には2つの重要な特性があります。第一に、それは凸関数であるため、凸関数専用に設計されたニューラルネットワークのクラスである入力凸ニューラルネットワーク（ICNN）によって近似できます。第二に、そのサポート関数の勾配は、任意のクエリで最適な一致キー自体を返します。

この作業の背後にある広範なアイデアは、計算コストをはるかに低く抑えながら、グラウンドトゥルースのサポート関数（またはその勾配を直接）を模倣できるニューラルネットワークを学習することです。これは、回帰として知られる古典的なML技術を通じて達成できます。入力ベクトル（ユーザーのクエリ）をサンプリングし、MIPSによって生成される関連出力信号（最適なスコアまたは最適なキーベクトルのいずれか）とペアにして、多数のトレーニング例を事前に計算するだけで済みます。その後、適切にパラメータ化されたニューラルネットワークをトレーニングして、入力が与えられたときの出力を再現するだけです。

最大内積検索を学習する2つの方法

上記の数学的アプローチに従って回帰問題を完全に解決できる場合、サポート関数を模倣するようにトレーニングされた同じニューラルネットワークは、その勾配（自動微分を介して）を評価することによって最適なキーを出力することもできます。実際には、これらの2つのタスクそれぞれに、個別の専門化されたネットワークをトレーニングする方がうまくいくことがわかっています。

SupportNetニューラルネットワークアーキテクチャは、（緩やかに制約された）ICNNを使用して、サポート関数を直接学習できます。このアプローチは数学的に原理的であり、問題の凸構造に一致します。このアプローチは、クラスタリングの品質を向上させるために、データクラスタごとに1つのサポート関数を計算するために使用する場合に輝くことがわかっています。

KeyNetモデルは、クエリが与えられたときに最適な一致キーベクトルを直接出力するように学習します。同次関数のオイラーの定理から導出された一貫性損失は、明示的な凸性は強制されていませんが、予測を基盤となるサポート関数に幾何学的に接地させます。

どちらのモデルも、関連するクエリ分布からサンプリングされた例でトレーニングされます。各サンプリングされたクエリについて、グラウンドトゥルースの最適な一致が網羅的な検索によって一度計算され、保存され、実際に重要なクエリの分布のMIPSソリューションランドスケープをキャプチャする教師ありデータセットが作成されます。

SupportNetによる学習済みルーティング

図1に示すように、大規模化を処理する最も簡単な方法は、キーをクラスタに分割し、クエリを各クラスタのセントロイドと比較し、これらのkスコアに基づいて最も有望なクラスタにクエリをルーティングすることです。これは、セントロイドだけではクラスタの広がりを考慮できないため、間違いを起こしやすい安価で粗いルールを提供します。実際、クラスタには、その自身のセントロイドから遠く離れているがクエリには近いキーが含まれている可能性があり、クエリがそのような外れ値の近くにある場合、セントロイドルーティングはそのクラスタのスコアを低くしてスキップし、真の最適な一致を見逃します。

各クラスタのサポート関数にアクセスできる場合、各クラスタ内の任意のキーがクエリにどれだけうまく一致するかを正確に定量化するため、正確なルーティングを行うことができます。これは、図2に示されており、埋め込み空間で変化するクエリに対して4つのクラスタのサポート関数がプロットされています。

［動画］

図2: セントロイドへの距離に依存することは誤解を招く可能性があります。なぜなら、それらの距離はクラスタ内のすべてのキーを考慮していないからです。各クラスタのサポート関数にアクセスできる場合、割り当ては、与えられたクエリがクラスタ内の任意のキーとどれだけうまく一致するかを基準に行うことができます。

この理論的な洞察に基づき、私たちはこれらのk個のサポート関数を、一度にk個のスコアを出力できるトレーニング済みSupportNetと共同で学習することを提案します。その後、これらのk個のスコアのうち最大のものを新しいクエリのルーティングに使用できます。SupportNetモデルは、単純なセントロイド比較よりも有用な質問に答えることができます。「このクラスタ内の最適なキーは、このクエリに対してどれくらい良いか？」SupportNetsは、古典的な機械学習の方法、すなわち回帰によってトレーニングされます。つまり、分布からサンプリングされたクエリ（通常はテキストクエリの埋め込み）を生成し、それらのクエリに対するグラウンドトゥルースのサポート関数を計算し、ニューラルネットワークをトレーニングしてこれらの値を再現します。

［動画］

図3: SupportNetsは各クラスタのサポート関数を近似し、典型的な幾何学的アプローチのようにセントロイドだけでなく、クラスタ内の最適な一致に基づいてクエリをクラスタに割り当てます。

私たちは、さまざまなサイズのSupportNet（XS、S、M、L → データベースサイズの1%、5%、10%、20%）を、128個のクラスタに分割されたNatural Questionsでトレーニングしました。Natural Questionsは、一般的な知識の質問と大規模なドキュメントデータベースからのパッセージをペアにした検索ベンチマークであり、私たちのセットアップでは約260万の一意のパッセージが含まれています。

以下の図4は、さまざまなコストの概念に対するリコール（検索精度）の関数としての検索パフォーマンスを示しています。ルーティングリコールは、真のパーティション（正確な最適なキーが含まれるパーティション）がルーターがプローブすることを選択したパーティションに含まれるクエリの割合です。両方のプロットで、クエリと一致する上位sクラスタ（sは2のべき乗（1、2、4、8、…、128）として増加）で網羅的な検索が実行されます。当然、すべての方法が、すべての128クラスタを調査するとリコール1に達します。

最初のタブは、固定された検索パフォーマンスに対して、学習済みルーターがセントロイドルーティングよりも低いレイテンシでそれを達成することを示しています。小さなSupportNetモデルは、セントロイドルーティングよりも約1.9倍高速（1.74 ms対3.36 ms）で95%のルーティングリコールに達し、大きなモデルではさらに改善され、1.74 msで97%に達します。このプロットは、セントロイドと比較する追加の計算コストを明示的に考慮に入れています。2番目のタブは、NNのオーバーヘッドを無視して、ルーティング品質のみに焦点を当てています。固定されたクラスタ探索予算では、それらは正しいパーティションをより頻繁に見つけます。データベースのわずか5%をスキャンするだけで、学習済みルーターは、セントロイドベースラインの87.5%と比較して、真のパーティションを約90〜93%の時間（中サイズ（M）では最大92.6%）で復元します。

図4: さまざまなサイズのSupportNetによるルーティングパフォーマンス（リコール）とレイテンシまたはルーターがスキャンできるデータベースの割合とのトレードオフ。すべてのルーターは深度16層で、128クラスタのNatural Questionsデータセットでトレーニングされています。SupportNetはセントロイドルーターを凌駕します。レイテンシの目標が与えられた場合、SupportNetを使用するとより高いリコールを得ることができます。逆に、セントロイドルーターよりも低いレイテンシで同じリコールを達成できます。

既存のインデックスへの接続

SupportNetルーティングは、既存の転置ファイルインデックス（IVF）システムに見られる単純なセントロイドベースのルーティングを置き換えます。KeyNetを使用すると、IVFシステムをそのまま使用し、クエリの変更バージョンをフィードすることもできます。そのシナリオでは、KeyNetは標準的な近似検索インデックスの前に「クエリリライター」として自然に機能します。最適な一致が埋め込み空間のどこにある可能性が高いかを予測し、元のクエリから開始するのではなく、お気に入りのIVFシステムにそこを見るように依頼します。IVFインデックスは高度に最適化されているため、KeyNetの支援を受けると結果をより速く見つけることができます。KeyNetをトレーニングするためにも回帰に依存しており、これは図5に要約されています。サンプリングされたクエリは、グラウンドトゥルースの最適な一致キーと照合され、KeyNetニューラルネットワークは、その関連付けを再現するようにトレーニングされます。

［動画］

図5: KeyNetは、クエリをデータベース内の最適な一致にマッピングするようにトレーニングされます。IVFインデックスと組み合わせると、このようなマッピングによるクエリの前処理は検索を大幅に容易にします。結果として得られる変換されたクエリは、取得対象に非常に近くなります。

さまざまなサイズのKeyNet（XS、S、M、L → データベースサイズの1%、5%、10%、20%）を評価し、以下の図6で、これらのさまざまなサイズのリコールとコストのトレードオフを示します。最初のタブでは、検索予算は壁時計レイテンシです。最大のモデル（0.23 ms対生のクエリの1.8 ms）で最大約8倍速く99%のリコールに達します。前述のように、2番目のタブは、元のクエリをKeyNet変換に置き換えることによって提供される「品質」の向上にのみ焦点を当てています。X軸は、インデックスが実際にスキャンするデータベースの割合です。データベースのわずか0.001%をスキャンするだけで、KeyNetマッピングされたクエリは、生のクエリの86%と比較して、89〜97%の時間（最大のモデル（L）では最大96.7%）で真のキーを取得します。同じコストでより良いリコール、またはその一部のコストで同じリコールです。

図6: IVFインデックスを使用してNatural Questionsデータセットから作成されたデータベースを検索する検索パフォーマンス（リコール）。プロットは、生のクエリでIVFインデックスをクエリする結果と、さまざまなサイズのKeyNetでクエリをマッピングした後に実行される検索結果を比較します。すべてのモデルは深度16層です。プロットのレイテンシは、クエリのマッピングと検索に必要な合計時間に対応しますが、生のクエリの場合は検索のみを考慮します。KeyNetによるクエリの前処理は、サイズに関係なく価値があります。前処理のオーバーヘッドは、後続の高速検索によって補償されます。

以前に議論したNatural Questionsデータセットから取得された例を示しており、FAISS IVFインデックスを介して提供されます。これは私たちが研究したいくつかのデータセットのうちの1つであり、傾向はそれら全体で一貫しています。KeyNetの予測キーを使用してインデックスをクエリすることは、生の質問を使用してクエリすることよりも常に優れており、検索予算全体にわたって一貫しています。

具体的には、以下の図7は、生のクエリとKeyNetマッピングされたクエリのトップ3ヒットを5つのNatural Questionsの例で示しています。

［画像: 図7: 生のクエリとKeyNetマッピングされたクエリのトップ3ヒット: 'Matchbox Twentyとは？'］

図7: 生のクエリとKeyNetマッピングされたクエリのトップ3ヒット: 'Matchbox Twentyとは？'

［画像: 図7: 生のクエリとKeyNetマッピングされたクエリのトップ3ヒット: 'The Eye of the Worldを書いたのは誰？'］

図7: 生のクエリとKeyNetマッピングされたクエリのトップ3ヒット: 'The Eye of the Worldを書いたのは誰？'

［画像: 図7: 生のクエリとKeyNetマッピングされたクエリのトップ3ヒット: 'ウィンドコントローラーは何に使用されますか？'］

図7: 生のクエリとKeyNetマッピングされたクエリのトップ3ヒット: 'ウィンドコントローラーは何に使用されますか？'

［画像: 図7: 生のクエリとKeyNetマッピングされたクエリのトップ3ヒット: 'Summertime Ballはいつ開催されますか？'］

図7: 生のクエリとKeyNetマッピングされたクエリのトップ3ヒット: 'Summertime Ballはいつ開催されますか？'

［画像: 図7: 生のクエリとKeyNetマッピングされたクエリのトップ3ヒット: 'Victoriousという番組はいつNickelodeonで放送されましたか？'］

図7: 生のクエリとKeyNetマッピングされたクエリのトップ3ヒット: 'Victoriousという番組はいつNickelodeonで放送されましたか？'

どのネットワークを使用すべきか？

両者の選択は、最終的には実用的なものです。推論速度がボトルネックであり、最適なキーを近似するだけでよい場合は、KeyNetを使用します。キャリブレーションされたスコア（再ランキング、閾値処理、または下流での使用のため）も必要な場合、またはモデルサイズが小さい場合に計算量がタイトな場合（このレジームではSupportNetが競争力を維持します）は、SupportNetを使用します。どちらも同様にトレーニングされ、既存の近似検索インフラストラクチャと統合され、クエリパターンの軽微な分布シフトに対して比較的堅牢です。最終的には、キーセットが巨大な非常に大規模なMIPSシステムでは、両方が必要になる可能性があります。すべてのクラスタ全体でSupportNetルーターをトレーニングし、次に各クラスタでKeyNetを使用して、その通常のインデックスが最適なキーをより速く絞り込めるようにします。

学習済み検索の新しい方向性

償却MIPSは、クエリ分布が事前にわかっている場合の検索のアプローチ方法の変化を反映しています。これは、本番MLシステムで一般的な条件です。MIPSソリューションを直接ネットワークの重みにエンコードすることにより、このアプローチは、一度のトレーニング（いわゆる償却）コストを、高速で軽量な推論と引き換えにし、検索対象に明示的に調整された学習済み検索システムへの実用的な道を開きます。

より良く、より効率的なRAGのための高速な検索

要約

全文翻訳

より良く、より効率的なRAGのための高速な検索

最大内積検索の基本

MIPSをサポート関数として

最大内積検索を学習する2つの方法

SupportNetによる学習済みルーティング

既存のインデックスへの接続

どのネットワークを使用すべきか？

学習済み検索の新しい方向性

関連文献と更新情報。