Vector Sets are part of Redis

Salvatore Sanfilippo

Vector SetsがRedisの一部になった

原文は Salvatore Sanfilippo により に公開されました。 このブログを購読する

昨日、ついにVector SetsをRedisにマージした。詳細はこちらのREADMEで確認できる:

https://github.com/redis/redis/blob/unstable/modules/vector-sets/README.md

新しいデータ構造の目的を一言で言えば、Sorted Setsに似た、しかしスコアがスカラー値ではなくベクトルであるような、新しい「Set系」のデータ型を作ることだ。Redisらしいやり方で要素を追加・削除でき、Redisが実装する抽象データ構造の性質以外を気にする必要はなく、クエリベクトル(あるいはすでにセット内にある要素に紐づいたベクトル)に似た要素を問い合わせるといったことができる。詳しくは後で述べるとして、まずは少し背景から:

READMEのパスからも分かるように、実装は「modules」以下にあるが、実際にはVector Setsはモジュールではない。Redisコアの一部だ。経緯としては、最初はモジュールとして開発を始め、後に内部のモジュール性を高めるために、あえてModules APIを使ったまま実装することを提案したのだ。これにより両方の利点が得られる。Redis 8以降のすべてのRedisインスタンスはVector Setsをネイティブなデータ型として備えることになるし、同時にコアと実装の間には明確な境界が保たれる。

久しぶりのRedisの新たな主要データ型

Redisで最後に登場した大きなデータ構造はStreamsだったと思うが、あれも僕が開発したものだ。辞職し、復帰し、その間にフォークも起こり、それでもRedisに新しいデータ型を導入する役目はやはり僕の肩にかかってきたようだ :D 正直、それはそれで構わない。プログラミングが好きなのと同じくらい、いやそれ以上に設計が大好きだし、ベクトルやベクトル類似度というのは概念的にとてもシンプルで、だからこそシンプルなAPIにふさわしいと感じていた。そして、それを実現しようとしたのだ。Vector Setsはまだベータ機能だが、ひとつ保証できることがある。APIは3分で覚えられるということだ。

ベクトル類似度を実装するにあたって、HNSWを一から再実装することは必須だと決めていた(僕の実装はhnsw.cで見られる)。それがコアとなるデータ構造になるのだから、GitHubから適当なコードを拾ってきて満足するようなことはしたくなかったのだ。ところが、論文を読み進めるうちに、いくつか足りないピースがあることに気づき始めた。

