Sentence Transformers は、検索拡張生成、セマンティック検索、意味的テキスト類似度など、さまざまな用途で埋め込みモデルやリランキングモデルを利用・学習するための Python ライブラリです。v6.0 では、4 種類目のモデルとして MultiVectorEncoder が導入されました。これは ColBERT 形式の後期インタラクション（late interaction）検索に対応するもので、完全な学習手法も提供されています。この記事では、マルチベクトルモデルを微調整して、汎用検索モデルを自分のデータ上で上回る方法を紹介します。この手法を使えば、強力な新しいマルチベクトルモデルをゼロから学習することもできます。以下の内容はすべて、pip install -U "sentence-transformers[train]" で実行できます。
マルチベクトルモデルの微調整には、複数のコンポーネントが関わります。モデル本体、データセット、損失関数、学習パラメータ、評価器、そしてトレーナークラスです。ここからは、これらのコンポーネントを順に紹介し、実例を通して高性能なマルチベクトルモデルを微調整する方法を説明します。
最後に、評価のセクションでは、この記事で紹介する手法を使い、1 枚の RTX 3090 で 14.5 時間かけて学習した multi-vector-encoder/mLateOn-medical モデルが、私の医療検索ベンチマークにおいて、見つけられたすべての汎用検索モデルを容易に上回ったことを紹介します。対象には、密ベクトル、スパース、語彙ベース、マルチベクトルの各モデルが含まれます。
密な埋め込みモデル、スパース埋め込みモデル、またはリランキングモデルの微調整に関心がある場合は、以前公開した埋め込みモデルの学習と微調整、スパース埋め込みモデルの学習と微調整、およびリランキングモデルの学習と微調整の記事もご覧ください。
この記事では、マルチベクトルモデルを学習する方法を説明します。マルチベクトルモデルの利用方法（モデルのロード、エンコード、ベクトルデータベースでのインデックス構築など）については、関連する記事Sentence Transformers でマルチベクトル（後期インタラクション）埋め込みモデルを利用するをご覧ください。
マルチベクトルモデルとは？
密な埋め込みモデルは、テキスト全体を 1 つのベクトルに圧縮します。そのような 2 つの要約の類似度は、1 回の内積計算だけで求められます。マルチベクトルモデル（後期インタラクションモデル、または ColBERT 形式モデルとも呼ばれます）は、この圧縮処理を行いません。各トークンに対して小さなベクトルを 1 つ保持し、MaxSim 演算子を使ってクエリとドキュメントのスコアを計算します。各クエリトークンについて、最もよく一致するドキュメントトークンを見つけ、そのスコアを合計します。トークン単位のマッチングにより、単一ベクトルに平均化して圧縮せざるを得ない細粒度のシグナルを保持できます。そのため、通常はより高い検索性能が得られますが、より大きなインデックスが必要になります。
関連する記事マルチベクトル埋め込みモデルでは、アーキテクチャ、エンコード、スコアリング、インデックス構築について詳しく説明しているため、ここでは詳しく扱わず、さっそく学習の話に進みます。
なぜ微調整するのか？
特定の分野に合わせてマルチベクトルモデルを微調整すると、検索性能を大幅に向上させられます。ウェブ検索、法務調査、コード検索、科学文献レビューでは、語彙、クエリの形式、関連性の定義がそれぞれ異なります。クエリとドキュメントをトークン単位で照合するマルチベクトルモデルは、単一ベクトルモデルでは通常平均化されてしまう、分野固有の細粒度なシグナルを捉えられます。そのため、分野内の微調整データが適量しかない場合でも、優れた性能を得ることができます。
また、公開されている検索モデルの大半は、短いテキストチャンクを前提に設定されています。従来の ColBERT チェックポイントはドキュメントを 180 または 300 トークンに切り詰めます。多くの主要な密なモデルも 256 または 512 トークンで切り詰められます。これは、MS MARCO 形式の学習データの多くが、この長さを超えることがほとんどないためです。ドキュメントが長い場合、これらのモデルはスコアリング前にドキュメントの大部分をひそかに破棄してしまいます。私の医療ベンチマークでは、テキストチャンクの平均長は 941 トークンでした。測定したところ、この切り詰め処理だけで NDCG@10 が最大 0.24 低下しました。これは、モデルアーキテクチャの違いによる影響をはるかに上回ります。自分でモデルを学習すれば、自分のデータの要件に合わせてドキュメント長を設定できます。
LightOn もコード検索の分野で同じ問題に直面しました。汎用モデルの LateOn では十分な性能が得られなかったため、LateOn-Code を学習しました。分野が医療、法律、金融、あるいは社内ドキュメントのいずれであっても、自分の用途に特化した公式モデルが用意されているとは限りません。この記事では、数時間と民生用 GPU 1 枚だけで、自分の分野に適したモデルを構築する方法を紹介します。
学習コンポーネント
MultiVectorEncoder モデルの学習には、以下のコンポーネントが関わります。
-
モデル：微調整するモデル、またはゼロから構築するアーキテクチャ。
-
データセット：学習と評価に使用するデータ。
-
損失関数：モデルの性能を測定し、最適化プロセスを導くための関数。
-
学習パラメータ（任意）：学習性能、追跡、デバッグに影響するパラメータ。
-
評価器（任意）：学習前、学習中、または学習後にモデルを評価するクラス。
-
トレーナー：すべての学習コンポーネントを統合するもの。
ここからは、各コンポーネントについて詳しく見ていきます。
モデル
マルチベクトルの学習では、複数の出発点から選択できます。どの出発点を選ぶかは、想像以上に重要になる場合があります。
既存のマルチベクトルモデルを微調整する
既存のマルチベクトルモデルをさらに微調整したい場合、アーキテクチャについて心配する必要はまったくありません。
from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder( "lightonai/mLateOn-unsupervised", model_kwargs={"torch_dtype": "float32"}, processor_kwargs={"model_max_length": 8192}, # tokenizer 層でのトークン数の上限)このチェックポイントには、その