最大内積検索(MIPS)は機械学習における重要なサブルーチンであり、与えられたクエリに対してデータベース(キー)内で最も内積が大きくなるベクトルを特定する。多くの応用では、クエリは既知の分布から繰り返しサンプリングされ、キーデータベースは固定されている。この特性に着目し、MITとAppleの研究者らは「償却MIPS(Amortized MIPS)」と呼ばれる回帰ベースのアプローチを提案した。これは、ニューラルネットワークを訓練してMIPSの解を直接予測させることで、クエリごとに高コストな厳密検索を実行する代わりに、訓練段階でコストを償却するものである。
論文の中心的な理論的洞察は、MIPS値関数がキー集合の支持関数(support function)であるという点である。支持関数はよく研究された凸関数であり、その勾配は最適なキーを与える。この性質が、償却モデルの設計に堅牢な数学的基盤を提供する。著者らはこれに基づいて2つの相補的なニューラルネットワークアーキテクチャを提案した。第一はSupportNetで、入力凸ニューラルネットワーク(input-convex neural network)であり、支持関数を回帰するように訓練される。SupportNetの出力はクラスタルータとして機能し、クエリを関連するデータベースパーティションに誘導することで検索範囲を絞り込む。第二はKeyNetで、ベクトル値ネットワークが最適キーを直接回帰する。KeyNetが生成するキーは元のクエリの代替として機能し、既存のインデックスパイプラインにそのまま入力できるため、システムを変更する必要がない。
提案手法の有効性を検証するため、研究者らはBEIRベンチマーク上で、特にドキュメント埋め込みを対象に実験を行った。実験結果は、計算量(FLOPs、プローブ数、実時間のそれぞれで測定)を考慮した場合、学習されたSupportNetとKeyNetがIVF(転置ファイル)インデックスのマッチ率を大幅に改善することを示した。これは、同じ計算予算の下で、償却手法がより関連性の高いキーを見つけ、検索品質を向上させることを意味する。論文のコードはhttps://github.com/apple/ml-amipsで公開されており、コミュニティによる再現と拡張が可能である。
本研究は、凸解析における支持関数の概念とニューラルネットワーク回帰を組み合わせ、MIPS問題に対して新たな償却の視点を提供する。従来の厳密検索や近似検索とは異なり、償却MIPSはクエリ分布の事前知識を利用して訓練段階でクエリから解へのマッピングを学習し、推論段階で高速な予測を実現する。この手法は、クエリパターンが安定しデータベースが固定されているシナリオ、例えば推薦システムや情報検索における埋め込み検索に特に適している。将来的には、このアイデアを他の凸構造を持つ検索問題に拡張したり、より複雑なインデックス構造と組み合わせたりすることが考えられる。