これは PostgreSQL Advent Calendar 2023 2 日目の記事です。
昨日は noborus さんでした。
先日(11/24)、PostgreSQL Conference Japan 2023 に参加した際、
- pgvector なんもわからん
- Vector Store なにそれ?
という声が多かったので(幻聴かもしれませんが)、pgvector の初歩の初歩について記していきます。
当日、石井さんのセッションに参加された方は読み飛ばしてください。
(あのセッション以上の説明はしないので)
pgvector とは?
PostgreSQL の拡張機能(Extension)です。
PostgreSQL に、
- ベクトル(
vector
)データ型 - ベクトル関数(距離計算など)・オペレータ
- 近似最近傍探索用のインデックス
を追加するものです。
何に使うもの?
DB のテーブルにベクトルデータを保存し、主に最近傍探索をするために使います。
具体的には?
最近よく使われる例として、
- 文章をベクトル化して DB のテーブルに保存する
- 入力文に近い意味・文脈の文章をテーブルからベクトルを使って検索する
- テーブルの検索結果を「ヒント」「文脈」のような形で入力文に付け加え、ChatGPT などの LLM(大規模言語モデル)にプロンプトとして送信する
いわゆる RAG(Retrieval-Augmented Generation) があります。
ChatGPT などで、モデルが学習済みの範囲よりも新しい(あるいは一般公開されていない)話題・情報についてカバーしようとする場合、一般的に
- 追加学習させる(ファインチューニング)
- プロンプトに「ヒント」「文脈」のような形で関連情報(文章)を与える
という方法で対応することになりますが、ファインチューニングにはコストも時間もかかる(そして、現時点ではまだ API として一般提供されていない LLM サービスが多い)ため、後者の方法が採用されるケースが多くなっています(併用することもあります)。
この後者の方法で構築するのが RAG です。
文章以外のものをベクトル化して検索することもできます(画像など)。
全文検索とは何が違うの?
全文検索では キーワードが一致するかどうか(一致度) を基準に文章を検索します。
キーワードが完全一致しない場合、そのシノニム(類義語)辞書を使ってキーワードに近い意味の文章を検索することになりますが、辞書を作るのには限界があります。
文章をベクトル化してベクトル同士を比較する場合は、 似た意味の(文脈が近い)文章が「距離が近い文章」 となり、キーワードの一致がなくとも類似度の高い文書が抽出可能です。
そのため、全文検索と比べて、あいまい検索や文章による検索に適しています。
使ってみる
一番簡単な方法として Docker で起動してみます。
Docker コンテナ起動
pgvector 開発者の方が提供されているイメージを使うと、最初から pgvector が組み込まれているので便利です。
$ docker pull ankane/pgvector
$ docker run --net=host -e POSTGRES_PASSWORD='【パスワード】' ankane/pgvector
2023/11/29 現在、提供されているコンテナイメージでは、
- PostgreSQL 15.4
- pgvector 0.5.1
の構成になっています。
docker exec -it 【コンテナID】 sh
でコンテナの中に入り、psql
コマンドで PostgreSQL に接続します。
# psql -h localhost -U postgres -d postgres
psql (15.4 (Debian 15.4-2.pgdg120+1))
Type "help" for help.
Extension 有効化
まずはpgvector
を有効化します。
postgres=# CREATE EXTENSION vector;
CREATE EXTENSION
postgres=# SELECT extversion FROM pg_extension WHERE extname = 'vector';
extversion
------------
0.5.1
(1 row)
テーブル作成・データ投入
pgvector の README を参考に、3次元のベクトルデータを含むテーブルを作成してみます。
postgres=# CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));
CREATE TABLE
vector(3)
の3
はベクトルの次元数です(上限は16000
。ただしインデックス(後述)がサポートする上限は2000
)。
続いてデータを投入します。
postgres=# INSERT INTO items (embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,9]'), ('[0,1,2]'), ('[3,4,5]'), ('[6,7,8]'), ('[9,0,1]'), ('[2,3,4]'), ('[5,6,7]'), ('[8,9,0]');
INSERT 0 10
postgres=# SELECT * FROM items;
id | embedding
----+-----------
1 | [1,2,3]
2 | [4,5,6]
3 | [7,8,9]
4 | [0,1,2]
5 | [3,4,5]
6 | [6,7,8]
7 | [9,0,1]
8 | [2,3,4]
9 | [5,6,7]
10 | [8,9,0]
(10 rows)
README の Storing には Insert 以外の例も列挙されています。
- Upsert
INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]'), (2, '[4,5,6]')
ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding;
- Update
UPDATE items SET embedding = '[1,2,3]' WHERE id = 1;
- Delete
DELETE FROM items WHERE id = 1;
検索(最近傍探索)
ベクトル同士の距離比較には、
- L2 距離(ユークリッド距離)
- 内積
- コサイン類似度(またはコサイン距離)
が使えます。
距離計算用のベクトル関数としては、pgvector 0.5.0 でl1_distance()
(L1 距離・マンハッタン距離)が追加されています。
postgres=# SELECT *, embedding <-> '[3,1,2]' AS l2_distance FROM items ORDER BY l2_distance LIMIT 5;
id | embedding | l2_distance
----+-----------+-------------------
1 | [1,2,3] | 2.449489742783178
4 | [0,1,2] | 3
8 | [2,3,4] | 3
5 | [3,4,5] | 4.242640687119285
2 | [4,5,6] | 5.744562646538029
(5 rows)
postgres=# SELECT *, (embedding <#> '[3,1,2]') * -1 AS inner_product FROM items ORDER BY inner_product LIMIT 5;
id | embedding | inner_product
----+-----------+---------------
4 | [0,1,2] | 5
1 | [1,2,3] | 11
8 | [2,3,4] | 17
5 | [3,4,5] | 23
7 | [9,0,1] | 29
(5 rows)
postgres=# SELECT *, 1 - (embedding <=> '[3,1,2]') AS cosine_similarity FROM items ORDER BY cosine_similarity LIMIT 5;
id | embedding | cosine_similarity
----+-----------+--------------------
4 | [0,1,2] | 0.5976143046671968
10 | [8,9,0] | 0.7324296566704842
1 | [1,2,3] | 0.7857142857142857
8 | [2,3,4] | 0.8436958338752907
7 | [9,0,1] | 0.8559079373463852
(5 rows)
それぞれで「距離の近さ」の順位が微妙に違いますね。
なお、OpenAI の embeddings API(text-embedding-ada-002
などのモデル)を使うケースなど、ベクトル化の際にベクトル長が 1 に正規化される場合は、内積を選択すると最もパフォーマンスが高くなります。
インデックス作成(近似最近傍探索用)
インデックスを作成しない場合、最近傍探索をするにはテーブル内の全行に対して距離を計算し比較する必要があります。
普通のテーブル全行スキャンより負荷が高そうですね。
この負荷を下げるためのインデックスですが、pgvector では(2023/11/29(pgvector 0.5.1)時点で)
- IVFFlat
- HNSW
の 2 種類が作成可能です。
インデックス作成・使用時の注意としては、まず
- インデックスを使った検索は「近似」最近傍探索になるので、インデックスを使わない「完全」最近傍探索と違う結果になることがある
点が挙げられます。
他には、
- 作るインデックスの種類によって特性が違う(それぞれの注意点がある)
- L2 距離・内積・コサイン距離で別のインデックスを作る必要がある
- ベクトル次元数上限が 2,000 に制限される(行データのベクトル次元数上限は 16,000)
などがあります。
種別 | インデックス作成所要時間 | インデックス作成時メモリ使用量 | 検索速度 | インデックス作成タイミング |
---|---|---|---|---|
IVFFlat | ○短い | ○少ない | ×遅い | データ行がある程度蓄積されてからでないと NG |
HNSW | ×長い | ×多い | ○速い | データ行がない状態でも OK |
次に示す例は IVFFlat・L2 距離(ユークリッド距離)インデックスを作成する例です。
postgres=# CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);
NOTICE: ivfflat index created with little data
DETAIL: This will cause low recall.
HINT: Drop the index until the table has more data.
CREATE INDEX
postgres=# SELECT *, embedding <-> '[3,1,2]' AS dist FROM items ORDER BY dist LIMIT 5;
id | embedding | dist
----+-----------+-------------------
1 | [1,2,3] | 2.449489742783178
4 | [0,1,2] | 3
8 | [2,3,4] | 3
5 | [3,4,5] | 4.242640687119285
2 | [4,5,6] | 5.744562646538029
(5 rows)
これ以外のインデックスの作り方は、README の Indexing に記されています。
インデックス作成時に
ERROR: memory required is 202 MB, maintenance_work_mem is 67 MB
のようなエラーが出ることがあります。
この場合はmaintenance_work_mem
の設定を変更する必要があります。
詳細は以下の Neon のドキュメントを参照してください。
- Indexing vectors(pgvector - Neon Docs)
IVFFlat は ベクトルをlists
で指定した数のリストに振り分ける 方式のインデックスですが、インデックス作成時のテーブル内のデータ行数が少なすぎることが理由で警告(NOTICE:
・DETAIL:
・HINT:
)が出ていますね。
IVFFlat では、各リストへの振り分けの際に「近い場所に存在するベクトル同士を同じリストに入れる」処理を行いますが、インデックス作成時にデータ行が少なすぎると、振り分けの基準に使う「重心」を正しく導き出すことができず、インデックス作成後に挿入されたデータが、真の「重心」から離れたリストに振り分けられる確率が高くなります。
結果として、検索時に誤ったベクトルが「最近傍」として抽出される可能性があります。
そのため、データ行がある程度蓄積されてからインデックスを作成する必要があります。
検索時にSET
でivfflat.probes
パラメータを指定すると、最近傍の重心を持つリスト以外の複数のリストからベクトルを検索できるようになります(デフォルトは1
)。
数値を大きくするほど最近傍のベクトルを取りこぼす確率は下がりますが、一方で検索速度は遅くなります。
また、インデックス作成後にデータ行を大量に追加した場合は再インデックスが推奨されます。
一方、HNSW(階層化したグラフを使って最近傍を探す方式のインデックス)ではデータ行がない状態でインデックスが作成可能ですが、以下の2 つのパラメータ指定で検索精度とパフォーマンスのバランスを取ることになります(詳細は README の HNSW や、冒頭の スライド資料 53 ページ目~ を参照)。
-
m
:レイヤ(階層)ごとの最大接続(グラフのリンク)数(デフォルト:16
) -
ef_construction
:インデックス作成時に「最近傍」ベクトルの候補をいくつ持つか?(デフォルト:64
)
検索時はSET
でhnsw.ef_search
パラメータを指定して検索精度とパフォーマンスのバランスを取ります(検索時の「最近傍」ベクトルの候補数。デフォルトは40
)。
HNSW については、以下の Neon のブログ記事も参考になります。
後日公開予定の記事に続きます。
- 検索負荷を下げるには?
- フィルタリング
- ハイブリッド検索
- バッファの暖機(プレウォーム)
- 検索精度と性能のバランスを取る
- 結局、象使いはどうすれば良い?
2024/5/19 追記:
さらにその続編を追加公開しました。
明日(3 日)の担当は asahide さんです。