検索タスクにおけるBM25のコサイン類似度とスコアの精度比較

追記

比較する条件を整理した改良版を書きました。本記事は記録として残しておきます。（2024/11/28）

概要

以下の記事の疑問に自分なりに答えを出すために、実際にBM25スコアとBM25ベクトルのコサイン類似度で検索精度にどう違いがあるのか検証しました。

【疑問】BM25でもTFIDF同様にコサイン類似度に基づいてランキングしてよいのか

背景

上記別記事で抱いた疑問の概略は以下です。

• 検索タスク等において、ランキングの指標として、TFIDFではTFIDF重みベクトルのコサイン類似度を用いるが、BM25ではBM25スコアを用いることが多い
• BM25スコアはクエリに含まれる単語を検索対象文書におけるその単語のBM25の重みに変換して足し合わせた値である。
• BM25でもBM25の重みベクトルのコサイン類似度（BM25コサイン類似度）をランキングに用いたらだめなのか？

記事で書いていない内容も含めてもう少し上記の疑問を深堀りします。

BM25スコアは、検索対象文書のBM25重みと、クエリの単語頻度の積和として計算されるため、BM25重みの１次の関数です。
一方で、BM25コサイン類似度は、検索対象文書もクエリもBM25の重みを成分とするベクトルとなるため、BM25重みの2次の関数です。
よって、BM25コサイン類似度のほうが、単語の重要度をより強調するような指標になっているといえます。

例えば、"a","b","c","d"という4種類の単語があるコーパスを考えてみます。
各単語の重みは以下のとおりとします。

weights = {"a": 0.9, "b": 0.1, "c": 0.8, "d": 0.2}

このとき、query = "a b c d"に対する検索対象文書text_a = "a b"およびtext_b = "c d"のBM25スコアおよびコサイン類似度を計算してみます。
BM25スコアは重みの線形な和なので、text_aとtext_bの場合で同じ（どちらも1.0）ですが、コサイン類似度は2次の計算となるため、大きい重みの"a"がより重要視された結果、text_aのほうがtext_bより大きくなります。

query = "a b c d"
text_a = "a b"
text_b = "c d"

weights = {"a": 0.9, "b": 0.1, "c": 0.8, "d": 0.2}

def score(text1,text2):
  tokens1 = text1.split()
  tokens2 = text2.split()
  score = 0
  for token, weight in weights.items():
    if token in tokens1 and token in tokens2:
      score += weight
  return score

def cos_sim(text1, text2):
    tokens1 = text1.split()
    tokens2 = text2.split()
    
    numerator = 0
    for token, weight in weights.items():
        if token in tokens1 and token in tokens2:
            numerator += weight ** 2
    
    denominator1 = sum(weights.get(token, 0) ** 2 * tokens1.count(token) ** 2 for token in tokens1)
    denominator2 = sum(weights.get(token, 0) ** 2 * tokens2.count(token) ** 2 for token in tokens2)
    
    similarity = numerator / ((denominator1 * denominator2) ** 0.5)
    return similarity
  
# Calculate and print scores and cosine similarities
print("Scores:")
print(f"Score between query and text_a: {score(query, text_a)}")
print(f"Score between query and text_b: {score(query, text_b)}")

print("\nCosine Similarities:")
print(f"Cosine similarity between query and text_a: {cos_sim(query, text_a)}")
print(f"Cosine similarity between query and text_b: {cos_sim(query, text_b)}")

Scores:
Score between query and text_a: 1.0
Score between query and text_b: 1.0

Cosine Similarities:
Cosine similarity between query and text_a: 0.7393691004272944
Cosine similarity between query and text_b: 0.6733003292241385

BM25コサイン類似度が単語の重要度をより強調する指標であること自体は、単なる性質なので、検索タスクにおいて、BM25コサイン類似度とBM25スコアのどちらがいいのか、言い換えると、単語の重要度を強調するほうがいいのかしないほうがいいのか、については実験で確かめてみるしかありません。

