Cloud SQL for PostgreSQL: pgvector インデックスによる類似検索の高速化
Google Cloud Japan Team
※この投稿は米国時間 2023 年 12 月 2 日に、Google Cloud blog に投稿されたものの抄訳です。
近似最近傍（ANN）探索が急速に進化するなか、最近発表された最も重要なものの一つとして pgvector 拡張機能バージョン 0.5.0 が挙げられ、Hierarchical Navigable Small Worlds（HNSW）インデックスもサポートされるようになりました。効率、速度、性能の絶え間ない追求により、HNSW インデックスは、ANN クエリのレイテンシを大幅に短縮し、望ましい再現率を実現します。Cloud SQL for PostgreSQL の pgvector 0.5.0 がサポートされるようになったことをお知らせいたします。
この投稿では、pgvector インデックスについて説明し、さまざまな構成を明確化します。また、HNSW インデックスを使用して pgvector ベースのアプリケーションの性能と実現性を向上させるための実践的なコーディング例を提示します。
Cloud SQL for PostgreSQL の pgvector 拡張機能
pgvector 拡張機能は、Cloud SQL for PostgreSQL の既存のインスタンス内に、以下のように CREATE EXTENSION コマンドを使ってインストールできます。既存のインスタンスがない場合は、今すぐインスタンスを作成してください。Cloud SQL for PostgreSQL は pgvector 0.5.0 をサポートするようになり、Inverted File Flat（IVFFlat）インデックスに加えて pgvector に HNSW インデックスが追加されました。
古いバージョンの pgvector を Cloud SQL for PostgreSQL インスタンスにインストール済みの場合は、以下のコマンドで pgvector 0.5.0 にアップグレードできます。
この拡張機能を使用すると、「ベクトル」と呼ばれる新しいデータ型が PostgreSQL に登録され、ベクトルに対し以下の複数の演算子が定義されます。
- 要素単位の加算（+）
- 要素単位の減算（-）
- 要素単位の乗算（*）
- ユークリッド距離（<->）
- コサイン距離（<=>）
距離演算子は、意味的に類似したベクトルを見つけることができます。距離演算子がどのように動作するかについては、以前の pgvector に関するブログ記事で詳細をご確認いただけます。
以下のコマンドは、PostgreSQL のテーブルにベクトルを挿入し、行単位のコサイン距離を計算する方法を示します。
pgvector のインデックス
データベースのインデックスによってデータ検索を高速化できるのと同様に、ベクトル インデックスによって類似検索演算を大幅に高速化し、デフォルトで使用される総当たり厳密最近傍探索を回避できます。大規模なデータベースを総当たりする検索は 100% の再現率を保証しますが、パフォーマンスには大きな影響が発生します。
pgvector 拡張機能は、すでに IVFFLAT インデックスをサポートしていましたが、バージョン 0.5.0 では、追加で HNSW インデックスもサポートされるようになりました。b-tree インデックスなどの正確な検索に使用されるインデックスとは異なり、HNSW インデックスは、同じ問い合わせに対して異なる結果を導き出すことがあります。
HNSW インデックス
HNSW インデックスは、近似最近傍探索を行うために最適化されたグラフを構築する Hierarchical Navigable Small Worlds（HNSW）アルゴリズムを使用します。ベクトル データベースで使用されるベクトル インデックス作成アルゴリズムの中では最高度の性能を誇ります。HNSW は IVFFlat を含む他の多くのアルゴリズムよりも優れたクエリ性能を発揮しますが、その代償として構築時間が長くなり、メモリ使用量が増えます。
HNSW のもう一つの利点として、インデックスを作成するときに学習ステップがないことが挙げられます。つまり、テーブルにデータがなくてもインデックスを作成でき、データが追加されるにつれて段階的にインデックスが構築されていきます。これは、インデックス作成の対象となるデータが時間とともに変化するため、より良い性能を確保するためにインデックスを定期的に再作成する必要がある IVFFlat とは異なる点です。
以下のコード スニペットは、異なる距離関数とインデックス パラメータを使用して pgvector で HNSW インデックスを作成する方法を示しています。
このスニペットは、HNSW パラメータ m と ef_construction をそれぞれ 24 と 100 に設定し、コサイン距離関数を使用して product_embeddings テーブルの埋め込み列に HNSW インデックスを作成します。
自分で試す前に、以下の HNSW チューニング パラメータを理解しておく必要があります。
- m: グラフ内の近隣データポイントとの最大接続数。値を大きくするとグラフの密度が高まり、検索速度が上がる代わりにビルド時間とメモリ使用量が増加します。HNSW の使用に関する原論文によると、これは最も重要なパラメータです。m 値の妥当な範囲は 5 から 48 で、pgvector のデフォルト値は 16 です。値が大きいほど、特に高次元データに対して、良好な再現率でより良い性能が保証されます。
- ef_construction: インデックス作成時、グラフ走査中に最も近い候補を保持するリストのサイズ。値が大きいほど、アルゴリズムはより多くの候補を検討するようになり、より良いインデックスを作成しやすくなります。ただし、ある点を超えると、このパラメータを大きくしても効果は少なくなります。
- ef_search: インデックスの構築を構成する他のパラメータとは異なり、このパラメータはクエリの実行を構成し、リストに保持される最近傍の数を制限します。値を大きくすると、アルゴリズムがより多くのノードを考慮し返すため、クエリ性能は犠牲になりますが、再現率は向上します。このパラメータはインスタンス、セッション、あるいはトランザクションのレベルで設定できます。
IVFFlat インデックス
また、IVFFlat インデックスを使用して、ベクトルをリストに分割し、入力ベクトルに最も近いリストのサブセットを特定することで、ベクトルの列にインデックスを付与できます。IVFFlat インデックスは、HNSW インデックスよりもビルド時間が短く、メモリ使用量も少なく済みます。ただし、クエリ性能は HNSW インデックスの方が優れています。
以下のコード スニペットは、pgvector で IVFFlat インデックスを作成する方法を示しています。
インデックスを使用したサンプル アプリケーションの拡充
前回のブログ投稿では、pgvector と LLM を使用した対話型の Google Colab アプリケーションを作成しました。その一環として、Kaggle で公開されている大規模な小売データセットからサンプリングしたデータセットを読み込みました。このアプリケーションで使用したデータセットには、約 800 点の玩具が各商品について説明とともに掲載されています。これを products という PostgreSQL のテーブルに読み込みました。
これらの長い記述は LangChain を使用して処理され、Vertex AI Text Embedding モデルを使用して各チャンクに対してエンベディングが生成されました。これらのエンベディングは、product_embeddings と呼ばれるテーブルにベクトル型として格納されました。最後に、検索キーワードから生成されたエンベディングの最近傍を見つけるために、類似検索演算子が使用されました。
このブログ記事では、同じサンプル アプリケーションを拡張して pgvector インデックスを活用し、合理的な数のベクトルでアプリケーションのスケーリングを可能にします。前回のブログ記事と同じ設定を使用し、対話型の Google Colab ノートブックを基盤としています。
Colab ノートブックの指示に沿って、環境をセットアップします。必要な名前のインスタンスが存在しない場合、ノートブックによって Cloud SQL PostgreSQL インスタンスが作成されます。ノートブックを実行すると、Google Cloud の料金が発生することがありますが、その費用分のクレジットを提供する無料トライアルをご利用いただける場合があります。
ベクトル類似度検索を最適化するための HNSW インデックスの追加
前回のブログ記事の手順に沿って、商品データセット全体を自然言語で検索できるようにしました。
入力されたクエリに対してベクトル エンベディングを生成し、pgvector コサイン類似度検索演算子を使用して関連商品を検索するだけです。では、類似度検索を最適化してみましょう。
EXPLAIN ANALYZE を付加してクエリプランを分析すると、このクエリは 800 行を超える行を反復処理し、その処理には 12 ms 以上かかります。
代わりに、HNSW インデックスを作成してみましょう。データセットは 800 行しかないため、インデックス パラメータにはより小さな値を使用するよう検討できます。
同じクエリを再度実行すると、フルスキャンと同じレスポンスが得られます。ただし、同じクエリに対し再度 EXPLAIN ANALYZE を行うと、PostgreSQL は、product_embeddings_embedding_idx インデックスを使用し、クエリの費用が大幅に下がります。
インデックスを追加すると、クエリの性能が大幅に向上し、クエリがインデックスを使用することで、呼び出しにかかる費用を大幅に削減できます。この例では、実行時間が約 0.5 ms に減少しており、pgvector インデックスの威力が見て取れます。
この性能の向上は、800 行しかないデータセットで確認できるものですが、より大きなデータセットではより性能向上が顕著になります。場合によっては、pgvector を行数が数十万から数百万の実際のデータセットにスケールする唯一の現実的な方法となります。
たとえば、Kaggle の元のデータセットから 30,000 行を使用してみましょう。この規模では、HNSW インデックスを使用する方がシーケンシャル スキャンよりもはるかに高速です。さらに行数と項目数が増えると、シーケンシャル スキャンは完全に実行不可能になります。
シーケンシャル スキャンの場合:
HNSW インデックス スキャンの場合:
GitHub の ANN-benchmarks パッケージには、これらの近似最近傍（ANN）アルゴリズムのいくつかをベンチマークするツールが備わっています。このパッケージを使用することで、さまざまな実世界のデータセットを使用して、HNSW インデックス、IVFFlat インデックス、フルスキャンを使用した pgvector の性能特性と挙動を明確に確認できます。
これらのインデックスを活用することで、何百万ものベクトル エンベディングを持つ実世界のデータに対しても、相当な秒間クエリ数によって大規模な類似検索を行うことができます。
まとめ
この投稿で、pgvector におけるインデックスの威力と、HNSW のようなインデックスを追加することで、PostgreSQL の本番環境データを用いたベクトル類似検索をいかに高速化し、スケールできるかを示すことができたかと思います。Google のハンズオン アプリケーションとデータセットを使用すことで、段階的手順に沿ってより簡単に独自のアプリケーションを構築できるはずです。
このブログ記事で説明している pgvector インデックスの利点について説明するため、既存の Google Colab ノートブックを更新しました。ぜひ実際にお試しください。
ー ソフトウェア エンジニア Eeshan Gupta
