こんにちは!2024年新卒入社のJinyangです。 時の流れは本当に早いもので、気づけば株式会社インゲージに入社してから一年が経ちました。今回でこのブログも三度目の投稿となります。
さて、今回のテーマは何にしようかと悩んでいたのですが、やはり「RAG(Retrieval-Augmented Generation)」について書くことにしました。
RAGといえば一時期大きな注目を集めましたが、最近では当たり前の技術として定着しつつありますね。特に今回注目したいのは、RAGの「Retrieve」部分、つまりナレッジベースからの効率的なドキュメント取得です。
現在、ベクトル検索(文章の意味や文脈の類似度を用いた検索手法)が主流となっており、Weaviate、Qdrant、Pinecone、LanceDBなど、専用データベースベンダーが激しい競争を繰り広げています。
一方で、Postgresでもベクトル検索ができるのです。そんなベクトル検索を可能にしてくれるpgvector
という拡張機能についてほんの少し深ぼってみたいと思います。
目次:
はじめに: ベクトル検索って?
検索とはデータの海から自分にとって有用であるデータを見つけることを指します。 キーワードを用いた検索は我々の日常によく溶け込んでおり、馴染みがある手法だと思われます。 例えば、新しい住まいを探していて、不動産検索サイトで「駅 5分」と入力すれば、その条件に合う物件がヒットするでしょう。これは非常に馴染み深く、便利な方法です。
しかし、このキーワード検索には限界があります。例えばあなたが、「犬を飼うのに適したマンション」を探しているとしましょう。キーワード検索で「犬」「マンション」と入力しても、必ずしも求めている情報にたどり着けるとは限りません。「ペット可物件」や「動物と暮らせる家」といった、直接キーワードには含まれないけれど意味的に関連性の高い情報は漏れてしまう可能性があります。
ここで登場するのがベクトル検索です。ベクトル検索は、単なるキーワードの一致ではなく、文章やデータの「意味」を理解して関連性の高い情報を探し出すことが得意です。
文章のベクトル化(Embedding)とは?
このベクトル検索を可能にするには当然ながら、文章をベクトル化しなければなりません。わかりやすい記事がありますので以下の記事にその説明を委ねることといたします。
コサイン類似度とは?
さて、文章をベクトル化に成功した場合にとある文章とある文章がどれほど似ているかを判定するアルゴリズムが当然必要となります。その際に、よく使われるアルゴリズムがコサイン類似度となります。
コサイン類似度が取りうる値は -1 ~ 1であり、1に近づけばつくほど意味的に似ている(ベクトルが同じ方向を向いている)ということを意味しており、-1に近づけつくほど意味的に逆の関係にあるということになります。
詳しい説明は以下の記事に委ねることといたします。
pgvectorとは?
pgvectorはPostgreSQLの拡張機能で、ベクトルデータ型とベクトル類似度検索機能を提供します。主な特徴は:
- ベクトルデータ型:任意の次元数のベクトルを格納できる
vector
型 - 類似度演算子:
<->
:L2距離(ユークリッド距離)<=>
:コサイン距離(1 - コサイン類似度)<#>
:内積距離(負の内積)
- インデックス方式:
- IVFLat:中規模データセット向け
- HNSW:高速な近似最近傍検索向け
pgvectorの検証
検証環境
今回の検証では以下の環境を使用します:
chunksテーブルのembeddingカラムが vector型
であり、ベクトル検索に使うカラムとなります。
- Timescaleクラウド環境
- 0.5 CPU / 2 GiB Memory
- 標準ストレージタイプ(最大16TB、16000 IOPS)
- データセット:JaGovFaqs-22k(日本の官公庁のWebサイトに掲載されている「よくある質問」)
- テーブル構造:
- chunksテーブル: 55298行
- sourcesテーブル: 5904行
- copyright_holdersテーブル: 55行
-- 著作権者テーブル CREATE TABLE copyright_holders ( id SERIAL PRIMARY KEY, name VARCHAR(255) NOT NULL UNIQUE, created_at TIMESTAMPTZ DEFAULT NOW() ); -- ソーステーブル CREATE TABLE sources ( id SERIAL PRIMARY KEY, copyright_holder_id INTEGER REFERENCES copyright_holders(id), url TEXT NOT NULL UNIQUE, created_at TIMESTAMPTZ DEFAULT NOW() ); -- チャンクテーブル CREATE TABLE chunks ( id SERIAL PRIMARY KEY, source_id INTEGER REFERENCES sources(id), chunk_text TEXT NOT NULL, embedding VECTOR(1536), metadata JSONB NOT NULL, created_at TIMESTAMPTZ DEFAULT NOW() );
検証1:indexなしの検索とindexありの検索
検証1-1 chunksテーブルのembeddingカラムにindexを貼らない場合
- 処理内容
- chunks テーブル約55K行をすべて走査(Seq Scan)
- 各行について (embedding <=> :query_vector) を計算
- その距離をキーに「top‐N heapsort」で上位10件をソート・抽出
- パフォーマンス要点
- 全体実行時間:約1,163 ms
- ディスクI/O:13,544ページ読み込み(約31.6 ms)
- キャッシュヒット:254,933ページ
▶︎クリックして詳しい統計を見る
Limit (cost=10849.91..10849.94 rows=10 width=369) (actual time=1163.668..1163.671 rows=10 loops=1) Buffers: shared hit=254933 read=13544 I/O Timings: shared read=31.640 -> Sort (cost=10849.91..10987.41 rows=54997 width=369) (actual time=1163.666..1163.668 rows=10 loops=1) Sort Key: vector Sort Method: top-N heapsort Memory: 52kB Buffers: shared hit=254933 read=13544 I/O Timings: shared read=31.640 -> Seq Scan on chunks (cost=0.00..9661.45 rows=54997 width=369) (actual time=0.115..1146.244 rows=55298 loops=1) Buffers: shared hit=254930 read=13544 I/O Timings: shared read=31.640 Planning: Buffers: shared hit=43 read=7 I/O Timings: shared read=0.040 Planning Time: 0.598 ms Execution Time: 1163.713 ms Allocated Memory: allocated_by_plan=531kB allocated_by_exec=2805174kB base_allocation=3349kB base_allocation_increase=260kB
検証1-2 hnsw indexを貼った場合
以下のようなsql文でindexを貼ります。
CREATE INDEX chunks_embedding_hnsw_cosine ON chunks USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 256);
HNSWパラメータの詳細
m = 16(接続数パラメータ)
- 意味: 各ノードが持つ最大接続数
- 推奨値: 16(デフォルト)
- 調整の指針:
- 小さくする(8-12): メモリ使用量削減、構築速度向上、検索精度やや低下
- 大きくする(32-64): 検索精度向上、メモリ使用量増加、構築時間増加
ef_construction = 256(構築時探索パラメータ)
- 意味: インデックス構築時の候補リストサイズ
- 推奨値: 64-512(データ量に応じて)
- 調整の指針:
- 小さくする(64-128): 構築速度向上、検索精度やや低下
- 大きくする(512-1000): 検索精度向上、構築時間大幅増加
ef_constructionに関しては、amazonのdocumentで256と推奨されていたこともあり、そうしました。
indexを貼るのにはなかなか時間がかかりました。
冒頭で書いていたスペックのインスタンスを使っていた際には、以下の注意がはかれて途中から動かなくなったので、
NOTICE: hnsw graph no longer fits into maintenance_work_mem after 38709 tuples DETAIL: Building will take significantly more time. HINT: Increase maintenance_work_mem to speed up builds.
インスタンスを 以下のスペックに増強してindexを貼りました。56秒ほどかかりました。
- CPU: 8
- Mem: 32Gib
- work_mem: 5242kb
indexを貼ったのちに同じようにクエリを投げてみたところ
- 処理内容
- chunks_embedding_hnsw_cosine HNSWインデックスを用いて chunks テーブルをスキャン
- 内部で “距離が小さい順” に近傍ノードをたどり、上位10件を直接取得
- パフォーマンス要点
- 全体実行時間:約 4.897 ms
- ディスクI/O:868ページ読み込み(約 3.040 ms)
- キャッシュヒット:133ページ
実行時間だけでいえば、×250倍の向上となりました。 一方で、hnswは近似解であることから、真の解と比較してどのくらい取得できているかについては計測できておりません。(俗にいうrecall率)
▶︎クリックして詳しい統計を見る
Limit (cost=1255.16..1277.89 rows=10 width=369) (actual time=4.600..4.839 rows=10 loops=1) Buffers: shared hit=133 read=868 I/O Timings: shared read=3.040 -> Index Scan using chunks_embedding_hnsw_cosine on chunks (cost=1255.16..126922.45 rows=55298 width=369) (actual time=4.598..4.836 rows=10 loops=1) Order By: (embedding <=> [ ... ]) Buffers: shared hit=133 read=868 I/O Timings: shared read=3.040 Planning: Buffers: shared hit=40 read=14 dirtied=1 I/O Timings: shared read=3.220 Planning Time: 3.618 ms Execution Time: 4.897 ms Allocated Memory: allocated_by_plan=556kB allocated_by_exec=971kB base_allocation=3343kB base_allocation_increase=183kB
検証2:フィルタリングを伴うクエリ
これよりembedding
にhnsw
indexを貼った状態で、検証を行う
pgvectorでは、ベクトル検索とフィルタリング条件を組み合わせる際に「ポストフィルタリング」という手法が使われます。これは以下のような流れで処理されます:
- まずベクトル類似度検索を実行
- その後、フィルタリング条件を適用
この方式では、フィルタリング条件の選択性(絞り込み効果)によってパフォーマンスとリコール(検索精度)に大きな影響が出ます。
よく知られる影響としては、ベクトル類似度検索が
検証2-1: 単一テーブル内でのwhere句
まずは、消費者庁
(私の環境では、copyright_holder_id = 77)に限定して、chunksをselectします。
SELECT id, content, embedding <=> :query_vector as distance FROM chunks where source_id IN ( SELECT id FROM sources WHERE copyright_holder_id = 77 ) ORDER BY distance LIMIT 10
SELECT id, content, embedding <=> :query_vector as distance FROM chunks WHERE source_id IN (5086, 5087, 5088, 5089, 5090, 5091, 5092, 5093, 5094, 5095, 5096, 5097, 5098, 5099, 5100, 5101, 5102, 5103) ORDER BY distance LIMIT 10
▶︎クリックして詳しい統計を見る
Limit (cost=1255.46..8888.83 rows=10 width=337) (actual time=1.227..1.395 rows=4 loops=1) Buffers: shared hit=1053 -> Nested Loop (cost=1255.46..129496.14 rows=168 width=337) (actual time=1.226..1.394 rows=4 loops=1) Buffers: shared hit=1053 -> Index Scan using chunks_embedding_hnsw_cosine on chunks (cost=1255.16..126507.71 rows=55298 width=351) (actual time=1.085..1.179 rows=40 loops=1) Order By: (embedding <=> '[...]'::vector) Buffers: shared hit=955 -> Memoize (cost=0.29..0.32 rows=1 width=4) (actual time=0.004..0.004 rows=0 loops=40) Cache Key: chunks.source_id Cache Mode: logical Hits: 13 Misses: 27 Evictions: 0 Overflows: 0 Memory Usage: 2kB Buffers: shared hit=81 -> Index Scan using sources_pkey on sources (cost=0.28..0.31 rows=1 width=4) (actual time=0.002..0.002 rows=0 loops=27) Index Cond: (id = chunks.source_id) Filter: (copyright_holder_id = 77) Rows Removed by Filter: 1 Buffers: shared hit=81 Planning: Buffers: shared hit=127 Planning Time: 0.396 ms Execution Time: 1.464 ms Allocated Memory: allocated_by_plan=1093kB allocated_by_exec=677kB base_allocation=3415kB base_allocation_increase=210kB
以上のSQLを発行すると、4件のチャンクしかヒットしません。
上で述べたように、pgvector
はポストフィルタリング戦略をとっているからです。
したがって、挙動としてはhnsw indexを使って、クエリからもっとも近傍にあるchunksを取ってきた後に、filteを通過できたものだけが返されます。そして、ef_constructionで定められた候補数が尽きたため、結果として4件のみが返されたという結果になります。
検証2-1: 複数テーブルをjoinし、chunksテーブル以外のテーブルのカラムにfilterを設けた場合
全く同じような結果だったため省略
検証結果と考察
ポストフィルタリングでは、フィルタリング条件の選択性が高いほど(フィルターを通過する行が少ないほど)、リコール率が低下する傾向があります。これは、最適なベクトル検索結果がフィルタリングによって除外されてしまうためです。
それを避けるためにはなるべくef_constructionを大きい数値にする必要がありますが、index構築するためにかかる時間が大きく伸びたりするなどのデメリットがあり、トレードオフについて慎重に検討する必要があるようです。
pgvectorscaleで同じようなことを検証してみよう
Timescale社が開発するpgvectorscaleは、pgvectorを補完し、大規模なベクトル検索ワークロードに対して高性能かつコスト効率の良い検索を実現する拡張機能です。
主な特徴
StreamingDiskANNインデックス:
- Microsoft Researchが開発したDiskANNアルゴリズムにインスパイアされたインデックス型
- 大規模データセットでの高速な近似最近傍検索を実現
Statistical Binary Quantization:
- Timescale社の研究者によって開発された圧縮方法
- 標準的なBinary Quantizationを改良し、より効率的なストレージ使用を実現
効率的なメモリ戦略:
インデックスの大部分をSSD上に配置し、「探索に必要なノードやエッジのみをメモリにロード」するメカニズムを取るため、HNSWより数分の一~数十分の一程度のRAMで済むケースが多い。
検証環境
index構築用のsql
search_list_sizeやnum_neighborsはhnswと同様な値に設置しています。(全く同じような環境で検証できているかは不透明)
CREATE INDEX chuns_embedding_disann_idx ON chunks USING diskann (embedding) WITH(search_list_size=256,num_neighbors=16);
検証: フィルタリングを伴うクエリ
SELECT c.id, c.content, c.embedding <=> :query_vector as distance, s.url, ch.name as copyright_holder_name FROM chunks c JOIN sources s ON c.source_id = s.id JOIN copyright_holders ch ON s.copyright_holder_id = ch.id WHERE ch.name = '消費者庁' ORDER BY distance LIMIT 10;
同じようなクエリを投げたところ、しっかり10件の候補が返ってきました。 queryにかかった時間も増えた一方で、hnswよりも作業用のメモリ消費が10倍以上となりました。
hnswと同じような条件ではない(indexのチューニングが同じような条件にできていない)ことに起因するのでしょうが、候補がしっかり10件返ってきていることは評価できるポイントであります。
▶︎クリックして詳しい統計を見る
Limit (cost=56.44..257.17 rows=10 width=414) (actual time=2.796..4.634 rows=10 loops=1) Buffers: shared hit=3637 -> Nested Loop (cost=56.44..21815.83 rows=1084 width=414) (actual time=2.795..4.632 rows=10 loops=1) Join Filter: (ch.id = s.copyright_holder_id) Rows Removed by Join Filter: 66 Buffers: shared hit=3637 -> Nested Loop (cost=56.44..20982.01 rows=55298 width=411) (actual time=2.627..4.528 rows=76 loops=1) Buffers: shared hit=3596 -> Index Scan using chuns_embedding_disann_idx on chunks c (cost=56.15..18007.13 rows=55298 width=351) (actual time=2.535..4.279 rows=76 loops=1) Order By: (embedding <=> '[...]'::vector) Buffers: shared hit=3458 -> Memoize (cost=0.29..0.31 rows=1 width=68) (actual time=0.003..0.003 rows=1 loops=76) Cache Key: c.source_id Cache Mode: logical Hits: 30 Misses: 46 Evictions: 0 Overflows: 0 Memory Usage: 8kB Buffers: shared hit=138 -> Index Scan using sources_pkey on sources s (cost=0.28..0.30 rows=1 width=68) (actual time=0.002..0.002 rows=1 loops=46) Index Cond: (id = c.source_id) Buffers: shared hit=138 -> Materialize (cost=0.00..1.64 rows=1 width=21) (actual time=0.000..0.000 rows=1 loops=76) Buffers: shared hit=1 -> Seq Scan on copyright_holders ch (cost=0.00..1.64 rows=1 width=21) (actual time=0.019..0.020 rows=1 loops=1) Filter: (name = '消費者庁'::text) Rows Removed by Filter: 54 Buffers: shared hit=1 Planning: Buffers: shared hit=680 Planning Time: 1.874 ms Execution Time: 4.818 ms Allocated Memory: allocated_by_plan=3134kB allocated_by_exec=10575kB base_allocation=3306kB base_allocation_increase=340kB
結論
- pgvectorはポストフィルタリングがかなりきつく、ef_constructionを256以上の値に設定しなければ、chunksが数万しかない現状でも十分に結果が返ってこないことがわかりました。
- pgvectorscaleはdiskAnnという目新しいindexアルゴリズムをとっているが、細かいチューニングがわからずpgvectorと十分に対照実験を行う必要があります。