RAG質問応答システムに使うRetrieverの精度比較

TL;DR

今回はRAGのRetrieverの性能を比較しました。 その結果としては、

• Dense Retrieverの中でデフォルトのGPTのEmbeddingモデル+Cosine類似度の組み合わせるが一番良かったです。

• Sparse Retrieverの中でBM25は計算スピードが早くてそこそこ良いパフォーマンスを出せています。

• Hybridのやり方で、Dense RetrieverとSparse Retrieverを組み合わせると一番良い結果を出せています。

RAG(Retrieval Augmented Generation)とは？

RAG（Retrieval-Augmented Generation）は、自然言語処理（NLP）タスクのための最新の機械学習モデルの一つです。RAGは、質問応答、文章生成、要約作成などのタスクに適用されます。このモデルは、あらかじめ学習された情報を取得（retrieval）し、その情報を利用して文を生成（generation）することが特徴です。

RAGは、主に以下の2つのコンポーネントから構成されています。

• 検索器（Retriever）：質問や入力文に関連する情報をデータセットから見つけ出す役割を担います。検索された情報は、文書や段落といった形式で提供されます。

• 生成器（Generator）：検索器から提供された情報を基に、適切な応答や文章を生成します。生成器は、Transformerベースのモデル（例：GPT-3、BART）で構築されることが一般的です。最近はOpenAIのGPTのAPIのを利用することが多いです。

RAG

生成器の部分はLLMを使うため、RAGの性能は検索器の性能に依存します。今回は、検索器の性能を比較しました。そのまとめは以下の図になります。

Alt text

使用するデータセット

今回使用するデータは東京都立大学のeラーニングシステムのQ&Aデータです。このデータは、東京都立大学で導入されたeラーニングシステムのユーザーから2015年4月から2018年7月までに報告された問題点としてのQ&Aデータを収集したものです。427の質問と79の回答が含まれています。質問にどの回答に紐づくかのラベルがあります。

データの様子は下記の通りです。

import pandas as pd
# https://zenodo.org/record/2783642
q_df = pd.read_csv("https://zenodo.org/record/2783642/files/Questions.csv")
a_df = pd.read_csv("https://zenodo.org/record/2783642/files/Answers.csv")
print("q_df.shape:", q_df.shape)
print("a_df.shape:", a_df.shape)
q_df.columns = [c.strip() for c in q_df.columns]
a_df.columns = [c.strip() for c in a_df.columns]
df = q_df.merge(a_df, on="AID")
df.columns = ["query","AID","document"]

metadata = a_df[["AID"]].to_dict(orient="records")
documents = a_df["Text"].tolist()
query_list = list(zip(q_df["Text"], q_df["AID"]))
display(q_df.head(3))
display(a_df.head(3))

q_df.shape: (427, 2)
a_df.shape: (79, 2)

	Text	AID
0	履修している授業で先生が資料をアップロードしているはずだが、コース上に資料が見当たらない。	A001
1	資料をマイページに置いたが、学生からは見えなかった。	A001
2	前期の科目の「資料」を学生から見られないようにするにはどうしたら良いか？	A001

	AID	Text
0	A001	資料が見つからない場合は、以下の点を確認してください。<br><br><br>【受講生編】<...
1	A002	資料のアップロードやお知らせ作成時の電子メールでの通知の有無は、各授業の担当教員が設定できま...
2	A003	kibacoにはファイルへパスワードを設定する機能はありません。資料は受講生全員に開示されま...

評価対象と評価方法、評価指標

以前のポストは主にDense RetrieverのEmbeddingモデルを比較しました。今回は以下の角度で比較します。

• 類似度の計算方法
  • Cosine類似度
  • MMR（Maximal Marginal Relevance）類似度
  • SVM
• Sparse Retrieverのモデル
  • BM25
  • TF-IDF

評価方法は以下の3つのステップです。

1. 79のドキュメントをEmbeddingに変換し、FAISSのVectorstoreとして保存する。
2. 427の質問をEmbeddingに変換し、FAISSのVectorstoreを使用して、79のドキュメントを近い順に並べる。
3. 並んだ順番でEmbeddingの性能を評価する

評価指標は以下の3つです。

