RAGの処理で、リランクとベクトル検索でできることの違いを検証／解説してみる

環境・使用するモデル

今回の実験はcollaboratory上で実施します。
使用するモデルはベクトルにはmultilingual-e5-baseを使用します。
日本語に対応しており、精度も日本語では他モデルに対して比較的良いとされています。
huggingface.co
リランカーモデルではjapanese-reranker-cross-encoder-large-v1を使用します。
huggingface.co
このモデルは日本語に特化した数少ないリランク専用のモデルになります。

実験内容

今回は宿のレビュー検索を想定した簡単な実験を行います。
まず、以下のレビュー文章を用意します。

上記のレビュー文章は1~5は温泉について言及していて、6~10は食事について言及しているレビューです。
このレビュー記事から「お肉の料理がおいしい宿」というクエリが入力される場合を考えます。
想定としては上記の中だとNo. 7とNo. 10の記事が取得できることを期待したいですが、ベクトル検索とリランクの結果がそれぞれどうなるのかを比較してみます。

ベクトル検索

ベクトル検索ではe5-baseを使用して変換したベクトルを、faissを使ってベクトル検索します。
faissで検索する部分はIndexFlatL2を使用して検索するだけなので、コードは省略します。
IndexFlatL2はL2ノルムを利用した逐次検索のため、シンプルに全件に対してL2距離を算出します。

import torch.nn.functional as F

from torch import Tensor
from transformers import AutoTokenizer, AutoModel

def average_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

input_texts = [
    <各文章のリスト>
]

# モデルを定義
tokenizer = AutoTokenizer.from_pretrained('intfloat/multilingual-e5-base')
model = AutoModel.from_pretrained('intfloat/multilingual-e5-base')
# ベクトルに変換した結果を取得 
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')
outputs = model(**batch_dict)
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
embeddings = F.normalize(embeddings, p=2, dim=1)

検索の結果としてトップ5件を並べると以下のようになりました。
この場合L2距離がスコアになるため、低いものほど近しいと判定されたドキュメントになります。

おおむね食事について言及しているレビューを取得できていますが、本来はトップに来てほしい肉料理について言及しているNo.7は3番目の表示順になっています。
ベクトル検索ではこのようにクエリの意図からは異なる文章がトップになることがよくあります。

リランクの場合

リランクの場合を見てみましょう。
リランクはクエリ＋レビュー文章を入力にして、10件分それぞれのスコアを算出し、スコアが高いものほど正しいと判断されたデータになります。

from sentence_transformers import CrossEncoder
import torch

# モデルを定義する
MODEL_NAME = "hotchpotch/japanese-reranker-cross-encoder-large-v1"
device = "cuda" if torch.cuda.is_available() else "cpu"
model = CrossEncoder(MODEL_NAME, max_length=512, device=device)

if device == "cuda":
    model.model.half()

query = "お肉の料理がおいしい宿"
# 各レビュー文章に対してスコアを算出する
scores = model.predict([(query, passage) for passage in passages])

そうすると上位5件の結果は以下のようになります。
リランクの場合はベクトル検索とは異なり、スコアが高いものが上位のドキュメントを表します。

肉料理について言及しているNo. 7とNo. 10が上位に来ており、さらにスコアでもかなり他と差がついていることがわかります。

結果

結果としてはベクトル検索を使用した場合はざっくり同じトピックの文章は上位に持って来れていますが、ほしい文章ではないものがトップに来ていたのに対して、リランクではほしい文章をトップに持ってくることができました。
ただし、リランクは前述した通り、対象のドキュメントごとにモデルを使ってスコアを計算するため、メモリなどのリソースが必要でなおかつ時間がかかります。
そのため、大量のドキュメントからおおむね同じトピックについて言及しているドキュメントをベクトル検索で広く取って、取得した結果をリランクするという手法がRAGにおける検索の主流になっています。