Elasticsearchの bbq_hnsw を使ったベクトル検索(実践編)

BLOG

はじめに

前回は、Elasticsearchで bbq_hnsw を利用したベクトルの量子化について、その基礎を解説しました。今回は、bbq_hnsw を使って実際にベクトル検索を行う手順を解説します。

対象読者

  • Elasticsearchでベクトル検索の導入を検討している方
  • Elasticsearchの初心者〜中級者

対象バージョン

  • Elasticsearch 8.19.0以上
  • Elasticsearch 9.1.0以上 1

この記事のサンプルでは rescore_vector の oversample に0を指定していますが、これがサポートされているのは上記のバージョン以降です。筆者はElasticsearch 8.19.0で検証しました。

検索の準備

1. インデックスの作成

今回の検証に利用するインデックスを作成します。今回はベクトル検索のみを行うため、形態素解析の設定は省略します。

以下のリクエストをKibanaのDev Toolsから実行してください。

(インデックス名 = waganeko_with_bbq_hnsw)

PUT /waganeko_with_bbq_hnsw/
{
  "settings": {
    "index": {
      "number_of_shards": 1,
      "number_of_replicas": 1,
      "refresh_interval": "3600s"
    }
  }
}

2. インデックスのマッピング設定

インデックスにフィールドを作成します。ここでは、以下の3つのフィールドを定義します。

フィールド名タイプ説明
chunk_nolongチャンク番号
contenttext本文
text_embeddingdense_vector384次元のfloat32密ベクトル

上記には bbq_hnsw で量子化された項目を記載していません。bbq_hnsw による量子化は、厳密にはフィールドではなく、dense_vector タイプのフィールドのindex_optionsとして設定されます。

なお、bbq_hnsw を使ったベクトル検索時のオーバーサンプリング係数は5としています。

現在の最新バージョンでのデフォルト値は3のようですが、ドキュメントに明記されているわけではないため、ご自身のバージョンでご確認ください。

以下のリクエストをDev Toolsから実行します。

PUT /waganeko_with_bbq_hnsw/_mapping
{
  "dynamic": false,
  "properties": {
    "chunk_no": {
      "type": "long"
    },
    "content": {
      "type": "text"
    },
    "text_embedding": {
      "type": "dense_vector",
      "dims": 384,
      "index_options": {
        "type": "bbq_hnsw",
        "rescore_vector": {
          "oversample": 5
        }
      }
    }
  }
}

dense_vector 型のフィールドの設定については、公式ドキュメントを参照してください。

3. モデルの準備

テキストから密ベクトルを生成するためのモデルを準備します。

  • Elasticsearchに内包されているモデルを利用する(例: .multilingual-e5-small_linux-x86_64)

または

  • Elasticsearchの外部にあるモデルを利用する(例: OpenAIのtext-embedding-3-small)

筆者は、前者の.multilingual-e5-small_linux-x86_64を利用しました。モデルの準備の詳細な手順は割愛します。2

4. インジェストパイプラインの作成

準備したモデルを利用するインジェストパイプラインを作成します。

.multilingual-e5-small_linux-x86_64を使う場合の例は以下のとおりです。

PUT /_ingest/pipeline/my-e5-small-pipeline
{
  "description" : "Text embedding pipeline",
  "processors" : [
    {
      "inference": {
        "model_id": ".multilingual-e5-small_linux-x86_64",
        "input_output": [
          {
            "input_field": "content",
            "output_field": "text_embedding"
          }
        ]
      }
    }
  ]
}

“content”“text_embedding” は、上で作成したインデックスのフィールド名と一致させてください。

このリクエストをDev Toolsから実行します。

5. インジェストパイプラインの確認

Embeddingモデルとパイプラインが正しく登録できたか確認します。

以下のリクエストをDev Toolsから実行します。

POST /_ingest/pipeline/my-e5-small-pipeline/_simulate
{
  "docs": [
    {
      "_index": "index",
      "_id": "id",
      "_source": {
        "content": "吾輩は猫である"
      }
    }
  ]
}

正しく登録できていれば、以下のように384次元のfloat32のベクトルが返却されます。

{
  "docs": [
    {
      "doc": {
        "_index": "index",
        "_version": "-3",
        "_id": "id",
        "_source": {
          "text_embedding": [
            0.06494276970624924,
            0.0074624731205403805,
            ...
          ]
        }
      }
    }
  ]
}

6. データの登録

今回は、青空文庫で公開されている「吾輩は猫である」のデータを検証に利用しました。データ登録方法はいくつかありますが、ここでは詳細な手順は割愛し、筆者が行った方法を紹介します。

6.1 NDJSONの用意

検証用のため、青空文庫のデータを1行1チャンクとしました。RAGなどの用途では、より適切なチャンキングを検討してください。

{"chunk_no":1, "content":"一\n吾輩は猫である。名前はまだ無い。"}
...
{"chunk_no":2207, "content":"次第に楽になってくる。...ありがたいありがたい。"}

