Elasticsearchでコサイン類似度を求める

M2のmarutakuです。カレンダーが全然埋まらなくて困ってます。どんどん参加登録してくださいお願いします。
今回は、過去に自分が作ろうとした画像検索のシステムでElasticsearchを使ってコサイン類似度を求めた話をします。

コサイン類似度について

コサイン類似度は、ベクトル同士の類似度を比較するために用いられます。数式は他のサイトでよく開設されているので割愛します。Word2Vecで、単語の類似度を測るときに使われています。gensimを使っているときは、gensim.models.Word2Vec.most_similarでコサイン類似度を求めると思います。しかし、Pythonでコサイン類似度を比較する処理を書くと遅くなりがちな気がしています。
そんな時にElasticsearchの登場です。

Elasticsearchについて

Elasticsearchはオープンソースの検索エンジンです。 大量のデータから高速に検索を行うことができます。全文検索や、ログ収集・可視化などに用いられています。Dockerの環境も用意されていて、手軽に試せるので是非試してみてください。

https://www.elastic.co/jp/blog/text-similarity-search-with-vectors-in-elasticsearch
上記に記事にもあるように、Elasticsearch7.0からベクトルを格納することができるようになり、Elasticsearch7.3からベクトルのドキュメントスコアリング向けの使用をサポートはじめました。これにより、Elasticsearch内でコサイン類似度を算出できるようになりました。

Elasticsearchを用意する

ElasticsearhはDockerイメージが用意されているので、今回はそのイメージを使用したいと思います。以下はdocker-composeで使用するdocker-compose.ymlの内容です。

version: "2.3"
services:
    elasticsearch:
        image: "elasticsearch:7.5.1"
        ports:
            - "9200:9200"
        volumes:
            - es-data:/usr/share/elasticsearch/data
        tty: true
        environment:
            discovery.type: single-node
volumes:
    es-data:
        driver: local

Elasticsearchでは、MySQLでテーブルを定義するように、あらかじめデータの構造を定義する必要があります。この定義のことをMappingと呼びます。以下は、過去に自分が画像の検索システムを作る時に使ったMappingです。
ベクトルはdense_vectorというフィールドタイプで、ベクトル長はあらかじめ指定しておく必要があります。

{
    "mappings": {
        "properties": {
            "item_url": {
                "type": "text"
            },
            "item_description": {
                "type": "text"
            },
            "image_vector": {
                "type": "dense_vector",
                "dims": 512
            }
        }
    }
}

Pythonでデータを挿入する

Pythonを使ってデータを入れてみようと思います。
ElasticsearchをPythonで操作するためのモジュールが提供されているので、まずはそれをインストールします。

pip install elasticsearch

Elasticsearchのなかにデータを入れます。

from elasticsearch import Elasticsearch

INDEX_NAME = 'test_index'
es_client = Elasticsearch(host=es_host, port=es_port)

def insert(tensor, item_url, item_description):
    es_client.index(index=INDEX_NAME, body={
        "item_url": item_url, "item_description": item_description, "image_vector": tensor
    })

検索してみる

データを格納することができたら実際に検索をしてみます。

def search(vector):
        query = {
            "script_score": {
                "query": {"match_all": {}},
                "script": {
                    "source": "cosineSimilarity(params.image_vector, doc['image_vector']) + 1.0",
                    "params": {"image_vector": vector}
                }
            }
        }
        response = es_client.search(
            index=INDEX_NAME,
            body={
                "size": 5,
                "query": query,
                "_source": {"includes": ["item_url", "item_description"]}
            }
        )
        return response

これで検索することができます。
自分の環境では、約1万件の512次元のベクトルから類似したものを検索するのに40ms程度でした。多分とっても速いと思います。

終わりに

今回はElasticsearchでコサイン類似度を求めてみました。
Pythonでやるよりも高速なので是非試してみてください。