pgvector 0.7.0 の新機能と、これから来る（かもしれない）pgvectorscale

JAWS-UG 名古屋 生成 AI 夏祭り 2024/7/29 LT

Transcript

1. pgvector 0.7.0 の新機能と、 これから来る（かもしれない）pgvectorscale JAWS-UG 名古屋 生成 AI 夏祭り　2024/7/29 まつひさ（hmatsu47）
2. 本日のネタ • 以下の 3 つの記事より ◦ 象使いのための pgvector 入門 (3) 0.6.0 〜 0.7.0 追加機能について https://qiita.com/hmatsu47/items/8cdafc093fcda793e863 ◦ pgvector 0.7.0 でバイナリインデックス＆量子化を試してみた https://zenn.dev/hmatsu47/articles/pgvector070-binaryvector ◦ pgvectorscale を少しだけ試してみた（失敗編） https://zenn.dev/hmatsu47/articles/pgvectorscale 2
3. pgvector おさらい（〜 0.6.x） • PostgreSQL をベクトル対応するための拡張機能 https://github.com/pgvector/pgvector ◦ RAG やセマンティック検索・推薦システムでよく使用される • インデックスとして IVFFlat・HNSW をサポート ◦ HNSW のほうが高性能 ▪ 0.5.0 で HNSW をサポート ▪ 0.6.0 で HNSW のパラレルビルドに対応 3
4. pgvector 0.7.0（主な機能追加） • 半精度・バイナリ（ビット）ベクトルに対応 ◦ 16 ビット浮動小数点数と 1 ビットのベクトルに対応 ▪ （それぞれ）インデックスも対応 ◦ ビット量子化とハミング距離・ジャッカード距離に対応 • 疎ベクトルに対応（説明は省略） • サブベクトルインデックス化に対応 ◦ ベクトルの一部だけを抜き出してインデックス化 4
5. いずれも • ベクトルデータの容量を削減するための機能追加 ◦ 1 次元あたりのビット数を削減 or ベクトルの次元数を削減 ◦ 特にインデックスの容量をできるだけ減らすのが大事 ▪ インデックスをメモリ上にキャッシュさせる ▪ 検索時のディスクアクセス頻度を減らす 5
6. 一方で • インデックスのビット数／次元数を削減しても検索精度 を下げない工夫が必要 ◦ Re-rank 検索（後述） 6
7. ビット量子化インデックスを試してみた（1/2） • AWS オレゴンリージョンで以下を使ってテスト ◦ RDS for PostgreSQL 16.3（db.m6g.large） ◦ Titan Text Embeddings V2（1024 次元・正規化あり） • 使ったデータはこちら ◦ https://huggingface.co/datasets/takaaki-inada/databricks-dolly-15k-ja-zundamon ▪ 2024/5/22 時点・トークンサイズ等制限に掛からない約 4,600 件 ▪ 質問文（instruction）と回答（output）をそれぞれ個別の行として登録 ▪ 埋め込みベクトルは 32 ビット float 値でデータ行に登録 7
8. ビット量子化インデックスを試してみた（2/2） • インデックスの作成時間と検索時間・検索結果を比較 ◦ インデックスなし（内積） ◦ 32 ビット float 値の HNSW インデックスを作成して検索（同上） ◦ ビット量子化して作成したインデックスを Re-rank 検索 8
9. Re-rank 検索とは？（ざっくり） • 一度検索して得た「候補」を別手法で順位づけ ◦ 例）低負荷・低精度の手法で粗く絞り込み→その中から高精度の 手法で順位づけ or 最適値選択する ▪ 例 1 : ビット量子化した近似最近傍探索用インデックスから 20 件絞り込み →20 件の中から 32 ビット精度で最近傍の 5 件を選択 ▪ 例 2 : 次元数を減らした近似最近傍探索用インデックス（32 ビット精度） から 20 件絞り込み→20 件の中から最近傍の 5 件を選択（元の次元数で） 9
10. 0. データ投入 • GitHub リポジトリにある Python スクリプトを使用 https://github.com/hmatsu47/pgvector070-test/blob/main/app_titan2_re_rank_ bin_index_text.py 10
11. 1. インデックスがない状態でベクトル検索（1/2） • ↓の質問文（instruction）に近い回答（output）を抽出 11 postgres=> SELECT id, data FROM rerank_test WHERE id = 100000; -[ RECORD 1 ]---------------------------------------------------------- id | 100000 data | ヴァージン・オーストラリア航空はいつから運航を開始したのですか？ Time: 139.100 ms
12. 1. インデックスがない状態でベクトル検索（2/2） 12 postgres=> SELECT id, (embedding <#> (SELECT embedding FROM rerank_test WHERE id = 100000)) * -1 AS inner_product, data FROM rerank_test WHERE id < 100000 ORDER BY inner_product DESC LIMIT 5; -[ RECORD 1 ]-+--------------------------------------------------------------------------------------… id | 0 inner_product | 0.7999640703201294 data | ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）はオーストラリアを拠点とする… -[ RECORD 2 ]-+--------------------------------------------------------------------------------------… id | 3355 inner_product | 0.4171184301376343 data | N.V. Virgin Express S.A.は、ヴァージン・グループ内に誕生したベルギーの航空会社なのだ。ブリュッセ … -[ RECORD 3 ]-+--------------------------------------------------------------------------------------… id | 6819 inner_product | 0.2454237937927246 data | OA形は、ボールドウィン機関車製作所がニュージーランドのウェリントン＆マナワツ鉄道（ WMR）のために製作し… -[ RECORD 4 ]-+--------------------------------------------------------------------------------------… id | 1540 inner_product | 0.22488926351070404 data | 橋が架かる前、サンフランシスコと現在のマリン郡を結ぶ唯一の実用的な近道は、サンフランシスコ湾の一角を船 … -[ RECORD 5 ]-+--------------------------------------------------------------------------------------… id | 3120 inner_product | 0.22459137439727783 data | Oy Air Finlandは、フィンランドのヴァンターにあるヘルシンキ空港を本社・拠点とする航空会社で、休暇先 … Time: 173.794 ms
13. 2. 32 ビットインデックス（HNSW）を作成 13 postgres=> CREATE INDEX ON rerank_test USING hnsw (embedding vector_ip_ops); CREATE INDEX Time: 9103.665 ms (00:09.104)
14. 3. 32 ビットインデックス（HNSW）で検索 14 postgres=> SELECT id, (embedding <#> (SELECT embedding FROM rerank_test WHERE id = 100000)) * -1 AS inner_product, data FROM rerank_test WHERE id < 100000 ORDER BY inner_product DESC LIMIT 5; -[ RECORD 1 ]-+--------------------------------------------------------------------------------------… id | 0 inner_product | 0.7999640703201294 data | ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）はオーストラリアを拠点とする… -[ RECORD 2 ]-+--------------------------------------------------------------------------------------… id | 3355 inner_product | 0.4171184301376343 data | N.V. Virgin Express S.A.は、ヴァージン・グループ内に誕生したベルギーの航空会社なのだ。ブリュッセ … -[ RECORD 3 ]-+--------------------------------------------------------------------------------------… id | 6819 inner_product | 0.2454237937927246 data | OA形は、ボールドウィン機関車製作所がニュージーランドのウェリントン＆マナワツ鉄道（ WMR）のために製作し… -[ RECORD 4 ]-+--------------------------------------------------------------------------------------… id | 1540 inner_product | 0.22488926351070404 data | 橋が架かる前、サンフランシスコと現在のマリン郡を結ぶ唯一の実用的な近道は、サンフランシスコ湾の一角を船 … -[ RECORD 5 ]-+--------------------------------------------------------------------------------------… id | 3120 inner_product | 0.22459137439727783 data | Oy Air Finlandは、フィンランドのヴァンターにあるヘルシンキ空港を本社・拠点とする航空会社で、休暇先 … Time: 175.982 ms インデックスなしのケースとほぼ変わらず（データ行数が少ないため？）
15. 4. ビット量子化インデックス（HNSW）を作成 15 • 先に作成した 32 ビットインデックスを削除してから作成 postgres=> DROP INDEX rerank_test_embedding_idx; DROP INDEX Time: 152.619 ms postgres=> CREATE INDEX ON rerank_test USING hnsw ((binary_quantize(embedding)::bit(1024)) bit_hamming_ops); CREATE INDEX Time: 3356.522 ms (00:03.357) インデックス作成所要時間が半分以下に
16. 5. ビット量子化インデックス使用の Re-rank 検索 16 postgres=> SELECT id, (embedding <#> (SELECT embedding FROM rerank_test WHERE id = 100000)) * -1 AS inner_product, data FROM (SELECT * FROM rerank_test WHERE id < 100000 ORDER BY binary_quantize(embedding)::bit(1024) <~> (SELECT binary_quantize(embedding)::bit(1024) FROM rerank_test WHERE id = 100000) LIMIT 20) ORDER BY inner_product DESC LIMIT 5; -[ RECORD 1 ]-+--------------------------------------------------------------------------------------… id | 0 inner_product | 0.7999640703201294 data | ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）はオーストラリアを拠点とする… -[ RECORD 2 ]-+--------------------------------------------------------------------------------------… id | 3355 inner_product | 0.4171184301376343 data | N.V. Virgin Express S.A.は、ヴァージン・グループ内に誕生したベルギーの航空会社なのだ。ブリュッセ … -[ RECORD 3 ]-+--------------------------------------------------------------------------------------… id | 6819 inner_product | 0.2454237937927246 data | OA形は、ボールドウィン機関車製作所がニュージーランドのウェリントン＆マナワツ鉄道（ WMR）のために製作し… -[ RECORD 4 ]-+--------------------------------------------------------------------------------------… id | 1540 （中略） -[ RECORD 5 ]-+--------------------------------------------------------------------------------------… id | 3120 （中略） Time: 140.448 ms 少しだけ検索所要時間が短縮
17. 6. ビット量子化インデックスで検索（非 Re-rank） 17 postgres=> SELECT id, (binary_quantize(embedding)::bit(1024) <~> (SELECT binary_quantize(embedding)::bit(1024) FROM rerank_test WHERE id = 100000)) AS humming_distance, data FROM rerank_test WHERE id < 100000 ORDER BY humming_distance LIMIT 5; -[ RECORD 1 ]----+-----------------------------------------------------------------------------------… id | 0 humming_distance | 224 data | ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）はオーストラリアを拠点と… -[ RECORD 2 ]----+-----------------------------------------------------------------------------------… id | 3355 humming_distance | 385 data | N.V. Virgin Express S.A.は、ヴァージン・グループ内に誕生したベルギーの航空会社なのだ。ブリュ … -[ RECORD 3 ]----+-----------------------------------------------------------------------------------… id | 5003 humming_distance | 419 data | エアバスA380は、エアバス社が開発・製造した大型ワイドボディー旅客機なのだ。世界最大の旅客機であり … -[ RECORD 4 ]----+-----------------------------------------------------------------------------------… id | 65 humming_distance | 426 data | イギリス領ヴァージン諸島（BVI）、通称ヴァージン諸島は、カリブ海にあるイギリス海外領土で、プエルト … -[ RECORD 5 ]----+-----------------------------------------------------------------------------------… id | 5206 （中略） Time: 139.787 ms Re-rank 検索と順位が変化・検索所要時間の差はわずか
18. ビット量子化インデックス（pgvector）結果まとめ • データ量が少なくてもある程度効果あり ◦ インデックス作成時間短縮 ▪ 所要時間が半分以下（9.1s → 3.4s） ◦ ビット量子化インデックス＋Re-rank 検索による検索時間短縮 ▪ 所要時間 2 割弱低減（0.17s → 0.14s） 18
19. pgvectorscale（概要） • Timescale 社による pgvector の拡張 ◦ https://github.com/timescale/pgvectorscale ◦ StreamingDiskANN インデックスをサポート ▪ マイクロソフトの DiskANN にインスパイアされた ▪ 効率的にメモリにキャッシュしてディスク（SSD）アクセス頻度を低減 ◦ 統計的バイナリ（ビット）量子化をサポート ▪ 通常のビット量子化よりも高い精度を保ったままビット数を削減可能？ ◦ OSS として提供されているが、RDS / Aurora では未サポート 19
20. TimescaleDB-HA コンテナで少しだけ触ってみた 20 postgres=# CREATE INDEX ON rerank_test USING diskann (embedding); CREATE INDEX 先ほどとは別環境なので所要時間は非表示・デフォルトで 1 ビットのインデックス （データ量が少なくメモリ上に全データがキャッシュされるせいか、同じ環境では pgvector のビット量子化インデックスのほうが速かった） postgres=# SELECT id, (embedding <#> (SELECT embedding FROM rerank_test WHERE id = 100000)) * -1 AS inner_product FROM (SELECT * FROM rerank_test WHERE id < 100000 ORDER BY embedding <=> (SELECT embedding FROM rerank_test WHERE id = 100000) LIMIT 20) ORDER BY inner_product DESC LIMIT 5; id | inner_product ------+--------------------- 0 | 0.7999640703201294 3355 | 0.4171184301376343 6819 | 0.2454237937927246 1540 | 0.22488926351070404 3120 | 0.22459137439727783 (5 rows) StreamingDiskANN インデックスは現時点でコサイン距離のみ対応