だから、かつてHyperLogLogでいくつかのギャップを埋めなければならなかったときのように(こちら:https://antirez.com/news/75)、今回も新たなアルゴリズム上の課題があった。特に実現したかったのは次の2点だ:

  1. ノードを真に削除できること。Vector SetsではVADDで新しい要素を追加し、VREMで要素を削除できるが、削除した際にはできるだけ早くメモリを解放したかった。
  2. 要素を削除しても、HNSWグラフの連結性が保たれるようにすること。

そのため、僕の実装は他のHNSW実装とはいくつか異なる点がある。tombstone(墓標)による削除は使わず、削除の瞬間にノードを実際にリンクから外し、他に有望な近傍ノードと再リンクする。そのために、リンクは相互に張られることを強制するように設計している。もはやベストエフォートな性質ではないのだ。そしてこれは挿入時にやるべきことをかなり変えることになる。

HNSWに加えたもう一つの変更は、述語関数を使ってグラフを走査できるようにしたことだ。これにより、与えられた式にマッチするノードを問い合わせることができる。そのためには貪欲なグラフ走査アルゴリズムをある程度修正する必要がある。訪問候補のノードを集め、結果セットを収集しつつ、クエリの選択性が高すぎる場合には早期に停止する条件も持たせる必要がある。もちろん、グラフ全体のフルスキャンを引き起こしたくはないからだ。

HNSW自体の変更以外にも、もう少し実務的で分かりやすいことがらを実現したかった:

  1. すべてのベクトル類似度リクエストのスレッド化だ。そう、これはRedisの世界では新しいことだが、シングルスレッドでシェアードナッシングが一般的には良い設計だと信じている一方で、ベクトルは特別だと思っている。他のRedisが扱うデータ構造に比べて圧倒的に遅いのだ。おまけとして、スレッド化されたVSIM(ベクトル類似度クエリを実行するコマンド)を実装している際に、ちょっとした工夫で書き込みの読み取り部分と書き込み部分を2つに分割できることにも気づいた。近傍候補の収集はバックグラウンドで行い、実際の挿入はフォアグラウンドで行うのだ。ただしこの分割はデフォルトではなく、VADDのCASオプションで明示的に有効にする必要がある。
  2. 量子化をサポートし、さらにそれをデフォルトにすること。Vector Setsは8ビット量子化とバイナリ量子化の両方を備えている。次元削減のためのランダムプロジェクションにも対応している。僕自身はRPもバイナリ量子化も好きだが、正直なところ「キラー」なのはint8量子化だ。とてつもなく高速で、メモリ使用量はFP32の25%で済み、埋め込みAIモデルによって生成されたほとんどのベクトルでは、フルベクトルの結果とほぼ同じ結果が得られる。

ちなみに、最終的な実装はかなり高速だと自負している。たとえば僕のマシンでは、300次元のベクトルを300万件格納したVector Setに対して、VSIM(上位10件取得)を毎秒5万〜6万回実行できている。手元のノートPCでの話だ。ぜひ皆さん自身でもベンチマークしてみてほしい。

また、Vector Setsはディスク上ではグラフとしてシリアライズされるため、Redis再起動後にメモリへロードする際に挿入時間を再び支払う必要がない。100万要素あたり数秒でロードでき、そうでなければインメモリのHNSWに再投入するのに何分もかかるようなケースと比べて大きな違いだ。

インデックスではなくデータ構造

ここまで述べてきたのはすべて低レイヤーの話だ。しかし僕にとってVector Setsで最も面白いのは、データモデルとそれを支えるAPIだ。多くのデータベースはベクトル類似度を一種のインデックスとして提供するが、ここはRedisだ。Redisにあるものはデータ構造であり、今回も例外ではない。追加はこんな感じでできる:

VADD mykey FP32 …blob of data… item1

といった具合だ。だから、キーごとに小さなVector Setをたくさん持つこともできる。そして重要なのは、ベクトルをN個の異なるキーに分割した場合(挿入するアイテムをハッシュしてどのキーを選ぶかを決めるなど)、異なるキーに対する複数のVSIM呼び出しを1つのレスポンスにマージできるということだ:

VSIM word_embeddings_int8 ele "banana" WITHSCORES COUNT 4
1) "banana"
2) "0.9997616112232208"
3) "bananas"
4) "0.8758847117424011"
5) "pineapple"
6) "0.8288004100322723"
7) "mango"
8) "0.8179697692394257"

異なるキーやインスタンスからこうした結果をいくつか取得すれば、スコア(1が同一、0が逆ベクトルを意味する)でソートするだけで済む。

だから、Vector Setsは異なるパターンに組み合わせて、大量のベクトル(かなりRAMを食う)を複数のインスタンスに分散させるといった使い方もできると感じている。また、分割することで書き込みが線形にスケールするのも面白い。サブセットごとに特定のキーにヒットし、複数の挿入を並列に実行できるからだ。

例によって、Redisコミュニティは今はまだ自明ではないような、さまざまな使い方のパターンをきっと見つけ出してくれるだろう。

フィルタリングの仕組み

スレッド化でさえRedisでは一般的ではなかったのに、ましてやJSONなど!だが今回初めて、JSONをRedisのAPIで、ユーザーに直接見える形で公開する十分な理由が見つかった:

VGETATTR word_embeddings_int8 banana
{"len": 6}

基本的に、VSETATTR / VGETATTR(およびVADD時にJSON属性を直接設定する同等のオプション)で、アイテムに文字列を紐づけることができる。

そうすると、こんなことができるようになる:

VSIM word_embeddings_int8 ele "banana" FILTER ".len == 3"
 1) "yam"
 2) "pea"
 3) "fig"
 4) "rum"
 5) "ube"
 6) "oat"
 7) "nut"
 8) "gum"
 9) "soy"
10) "pua"

フィルタ式はプログラミング言語ではなく、高水準プログラミング言語のif()文の中に書けるようなものだ。&&や||、その他当然あるべき演算子はすべて使える(もちろん、今後さらに追加していくつもりだ)。

詳しくはドキュメントに書いてある。メモリ使用量の例や、特定の機能に関する詳細な議論なども載っている。近いうちにドキュメントも拡充する予定だ。とにかく、Vector Setsを(本当に!)楽しんでもらえれば嬉しい。バグを見つけたら、ぜひ教えてほしい :)

この記事は「muse-spark-1.2-contributor」を使用して翻訳されました。

コメント