この記事は Elastic Stack (Elasticsearch) Advent Calendar 2023 18日目の記事です。
こんにちは。
Acroquestのデータサイエンスチーム「YAMALEX」に所属する@shin0higuchiです😊
YAMALEXチームでは、コンペティションへの参加や自社製品開発、技術研究などに日々取り組んでいます。
さて、最近はLLMの発展に伴ってRAG（Retrieval-Augumented Generation）が盛んに活用されています。
その中で、キーワードベースの検索だけでなくベクトル検索を併用するケースが多く見られ、実際にElasticsearchが利用されているケースも多く目にします。そのため、Elasticsearchのベクトル検索に興味を持っている方も多いと思います。今回の記事では、Elasticsearchのベクトル検索の速度などを、簡単に計測してみました。
概要
検証環境のOS／ミドルウェアバージョン
| OS/ミドルウェア | バージョン | 備考 |
|---|---|---|
| Elasticsearch | 8.11.3 | |
| Kibana | 8.11.3 | |
| OS | Ubuntu 22.04.3LTS | CPU 32コア, RAM 128GB。データ投入は別のWindowsマシンから実施 |
今回は以下の想定で検証を実施してみました。
RAGでよく用いるOpenAIのEmbeddingを用い、100万件の文章のベクトルをESに投入したと仮定して、その検索速度を計ってみたいと思います。
なお、OpenAIのEmbeddingを利用してベクトル化すると、各文章は1536次元のベクトルになります。
検証
インデックスのセットアップ
まずはデータを入れるためのインデックスにデータ型の定義を実施しておきます。
ここではシンプルに `vector` というフィールドのみを定義します。dense_vector型を指定し、1536次元を指定します。（上限はVer.8.11現在では4096次元です）
PUT single_vector_test { "mappings": { "properties": { "vector": { "type": "dense_vector", "dims": 1536 } } } }
データ投入
データ投入にはLogstashを用います。
以下のような設定ファイルを作成して利用します。
input {
file {
path => "//pc-shin-gpu1/share/data/single_vector.txt"
start_position => "beginning"
sincedb_path => nul
}
}
filter {
mutate {
rename => {"message" => "vector"}
}
json {
source => "vector"
target => "vector"
}
mutate {
remove_field => ["message", "host", "path", "@version"]
}
}
output {
elasticsearch {
hosts => ["https://pc-shin-gpu1:9200"]
index => "single_vector_test"
# ここでは検証を簡単化するためにSSLの検証をスキップさせたり、elasticユーザを利用しています。
# 本番で利用する際は、SSLの設定を適切に行い、API Keyの利用なども検討してください。
ssl_verification_mode => "none"
user => "elastic"
password => "XXXXX"
}
stdout { codec => dots }
}
読み込む対象のデータは single_vector.txt に以下のようなフォーマットで保存されている前提です。
[2,31,0,212,6,.....,23] #1536次元 [83,11,5,0,0,.....,98] ....
Elasticsearch上でのデータ確認
Cat Indices APIで件数を確認してみます。1,000,000件入っていることが確認できます。
(レプリカ数)
# GET _cat/indices/single_vector_test?v health status index uuid pri rep docs.count docs.deleted store.size pri.store.size dataset.size green open single_vector_test kNg4GFrKSzKrChb8EELNxA 1 0 1000000 0 33.4gb 33.4gb 33.4gb
検索処理
Elastisearchには正確に近傍検索をおこなう kNN と、近似計算で高速に処理をする ANN の２パターンがサポートされています。
kNN
GET single_vector_test/_search { "size": 10, "query": { "script_score": { "query": { "match_all": {} }, "script": { "source": "cosineSimilarity(params.queryVector, 'vector') + 1.0", "params": { "queryVector": [106, 142, 86, 87,...., 149] } } } } }
500ミリ秒くらいかかっていますね。
(なお、32シャードに分割して同検証を実施すると170ミリ秒程度までは短縮されました)
{ "took": 501, "timed_out": false, "_shards": { "total": 1, "successful": 1, "skipped": 0, "failed": 0 }, "hits": { "total": { "value": 10000, "relation": "gte" }, "max_score": 2, "hits": [ { "_index": "single_vector_test", "_id": "9D6UdowBpziQa9Jus0xV", "_score": 2, "_ignored": [ "event.original.keyword" ], "_source": { "@timestamp": "2023-12-17T07:01:43.527208800Z", "log": { "file": { "path": "//pc-shin-gpu1/share/data/single_vector.txt" } }, "vector": [.....] } } ] } }
ANN
一方ANNはこちら。
GET single_vector_test/_search { "knn": { "field": "vector", "k": 10, "num_candidates": 100, "query_vector": [....] } }
12ミリ秒と非常に高速に検索できていることがわかります。
ただし、あくまでkNNと違い近似の結果にはなるので、正確に類似度順に取得する必要がある場合は kNNを利用するようにしてください。
{ "took": 12, "timed_out": false, "_shards": { "total": 1, "successful": 1, "skipped": 0, "failed": 0 }, "hits": { "total": { "value": 10, "relation": "eq" }, "max_score": 1, "hits": [ .... ] } }
まとめ
1536次元という比較的高次元なベクトルも、ANNを用いることでかなり高速に検索することができました。RAGに利用するという観点で言えば、検索結果の厳密性よりも高速にある程度の精度で類似検索ができるメリットは大きいように思います。
冒頭で触れた通り、キーワード検索とベクトル検索を併用できるElasticsearchはLLMの後ろで使うには非常に便利ですね。
ということで、今回の記事は以上となります。
お読みいただきありがとうございました。
Acroquest Technologyでは、キャリア採用を行っています。
- ディープラーニング等を使った自然言語／画像／音声／動画解析の研究開発
- Elasticsearch等を使ったデータ収集／分析／可視化
- マイクロサービス、DevOps、最新のOSSを利用する開発プロジェクト
- 書籍・雑誌等の執筆や、社内外での技術の発信・共有によるエンジニアとしての成長
少しでも上記に興味を持たれた方は、是非以下のページをご覧ください。Kaggle Grandmasterと一緒に働きたエンジニアWanted！ - Acroquest Technology株式会社のデータサイエンティストの採用 - Wantedlywww.wantedly.com
