はじめに

前回の記事では、1,280次元の画像特徴ベクトルを約100万用意し、Amazon Elasticsearch Serviceに投入したが、レスポンス時間が15秒/クエリという実用からは程遠い結果が得られた。ありがたいことに、Amazon ES Teamからアドバイスを頂いたので、それに沿って再度検証を行った。

この手順にについてはElasticsearchに長けている人であれば、ささいなことなのかもしれないが、実サービスでElasticsearchを運用した経験がない私にとっては有用だった。

Segments

前回のElasticsearchからのレスポンスを見てみると、hitsが1,203であることがわかる。このとき、近傍10ベクトルを検索していたので、k=10であったので、少なくとも120のElasticsearch Indexから結果が返ってきていたことがわかる。

{
  "took": 15471,
  "timed_out": false,
  "_shards": {
    "total": 5,
    "successful": 5,
    "skipped": 0,
    "failed": 0
  },
  "hits": {
    "total": {
      "value": 1203,
      "relation": "eq"
...

Elasticsearch Indexは、Shardという単位で分割されており、それぞれがLucene Indexである。Lucene Indexは内部的には複数のファイルに分割されており、それがSegmentと言われるものである。Segmentはシーケンシャルに検索されるので、Segmentの数が少なければ少ないほど検索効率は高くなる。

Amazon ESのデフォルトでは、Shard数は5であるので、検索効率を考えた場合、Segment数も5であることが望ましい。

設定

検索効率と改善するために、以下の設定が提案された。

index.refresh_interval = -1 (default: 1 sec)
index.translog.flush_threshold_size = ‘10gb’ (default: 512mb)
index.number_of_replicas = 0 (default: 1)

さらに、インデックス効率を改善するために、インデックス時のスレッド数についても提案があった。

下記が、改善後のデータ挿入のためのコード

import time
import math
import numpy as np
import json
import certifi
from elasticsearch import Elasticsearch, helpers
from sklearn.preprocessing import normalize

dim = 1280
fvecs = np.memmap('fvecs.bin', dtype='float32', mode='r').view('float32').reshape(-1, dim)

idx_name = 'imsearch'
es = Elasticsearch(hosts=['https://vpc-xxxxxxxxxxx.us-west-2.es.amazonaws.com'],
                   ca_certs=certifi.where())

res = es.cluster.put_settings({'persistent': {'knn.algo_param.index_thread_qty': 2}})
print(res)

mapping = {
    'settings' : {
        'index' : {
            'knn': True,
            'knn.algo_param' : {
                'ef_search' : 256,
                'ef_construction' : 128,
                'm' : 48
            },
            'refresh_interval': -1,
            'translog.flush_threshold_size': '10gb',
            'number_of_replicas': 0
        },
    },
    'mappings': {
        'properties': {
            'fvec': {
                'type': 'knn_vector',
                'dimension': dim
            }
        }
    }
}

res = es.indices.create(index=idx_name, body=mapping, ignore=400)
print(res)

bs = 200
nloop = math.ceil(fvecs.shape[0] / bs)
for k in range(nloop):
    rows = [{'_index': idx_name, '_id': f'{i}',
             '_source': {'fvec': normalize(fvecs[i:i+1])[0].tolist()}}
             for i in range(k * bs, min((k + 1) * bs, fvecs.shape[0]))]
    s = time.time()
    helpers.bulk(es, rows, request_timeout=30)
    print(k, time.time() - s)

Merge Segments

上記コードであっても、Segment数は5にはならないので、下記エンドポイントを呼んで、能動的にマージする必要がある。

POST /imsearch/_forcemerge?max_num_segments=1

この実行には時間がかかるので、ステータスコード:200が返ってくるまで数分おきくらいに呼び、完了したかどうかを確認する必要がある。もちろん、定期的にポーリングせずにしばらく放置しててもいい。最終的には、Segment数は5になるはず。

Refresh

index.refresh_interval = -1 によってリフレッシュが行われていないので、下記エンドポイントを呼ぶ必要がある。 (デフォルトの1sに戻してもいい)

POST /imsearch/_refresh

リフレッシュが完了すると、挿入したベクトルの検索が可能となる。

で、実際に実行してみると、その結果は、、、7秒。。。結構手間かかってるにも関わらず、これはショック。もちろん、前回の15秒に比べると圧倒的に改善しているが、7秒は実用的とは言えない。また、ウォームアップも特に効果はなかった。ここで、検索時間以外では、hitsが50になっているのは、Shardあたり1つのSegmentになっているので、これは理想的。

{
  "took": 7269,
  "timed_out": false,
  "_shards": {
    "total": 5,
    "successful": 5,
    "skipped": 0,
    "failed": 0
  },
  "hits": {
    "total": {
      "value": 50,
      "relation": "eq"
    },
...

Memory

今回の実験で利用しているのは、r5.largeインスタンスで、2コアCPUと16GBメモリが利用可能である。Amazon ESのデフォルトでは、32GBを上限としてサーバメモリの50%がElasticsearchに割り当てられ、残りのメモリの一部がkNNに割り当て可能となる。余ったメモリすべてをkNNに割り当てるのは危険であるため、circuit breakerが実装されており、メモリ使用量に上限が設けられている。これは、knn.memory.circuit_breaker.limitで設定可能で、デフォルト値は60%である。したがって、16GB * 50% * 60% = 4.8GBがkNNへの割当上限となる。

HNSWではグラフ/インデックスを保持するのに4 * d + 8 * Mバイト必要。実験に使ったのは、1,280次元の100万ベクトルで、さらに、Amazon ES Teamからのアドバイスによると、グラフ構築などを考えると1.5倍が必要とのことで、今回は、M=48なので、(4 * 1,280 + 8 * 48) * 1.5 * 1M = 7.7GB必要となる。

r5.largeでは足りないので、r5.xlarge(4コアCPU, 32GBメモリ)でクラスタを作り直して再度試した。このインスタンスタイプであれば、32GB * 50% * 60% = 9.6GB割当可能なので、7.7GBを満たしている。CPUコア数が増えたことにより、ついでに、knn.algo_param.index_thread_qtyも4にして、上記手順でインデックスを生成した。

その後、検索リクエストを1000回送って、検索時間を計測すると平均14msという結果が得られ、実用レベルにまで改善した。

ANN with nmslib

ANNライブラリとして、普段Faissを利用しているため、前回の実験でも、HNSWについてもFaissを使ったが、Elasticsearchに使われているものは、nmslibであるため、念の為比較しておいた。インスタンスタイプは、クラスタに使った、r5.xlargeである。

Elasticsearch kNNを使った場合に比べて、2倍ほど高速であるが、ベクトル検索エンジンではなく、全文検索エンジンであることを考えると、14msという結果は十分に実用レベルであると言える。

まとめ

Amazon ES Teamのサポートのおかげで、Similarity SearchをElasticsearchで実用レベルの検索時間で実現できた。具体的には、Segment数とメモリ割り当てを適切にすることで、前回とは比べ物にならない改善が見られた。HNSWは素晴らしいANNアルゴリズムではあるものの、実データを保持する必要がありメモリ効率に関しては、Inverted File with Product Quantization (IVFPQ)などと比べるとよいとは言えない。高次元のベクトルをElasticsearchで扱う場合には、やはり可能な範囲で次元圧縮を行ったほうがいいという印象を受けた。HNSWは、million-scaleのデータに関してはよくできたアルゴリズムで、それをElasticsearchによってクラスタ化することで、billion-scaleのデータも扱えるようなスケーラビリティが得られたと言える。

続) Elasticsearchで類似ベクトル探索 / 類似画像検索