FAISS Hands-On

テストコードを動かして以下の4つの検索手法をなんとなく理解することを目的とします

IndexFlatL2
IndexFlatIP
IndexIVFFlat
IndexHNSW

事前準備

pip install faiss-cpu

Embedding

FAISSに格納するデータをembeddingする

embedding用のコードを以下にしめす。後にFAISSに格納するembedding対象の日本語は以下の5つである

"好きな食べ物は何ですか?",
"どこにお住まいですか?",
"朝の電車は混みますね",
"今日は良いお天気ですね",
"最近景気悪いですね",

# 必要なライブラリをインポートします。
import torch.nn.functional as F
from torch import Tensor
import faiss  # FAISSライブラリをインポート
import numpy as np  # NumPyライブラリをインポート
from transformers import AutoTokenizer, AutoModel

# 最後の隠れ層の状態を平均プーリングする関数を定義します。
def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    # print(attention_mask) # 情報があるところが1, paddingが0で表示
    # print(attention_mask[..., None]) # 次元を一つ追加して3次元にしている
    # print(attention_mask[..., None].bool()) # 1をTrue, 0をFalseへ変換
    # print(~attention_mask[..., None].bool()) # TrueとFalseを入れ替え (paddingがTrueになる)
    
    # last_hidden_statesの値におけるpaddingの値が0.0に置き換えられる
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)

    # last_hidden.sum(dim=1)でシーケンシャル方向に和を取る
    # つまり、[<batch_size>, <sequence_length>, <隠れ層の次元数>] => [<batch_size>, <隠れ層の次元数>] に変換
    # attention_mask.sum(dim=1)[..., None]は、padding除いたToken数の和。この値で割って平均化する
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

# 解析したいテキストのリストを定義します。
input_texts = [
    "好きな食べ物は何ですか?",
    "どこにお住まいですか?",
    "朝の電車は混みますね",
    "今日は良いお天気ですね",
    "最近景気悪いですね",
    #"最近、出かけていないので、たまには外で食事でもどうですか？"
]

# 使用するモデルとトークナイザーの事前学習済みのパスを指定してロードします。
tokenizer = AutoTokenizer.from_pretrained('intfloat/multilingual-e5-large')
model = AutoModel.from_pretrained('intfloat/multilingual-e5-large')

# テキストをモデルが処理できる形式にトークン化します。
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

# batch_dict内の各テンソルのサイズを確認する
for i, (key, tensor) in enumerate(batch_dict.items()):
    print(f"[{i}][batch_dict] {key}: {tensor.size()}")

# モデルを実行して出力を取得します。
outputs = model(**batch_dict)

# batch_dict内の各テンソルのサイズを確認する
for i, (key, tensor) in enumerate(outputs.items()):
    print(f"[{i}][outputs] {key}: {tensor.size()}")

# `average_pool`関数を使って、最終的な埋め込みベクトルを計算します。
# print("attention_mask:", batch_dict['attention_mask'].sum(dim=1))
# print("outputs.last_hidden_state.sum(dim=1).shape:", outputs.last_hidden_state.sum(dim=1).shape)
# print("last_hidden_states:", outputs.last_hidden_state.sum(dim=1))
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

# PyTorchのTensorからNumPy配列に変換します。
embeddings_np = embeddings.cpu().detach().numpy()
print("embeddings_np:", embeddings_np.shape)

# わかりやすさのための表示
for i, vec in enumerate(embeddings_np[:]):
    print(f"----{i}----")
    print(input_texts[i])
    print(vec)

以下のような結果を得て、Inputした文章が1024次元のベクトル化されていることが分かる

[0][batch_dict] input_ids: torch.Size([5, 10])
[1][batch_dict] attention_mask: torch.Size([5, 10])
[0][outputs] last_hidden_state: torch.Size([5, 10, 1024])
[1][outputs] pooler_output: torch.Size([5, 1024])
embeddings_np: (5, 1024)
----0----
好きな食べ物は何ですか?
[ 0.6108155  -0.2260578  -0.10909441 ... -0.62022376 -0.83260953
  0.28850102]
