📏
pgvectorscale を少しだけ試してみた（失敗編）
前回 ↓ の記事を書いた後に、Timescale 社が pgvectorscale をリリースしたので、少しだけ試してみました。
pgvectorscale とは？
TimescaleDB でおなじみの Timescale 社が OSS としてリリースした PostgreSQL 拡張機能（extension）です。
pgvector に上乗せする形で使用します。
GitHub リポジトリの README にも書かれているとおり、現時点では以下の 2 つの機能をサポートしています。
- StreamingDiskANN インデックス
- Microsoft の研究から生まれた DiskANN にインスパイアされたインデックス
効率的にメモリにキャッシュしてディスク（SSD）アクセス頻度を低減する- ディスク（SSD）に効率的にアクセスすることで、少ないメモリ容量で大規模ベクトルデータを扱うケースでの性能と精度を可能な限り確保する
- 統計的バイナリ量子化
- 標準的なバイナリ量子化を改良したもの
なぜ「失敗編」？
簡単にいうと pgvectorscale に合わない環境・シチュエーションで試してしまったから です（後述）。
今回試したこと
RDS for PostgreSQL や Aurora PostgreSQL では pgvectorscale をサポートしていません。
そのため、普通の EC2 インスタンス（m6g.large）を用意し、
- pgvector + pgvectorscale をあらかじめ内包する TimescaleDB-HA コンテナを Docker で起動
- 前回の記事と同じデータを投入
して試してみました。
Docker コンテナ起動
公式ドキュメントどおりの手順で起動しました。
Dockerコンテナ起動
docker pull timescale/timescaledb-ha:pg16
docker run -d --name timescaledb -p 5432:5432 -e POSTGRES_PASSWORD=password timescale/timescaledb-ha:pg16
データ投入
前回の記事と同じなので省略します。
pgvectorscale で StreamingDiskANN インデックスを作成して Re-rank 検索
以下のとおりデフォルト設定で作成しました。
StreamingDiskANNインデックス作成
postgres=# CREATE INDEX ON rerank_test USING diskann (embedding);
NOTICE: Starting index build. num_neighbors=-1 search_list_size=100, max_alpha=1.2, storage_layout=SbqCompression
WARNING: Indexed 9188 tuples
CREATE INDEX
Time: 12289.672 ms (00:12.290)
このときembeddingは 900 次元以上なので、インデックスの 1 次元あたりのビット数は 1 です（← バイナリインデックス）。
pgvector でのバイナリ量子化インデックス検索とは違い、pgvectorscale（0.2.0）ではコサイン距離（コサイン類似度）を使って検索します（他は未対応）。
バイナリインデックスを使ってコサイン類似度が高いほうから 20 件抽出して、さらに 32 ビット float 値のベクトルの内積で近いものから 5 件抽出しています。
pgvectorscaleでRe-rank検索
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 embedding <=> (SELECT embedding FROM rerank_test WHERE id = 100000) LIMIT 20) ORDER BY inner_product DESC LIMIT 5;
id | inner_product | data
------+---------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
0 | 0.7999640703201294 | ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）はオーストラリアを拠点とするヴァージン・ブランドを冠する最大の船団規模を持つ航空会社なのだ。2000年8月31日に、ヴァージン・ブルー空港として、2機の航空機、1つの空路を運行してサービスを開始しました。2001年9月のアンセット・オーストラリア空港の崩壊後、オーストラリアの国内市場で急速に地位を確立しました。その後はブリスベン、メルボルン、シドニーをハブとして、オーストラリア国内の32都市に直接乗り入れるまでに成長しました。
3355 | 0.4171184301376343 | N.V. Virgin Express S.A.は、ヴァージン・グループ内に誕生したベルギーの航空会社なのだ。ブリュッセル空港をハブ空港として、主に南欧方面へのフライトを運航していました。航空券は主にインターットを通じて販売されました。SNブリュッセル航空と合併してブリュッセル航空となり、2007年3月25日に運航を開始した。 バージン・エクスプレスの本社はブリュッセル近郊のベルギー・ザベンタムのブリュッセル空港の116号館にある。 +
| | +
| | 沿革 +
| | 1996年4月23日、ヴァージン・グループ（会長リチャード・ブランソン）は、ビクター・ハッソンとジョルジュ・グーテルマンが設立したベルギーのレジャー航空会社EBA - ユーロベルジアン・エアラインズを買収し、ヴァージン・エクスプレスのブランド名に改名しました またEBAのボーイング737を引き継ぎ、以降このタイプの航空機で運行しました。すぐにブリュッセルのハブ空港から低予算の定期便に集中し、サベナや後のSNブリュッセル航空の主要な競争相手となった。 +
| | +
| | 2004年10月、ヴァージン・グループはSNブリュッセル航空に資産売却し、両航空会社はエチエンヌ・ダヴィニョン子爵が会長を務める親会社の持ち株会社SNエアホールディングに統合されました。 +
| | +
| | 2006年3月31日、SNブリュッセル航空とヴァージン・エクスプレスは合併し、ブリュッセル航空と命名したことを発表しました。統合後の航空会社は長距離路線を追加し、アフリカでの地位を強化した。
6819 | 0.2454237937927246 | OA形は、ボールドウィン機関車製作所がニュージーランドのウェリントン＆マナワツ鉄道（WMR）のために製作した孤高の蒸気機関車なのだ。1894年に発注され、同年8月に13号機として運行を開始し、世界初のナローゲージのヴォークレイン配合となりました。1908年、WMRとその機関車群はニュージーランド鉄道局（NZR）に買収され、全国鉄道網に組み込まれた。13号機はOクラスのメンバーとは似ているが、十分に異なるため、別の分類が必要なのだ。13号車はO型に似ていたが、十分な違いがあり、別の分類が必要であったため、OAという名称が作られ、OA 457という番号が付けられた。1929年12月にオークランドで撤去されるまで、さらに20年間運行された。この機関車はWMRのスタッフには「The Lady」と呼ばれていた。
1540 | 0.22488926351070404 | 橋が架かる前、サンフランシスコと現在のマリン郡を結ぶ唯一の実用的な近道は、サンフランシスコ湾の一角を船で渡ることであった。1820年には早くもフェリーの運航が始まり、1840年代にはサンフランシスコへの水運を目的とした定期運航が開始された。 +
| | +
| | 1867年に開始されたサウサリート・ランド・アンド・フェリー・カンパニーのサービスは、やがてサザン・パシフィック鉄道の子会社であるゴールデン・ゲート・フェリー・カンパニーとなり、1920年代後半には世界最大のフェリー運航会社となりました。かつては鉄道の乗客や顧客だけのものだったサザン・パシフィックの自動車フェリーは、非常に収益性が高く、地域経済にとって重要なものとなりました。サンフランシスコのハイド・ストリート・ピアとマリン郡のサウサリート・フェリー・ターミナルを結ぶフェリーの所要時間は約20分、料金は車1台につき1ドルだった[when?]が、後に新しい橋に対抗するために値下げされた。サンフランシスコ・フェリー・ビルディングからの所要時間は27分であった。
5206 | 0.2185841202735901 | オーストラリア太平洋空港株式会社（APAC）は、非上場企業であり、オーストラリアの2つの空港の所有者である：メルボルン空港とローンセストン空港の2つの空港を所有しているのだ。各空港は、APACが支配権を持つ空港賃借人会社によって運営されています：オーストラリア・パシフィック・エアポーツ（メルボルン）Pty Ltdとオーストラリア・パシフィック・エアポーツ（ローンセストン）Pty Ltdがそれぞれ支配的な株式を保有している。
(5 rows)
Time: 4.919 ms
前回と比べると、5 番目だけ結果が違いました。
バイナリインデックス単独での検索結果もあげておきます。
バイナリインデックス単独での検索結果
postgres=# SELECT id, (embedding <=> (SELECT embedding FROM rerank_test WHERE id = 100000)) AS cos_dist, data FROM rerank_test WHERE id < 100000 ORDER BY cos_dist LIMIT 5;
id | cos_dist | data
------+--------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
0 | 0.2000359773614292 | ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）はオーストラリアを拠点とするヴァージン・ブランドを冠する最大の船団規模を持つ航空会社なのだ。2000年8月31日に、ヴァージン・ブルー空港として、2機の航空機、1つの空路を運行してサービスを開始しました。2001年9月のアンセット・オーストラリア空港の崩壊後、オーストラリアの国内市場で急速に地位を確立しました。その後はブリスベン、メルボルン、シドニーをハブとして、オーストラリア国内の32都市に直接乗り入れるでに成長しました。
3355 | 0.5828816071556515 | N.V. Virgin Express S.A.は、ヴァージン・グループ内に誕生したベルギーの航空会社なのだ。ブリュッセル空港をハブ空港として、主に南欧方面へのフライトを運航していました。航空券は主にインターネットを通じて販売されました。SNブリュッセル航空と合併してブリュッセル航空となり、2007年3月25日に運航を開始した。 バージン・エクスプレスの本社はブリュッセル近郊のベルギー・ザベンタムのブリュッセル空港の116号館にある。 +
| | +
| | 沿革 +
| | 1996年4月23日、ヴァージン・グループ（会長リチャード・ブランソン）は、ビクター・ハッソンとジョルジュ・グーテルマンが設立したベルギーのレジャー航空会社EBA - ユーロベルジアン・エアラインズを買収し、ヴァージン・エクスプレスのブランド名に改名しました またEBAのボーイング737を引き継ぎ、以降このタイプの航空機で運行しました。すぐにブリュッセルのハブ空港から低予算の定期便に集中し、サベナや後のSNブリュッセル航空の主要な競争相手となった。 +
| | +
| | 2004年10月、ヴァージン・グループはSNブリュッセル航空に資産を売却し、両航空会社はエチエンヌ・ダヴィニョン子爵が会長を務める親会社の持ち株会社SNエアホールディングに統合されました。 +
| | +
| | 2006年3月31日、SNブリュッセル航空とヴァージン・エクスプレスは合併し、ブリュッセル航空と命名したことを発表しました。統合後の航空会社は長距離路線を追加し、アフリカでの地位を強化した。
1540 | 0.7751107632981805 | 橋が架かる前、サンフランシスコと現在のマリン郡を結ぶ唯一の実用的な近道は、サンフランシスコ湾の一角を船で渡ることであった。1820年には早くもフェリーの運航が始まり、1840年代にはサンフランシスコへの水運を目的とした定期運航が開始された。 +
| | +
| | 1867年に開始されたサウサリート・ランド・アンド・フェリー・カンパニーのサービスは、やがてサザン・パシフィック鉄道の子会社であるゴールデン・ゲート・フェリー・カンパニーとなり、1920年代後半には世界最大のフェリー運航会社となりました。かつては鉄道の乗客や顧客だけのものだったサザン・パシフィックの自動車フェリーは、非常に収益性が高く、地域経済にとって重要なものとなりました。サンフランシスコのハイド・ストリート・ピアとマリン郡のサウサリート・フェリー・ターミナルを結ぶフェリーの所要時間は約20分、料金は車1台につき1ドルだった[when?]が、後に新しい橋に対抗するために値下げされた。サンフランシスコ・フェリー・ビルディングからの所要時間は27分であった。
5206 | 0.7814159188122891 | オーストラリア太平洋空港株式会社（APAC）は、非上場企業であり、オーストラリアの2つの空港の所有者である：メルボルン空港とローンセストン空港の2つの空港を所有しているのだ。各空港は、APACが支配権を持つ空港賃借人会社によって運営されています：オーストラリア・パシフィック・エアポーツ（メルボルン）Pty Ltdとオーストラリア・パシフィック・エアポーツ（ローンセストン）Pty Ltdがそれぞれ支配的な株式を保有している。
344 | 0.7876815687702661 | ラム・エア社は、1934年にマニトバ州ザ・パスで運航を開始し、1981年に廃業したカナダの航空会社なのだ。 +
| | +
| | 沿革 +
| | トム・ラムは、19世紀後半にイギリスから移住してきたトーマス・ヘンリー・ピーコック（THP）・ラムの息子なのだ。THPラムは学校の先生から毛皮商に転身し、1900年、マニトバ州ムースレイクでラムズ・ストアを始めた。トムとその兄弟姉妹はマニトバ州北部で育ち、父のもとで働きました。 +
| | +
| | トム・ラムは、3年生を終える前に学校を去りました。後年、マニトバ大学から名誉法学博士号を授与されたときのスピーチで、「もし4年生までだったら」というコメントを残しているのだ。10歳のとき、自分の馬とそりをもっていたトムは、魚の運搬の仕事で大人の男たちと競争していた。ソリに荷物を載せるには、魚の箱で足場を固めなければならなかった。 +
| | +
| | ラム家の事業のひとつに「物流・輸送」がある。魚、材木、木、毛皮、物資は、あらゆる手段で運ぶ必要があった。犬チーム、馬、ボート、トラック、トラクターが使われた。1930年代には、輸送革命が起きていた。飛行機がカナダ北部に進出してきたのなのだ。トム・ラムは初めて飛行機を見たとき、その可能性に気づきました。1930年、トムは最初の航空機、スティンソンSR8購入しました。1930年、トム・ラムは飛行を学ぶためにウィニペグへ旅立ちました。費用を最小限に抑えるため、ウィニペグ・フライング・クラブの裏手にあるテントで生活していた。The Pasに戻った彼は、スティンソンで自分をチェックアウトした。 +
| | +
| | 1935年、トムはラム・エアウェイズ・リミテッドを設立しました。航空会社は何度か名前を変えながら、トムはカナダ人でなくとも、世界で最もよく知られたマニトバ人の一人になった。彼の冒険は、本やテレビのドキュメンタリー番組、そして歌にもなっている。
(5 rows)
Time: 3.795 ms
所要時間比較（単位：ミリ秒）
| extension / INDEX | INDEX 作成 | 検索（3 回平均） |
|---|---|---|
| pgvector（INDEX なし） | - | 41.388 |
| pgvector（HNSW 32 ビット） | 10,108.801 | 39.227 |
| pgvector（HNSW バイナリ Re-rank） | 3,514.146 | 2.404 |
| pgvectorscale（StreamingDiskANN バイナリ Re-rank） | 12,289.672 | 4.810 |
結果としては pgvector によるバイナリ Re-rank 検索よりも遅くなってしまいました。
データの量が少なく全体がメモリにキャッシュされる容量だったのが大きい と思います。
すべてがメモリにキャッシュされるのであれば、データ容量に対して少ない容量のメモリを使ってキャッシュ効率が良くなるようにディスクから読み出す処理がかえってオーバーヘッドになってしまっても不思議ではありませんので。
2024/8/16 追記：
このあたり ↓ の記事を見て、pgvectorscale 実装の経緯がわかりました。
今回実験してみたデータセットとユースケースは確かに目的に合っていませんね。
（Pinecone 社による pgvector の欠点の指摘）
（それを受けて Timescale 社 が pgvectorscale を実装したこと、および Pinecone の DiskANN よりも pgvectorscale が優れていることを説明する記事）
余計なお世話かもしれませんが、Timescale 社のブログ記事、似たようなテーマの記事が複数ポストされていて、GitHub リポジトリなどからはその一方の記事にしかリンクされていなかったりするので、目的の記事が見つけづらいんですよね…。
2024/8/18 追記：
先の追記で挙げた記事からリンクされているこちら ↓ の記事に、
- DiskANN アルゴリズムにストリーミング取得のサポートを追加して StreamingDiskANN インデックスを実装した理由
- HNSW ではインデックスを探索する際にランダムアクセスが発生しがち
- pgvector での
hnsw.ef_searchによる探索のカットオフが、セカンダリフィルター適用時に悪影響を及ぼす可能性がある
- pgvector に実装されたバイナリ量子化（BQ）インデックスとは別に統計的バイナリ量子化（SBQ）インデックスを実装した理由
- BQ による「0」「1」のビット化境界は、実際の値を正しく二分する境界と異なる
- 次元数が少ないケース（768 次元など）では次元あたりのビット数を「1」ではなく「2」としたほうが精度が向上する
が記されています。
なお、この Zenn 記事で試したのとは別のデータセットで軽く試した範囲では、データ量が圧倒的に足りなかったためか、hnsw.ef_searchの問題は再現できませんでした。
Discussion