1. Mean Reciprocal Rank（MRR）: 正解ドキュメントの順位の平均の逆数で、ランク全体を評価する指標。
2. Recall@1: 正解ドキュメントが1番目に並んでいるかどうかを評価する指標。
3. Recall@5: 正解ドキュメントが上位5位以内に入っているかどうかを評価する指標。

類似度の計算方法

Embeddingを得た後、類似度の計算によく使う方法はCosine類似度です。それを今回のベースラインとします。

MMR（Maximal Marginal Relevance）とは、は検索クエリとの関連性を維持しつつも、検索結果多様性を持たすように検索結果の順位を並べ替える手法。推薦システムで使うことが多いです。(https://yolo-kiyoshi.com/2020/05/08/post-1781/)

SVMはもともと分類の機械学習モデルですが、それを検索に使うことをOpenAIの創立者karpathyがTwitterで提案しました。 「感覚」としてはCosine Similarlyより良いだけでLangchainに取り込まされました。 (https://github.com/karpathy/randomfun/blob/master/knn_vs_svm.ipynb)

評価用コードがないので折り畳みしました。興味がある方は下の「Show the code」をクリックしてください。

Show the code

from dataclasses import dataclass
from tqdm.auto import tqdm

DOC_NUM = len(a_df)
@dataclass
class EvaluationResults:
    result_df: pd.DataFrame
    mrr: float
    recall_at_1: float
    recall_at_5: float

def mrr(rank_array):
    return (1 / rank_array).mean()

def recall_at_k(rank_array, k):
    return (rank_array <= k).mean()

def evaluate(query_list, search_func):
    result_list = []
    for query, aid in tqdm(query_list):
        rank_result = get_rank(query, search_func=search_func)
        if aid not in rank_result:
            rank = DOC_NUM + 1
        else:
            rank = rank_result.index(aid) + 1
        
        result_list.append((query, rank, rank_result))

    result_df = pd.DataFrame(result_list, columns=["query", "rank", "rank_result"])
    return EvaluationResults(result_df, mrr(result_df["rank"]), recall_at_k(result_df["rank"], 1), recall_at_k(result_df["rank"], 5))

def get_rank(query, search_func):
    results = search_func(query)
    aid_list = []
    for doc in results:
        aid = metadata[documents.index(doc.page_content)]["AID"]
        aid_list.append(aid)
    return aid_list

Show the code

from langchain.vectorstores import FAISS
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.retrievers import KNNRetriever
from langchain.retrievers import SVMRetriever
from langchain.retrievers import TFIDFRetriever
from langchain.retrievers import ElasticSearchBM25Retriever

embedding_openai = OpenAIEmbeddings()

faiss_vectorstore = FAISS.from_texts(documents, embedding_openai)

svm_retriever = SVMRetriever.from_texts(documents, embedding_openai)
svm_retriever.k = DOC_NUM

faiss_similarity_result = evaluate(query_list, lambda q: faiss_vectorstore.similarity_search(q, k=DOC_NUM))
faiss_mmr_result = evaluate(query_list, lambda q: faiss_vectorstore.max_marginal_relevance_search(q, k=DOC_NUM))
svm_result = evaluate(query_list, lambda q: svm_retriever.get_relevant_documents(q))

result_df = pd.DataFrame(
    [
        ["faiss_similarity", faiss_similarity_result.mrr, faiss_similarity_result.recall_at_1, faiss_similarity_result.recall_at_5],
        ["faiss_mmr", faiss_mmr_result.mrr, faiss_mmr_result.recall_at_1, faiss_mmr_result.recall_at_5],
        ["svm", svm_result.mrr, svm_result.recall_at_1, svm_result.recall_at_5],
        ], 
    columns = ["model_id","mrr","recall_at_1","recall_at_5"]
    ).sort_values("mrr", ascending=False)

結果はこちらです。ご覧の通り、MMRとSVM両方ともベースラインより弱いです。特にSVMの性能がひどいです。またくCosine類似度の方法に比べにならないです。

result_df

	model_id	mrr	recall_at_1	recall_at_5
0	faiss_similarity	0.685327	0.550351	0.868852
1	faiss_mmr	0.622978	0.550351	0.679157
2	svm	0.520898	0.388759	0.683841

Sparse Retrieverのモデル

次にSparse RetriewerのBM25とTF-IDFを見ます。

TF-IDFとBM25は、情報検索において、ドキュメントと検索クエリの関連性を評価するために使用される代表的な手法です。単語のカウントをベースとしているため、得たドキュメントのベクトルはsparseになります。そのため、この手法をSparse Retrieverと呼びます。

このポストはその効果の紹介を目的しているので、TF-IDFとBM25の詳細な説明は折り畳みしました。興味がある方は下の青いバーをクリックください。

TF-IDFとBM25とは

TF-IDF: TF-IDFは、単語の重要度を評価するために使用される統計的手法です。TF-IDFは、2つの主要な要素で構成されています。 - TF（Term Frequency）：文書内での単語の出現頻度。単語が文書内で頻繁に出現するほど、その単語は文書内で重要であると考えられます。

• IDF（Inverse Document Frequency）：文書全体のセット（コーパス）における単語の希少性を測定する指標。単語が多くの文書で出現するほど、その単語は一般的であると考えられ、IDFの値は低くなります。逆に、単語が少数の文書でのみ出現する場合、その単語は特定の文書に特有であり、IDFの値が高くなります。

TF-IDFスコアは、これら2つの要素の積として計算され、このスコアが高い単語ほど、検索クエリと関連性が高いと考えられます。

BM25: BM25（Best Matching 25）は、TF-IDFを拡張した検索アルゴリズムで、検索クエリとドキュメントの関連性をより正確に評価することができます。BM25は、以下の特徴を持っています。

• 長さ正規化：長いドキュメントは、短いドキュメントよりも多くの単語を持つため、TF-IDFでは不利になる可能性があります。BM25では、ドキュメントの長さを正規化することで、この問題を解決しています。

• 単語の出現頻度の飽和：単語がある一定の出現回数を超えると、その単語の重要度が飽和し、それ以上の出現回数が重要度に大きな影響を与えなくなります。これにより、ある単語が特定のドキュメントで極端に多く出現する場合でも、適切な関連性評価が可能になります。

BM25スコアは、TF-IDFスコアを改善したものであり、検索クエリとドキュメントの関連性をより正確に評価することができます。多くの情報検索システムや検索エンジンでは、BM25が関連性スコアとして使用されています。

実装について、LangchainにあるTF-IDFモジュールには前処理がないため、それを追加しました。また、TF-IDFのモジュールに模倣してBM25を実装しました。

一点注意すべきところとしては、Sparse Retriewerは前処理が非常に重要です。今回の前処理は簡単にMecabでテキストを単語単位に分割し、ストップワードを除去しました。

Show the code

from __future__ import annotations

from typing import Any, Dict, Iterable, List, Optional, Callable

from pydantic import BaseModel

from langchain.schema import BaseRetriever, Document


class TFIDFRetriever(BaseRetriever, BaseModel):
    vectorizer: Any
    docs: List[Document]
    tfidf_array: Any
    k: int = 4
    preprocess_func: Callable[[str], str] = None

    class Config:
        """Configuration for this pydantic object."""

        arbitrary_types_allowed = True

    @classmethod
    def from_texts(
        cls,
        texts: Iterable[str],
        metadatas: Optional[Iterable[dict]] = None,
        tfidf_params: Optional[Dict[str, Any]] = None,
        preprocess_func: Optional[Callable[[str], str]] = None,
        **kwargs: Any,
    ) -> TFIDFRetriever:
        try:
            from sklearn.feature_extraction.text import TfidfVectorizer
        except ImportError:
            raise ImportError(
                "Could not import scikit-learn, please install with `pip install "
                "scikit-learn`."
            )

        tfidf_params = tfidf_params or {}
        vectorizer = TfidfVectorizer(**tfidf_params)
        if preprocess_func:
            processed_texts = [preprocess_func(t) for t in texts]
            tfidf_array = vectorizer.fit_transform(processed_texts)
        else:
            tfidf_array = vectorizer.fit_transform(texts)
        metadatas = metadatas or ({} for _ in texts)
        docs = [Document(page_content=t, metadata=m) for t, m in zip(texts, metadatas)]
        return cls(vectorizer=vectorizer, docs=docs, tfidf_array=tfidf_array,preprocess_func=preprocess_func,  **kwargs)

    @classmethod
    def from_documents(
        cls,
        documents: Iterable[Document],
        *,
        tfidf_params: Optional[Dict[str, Any]] = None,
        preprocess_func: Optional[Callable[[str], str]] = None,
        **kwargs: Any,
    ) -> TFIDFRetriever:
        texts, metadatas = zip(*((d.page_content, d.metadata) for d in documents))
        return cls.from_texts(
            texts=texts, tfidf_params=tfidf_params, metadatas=metadatas, preprocess_func=preprocess_func, **kwargs
        )

    def get_relevant_documents(self, query: str) -> List[Document]:
        from sklearn.metrics.pairwise import cosine_similarity
        if self.preprocess_func:
            query = self.preprocess_func(query)
        query_vec = self.vectorizer.transform(
            [query]
        )  # Ip -- (n_docs,x), Op -- (n_docs,n_Feats)
        results = cosine_similarity(self.tfidf_array, query_vec).reshape(
            (-1,)
        )  # Op -- (n_docs,1) -- Cosine Sim with each doc
        return_docs = []
        for i in results.argsort()[-self.k :][::-1]:
            return_docs.append(self.docs[i])
        return return_docs

    async def aget_relevant_documents(self, query: str) -> List[Document]:
        raise NotImplementedError
    
    

class BM25Retriever(BaseRetriever, BaseModel):
    vectorizer: Any
    docs: List[Document]
    k: int = 4
    preprocess_func: Callable[[str], str] = None
    tokenizer: Callable[[str], List[str]] = None

    class Config:
        """Configuration for this pydantic object."""

        arbitrary_types_allowed = True

    @classmethod
    def from_texts(
        cls,
        texts: Iterable[str],
        metadatas: Optional[Iterable[dict]] = None,
        bm25_params: Optional[Dict[str, Any]] = None,
        preprocess_func: Optional[Callable[[str], str]] = None,
        tokenizer : Optional[Callable[[str], List[str]]] = None,
        **kwargs: Any,
    ) -> BM25Retriever:
        try:
            from rank_bm25 import BM25Okapi
        except ImportError:
            raise ImportError(
                "Could not import rank_bm25, please install with `pip install "
                "rank_bm25`."
            )
            
        if preprocess_func:
            texts_processed = [preprocess_func(t) for t in texts]
        else:
            texts_processed = texts
            
        if tokenizer:
            tokenized_texts = [tokenizer(t) for t in texts_processed]
        else:   
            tokenized_texts = [t.split() for t in texts_processed]
        
        bm25_params = bm25_params or {}
        vectorizer = BM25Okapi(tokenized_texts, **bm25_params)
        metadatas = metadatas or ({} for _ in texts)
        docs = [Document(page_content=t, metadata=m) for t, m in zip(texts, metadatas)]
        return cls(vectorizer=vectorizer, docs=docs, preprocess_func=preprocess_func,  **kwargs)

    @classmethod
    def from_documents(
        cls,
        documents: Iterable[Document],
        *,
        bm25_params: Optional[Dict[str, Any]] = None,
        preprocess_func: Optional[Callable[[str], str]] = None,
        tokenizer : Optional[Callable[[str], List[str]]] = None,        
        **kwargs: Any,
    ) -> BM25Retriever:
        texts, metadatas = zip(*((d.page_content, d.metadata) for d in documents))
        return cls.from_texts(
            texts=texts, tfidf_params=bm25_params, metadatas=metadatas,preprocess_func=preprocess_func, tokenizer=tokenizer, **kwargs
        )

    def get_relevant_documents(self, query: str) -> List[Document]:
        if self.preprocess_func:
            query = self.preprocess_func(query)
        tokenized_query = query.split()
        return_docs = self.vectorizer.get_top_n(tokenized_query,self.docs, n=DOC_NUM)
        return return_docs

    async def aget_relevant_documents(self, query: str) -> List[Document]:
        raise NotImplementedError

Show the code

import requests
# download japanese stopwords
url = "https://raw.githubusercontent.com/stopwords-iso/stopwords-ja/master/stopwords-ja.txt"
stopwords = requests.get(url).text.split("\n")

import MeCab
import ipadic


# parser = MeCab.Tagger("-Owakati")
def extract_nouns_verbs(text):
    parser = MeCab.Tagger(ipadic.MECAB_ARGS)
    parsed_text = parser.parse(text)
    lines = parsed_text.split('\n')
    nouns_verbs = []

    for line in lines:
        if '名詞' in line or '動詞' in line or "形状詞" in line:
            parts = line.split('\t')
            word = parts[0]
            if not word.isascii():
                nouns_verbs.append(word)
    return nouns_verbs

def preprocess(text):
    token_list = [token for token in extract_nouns_verbs(text) if token not in stopwords]
    return " ".join(token_list)

Show the code

tfidf_search = TFIDFRetriever.from_texts(a_df["Text"].tolist(), preprocess_func=preprocess)
tfidf_search.k = DOC_NUM

bm25_search = BM25Retriever.from_texts(a_df["Text"].tolist(), preprocess_func=preprocess)
bm25_search.k = DOC_NUM

tfidf_result = evaluate(query_list, lambda query: tfidf_search.get_relevant_documents(query))
bm25_result = evaluate(query_list, lambda query: bm25_search.get_relevant_documents(query))

result_df = pd.DataFrame(
    [
    ["bm25", bm25_result.mrr, bm25_result.recall_at_1, bm25_result.recall_at_5],
    ["tfidf", tfidf_result.mrr, tfidf_result.recall_at_1, tfidf_result.recall_at_5]],
    columns = ["model_id","mrr","recall_at_1","recall_at_5"]
    ).sort_values("mrr", ascending=False)

結果を見ましょう。Sparse Retriewerのほうは意外に良い結果がを得ています。計算の速さやコスト等を考えると相当試すべき手法です。

result_df

	model_id	mrr	recall_at_1	recall_at_5
0	bm25	0.608793	0.475410	0.798595
1	tfidf	0.592323	0.454333	0.763466

Hybrid Retriever

これまでDense RetrieverとSparse Retrieverを比較しました。機械学習の中でEnsembleという概念がありまして、つまり複数のモデルを組み合わせて、より良い結果を得ることができます。それと似たような概念で、Dense RetrieverとSparse Retrieverを組み合わせて、Hybrid Retrieverを作るとより良い結果を得ることができます。

Hybridのほうほうとしては、Reciprocal Rank Fusion（RRF）という手法を使います。RRFは、複数の検索結果を組み合わせて、より良い検索結果を得るための手法です。RRFは、検索結果のランキングを組み合わせることで、検索結果のランキングを改善します。RRFは、以下の手順で実行されます。

\[\operatorname{RRF}(d) = \sum_{i=1}^{n} \frac{1}{k + r_i(d)}\]

ここで、

\(d\) は検索結果のドキュメントです。 \(n\) は統合する検索結果の数です。 \(r_i(d)\) は、\(i\)番目の検索結果におけるドキュメント\(d\)の順位です。 \(k\) は、RRFのパラメータです。この値を大きくすることで、検索結果の順位に対するペナルティを調整できます。通常、\(k\)は60などの定数値をとります。

実際の例で語りましょう。

BM25 Ranking	Dense Ranking	Results
A	B	A: 1/1 + 1/3 = 1.3
B	C	B: 1/2 + 1/1 = 1.5
C	A	C: 1/3 + 1/2 = 0.83

(https://weaviate.io/blog/hybrid-search-explained)

上のように、BM25とDense Retriever両方の結果が出た後、それぞれのドキュメントのランクの逆数を足して、その結果をもとに並べ替えます。例えば、AはBM25の結果で1位、Dense Retrieverの結果で3位なので、Kが0の場合はAの最終のスコアは1/1 + 1/3 = 1.3となります。

パラメーター\(k\)の影響

\(k\)パラメータは、RRFの式において検索結果の順位に対するペナルティを調整する役割を持っています。\(k\)が大きいほど、順位が低いドキュメントへのペナルティが緩やかになります。逆に、\(k\)が小さいほど、順位が低いドキュメントへのペナルティが厳しくなります。

具体的には、\(k\)が大きい場合、異なる検索結果間で順位が低いドキュメントでも、それらの組み合わせによってRRFスコアが上がる可能性があります。これにより、検索結果の多様性が増すことが期待されます。

一方で、\(k\)が小さい場合、順位が高いドキュメントがより重視されるため、検索結果の集中度が高まることが期待されます。ただし、あまりにも\(k\)が小さいと、異なる検索結果間でのバランスが悪くなり、検索結果の統合効果が十分に発揮されない可能性があります。

通常、\(k\)は60などの定数値をとりますが、実際の検索タスクや評価指標によって最適な\(k\)の値は異なる場合があります。実際の応用においては、パラメータチューニングや交差検証を用いて適切な\(k\)の値を決定することが望ましいです。

その実装は以下の通りです。

def weighted_reciprocal_rank_fusion(rank_lists, weights, k=-1):
    """
    Perform weighted Reciprocal Rank Fusion on multiple rank lists.

    Args:
        rank_lists (list of lists): A list of rank lists, where each rank list contains unique items.
        weights (list of float): A list of weights corresponding to the rank lists.
        k (float, optional): A constant added to the rank, controlling the balance between the importance
            of high-ranked items and the consideration given to lower-ranked items. Default is 0.

    Returns:
        list: The final aggregated list of items sorted by their weighted RRF scores in descending order.
    """
    if k == -1:
        k = 0.5 * len(rank_lists[0])

    if len(rank_lists) != len(weights):
        raise ValueError("Number of rank lists must be equal to the number of weights.")
    
    rrf_scores = {}
    
    for rank_list, weight in zip(rank_lists, weights):
        for rank, item in enumerate(rank_list, start=1):
            rrf_score = weight * (1 / (rank + k))
            
            if item in rrf_scores:
                rrf_scores[item] += rrf_score
            else:
                rrf_scores[item] = rrf_score

    # Sort items by their RRF scores in descending order
    sorted_items = sorted(rrf_scores.keys(), key=lambda x: rrf_scores[x], reverse=True)

    return sorted_items

Show the code

bm25_rank_result = bm25_result.result_df["rank_result"]
tfidf_rank_result = tfidf_result.result_df["rank_result"]
faiss_rank_result = faiss_similarity_result.result_df["rank_result"]

fused_rank_result = [
    weighted_reciprocal_rank_fusion(
        [bm25_rank_result[i],  faiss_rank_result[i]],
        [0.2, 0.8]
)
    for i in range(len(bm25_rank_result))
]

fused_rank_s = [fused_rank_result[i].index(query_list[i][1]) + 1 for i in range(len(query_list))]
fused_df = bm25_result.result_df.copy()
fused_df["rank"] = fused_rank_s

fused_result = EvaluationResults(fused_df,  mrr=mrr(fused_df["rank"]), recall_at_1=recall_at_k(fused_df["rank"], 1), recall_at_5=recall_at_k(fused_df["rank"], 5))

result_dict = {
    "faiss__cosine_similarity": faiss_similarity_result,
    "faiss_mmr": faiss_mmr_result,
    "svm": svm_result,
    "tfidf": tfidf_result,
    "bm25": bm25_result,
    "hybrid": fused_result,
}
result_df = pd.DataFrame(
    [
        [k, v.mrr, v.recall_at_1, v.recall_at_5]
        for k, v in result_dict.items()
    ],
    columns=["model_id", "mrr", "recall_at_1", "recall_at_5"],
).sort_values("mrr", ascending=False)

これでテストすると、Hyperのやり方は一番精度が良いことがわかります。

result_df

	model_id	mrr	recall_at_1	recall_at_5
5	hybrid	0.703478	0.573770	0.882904
0	faiss__cosine_similarity	0.685327	0.550351	0.868852
1	faiss_mmr	0.622978	0.550351	0.679157
4	bm25	0.608793	0.475410	0.798595
3	tfidf	0.592323	0.454333	0.763466
2	svm	0.520898	0.388759	0.683841

まとめ

今回はRAGのRetrieverの性能を比較しました。 その結果としては、

• Dense Retrieverの中でデフォルトのGPTのEmbeddingモデル+Cosine類似度の組み合わせるが一番良かったです。

• Sparse Retrieverの中でBM25は計算スピードが早くてそこそこ良いパフォーマンスを出せています。

• Hybridのやり方で、Dense RetrieverとSparse Retrieverを組み合わせると一番良い結果を出せています。