6.2 一時インデックスへのアップロード

作成した NDJSON を Elastic の File Upload 機能を使って一時的なインデックス(waganeko_tmp)へアップロードします。(Data View は作成しません。)

6.3 _reindexの実行

作成したパイプラインを使って、一時インデックス(waganeko_tmp) からwaganeko_with_bbq_hnsw インデックスへ _reindex します。

POST /_reindex?wait_for_completion=false
{
  "source": {
    "index": "waganeko_tmp"
  },
  "dest": {
    "index": "waganeko_with_bbq_hnsw",
    "pipeline": "my-e5-small-pipeline"
  }
}

時間がかかるので、wait_for_completion=false を指定しておきます。

これにより、上記のリクエストを発行すると、taskid が返却されます。

6.4 タスクの完了確認

タスクが完了したか、以下のコマンドで確認します。

GET /_tasks/{taskid}

6.5 _refreshの実行

タスクが completed になったら、インデックスをリフレッシュします。

POST /waganeko_with_bbq_hnsw/_refresh

登録データのストレージ利用量確認

bbq_hnsw を利用する場合、量子化されたベクトルだけでなく、元のベクトルも保存されていると前回説明しました。実際にストレージ利用量を確認してみましょう。

まず、登録されたドキュメント数を確認します。

GET /waganeko_with_bbq_hnsw/_count

レスポンス

{
  "count": 2207,
  ...
}

2207件のドキュメント、つまり2207個のベクトルが登録されています。(今回は、1ドキュメント = 1ベクトル としているため。)

次に、インデックスのストレージ利用量を調べます。

POST /waganeko_with_bbq_hnsw/_disk_usage?run_expensive_tasks=true

レスポンス

{
  "_shards": {
    "total": 1,
    "successful": 1,
    "failed": 0
  },
  "waganeko_with_bbq_hnsw": {
    ...
    "text_embedding": {
      ...
      "knn_vectors_in_bytes": 3577380
    }
  }
}

knn_vectors_in_bytes に表示された3577380バイトが、ベクトルデータのサイズです。

この値を理論値と比較してみましょう。

項目
ベクトルの総数2207個
次元数384次元
float32ベクトルの1要素のサイズ4バイト
bbq量子化後のベクトルの1要素のサイズ1/8バイト

量子化前のサイズと量子化後のサイズを合計すると…

量子化前のサイズ: 2207 * 384 * 4 = 3,389,952バイト

量子化後のサイズ: 2207 * 384 * 1/8 = 105,936バイト

合計: 3,389,952+105,936=3,495,888バイト

実測値の3,577,380バイトと理論値の3,495,888バイトは非常に近い値です。この結果から、インデックス内には量子化前のベクトルと量子化後のベクトルの両方が保存されていると推測されます。

参考までに、他の量子化方法でも計測した結果を記載しておきます。

量子化方法knn_vectors_in_bytes (実測値)理論値
量子化なし3,398,7803,389,952
bbq_hnsw3,577,3803,495,888
int8_hnsw4,310,9124,237,440

いずれも、理論値に近い実測値となっています。

ベクトル検索の実行

それでは、いよいよbbq_hnswを使ったベクトル検索を実行してみましょう。

rescore_vector を行わないベクトル検索

まずは、rescore_vector を行わない場合の検索結果を確認します。

検索クエリ: “鏡が持ってこられた場所”

取得件数: 10件

Elasticsearch 8.19.0ではデフォルトで rescore_vector が行われるため、”oversample”: 0を指定して意図的に無効化しています。これはあくまで検証目的であり、通常は推奨されません。

GET /waganeko_with_bbq_hnsw/_search
{
    "_source": false,
    "fields": [
        "chunk_no", "content"
    ],
    "knn": {
        "field": "text_embedding",
        "query_vector_builder": {
            "text_embedding": {
                "model_id": ".multilingual-e5-small_linux-x86_64",
                "model_text": "鏡が持ってこられた場所"
            }
        },
        "k": 10,
        "num_candidates": 100,
        "rescore_vector": {
            "oversample": 0
        }
    }
}

このリクエストを実行した結果が以下です。

rankscorechunk_nocontent
10.93101472鏡と云えば風呂場にあるに極まっている。現に吾輩は今朝風呂場でこの鏡を見たのだ。…
20.92831477風呂場にあるべき鏡が、しかも一つしかない鏡が書斎に来ている以上は…
30.92521493鏡は己惚の醸造器であるごとく、…
40.92341471…しかし主人は何のために書斎で鏡などを振り舞わしているのであろう。鏡と云えば風呂場にあるに極まっている。
50.92301786「誰が警察から油壺を貰ってくるものか。…
60.92281458…図書館もしくは博物館へ馳けつけて、…
70.92262009「…僕が昔し姥子の温泉に行って、一人のじじいと相宿になった事がある。…
80.92231478…なにさ今鏡を造ろうと思うて一生懸命にやっておるところじゃと答えた。
90.92151480かくとも知らぬ主人ははなはだ熱心なる容子をもって一張来の鏡を見つめている。…
100.92061263…不用心だって金のないところに盗難のあるはずはない。…

