この記事はスタンバイAdvent Calaendar 2021の7日目の記事です。
はじめに
ANNとは
与えられた集合の中で与えられた点に最も近い点を見つけるという最適化問題である最近傍探索(Nearest Neighbor Search)というものがあります。また、近い順にk個の解をみつける場合は、kNN Searchと言います。ANN(Approximate Nearest Neighbor: 近似最近傍探索)とは、kNN Searchの一種で、厳密な最近傍を求めず、ある程度の精度で解を見つけることを許容することで、高速に解を求める手法です。
LuceneによるANN Searchのサポート
最近になってANN SearchがApache Luceneでサポートされました1。
LuceneはANNを今後のリリース予定のバージョン9.0で提供する予定です。
なぜ、全文検索エンジンであるApache LuceneでANN Searchがサポートされるようになったのでしょうか?Luceneのissueには以下のような記載があります1。
"Semantic" search based on machine-learned vector "embeddings" representing terms, queries and documents is becoming a must-have feature for a modern search engine
自然言語処理における、Embedding(埋め込み)とは文章や単語などをベクトル表現に変換することです。具体的にはニューラルネットワークなどの機械学習によりベクトル空間にマッピングします。このようにして得られたベクトルは、単語や文章の意味を反映していることが期待されます。すなわち、似た意味をもつ文章や単語のベクトルは類似度が高くなります。この性質を用いることでSemantic Searchを実現することができます。
EmbeddingによるSemantic SearchとANN Search
Semantic Searchとはユーザの検索クエリの意図や意味を理解した上で結果を提供する検索です。通常の検索では文字列のマッチングで検索結果を返します。この場合、検索対象のドキュメントにクエリで指定された単語が含まれていなければマッチしません。例えば、検索クエリがバイト
の場合、同義語辞書などの工夫をしないと似た意味のパート
を含むドキュメントにマッチさせることができません。Semantic Searchにより検索クエリの意図を理解できれば、クエリで表現されていないユーザの意図にあった検索結果を返すことができます。
先ほど、Embeddingにより得られたベクトル表現は、似た意味のものであれば類似度が高くなると説明しました。上の例で言えば、意味的に近いバイト
とパート
はベクトル表現でも近くなると考えられます。そこで、検索対象のドキュメントとクエリをそれぞれベクトル表現にし、その類似度を計算することでクエリとドキュメントが意味的にマッチしているものを取得することができます。
ここでANNの話に戻りますが、ドキュメントのベクトルごとにクエリとの類似度を計算していくと、全てのドキュメントに対して計算をする必要があるため時間がかかります。検索エンジンの多くのユースケースでは検索にかかかる時間は短い方が好ましいです。そのため、高速に検索結果を取得できるようにANNが用いられることになります。
実際に、このような仕組みでSemantic Searchを実現している企業が複数存在し、論文も公開されています23。
Elasticsearchにおける実装
現在Luceneで実装されたANN SearchをElasticsearchでも利用できるように開発が進められています4。Luceneは次期9.0でANNをリリースするため、この機能はElasticsearch 8.xのみが対象となっています。
ANN Searchにはさまざまなアルゴリズムが存在します5が、LuceneではHNSW
6というアルゴリズムを使用されているため、ElasticsearchでもHNSW
での実装が進められています。
ANN Searchでは高速化のために、検索に適した形でデータを保持することがあります。HNSW
ではデータを階層的なグラフとして表現することで高層化を実現しています。
HNSWの詳細や他のアルゴリズムを知りたい方は、以下の資料がわかりやすいです。
https://speakerdeck.com/matsui_528/jin-si-zui-jin-bang-tan-suo-falsezui-qian-xian
ここからは、実際にElasticsearchでANN Searchを行う方法を紹介します。
以下の内容は主にこちらのissueの内容を参考にして作成しました。
https://github.com/elastic/elasticsearch/issues/78473
Elasticsearchの設定をする
はじめに、masterブランチのElasticsearchをビルドします。
2021/11月現在、ElasticsearchのANN Searchはリリース前なので、機能を試したい場合、masterブランチで動かす必要があります。
git clone https://github.com/elastic/elasticsearch.git
cd elasticsearch
./gradlew localDistro
cd distribution/local/elasticsearch-8.1.0-SNAPSHOT/bin
./elasticseaerch
Elasticsearchが起動したら、インデックスの設定をおこないます。
以下のようなHTTPリクエストをElasticsearchに送ることで設定ができます。
PUT http://127.0.0.1:9200/ann_sample
Content-Type: application/json
{
"mappings": {
"properties": {
"text_vector": {
"type": "dense_vector",
"dims": 768,
"index": true,
"similarity": "l2_norm",
"index_options": {
"type" : "hnsw",
"m" : 15,
"ef_construction" : 50
}
},
"text" : {
"type" : "keyword"
}
}
}
}
このインデックスはtext
とtext_vector
という二つのフィールドを持ちます。
text
は検索したい文章をそのまま保持するフィールドです。text_vector
は文書ベクトルを保持するフィールドです。
text_vector
フィールドはdense_vector
として表現されます。
dense_vector
フィールドは、現在のバージョンのElasticsearchで既にリリースされているフィールドです。
しかしながら、これまでのdense_vector
フィールドは、単にベクトルを保存できるだけであり、クエリ、ソートをサポートされておらず、スクリプト内で専用のベクトル関数を使ってのみアクセスが許可されていたため、cosine類似度をスクリプトで計算するといった形で使用することしかできませんでした。
https://www.elastic.co/guide/en/elasticsearch/reference/7.15/dense-vector.html
ANN Searchの機能はdense_vector
の拡張という形で提供されています。
既存のdense_vector
フィールドではdims
というベクトルの次元数を表すパラメータのみが存在していましたが、
ANNをサポートするためにdense_vector
フィールドに新たなパラメータが追加されています。
-
index
-
true
に設定すると、そのフィールドをANNインデックスに追加することになります - デフォルトでは互換勢担保のため
false
になっています
-
-
similarity
-
index
がtrue
の場合、ベクトルのインデックスを作成する際に使用する類似性を指定する必要があります。今のところ、l2_norm
とdot_product
が指定できます。これはLucene
のオプションと一致しています。
-
-
index_options
-
type
- 使用するkNNアルゴリズム
- 現状は
hnsw
のみサポート
-
m
- HNSWグラフにおける各ノードの最大接続数
-
ef_construction
-
HNSW
用のパラメータ - 新たに挿入された各ノードのグラフを検索する際に追跡する候補数
- デフォルトは100
-
-
データをElasticsearchにインデックスする
データセットのダウンロード
今回はデータセットとしてlivedoorニュースコーパスを使用しました。
livedoorニュースコーパスはlivedoor NEWSに掲載されたニュース記事を収集したもので、株式会社ロンウィットにより公開されています7。
# ダウンロード
wget https://www.rondhuit.com/download/ldcc-20140209.targ.gz
# ファイルの解答
tar -zxf ldcc-20140209.tar.gz
文書ベクトルの生成とインデクシング
今回文章をベクトル表現にするに当たってBERTを使用しました。BERTはGoogleによって提案されたモデルで、さまざまな言語タスクで高い性能を発揮しています8。
まず、BERTを利用するにあたって必要なライブラリをインストールします。
pip install -r requirements.txt
transformers==4.5.0
fugashi==1.1.0
ipadic==1.0.0
torch==1.8.0
elasticsearch==7.15.2
今回はPyTorch
とTransformers
を使用します。
PyTorch
はFacebookが開発した深層学習のフレームワークで、Transformers
はさまざまなニューラルネットワークを用いた言語モデルが実装されているライブラリです。Transformers
では、さまざまな言語の事前学習モデルが利用でき、日本語のモデルもあります。
以下のindexer.py
は文書ベクトルの生成とインデクシングを行うコードです。
import glob
from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk
import torch
from transformers import BertJapaneseTokenizer, BertModel
from joblib import Parallel, delayed
model_name = 'cl-tohoku/bert-base-japanese-whole-word-masking'
bert = BertModel.from_pretrained(model_name)
tokenizer = BertJapaneseTokenizer.from_pretrained(model_name)
max_length = 256
client = Elasticsearch()
INDEX_NAME = "ann_sample"
def index_batch(doc):
requests = get_request(doc)
print(requests)
bulk(client, requests)
def get_request(doc):
return [{
"_op_type": "index",
"_index": INDEX_NAME,
"text": doc,
"text_vector": embedding(doc)
}]
def embedding(text):
encoding = tokenizer(
text,
max_length=max_length,
padding='max_length',
truncation=True,
return_tensors='pt'
)
attention_mask = encoding['attention_mask']
# 文章ベクトルを計算
# BERTの最終層の出力の平均を計算する。(ただし、[PAD]は除く)
with torch.no_grad():
output = bert(**encoding)
last_hidden_state = output.last_hidden_state
averaged_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(1) \
/ attention_mask.sum(1, keepdim=True)
return averaged_hidden_state[0].tolist()
if __name__ == '__main__':
BATCH_SIZE = 1000
docs = []
count = 0
category_list = [
'dokujo-tsushin',
'it-life-hack',
'kaden-channel',
'livedorr-homme',
'movie-enter',
'peachy',
'smax',
'sports-watch',
'topic-news'
]
for category in category_list:
for file in glob.glob(f"../data/text/{category}/{category}*"):
lines = open(file).read().splitlines()
print(lines[0])
text = '\n'.join(lines[3:])
index_batch(text)
def embedding(text):
で文章のベクトル表現を生成しています。
BERTの利用にあたっては、以下の本のコードを参考にさせていただきました。
https://github.com/stockmarkteam/bert-book/blob/master/Chapter10.ipynb
ANNで検索する
最後に、作成したインデックスを使ってANN Searchを行ってみます。
次のようなリクエストを/_knn_search
というエンドポイントに投げることで検索ができます。
GET http://127.0.0.1:9200/ann_sample/_knn_search
Content-Type: application/json
{
"knn": {
"field": "text_vector",
"query_vector": [0.22..., ....]
"k": 10,
"num_candidates": 100
},
"_source": false,
"fields": ["text"]
}
各パラメータの説明は以下になります。
-
query_vector
- クエリのベクトル
-
k
- 取得する結果の件数
-
num_candidates
- 各シャード上ごとに候補ベクトルを何件集めるか
k
とnum_candidates
の関係ですが、まず結果を収集するために、
各シャード上の近似最近傍候補のベクトルをnum_candidates
個見つけます。これらの候補ベクトルとクエリベクトルの類似性を計算し、各シャードからk
このもっとも類似した結果を選択します。最後に各シャードからの結果をマージしてグローバルなトップk個の候補を返却します。
以下のsearcher.pyはann_searchを行うコードです。
def handle_query():
query = input("Enter query: ")
# script queryでdense_vectorを使う場合のクエリ
# script_query = {
# "script_score": {
# "query": {"match_all": {}},
# "script": {
# "source":"cosineSimilarity(params.query_vector, doc['text_vector']) + 1.0",
# "params":{"query_vector": query_vector}
# }
# }
# }
es_query = {
"knn": {
"field": "text_vector",
"query_vector": embedding(query),
"k": 10,
"num_candidates": 100
},
"_source": False,
"fields": ["text"]
}
url = 'http://127.0.0.1:9200/livedoor/_knn_search'
headers = {
'Content-Type': 'application/json',
}
req = urllib.request.Request(url, json.dumps(es_query).encode(), headers)
with urllib.request.urlopen(req) as res:
body = res.read().decode()
print(body)
検索結果のレスポンスのフォーマットは通常の_search
と同じになります。
まとめ
Elasticsearchに新しく追加された ANN Searchを試してみました。
ここで紹介した機能は、まだ開発中であり変わる可能性があるため、ご注意ください。
今回は単純に機能を使ってみるだけでしたが、今後はEmbeddingの方法を変更したり、パラメータを調整したりして精度を改善できないか試してみたいです。また、性能比較をしてANNによりどれだけ高速化されたかをみてみようと思います。