【langchain】FAISSの距離尺度をコサイン類似度にしたい場合はdistance_strategyにMAX_INNER_PRODUCTを指定する

概要

langchainの埋め込み類似度計算を行うクラスの一つであるFAISSでは、デフォルトの距離尺度がL2となっています。距離尺度をコサイン類似度にする方法がよくわからなかったので調べました。結論としては、インスタンス生成時の引数でdistance_strategy="MAX_INNER_PRODUCT"およびnormalize_L2=Trueを指定するとよいようです。

詳細

langchain==0.0.304で検証します。

langchainのFAISS.similarity_search_with_scoreで類似度検索を実施してみます。埋め込みモデルはoshizoさんの日本語lukeモデルを使わせていただきました。

類似度の指標は、特に指定しない場合は、L2距離が使われます。「こんにちは」同士の類似度が0でないのが少し気になりますが、浮動小数点の丸め誤差と思われます。スコアだけではL2距離かどうかは本当はわからないのですが、numpyなどで直接計算した結果と比較すると、L2距離だとわかります。

from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.vectorstores.utils import DistanceStrategy

embedding = HuggingFaceEmbeddings(model_name = "oshizo/sbert-jsnli-luke-japanese-base-lite", encode_kwargs={"normalize_embeddings":True })
vectorstore = FAISS.from_texts(
    ["こんにちは","こんばんは","さようなら"]
    , embedding
)
vectorstore.similarity_search_with_score("こんにちは")

[(Document(page_content='こんにちは', metadata={}), 2.1613458e-12),
 (Document(page_content='こんばんは', metadata={}), 0.41432357),
 (Document(page_content='さようなら', metadata={}), 1.0592083)]

類似度の指標をコサイン類似度にすることを試みます。
FAISSはインスタンス生成時にdistance_strategyという引数を与えることができます。ここにコサインを指定してみます。

vectorstore = FAISS.from_texts(
    ["こんにちは","こんばんは","さようなら"]
    , embedding
    , distance_strategy = DistanceStrategy.COSINE
)
vectorstore.similarity_search_with_score("こんにちは")

[(Document(page_content='こんにちは', metadata={}), 2.1613458e-12),
 (Document(page_content='こんばんは', metadata={}), 0.41432357),
 (Document(page_content='さようなら', metadata={}), 1.0592083)]

指定しない場合と全く同じ結果になりました。つまりL2距離のままです。

試しに、コサイン類似度の仲間である内積をdistance_strategyに指定してみます。

vectorstore = FAISS.from_texts(
    ["こんにちは","こんばんは","さようなら"]
    , embedding
    , distance_strategy = DistanceStrategy.MAX_INNER_PRODUCT
)
vectorstore.similarity_search_with_score("こんにちは")

[(Document(page_content='こんにちは', metadata={}), 1.0000001),
 (Document(page_content='こんばんは', metadata={}), 0.7928383),
 (Document(page_content='さようなら', metadata={}), 0.47039595)]

今度はコサイン類似度っぽい値になりました。numpyによる計算結果と比べても、確かにこれがコサイン類似度のようです（埋め込みの大きさを1に正規化しているので、コサイン類似度と内積は同じになります。）。

ということで、類似度の指標にコサインを使いたい場合は、COSINEではなくMAX_INNTER_PRODUCTを指定する必要があるようです。
釈然としません（langchainのバグ？）が、一応、解決策はわかりました。

付録

原因の調査

ここからはオマケみたいなものですが、langchainの実装がどうなっているのか深掘りしました。

まず、langchainのFAISSは、facebookが開発したベクトル検索ライブラリfaissのラッパークラスであり、内部ではfaissを使っています。内部のfaissの定義がdistance_strategyに応じてどう変わるか確認しておきます。

vectorstore = FAISS.from_texts(
    ["こんにちは","こんばんは","さようなら"]
    , embedding
)
print("DEFAULT", type(vectorstore.index))

vectorstore = FAISS.from_texts(
    ["こんにちは","こんばんは","さようなら"]
    , embedding
    , distance_strategy = DistanceStrategy.COSINE
)
print("COSINE", type(vectorstore.index))
vectorstore = FAISS.from_texts(
    ["こんにちは","こんばんは","さようなら"]
    , embedding
    , distance_strategy = DistanceStrategy.MAX_INNER_PRODUCT
)
print("MAX_INNER_PRODUCT", type(vectorstore.index))

DEFAULT <class 'faiss.swigfaiss.IndexFlatL2'>
COSINE <class 'faiss.swigfaiss.IndexFlatL2'>
MAX_INNER_PRODUCT <class 'faiss.swigfaiss.IndexFlatIP'>

IndexFlatL2はL2、IndexFlatIPはコサイン類似度に対応するクラスですので(参考)、やはりCOSINEを指定してもL2のままになっているようです。

次にlangchainのFAISSのソースコードを見てみます。