そこで、BM25スコアとBM25コサイン類似度をはじめ、いくつかのキーワード検索の比較を実施してみます。

条件

以下の条件を比較します。

• count-count: 単語の頻度ベクトル同士のコサイン類似度
• count-tfidf: TFIDFスコア。クエリの単語頻度ベクトルと、検索対象文書のTFIDFベクトルの内積。
• count-bm25: scikit-learnで計算したBM25スコア。クエリの単語頻度ベクトルと、検索対象文書のBM25ベクトルの内積。
• count-rankbm25: rank_bm25ライブラリで計算したBM25スコア。クエリの単語頻度ベクトルと、検索対象文書のBM25ベクトルの内積。
• tfidf-tfidf: TFIDFベクトル同士のコサイン類似度。
• bm25-bm25: scikit-learnで計算したBM25ベクトル同士のコサイン類似度。
• rankbm25-rankbm25: rank_bm25ライブラリで計算したBM25ベクトル同士のコサイン類似度。

条件設定の意図は次のとおりです。

• キーワード検索の代表的な方法であるカウント、TFIDF、BM25を比較
• それぞれの方法について、スコアとコサイン類似度の差を見たい。ただしカウントは重みの概念がないため１条件のみ
• bm25はライブラリが異なっても傾向があるか見ておきたい。scikit-learnベースのものと、rank_bm25を比較（ちなみにTFIDFの計算はtfidfを使用）

また、コーパスが小規模（数千件）と大規模（１０万件）でそれぞれ比較します。

実装

検索コーパス

検索精度を評価するためのコーパスを準備します。
miraclの日本語データセット、devスプリットを使います。

miraclをダウンロードするため、huggingfaceのアクセストークンを取得し、環境変数ファイルに記載します。

.env

HF_ACCESS_TOKEN="トークン"

miraclのコーパスとクエリをダウンロードします。

import os
import dotenv
import datasets
dotenv.load_dotenv()

docs = datasets.load_dataset('miracl/miracl-corpus', "ja")
queries = datasets.load_dataset("miracl/miracl", "ja", token = os.environ["HF_ACCESS_TOKEN"], split="dev")

クエリとコーパスから必要な情報を取得します。

• positive_id_set: クエリの正解となる文書idの集合。検索対象コーパスの作成に使用。
• query_texts: クエリのテキスト
• query_positive_ids: クエリごとの正解idのリスト

positive_id_set = set()
query_texts = []
query_positive_ids = []

for data in queries:
  query_id = data['query_id']
  query = data['query']
  positive_passages = data['positive_passages']
  negative_passages = data['negative_passages']
  
  query_texts.append(query)
  query_positive_ids.append([entry["docid"] for entry in positive_passages])
  
  for entry in positive_passages:
    docid = entry['docid']
    title = entry['title']
    text = entry['text']
    if docid not in positive_id_set:
      positive_id_set.add(docid)

正解文書とそれ以外にcorpusを分けます。後ほど最終的に評価に用いるコーパスを作成するときに、確実に正解文書が含まれるようにしつつコーパスサイズの調整を容易にするためです。

# Filter the documents based on corpus_ids
positive_docs = docs['train'].filter(lambda example: example['docid'] in positive_id_set)
docs_without_positive = docs['train'].filter(lambda example: example['docid'] not in positive_id_set)

分かち書き関数

日本語テキストを分かち書きするための関数を作成しておきます。

import MeCab
mecab = MeCab.Tagger('-Owakati')
def tokenize(text: str):
  return mecab.parse(text).strip()

rank_bm25の拡張

rank_bm25はBM25の文章を重みベクトルに変換するメソッドをもたないので、rank_bm25を拡張して自作します。
rank_bm25のBM25Okapiを継承したクラスをつくり、以下の変更を実施しています。

