Apple の学習型ルーティング、重心検索を置き換えて RAG の遅延を削減

Apple Machine Learning Research の研究者たちは、推薦システムや検索拡張生成（RAG）の中核エンジンである最大内積検索（MIPS）を加速するための新しい手法を導入しました。ICML 2026 で発表された論文によると、この研究では従来の幾何学的インデックスを、ユーザークエリの特定の分布を学習するオフライン訓練済みの小さなニューラルネットワークに置き換えています。このアプローチにより、標準的な重心ベースのルーティングと比較して、リコール精度を向上させつつ、レイテンシを大幅に削減しています。

何が起きたか

データベースに数百万もの高次元ベクトルが含まれている場合、最大内積検索は計算コストが高くなります。従来の近似手法は、キーをクラスタリングし、重心を使ってクエリをルーティングすることで処理を高速化しますが、これらの幾何学的な仮定はしばしば失敗します。これは、ユーザーのクエリが、探している精心されたドキュメントと構造的に異なるためです。Apple のチームは、汎用的でクエリに依存しない手法が、本番環境における安定かつ特徴づけられたクエリ分布を無視しており、性能向上の可能性を活かしきれていないことを観察しました。

これに対応するため、研究者たちは固定されたキーのデータベース上で小さなニューラルネットワークをオフラインで訓練しました。これらのネットワークは、網羅的な比較や単純な重心ルーティングよりも効率的に、クエリを最適な一致項目にマッピングする方法を学習します。チームは SupportNet と KeyNet という2つのアーキテクチャを、Natural Questions ベンチマークを使用して評価しました。このベンチマークは、一般知識に関する質問と約260万個のユニークなパッセージを組み合わせたものです。結果は、広範な幾何学的近似に頼るのではなく、クエリ分布の特定の幾何学構造を活用することで、学習型ルーティングがより低いレイテンシで高いリコールを達成できることを示しました。

仕組み

この核心的な洞察は、凸解析の観点から MIPS を捉えることにあります。クエリをセット内の任意のキーとの最大の内積に写像する関数は、支持関数として知られています。この関数は凸関数であり、入力凸ニューラルネットワーク（ICNN）によって近似可能です。さらに、任意のクエリ点におけるこの支持関数の勾配は、最もよく一致するキーベクトルを返します。ニューラルネットワークがこの支持関数またはその勾配を模倣するように訓練されることで、システムはデータベース全体を検索するよりもはるかに低い計算コストで最適な一致項目を予測できます。

研究者たちは、この数学的構造を活用するために2つの具体的なモデルを開発しました。SupportNet は各データクラスタに対して支持関数を直接学習し、クラスタ内の最良キーがクエリにどれだけよく一致するかを回答できるようにします。これにより、重心距離だけに基づくよりも正確なルーティングが可能になります。一方、KeyNet は最もよく一致するキーベクトルを直接出力することを学習します。これはクエリライターとして機能し、元のクエリを埋め込み空間内でターゲットに近いベクトルに変換してから、標準的な検索インデックスに渡します。

主要な詳細

• この研究は Apple Machine Learning Research により ICML 2026 で発表されました。
• SupportNet モデルは、重心ルーティングと比較して1.9倍速く（1.74 ms vs 3.36 ms）、95% のルーティングリコールを達成しました。
• KeyNet は、生のクエリ検索と比較して99% のリコール到達時間を最大8倍短縮しました（0.23 ms vs 1.8 ms）。
• モデルは、約260万個のユニークなパッセージを含む Natural Questions データセットで訓練されました。
• SupportNet は、データベースサイズの1%から20%までの範囲でテストされ、深さは16層でした。
• KeyNet によって変換されたクエリは、データベースのわずか0.009%をスキャンするだけで99%のリコールに達しました。これは生クエリの0.10%と比較されます。

なぜ重要なのか

RAG システムを構築するエンジニアにとって、短いユーザー質問と長いドキュメント埋め込み間のギャップは、持続的なボトルネックを生み出します。標準的な転置ファイルインデックスは、クラスタ内のキーの広がりを考慮しない重心に依存しており、これが近傍の見逃しや不要な計算資源の浪費につながります。特定のクエリ分布を学習することにより、開発者は計算資源をより効率的に再割り当てでき、網羅的な検索を行わずとも最も関連性の高いドキュメントを取得することが保証されます。

この手法はまた、既存のインフラストラクチャへの統合に向けた現実的な道も提供します。KeyNet は FAISS などの標準的な近似検索インデックスの前処理ステップとして機能でき、基盤となるストレージシステムに変更を加える必要はありません。つまり、チームは現在の検索パイプラインの前に軽量なニューラルネットワーク層を追加するだけで、ベクトル検索アーキテクチャ全体を再構築せずに、レイテンシの大幅な削減と精度の向上を達成できます。

実施可能なこと

• この手法は代表データに対するオフライン訓練に依存するため、本番環境のクエリ分布が安定しており、十分に特徴づけられているかどうかを評価してください。
• ルーティングエラーがリコールに影響を与えている場合、転置ファイルインデックスでの重心ベースのルーティングを置き換えるために SupportNet の実装を検討してください。
• 取得時にスキャンされるキーの数を減らすために、既存のベクトル検索インデックスの前にクエリライトアップ層として KeyNet を使用してください。
• サンプルクエリとその正解ベストマッチに対する回帰分析を用いて、小さなニューラルネットワークをオフラインで訓練し、計算コストを償却してください。
• モデルサイズとレイテンシのトレードオフをベンチマークしてください。小さなモデル（データベースサイズの1%）でも、ベースライン手法に対して顕著な速度向上を提供することに注意してください。
• KeyNet モデルの訓練時、予測を幾何学的に根拠づけるために、オイラーの定理から導かれる一貫性損失をテストしてください。