----1----
どこにお住まいですか?
[ 0.99075234  0.6782665  -0.68046105 ... -0.7541518  -1.3915895
 -0.14331307]
----2----
朝の電車は混みますね
[ 1.1466048   0.16233358 -0.02154386 ...  0.07610887 -0.13563985
  0.2248166 ]
----3----
今日は良いお天気ですね
[ 0.9829717   0.06835324 -0.34694287 ...  0.00633434 -0.7837953
 -0.28379193]
----4----
最近景気悪いですね
[ 0.26057956  0.5884724  -0.49848753 ... -0.19543631 -0.8918747
 -1.4221884 ]

query用文章のembedding

同様にquery用の文章をembeddingしておきます。クエリ用の日本語は "今日は雨が振らなくてよかった"です。

# クエリテキストを定義します。
query_text = ["今日は雨が振らなくてよかった"]

# クエリテキストをトークナイズし、モデルが処理できる形式にします。
query_dict = tokenizer(query_text, max_length=512, padding=True, truncation=True, return_tensors='pt')

# モデルを実行して出力を取得します。
query_output = model(**query_dict)

# `average_pool`関数を使って、クエリの埋め込みベクトルを計算します。
query_embeddings = average_pool(query_output.last_hidden_state, query_dict['attention_mask'])

# クエリの埋め込みベクトルをNumPy配列に変換します。
query_embeddings_np = query_embeddings.cpu().detach().numpy()

FAISSへのデータ格納

今回は以下の4つの方法でデータを格納した。

IndexFlatL2
- ベクトル間のユークリッド距離（L2距離）を使用して類似性を計測します。全データ点に対して線形探索を行うため、計算は正確ですが、大規模データセットでは計算コストが高くなります。
IndexFlatIP
- いわゆるcos類似度でありベクトル間の内積を用いて類似性を評価します。この指標は類似度が高いほど大きな値となるため、類似性の高いベクトルを効果的に特定できますが、大きなデータセットには適していません。
IndexIVFFlat
- 先にデータをクラスタリングして索引を作成し、検索時には近似的な検索を行います。これにより、IndexFlatL2やIndexFlatIPと比較して高速な検索が可能ですが、精度は若干落ちます。
IndexHNSW
- 階層的ナビゲータブル小世界（Hierarchical Navigable Small World）グラフを使用し、近似最近傍探索を非常に高速に行います。大規模なデータセットに適しており、距離計算の回数を劇的に削減することができます。

# ベクトルの次元数を設定
dimension = embeddings_np.shape[1]

# ベクトルを正規化
normalized_embeddings_np = embeddings_np / np.linalg.norm(embeddings_np, axis=1, keepdims=True)

# IndexFlatL2インデックスの初期化とデータ追加
index_flat_l2 = faiss.IndexFlatL2(dimension)
index_flat_l2.add(embeddings_np)  # 正規化なしで追加
print("Number of vectors in the IndexFlatL2:", index_flat_l2.ntotal)

# IndexFlatIP（内積を使用するインデックス）の初期化と正規化されたデータの追加
index_flat_ip = faiss.IndexFlatIP(dimension)
index_flat_ip.add(normalized_embeddings_np)  # 正規化されたデータを追加
print("Number of vectors in the IndexFlatIP:", index_flat_ip.ntotal)

# IndexIVFFlat（量子化されたインデックス）の初期化
nlist = 4  # クラスタ数
quantizer = faiss.IndexFlatL2(dimension)  # 量子化器としてIndexFlatL2を使用
index_ivf_flat = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_L2)
index_ivf_flat.train(embeddings_np)
index_ivf_flat.add(embeddings_np)
print("Number of vectors in the IndexIVFFlat:", index_ivf_flat.ntotal)

# IndexHNSW（階層的ナビゲーティブ小世界グラフを使用するインデックス）の初期化とデータ追加
index_hnsw = faiss.IndexHNSWFlat(dimension, 10)  # ここで10はHNSWのグラフの近隣接続数
index_hnsw.add(embeddings_np)
print("Number of vectors in the IndexHNSW:", index_hnsw.ntotal)