• __init__の変更: transformメソッド等で使用するプラパティを保存するようにしています。
• transformメソッド追加： クエリをBM25の重みベクトルに変換します。
• count_transformメソッド追加: クエリを単語頻度ベクトルに変換します。scikit-learnのCountVectorizerでも同様の処理ができますが、vocabularyを同じにするのが手間だったので、rank_bm25側でも実装することにしました。

from rank_bm25 import BM25Okapi
import scipy
import numpy as np

class ExtendedBM25Okapi(BM25Okapi):
    def __init__(self, corpus, tokenizer=None, k1=1.5, b=0.75, epsilon=0.25):
        super().__init__(corpus, tokenizer=tokenizer, k1=k1, b=b, epsilon=epsilon)
        
        self.vocabulary = list(self.idf.keys())
        self.word_to_id = {word: i for i, word in enumerate(self.vocabulary)}
    
    def transform(self, queries: list[list[str]]) -> scipy.sparse.csr_matrix:
        from scipy.sparse import csr_matrix
        
        rows = []
        cols = []
        data = []
        
        for i, query in enumerate(queries):
            query_len = len(query)
            
            for word in set(query):
                if word in self.word_to_id:
                    word_id = self.word_to_id[word]
                    tf = query.count(word)
                    idf = self.idf.get(word, 0)
                    
                    # BM25 scoring formula
                    numerator = idf * tf * (self.k1 + 1)
                    denominator = tf + self.k1 * (1 - self.b + self.b * query_len / self.avgdl)
                    
                    score = numerator / denominator
                    
                    rows.append(i)
                    cols.append(word_id)
                    data.append(score)
        
        return csr_matrix((data, (rows, cols)), shape=(len(queries), len(self.vocabulary)))
    def count_transform(self, queries: list[list[str]]) -> scipy.sparse.csr_matrix:
        from scipy.sparse import csr_matrix
        
        rows = []
        cols = []
        data = []
        
        for i, query in enumerate(queries):
            for word in query:
                if word in self.word_to_id:
                    word_id = self.word_to_id[word]
                    rows.append(i)
                    cols.append(word_id)
                    data.append(1)  # Count is always 1 for each occurrence
        
        return csr_matrix((data, (rows, cols)), shape=(len(queries), len(self.vocabulary)))

ExtendedBM25の処理の検算

ExpectedBM25のtransformが正しく実装できてるか確認するために、rank_bm25がもともともっているget_scores関数の戻り値と、比較します。

positive_texts_tokenized = list(map(lambda x: tokenize(x).split(), positive_docs['text']))
query_texts_tokenized = list(map(lambda x: tokenize(x).split(), query_texts))


original_bm25 = BM25Okapi(positive_texts_tokenized)
extended_bm25 = ExtendedBM25Okapi(positive_texts_tokenized)

original_scores = []
for query in query_texts_tokenized:
  scores = original_bm25.get_scores(query)
  original_scores.append(scores)

extended_corpus_vectors = extended_bm25.transform(positive_texts_tokenized)
extended_count_query_vectors = extended_bm25.count_transform(query_texts_tokenized)
extended_scores = extended_count_query_vectors.dot(extended_corpus_vectors.T)

# Convert lists to numpy arrays for easier manipulation
original_scores_array = np.array(original_scores)
extended_scores_array = extended_scores

# Calculate the difference between extended and original scores
score_difference = extended_scores_array - original_scores_array

# Calculate statistics
mean_difference = np.mean(score_difference)
std_difference = np.std(score_difference)
max_difference = np.max(score_difference)
min_difference = np.min(score_difference)

print(f"Mean difference: {mean_difference}")
print(f"Standard deviation of difference: {std_difference}")
print(f"Max difference: {max_difference}")
print(f"Min difference: {min_difference}")

# Calculate correlation
correlation = np.corrcoef(original_scores_array.flatten(), extended_scores_array.toarray().flatten())[0, 1]
print(f"Correlation between original and extended scores: {correlation}")

Mean difference: -6.688436291558274e-16
Standard deviation of difference: 1.2349400906308533e-15
Max difference: 2.1316282072803006e-14
Min difference: -1.4210854715202004e-14
Correlation between original and extended scores: 1.0