FAISS.from_textsはdistance_strategy引数をcls.__fromに渡しているので、FAISS.__fromをみます。以下のような記述があり、MAX_INNER_PRODUCTが指定された場合のみIndexFlatIPを、それ以外はCOSINEも含めて、IndexFlatL2を使うように定義されていることがわかります。

        if distance_strategy == DistanceStrategy.MAX_INNER_PRODUCT:
            index = faiss.IndexFlatIP(len(embeddings[0]))
        else:
            # Default to L2, currently other metric types not initialized.
            index = faiss.IndexFlatL2(len(embeddings[0]))

ということで、やはりMAX_INNER_PRODUCTの場合のみ、コサイン類似度を使う、という指標になっているようです。

バグなのか仕様なのかは微妙なところです。
L2とMAX_INNER_PRODUCTの２種類のみしか想定しないのであればこのコードで特に違和感はないのですが、DistanceStrategy.COSINEという定数がわざわざlangchainのコードに用意されてしまっているので、少し混乱させる実装な気はしますね。

他の人、あるいは、公式はどう思っているのか気になってい、githubのディスカッションを検索してみると、FAISSでコサイン類似度を使いたい場合は、「MAX_INNER_PRODUCTを指定して、normalize_L2=Trueにしてくれ」という回答が書かれていました。この記事では、HuggingFaceEmbedding側でencode時にnormalizeすることで内積をコサインにしていましたが、FAISS側でnormalize_L2を指定するほうがコードとしては確かにシンプルですね。

一応やってみます。
HuggingFaceEmbeddingsの初期化時に、encode_kwargsを渡していたのをやめます。するとただの内積なので１よりも大きいスコアになります。

embedding = HuggingFaceEmbeddings(model_name = "oshizo/sbert-jsnli-luke-japanese-base-lite")
vectorstore = FAISS.from_texts(
    ["こんにちは","こんばんは","さようなら"]
    , embedding
    , distance_strategy = DistanceStrategy.MAX_INNER_PRODUCT
)
vectorstore.similarity_search_with_score("こんにちは")

[(Document(page_content='こんにちは', metadata={}), 113.14021),
 (Document(page_content='こんばんは', metadata={}), 89.48091),
 (Document(page_content='さようなら', metadata={}), 49.38093)]

normalize_L2 = Trueを指定すると0から1の範囲に収まるようになります。つまりコサイン類似度になります。

embedding = HuggingFaceEmbeddings(model_name = "oshizo/sbert-jsnli-luke-japanese-base-lite")
vectorstore = FAISS.from_texts(
    ["こんにちは","こんばんは","さようなら"]
    , embedding
    , distance_strategy = DistanceStrategy.MAX_INNER_PRODUCT
    , normalize_L2 = True
)
vectorstore.similarity_search_with_score("こんにちは")

[(Document(page_content='こんにちは', metadata={}), 1.0),
 (Document(page_content='こんばんは', metadata={}), 0.7928383),
 (Document(page_content='さようなら', metadata={}), 0.47039598)]

ちなみに前述のディスカッションは「MAX_INNER_PRODUCTとCOSINEの違いは何？」という質問が最後にされていて、その返答がない状態になっています。

DistanceStrategy.DOT_PRODUCT

MAX_INNER_PRODUCTの他にDOT_PRODUCTというDistanceStrategyもあるようです。
ただ上記で説明したコードを読む限り、やはりMAX_INNER_PRODUCT以外は全部L2で計算されることになりそうです。

numpy計算結果との照合

FAISS.similarity_search_with_scoreの戻り値が本当にL2またはコサイン類似度なのか、一応確かめておきます。

import numpy as np
from langchain.embeddings import HuggingFaceEmbeddings
embedding = HuggingFaceEmbeddings(model_name = "oshizo/sbert-jsnli-luke-japanese-base-lite", encode_kwargs={"normalize_embeddings":True })
docs = ["こんにちは","こんばんは","さようなら"]
query = "こんにちは"
doc_embs = embedding.embed_documents(docs)
query_emb = embedding.embed_query(query)

doc_emb_arrays = [np.array(emb) for emb in doc_embs]
query_emb_array = np.array(query_emb)


# l2距離    
for i, doc_emb_array in enumerate(doc_emb_arrays):
    l2_distance = np.linalg.norm(query_emb_array - doc_emb_array) ** 2
    print(f"l2: {docs[i]}, {l2_distance}")

# コサイン類似度    
for i, doc_emb_array in enumerate(doc_emb_arrays):
    cosine_similarity = np.dot(query_emb_array, doc_emb_array) / np.linalg.norm(query_emb_array) / np.linalg.norm(doc_emb_array)
    print(f"cosine similarity: {docs[i]}, {cosine_similarity}")

l2: こんにちは, 2.1613458169806644e-12
l2: こんばんは, 0.4143235983727727
l2: さようなら, 1.059208282091555
cosine similarity: こんにちは, 0.9999999999989246
cosine similarity: こんばんは, 0.7928382138498853
cosine similarity: さようなら, 0.47039589941306786

大体一致しました。L2は２乗すると値が近くなったので、2乗距離みたいですね。

おわりに

FAISSの使い方で混乱したので、整理してみました。
似たような混乱を生じた方のお役に立てれば幸いです。