現代の検索拡張生成 (RAG) では、大規模なベクトルデータセットから高次元クエリベクトルの k 近傍を検索することが一般的です。データセットが大きくなるにつれて、すべてのベクトルをメモリに保存するコストは非常に高くなります。ベクトルの保存に必要な総メモリ量は次の式で表されます。
total_bits = num_documents × data_dimensionality × bits_per_dimension
この式からわかるように、メモリは各ベクトルの各次元をエンコードするビット数に比例して増加します。ドキュメント数や次元数を削減できない場合、メモリ使用量とコストを抑える一般的な方法は、各ベクトルの表現を標準の float32 (32 ビット) からより少ないビット数に圧縮することです。
最も基本的な圧縮手法であるスカラー量子化は、別の浮動小数点表現を使用することで圧縮を実現します。たとえば float16 に切り替えると、メモリ使用量が 50% 削減されます (2 倍の圧縮率)。これで十分な場合もありますが、データセットを RAM に収めるためにさらに高い圧縮率が必要になることもあります。最も極端なスカラー量子化であるバイナリ量子化は、すべての浮動小数点数を 1 ビットに削減し、32 倍の圧縮を実現します。この場合、メモリの式は次のように簡略化されます。
total_bits = num_documents × data_dimensionality
具体例を挙げると、1 億個の 768 次元ベクトルのデータセットは、float32 では 300 GB 以上のメモリが必要ですが、バイナリ量子化を使用すれば 10 GB 未満に収まります。
ただし、このメモリ削減には大きなトレードオフがあります。それはリコール (再現率) の低下です。ベクトルを圧縮すればするほど、元の表現間の距離を正確に推定することが難しくなり、特定のクエリに対する真の k 近傍を見つけることが困難になります。この問題を完全に回避することはできませんが、圧縮による悪影響を最小限に抑えるさまざまな技術が開発されています。本記事では、そのような手法の 1 つである非対称距離計算 (ADC) に焦点を当てます。ADC により、圧縮されていないクエリで圧縮されたインデックスを検索でき、従来のバイナリ量子化と比べてリコールを改善しながら 32 倍の圧縮を実現できます。
バイナリ量子化の仕組み
各次元を 1 ビットでエンコードする場合、バイナリ量子化されたベクトルは各次元に 2 つの値 (通常 ±1 と解釈) のいずれかしか割り当てられません。最も簡単な量子化方法は、元のベクトルの各要素の符号を取ることです。負の値は -1 (バイナリで 0) に、正の値は +1 (バイナリで 1) にマッピングされます。
実際には、量子化の前にベクトルを平均中心化することが効果的です。これは、データセット全体の平均 (座標ごと) を各ベクトルから減算し、データセットを原点方向にシフトする処理です。この平均中心化されたベクトルに対して量子化を適用します。これにより、特に一部 (またはすべて) の座標が同じ符号を持つデータセットで、元のベクトルの情報をより多く保持できます。たとえば、SIFT データセットのすべてのベクトルは非負の値のみを持ちます。平均中心化なしでは、すべてのポイントがすべて 1 のベクトルにエンコードされてしまいます。平均中心化により、データはより多様なバイナリ表現で表されるようになります。
シフトと量子化のプロセスを次の図に示します。
図: ドキュメントコーパスのバイナリ量子化を計算するには、まずデータを平均中心化して平均ドキュメントベクトルが原点に位置するようにし、次に各座標をその符号に応じて ±1 の値に変換します。
2 次元では 4 つの量子化パターンしかないため、多くのベクトルが同じ値に量子化されます。しかし、より高い d 次元空間では、バイナリ量子化のパターン数は通常コーパスのサイズを大幅に超えるため、典型的な高次元データセットでは衝突が比較的少なくなります。各ドキュメントについて、完全精度のベクトルはディスクに保存し、バイナリ量子化されたベクトルのみを RAM に保存します。
バイナリ量子化の仕組みを理解したところで、これらの圧縮されたベクトルを使って k-NN 検索を実行する 2 つの方法、対称距離計算 (SDC) と非対称距離計算 (ADC) を見ていきましょう。
対称距離計算
バイナリ量子化されたインデックスで k-NN 検索を実行する従来の方法は、前述と同じ手順でクエリベクトルを量子化することから始まります。クエリベクトルを量子化することで、RAM に保存されている量子化されたデータセットと同じスケールとデータ型になり、クエリベクトルとドキュメントベクトル間の距離計算が大幅に簡素化されます。
例として、2 つの 4 次元ベクトル
しかし、この計算は別の方法でも表現できます。ユークリッド距離の二乗は、
最新のプロセッサは、64 ビットワード全体に対して xor (「これらのビット列はどこで異なるか?」) と popcount (「このビット列に 1 はいくつあるか?」) 演算を一度に実行できます (SIMD 拡張を使用すればさらに多く)。そのため、バイナリ量子化では多くの次元の計算を並列に処理でき、距離計算の複雑さを大幅に削減できます。
SDC の欠点は、リコール (検索アルゴリズムが返す結果に含まれる真の k 近傍の割合) への影響です。すべてのベクトルがすべての次元で同時に歪められると、常に最も近いベクトルを正しく特定できるとは限りません。
図: 中心化後のクエリベクトル (オレンジ色の四角) は、元々紫色のドキュメントベクトルよりも青色のドキュメントベクトルにはるかに近い位置にあります。しかし、対称量子化がすべてのベクトルを各座標で ±1 にマッピングした後、この情報は失われ、オレンジ色のクエリベクトルは両方の量子化されたドキュメントベクトルから等距離に見えます。
次のセクションでは、ADC がこの弱点にどのように対処するかを説明します。
非対称距離計算
SDC では、アルゴリズムが 2 つの異なるポイントで距離計算に誤差を導入していることに注目してください。まず、コーパス内のドキュメントベクトルが量子化され、バイナリ量子化による 32 倍のメモリ削減が実現されます。次に、クエリベクトルも量子化されたドキュメントのスケールとデータ型に合わせて量子化されます。
しかし、この 2 番目の量子化は k-NN 検索のメモリ消費を実質的に削減しません。非常に高次元のデータセットでも、量子化されていないクエリのサイズは通常数キロバイト程度です。ADC では、クエリベクトルを元の量子化されていない形式のまま保持し、バイナリ量子化されたドキュメントベクトルと比較することでリコールを改善します。この非対称性、つまり一方のベクトルを高精度に保ちながらもう一方をバイナリのままにすることで、元の距離関係に関するより多くの情報を保持できます。
ADC が解決する主な課題は、バイナリ量子化が平均中心化後に同じ符号を持つ値をすべて同一に扱ってしまうことです。元の大きさに大きな違いがあっても同様です。例えば、平均中心化後、+0.6 という大きな正の値と +0.1 という小さな正の値は、元の空間では大きく異なるにもかかわらず、どちらも +1 に量子化されます。逆に、-0.1 と +0.1 は元の空間ではほぼ同じ値なのに、反対の値に量子化されてしまいます。
ADC は、クエリベクトルを完全精度のまま保持しつつ、ドキュメントベクトルに合わせてスケーリングすることでこの制限に対処します。量子化されたクエリではなく、再スケーリングされたクエリをバイナリドキュメントベクトルと直接比較することで、クエリも量子化した場合に失われる元のクエリベクトルの微妙な変動をより適切に捉えることができます。
図: ADC では、ドキュメントベクトルは SDC と同様に量子化されます。ただし、クエリベクトルは完全な FP32 精度で保持され、±1 の値に量子化する必要はありません。
ADC の再スケーリング
クエリベクトルを元の精度で維持していても、量子化されたドキュメントベクトルと同じスケールで動作することを保証する必要があります。バイナリ量子化は、各次元の元の大きさに関係なくすべての値を ±1 にマッピングし、データを暗黙的に再スケーリングすることを思い出してください。距離計算を意味のあるものにするには、クエリベクトルもドキュメント量子化時に使用したのと同じデータセット平均で再中心化し、バイナリ量子化のスケールに合わせて再スケーリングする必要があります。
次の図は、再スケーリングしないと ADC が間違った最近傍を見つけてしまう例を示しています。元の空間で最近傍のドキュメント (左上の象限の赤い点) が近くにあっても、量子化後は遠くの点 (右上の象限の青い点) よりも遠くに見えてしまう可能性があります。再スケーリングは、相対的な近さができるだけ保持されるように完全精度のクエリベクトルを再配置します。
図: クエリの再スケーリングが最近傍検索に与える影響。ドキュメント量子化は、本質的にドキュメントベクトルに一種の座標再スケーリングを適用します。クエリの再スケーリングがないと、クエリの近くにあった一部のポイントが遠ざかり、クエリから遠かったポイントが近づく可能性があります。クエリの再スケーリングは、ドキュメントベクトルの再配置の影響を打ち消すようにクエリベクトルを更新します。
クエリの適切なスケールを特定するために、データセットの各次元で「典型的な」ドキュメントがどのように歪められたかを調べます。データインデックス作成時に、データセットの各次元
-
\vec{l}_i i -
\vec{h}_i i
図: 各次元について、-1 (
これらの重心は、その次元の 2 つの量子化バケットのいずれかに入る「平均的な」値を表します。バイナリ量子化により、量子化後の各次元における正に量子化された値と負に量子化された値の差は正確に 2 になります。ADC でこの特性を再現するには、スケーリング後の重心間の平均距離も 2 になるようにデータをスケーリングします。つまり、クエリベクトル v を新しい v' に再スケーリングします。
重心情報はデータセット全体に対して 1 回だけ計算すればよく、ドキュメントベクトル自体と比較して追加のストレージはごくわずかです (2 つの完全精度ベクトル分)。さらに、クエリ時には、コーパス内のすべてのドキュメントベクトルと比較する前に、クエリベクトルに 1 回だけ再スケーリングを適用すれば済みます。
本番環境での ADC の使用
Amazon OpenSearch Service で ADC を使用したインデックスを作成するには、インデックス作成時に enable_adc を true に設定します。
PUT /vector-index
{
"settings" : {
"index": {
"knn": true
}
},
"mappings": {
"properties": {
"vector_field": {
"type": "knn_vector",
"dimension": 8,
"method": {
"name": "hnsw",
"engine": "faiss",
"space_type": "l2",
"parameters": {
"encoder": {
"name": "binary",
"parameters": {
"bits": 1,
"enable_adc": true
}
}
}
}
}
}
}
}
ADC を有効にしてインデックスを作成すると、標準の OpenSearch クエリ構文を使用して k-NN 検索を実行することができます。ADC の最適化は検索時に自動的に適用されます。
まとめ
ADC は、ベクトル検索の最適化における重要な進歩です。完全精度のベクトルと比較してメモリ使用量を 32 分の 1 に削減しながら、SDC で通常発生するリコールの低下を大幅に軽減します。この技術は、大規模なベクトル検索における重要な課題、つまりメモリ効率と検索品質のバランスを取ることに対処します。OpenSearch 3.2 以降、ADC を使用して、結果の関連性を損なうことなくベクトル検索の実装を最適化することができます。
ぜひ OpenSearch デプロイメントで ADC をお試しいただき、OpenSearch コミュニティフォーラム で経験を共有してください。皆様のフィードバックは、ベクトル検索機能の継続的な改善に役立ちます。
著者について
-
Tal Wagner は、OpenSearch Project と Amazon OpenSearch Service に取り組んでいる AWS のシニア応用科学者です。2020 年に MIT の CSAIL でコンピュータサイエンスの博士号を取得しました。大規模データセットのアルゴリズム設計と大規模機械学習に興味を持っています。
-
Finn Roblin は、OpenSearch ベクトル検索に取り組んでいるソフトウェアエンジニアです。大規模機械学習、推薦システム、トライアスロンスポーツ科学に興味を持っています。
-
Yonatan Naamad は、OpenSearch Project に取り組んでいる AWS のシニア応用科学者です。2017 年にプリンストン大学でコンピュータサイエンスの博士号を取得しました。グラフアルゴリズムと機械学習に興味を持っています。
OpenSearch Project(OSS) の Publicationです。 OpenSearch Tokyo User Group : meetup.com/opensearch-project-tokyo/
Discussion