最も関連性の高いと思われるchunk_no = 1477が2位になっており、全体的に検索結果の精度が微妙な印象です。

rescore_vector を行ったベクトル検索

次に、検索結果の上位50件(k(=10) * oversample(=5))を元のベクトルでrescoreしてみます。rescore_vector を指定することで、元のベクトルを使った再評価が行われます。

GET /waganeko_with_bbq_hnsw/_search
{
    "_source": false,
    "fields": [
        "chunk_no", "content"
    ],
    "knn": {
        "field": "text_embedding",
        "query_vector_builder": {
            "text_embedding": {
                "model_id": ".multilingual-e5-small_linux-x86_64",
                "model_text": "鏡が持ってこられた場所"
            }
        },
        "k": 10,
        "num_candidates": 100,
        "rescore_vector": {
            "oversample": 5
        }
    }
}

※mappings作成時に “rescore_vector”: { “oversample”: 5 } を指定しているので、本来は省略可能ですが、ここでは、わかりやすいよう明示しています。

rescore_vector の動作:

公式ドキュメントによると、このリクエストを実行した場合、以下の流れで検索が行われます。

(1) 各シャードごとに上位 num_candidates 件(今回は100件)のドキュメントを bbq_hnsw を使って取得します。

(2) そのうち、上位(k × oversample)件(今回は10 * 5=50件)のドキュメントを、元のベクトルを使って再評価(rescore)します。

(3) 再評価されたドキュメントの中から、最終的な上位k件(今回は10件)を取得します。

(4) 複数のシャードがある場合は、(1) – (3) をシャードごとに行い、最終的にマージして上位10件を取得します。今回はシャードが1つなので、この手順は省略されます。

rescore_vector を行った結果が以下です。

rankscorechunk_nocontent
10.92881477風呂場にあるべき鏡が、しかも一つしかない鏡が書斎に来ている以上は…
20.92481493鏡は己惚の醸造器であるごとく、…
30.92191480かくとも知らぬ主人ははなはだ熱心なる容子をもって一張来の鏡を見つめている。…
40.92151472鏡と云えば風呂場にあるに極まっている。現に吾輩は今朝風呂場でこの鏡を見たのだ。…
50.92141786「誰が警察から油壺を貰ってくるものか。…
60.92121684…昨日は鏡の手前もある事だから、おとなしく独乙皇帝陛下の真似をして整列したのであるが、…
70.9206633「…今日は土曜ですからこれから廻ったら、もう帰っておりましょう。…
80.92051471…しかし主人は何のために書斎で鏡などを振り舞わしているのであろう。鏡と云えば風呂場にあるに極まっている。
90.91971505…例のごとく赤い手をぬっと書斎の中へ出した。右手に髯をつかみ、左手に鏡を持った主人は、そのまま入口の方を振りかえる。…
100.9184533行きたいところへ行って聞きたい話を聞いて、舌を出し尻尾を掉って、髭をぴんと立てて悠々と帰るのみである。…

検索結果を見ると、最も関連性の高いドキュメント(chunk_no = 1477)が1位に上がっています。他のドキュメントも、rescoreなしの場合と比べて、より関連性の高いものが上位に来ている印象です。

この結果から、bbq による量子化を行った場合でも rescore_vector を使用することで、良好な検索結果が得られることがわかります。

なお、ベクトル値を再評価する方法は、rescore_vector 以外にも rescore セクションを指定する方法などがあります。詳しくは公式ドキュメントを参照してください。

まとめ

  • bbq_hnsw で密ベクトルを格納した場合、インデックス内には量子化前のベクトルと量子化されたベクトルの両方が保持されていることを実測値と理論値の比較で確認しました。
  • bbq_hnsw を利用した密ベクトル検索では、rescore_vectorを使うことで、検索精度を向上させられることを実証しました。


  1. Elasticsearch 9.1.0 は、デフォルトで directio を使用する、という仕様になっています。bbq_hnsw を利用時でメモリが十分に足りている場合は directio を使用しない方がいいケースもあるため、directio を使用しないよう設定する、あるいは Elasticsearch 9.1.1 以上へバージョンアップする、なども検討してみてください。
    https://www.elastic.co/docs/release-notes/elasticsearch/known-issues#elasticsearch-9.1.0-known-issues
    https://www.elastic.co/search-labs/blog/knn-vector-search-rescoring-direct-io
    ↩︎
  2. .multilingual-e5-small_linux-x86_64 を利用する準備の手順は、下記を参考にしてください。
    https://elastic.sios.jp/blog/preparing-for-vector-search/
    ↩︎