ほぼ0に近い差なので、うまく計算できていそうです。丸め誤差はあるので、完全なゼロにはならないと思われます。

検索スコアの計算

比較用関数

各種比較を実施し、検索精度を出力する関数を作ります。
処理の流れはだいたい以下のような感じです。

• 入力として、スペースで分かち書きされたテキスト（corpus_texts_tokenized）とそのdocid（corpus_ids）、スペースで分かち書きされたクエリのテキスト（query_texts_tokenized）、クエリごとの正解文書のdocid（query_positive_ids）を受け取ります。
• CountVectorizer, TfidfVectorizer, BM25Vectorizer、RankBM25Vectorizerを初期化し、corpusでfitします。BM25Vectorizerはnocchi1さんが作成したscikit-learnベースのBM25Vectorizerを使います。その後、コーパス、クエリを各ベクトライザでベクトル化した行列をえます。
• 比較条件となる指標での検索結果文書ランキングを作成します。
• ランキングからhit_at_nを計算します。hit_at_nは上位n件に一つ以上の正解が含まれたクエリの割合です。

from sklearn.feature_extraction.text import CountVectorizer,TfidfVectorizer
from bm25 import BM25Vectorizer
from sklearn.metrics.pairwise import cosine_similarity

def eval_corpus(corpus_texts_tokenized: list[str], corpus_ids: list[str], query_texts_tokenized: list[str], query_positive_ids: list[list])->dict:
    """
    Evaluate the performance of different vectorization and similarity methods on a corpus.

    This function takes a corpus of texts, query texts, and their corresponding IDs, and evaluates
    the performance of various vectorization and similarity calculation methods using Hit@N metric.

    Parameters:
    - corpus_texts_tokenized (list[str]): List of tokenized corpus texts
    - corpus_ids (list[str]): List of corpus document IDs
    - query_texts_tokenized (list[str]): List of tokenized query texts
    - query_positive_ids (list[list]): List of lists containing positive document IDs for each query

    Returns:
    - dict: A dictionary containing Hit@N scores for different vectorization and similarity methods,
            where N is 1, 3, 5, and 10.
    """

    # Initialize vectorizers
    count_vectorizer = CountVectorizer()
    tfidf_vectorizer = TfidfVectorizer()
    bm25_vectorizer = BM25Vectorizer()
    rankbm25_vectorizer = ExtendedBM25Okapi([doc.split() for doc in corpus_texts_tokenized], k1=1.2)

    # Transform corpus texts
    count_corpus_vectors = count_vectorizer.fit_transform(corpus_texts_tokenized)
    tfidf_corpus_vectors = tfidf_vectorizer.fit_transform(corpus_texts_tokenized)
    bm25_corpus_vectors = bm25_vectorizer.fit_transform(corpus_texts_tokenized)
    rankbm25_corpus_vectors = rankbm25_vectorizer.transform([doc.split() for doc in corpus_texts_tokenized])

    # Transform query texts
    count_query_vectors = count_vectorizer.transform(query_texts_tokenized)
    tfidf_query_vectors = tfidf_vectorizer.transform(query_texts_tokenized)
    bm25_query_vectors = bm25_vectorizer.transform(query_texts_tokenized)
    rankbm25_query_vectors = rankbm25_vectorizer.transform([doc.split() for doc in query_texts_tokenized])
    rankbm25_count_query_vectors = rankbm25_vectorizer.count_transform([doc.split() for doc in query_texts_tokenized])

    # Calculate similarity matrices
    count_count_cosine_sim = cosine_similarity(count_query_vectors, count_corpus_vectors)
    count_tfidf_dot_product = count_query_vectors.dot(tfidf_corpus_vectors.T).toarray()
    count_bm25_dot_product = count_query_vectors.dot(bm25_corpus_vectors.T).toarray()
    count_rankbm25_dot_product = rankbm25_count_query_vectors.dot(rankbm25_corpus_vectors.T).toarray()
    tfidf_tfidf_cosine_sim = cosine_similarity(tfidf_query_vectors, tfidf_corpus_vectors)
    bm25_bm25_cosine_sim = cosine_similarity(bm25_query_vectors, bm25_corpus_vectors)
    rankbm25_rankbm25_cosine_sim = cosine_similarity(rankbm25_query_vectors, rankbm25_corpus_vectors)

    # Calculate rankings
    count_count_rank = count_count_cosine_sim.argsort(axis=1)[:, ::-1]
    count_tfidf_rank = count_tfidf_dot_product.argsort(axis=1)[:, ::-1]
    count_bm25_rank = count_bm25_dot_product.argsort(axis=1)[:, ::-1]
    count_rankbm25_rank = count_rankbm25_dot_product.argsort(axis=1)[:, ::-1]
    tfidf_tfidf_rank = tfidf_tfidf_cosine_sim.argsort(axis=1)[:, ::-1]
    bm25_bm25_rank = bm25_bm25_cosine_sim.argsort(axis=1)[:, ::-1]
    rankbm25_rankbm25_rank = rankbm25_rankbm25_cosine_sim.argsort(axis=1)[:, ::-1]

    # Function to calculate hit@n
    def count_hit_at_n(rank_matrix, query_positive_ids, corpus_ids, n=3):
        hit_at_n = 0
        for i, pos_ids in enumerate(query_positive_ids):
            top_n_indices = rank_matrix[i][:n]
            top_n_docids = [corpus_ids[idx] for idx in top_n_indices]
            if any(docid in pos_ids for docid in top_n_docids):
                hit_at_n += 1
        return hit_at_n / len(query_positive_ids)

    # Calculate hit@n for different n values (1, 3, 5, 10)
    n_values = [1, 3, 5, 10]
    hit_at_ns = {}

    for n in n_values:
        hit_at_ns[n] = {
            "Count-Count": count_hit_at_n(count_count_rank, query_positive_ids, corpus_ids, n=n),
            "Count-TFIDF": count_hit_at_n(count_tfidf_rank, query_positive_ids, corpus_ids, n=n),
            "Count-BM25": count_hit_at_n(count_bm25_rank, query_positive_ids, corpus_ids, n=n),
            "Count-RankBM25": count_hit_at_n(count_rankbm25_rank, query_positive_ids, corpus_ids, n=n),
            "TFIDF-TFIDF": count_hit_at_n(tfidf_tfidf_rank, query_positive_ids, corpus_ids, n=n),
            "BM25-BM25": count_hit_at_n(bm25_bm25_rank, query_positive_ids, corpus_ids, n=n),
            "RankBM25-RankBM25": count_hit_at_n(rankbm25_rankbm25_rank, query_positive_ids, corpus_ids, n=n)
        }

        # Print results for each n
        print(f"Hit@{n}: Count-Count: {hit_at_ns[n]['Count-Count']:.4f}, Count-TFIDF: {hit_at_ns[n]['Count-TFIDF']:.4f}, Count-BM25: {hit_at_ns[n]['Count-BM25']:.4f}, Count-RankBM25: {hit_at_ns[n]['Count-RankBM25']:.4f}, TFIDF-TFIDF: {hit_at_ns[n]['TFIDF-TFIDF']:.4f}, BM25-BM25: {hit_at_ns[n]['BM25-BM25']:.4f}, RankBM25-RankBM25: {hit_at_ns[n]['RankBM25-RankBM25']:.4f}")
    
    return hit_at_ns

