多言語E5をファインチューニングして検索性能を向上させる

多言語のテキスト埋め込み用のモデルであるMultilingual-E5[1]をファインチューニングして、検索性能が向上するか否かを日本語のデータセット検証してみました。結果としては、ファインチューニングすることで、OpenAIのtext-embedding-ada-002を上回る性能を確認できました。なお、E5については以下の記事で解説しているので、必要に応じて参照してください。

hironsan.hatenablog.com

本記事の構成は次のとおりです。

実験設定
実験結果
参考資料

実験設定

今回の実験では、多言語E5をファインチューニングして得られたモデルをベクトル検索に使うことで、検索性能がどのように変化するかを検証します。多言語E5にはbaseとlargeの2つのモデルがあるので、これらをファインチューニングする前後の検索性能を測定することにします。また、比較用のモデルとしてOpenAIのtext-embedding-ada-002での性能も測定します。

評価用のデータセットとしては、尼崎市のQAデータ[2]を使用します。このデータセットには、784の質問に対して対応する回答がAからCの3つのカテゴリでラベル付けされています。Aの場合は正しい情報を含み、Bであれば関連する情報を含み、Cであればトピックが同じであることを意味します。今回はこれらのカテゴリを関連文書として扱うことにします。

学習用のデータセットとしては、同じく尼崎市のQAデータに含まれるQAペアを使用します。このQAペアは、尼崎市のWebサイトから収集されており、1786件からなります。今回は、この内の80%を学習用、20%を検証用として利用することにします。

E5の学習には、SentenceTransformersのMultipleNegativesRankingLossを利用します。この損失関数は、バッチ内の正しいQAペアを正例、それ以外の組み合わせを負例として使うことで学習を進めます。損失関数の詳細については論文や実装が参考になると思います。以下は、学習に用いたコードからの抜粋です。

from sentence_transformers import InputExample
from sentence_transformers import SentenceTransformer
from sentence_transformers import losses
from sentence_transformers.evaluation import InformationRetrievalEvaluator
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader

examples = [
    InputExample(texts=[f"query: {q}", f"passage: {a}"])
    for q, a in zip(questions, answers)
]
train_examples, test_examples = train_test_split(
    examples,
    test_size=0.2,
    random_state=42
)
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=8)
test_dataloader = DataLoader(test_examples, batch_size=8)

model = SentenceTransformer("intfloat/multilingual-e5-base")
loss = losses.MultipleNegativesRankingLoss(model)
EPOCHS = 2
warmup_steps = int(len(train_dataloader) * EPOCHS * 0.1)
model.fit(
    train_objectives=[(train_dataloader, loss)],
    epochs=EPOCHS,
    warmup_steps=warmup_steps,
    output_path="results",
    show_progress_bar=True,
    evaluator=evaluator, 
    evaluation_steps=50,
)

評価については上位10件のヒット率とMRR（Hit Rate@10、MRR@10）でします。

実験結果

評価結果は以下のとおりです。学習前のモデルはtext-embedding-ada-002と同程度かそれより低い性能ですが、学習をすることで性能が大きく改善することを確認できました。とはいえ、学習と評価でQAペアのうち回答が同じ文書集合を対象にしているので、そこは差し引いて見る必要があるかなと思います。

モデル	Hit Rate@10	MRR@10
`text-embedding-ada-002`	0.8355	0.6436
`multilingual-e5-base`	0.7679	0.5352
`multilingual-e5-large`	0.8457	0.6364
`multilingual-e5-base-finetuned`	0.8648	0.6456
`multilingual-e5-large-finetuned`	0.8980	0.6660

気になったので、GPT-3.5 Turboを用いて質問文から回答を生成し、それらのペアを用いて学習した結果を以下に示します。結果を見ると、チューニングしない場合よりは性能が向上していますが、あらかじめ用意されていたデータを用いた場合と比べると低い性能になっています。回答の生成にはHyDEのプロンプトを使用したので、回答が事実に基づいていない点や、回答の文字数の分布が生成したものとそうでないもので大きく異なっていることが原因として考えられます。