Queryの実行

# 上位3つの近いベクトルを検索するための設定（k=3）
k = 2

# クエリの処理（以前のコードからのクエリテキストの部分を続ける）
query_embeddings_np = query_embeddings.cpu().detach().numpy()
# クエリの埋め込みベクトルも正規化
normalized_query_embeddings_np = query_embeddings_np / np.linalg.norm(query_embeddings_np, axis=1, keepdims=True)

# 各インデックスに対して検索を実行し、結果を取得
D_l2, I_l2 = index_flat_l2.search(query_embeddings_np, k)
D_ip, I_ip = index_flat_ip.search(normalized_query_embeddings_np, k) # 正規化されたクエリを使用
D_ivf, I_ivf = index_ivf_flat.search(query_embeddings_np, k)
D_hnsw, I_hnsw = index_hnsw.search(query_embeddings_np, k)

# 結果を表示する関数
def display_results(D, I, input_texts, index_type):
    print(f"Results for {index_type}:")
    for i in range(k):
        print(f"  Rank: {i+1}, Index: {I[0][i]}, Distance: {D[0][i]}, Text: '{input_texts[I[0][i]]}")

# 各インデックスタイプの結果を表示
display_results(D_l2, I_l2, input_texts, "IndexFlatL2")
display_results(D_ip, I_ip, input_texts, "IndexFlatIP")
display_results(D_ivf, I_ivf, input_texts, "IndexIVFFlat")
display_results(D_hnsw, I_hnsw, input_texts, "IndexHNSW")

以下のような結果をえます。この例では"今日は雨が振らなくてよかった"に文章として一番近く感じる、"今日は良いお天気ですね"が、どの手法でも選ばれていることが分かります。

IndexFlatIPに関してはベクトルの長さを正規化しているため必ず0から1になる値が返ってきています。

Results for IndexFlatL2:
  Rank: 1, Index: 3, Distance: 202.58700561523438, Text: '今日は良いお天気ですね
  Rank: 2, Index: 4, Distance: 267.9387512207031, Text: '最近景気悪いですね
Results for IndexFlatIP:
  Rank: 1, Index: 3, Distance: 0.8827050924301147, Text: '今日は良いお天気ですね
  Rank: 2, Index: 4, Distance: 0.8428647518157959, Text: '最近景気悪いですね
Results for IndexIVFFlat:
  Rank: 1, Index: 3, Distance: 202.58700561523438, Text: '今日は良いお天気ですね
  Rank: 2, Index: -1, Distance: 3.4028234663852886e+38, Text: '最近景気悪いですね
Results for IndexHNSW:
  Rank: 1, Index: 3, Distance: 202.58700561523438, Text: '今日は良いお天気ですね
  Rank: 2, Index: 4, Distance: 267.9387512207031, Text: '最近景気悪いですね

まとめ

いろいろ調べていると、以下のようにまとめることができそうです

小規模データで何も考えず、とりあえず高精度に使いたい場合は IndexFlatL2
データが正規化されているなら IndexFlatIP が最も精度が高い
大規模データで高速処理が必要なら IndexHNSWは最も期待できる

Appendix

`IndexHNSW`のチューニングパラメータ

IndexHNSWの検討すべきパラメータとしては以下があります。必要に応じてチューニングしましょう。

M(Default: 16):
- 各ノードが保持する近傍の数。大きいほど精度は向上しますが、構築とメモリコストが増加します
efConstruction (Default: 200):
- インデックス構築時の動的なパラメータで、探索の深さを定義します
ef(Default: 200):
- 検索時の動的なパラメータで、より広い範囲を探索します。大きいほど検索精度が向上しますが、検索時間が増加します
- 検討すべき処理: パフォーマンスと精度のバランスを取るためにMとefを調整します。大規模なデータセットでは、これらのパラメータが検索性能に大きな影響を与えるため、実際のデータ特性に応じたチューニングが必要になります