小規模コーパス

まず小規模なコーパスで実験します。
ここではクエリのいずれかの正解となる文書のみで構成された最小サイズのコーパスを小規模コーパスとします。

# positive_corpus
positive_texts_tokenized = list(map(tokenize, positive_docs['text']))
positive_ids = positive_docs['docid']
query_texts_tokenized = list(map(tokenize, query_texts))
eval_corpus(positive_texts_tokenized, positive_ids, query_texts_tokenized, query_positive_ids)

Hit@1: Count-Count: 0.5140, Count-TFIDF: 0.6442, Count-BM25: 0.6756, Count-RankBM25: 0.6663, TFIDF-TFIDF: 0.7012, BM25-BM25: 0.7233, RankBM25-RankBM25: 0.6884
Hit@3: Count-Count: 0.6593, Count-TFIDF: 0.7709, Count-BM25: 0.7988, Count-RankBM25: 0.7826, TFIDF-TFIDF: 0.8198, BM25-BM25: 0.8279, RankBM25-RankBM25: 0.7953
Hit@5: Count-Count: 0.7023, Count-TFIDF: 0.8035, Count-BM25: 0.8372, Count-RankBM25: 0.8093, TFIDF-TFIDF: 0.8442, BM25-BM25: 0.8547, RankBM25-RankBM25: 0.8209
Hit@10: Count-Count: 0.7605, Count-TFIDF: 0.8465, Count-BM25: 0.8616, Count-RankBM25: 0.8291, TFIDF-TFIDF: 0.8721, BM25-BM25: 0.8709, RankBM25-RankBM25: 0.8384

大規模コーパス

大規模コーパスはdocs_without_positiveから追加で10万件取得して、検索対象コーパスとします。

# large corpus
additional_docs_size = 100000
additional_docs = docs_without_positive.select(range(additional_docs_size))

additional_texts_tokenized = list(map(tokenize, additional_docs['text']))
additional_ids = additional_docs['docid']

corpus_texts_tokenized = positive_texts_tokenized + additional_texts_tokenized
corpus_ids = positive_ids + additional_ids

eval_corpus(corpus_texts_tokenized, corpus_ids, query_texts_tokenized, query_positive_ids)

Hit@1: Count-Count: 0.1663, Count-TFIDF: 0.2523, Count-BM25: 0.1721, Count-RankBM25: 0.4000, TFIDF-TFIDF: 0.4023, BM25-BM25: 0.3279, RankBM25-RankBM25: 0.3012
Hit@3: Count-Count: 0.2523, Count-TFIDF: 0.3860, Count-BM25: 0.3209, Count-RankBM25: 0.5442, TFIDF-TFIDF: 0.5407, BM25-BM25: 0.5116, RankBM25-RankBM25: 0.4779
Hit@5: Count-Count: 0.2988, Count-TFIDF: 0.4523, Count-BM25: 0.4035, Count-RankBM25: 0.5860, TFIDF-TFIDF: 0.5953, BM25-BM25: 0.5977, RankBM25-RankBM25: 0.5674
Hit@10: Count-Count: 0.3570, Count-TFIDF: 0.5186, Count-BM25: 0.5035, Count-RankBM25: 0.6395, TFIDF-TFIDF: 0.6663, BM25-BM25: 0.6826, RankBM25-RankBM25: 0.6465

考察

TFIDFについてはコサイン類似度（TFIDF-TFIDF）がスコア（Count-TFIDF）より精度が高い傾向があるため、コサイン類似度を使うと良さそうです。

BM25については、scikit-learnのスコア（Count-BM25）が精度が低いことを除けば、ほかは、スコアでもコサイン類似度でも似たような精度です。scikit-learnよりはrank-bm25のほうが内部的に素直な処理をしていると思われますので、rank-bm25の結果をより信用するのであれば、スコアでもコサイン類似度でもあまり差はないというのが今のところの感触です。

おわりに

一応、以下のように結果をまとめておきます。

• BM25検索においてスコアとコサイン類似度のどちらがいいのかは、なんともいえない。逆に言うと、はっきりどっちがいい、ということは今のところないので、好きな方法を使えば良い。
• TFIDFはscikit-learnの場合はコサイン類似度がよい。
• 今後の課題として、多様なデータセットで検証し、BM25スコアとコサイン類似度の使い分け方について知見を深めたい。