Sentence Transformers を用いたリランカーモデルのトレーニングとファインチューニング

Hugging Face は、Sentence Transformers ライブラリを用いたリランカー（Cross Encoder）モデルのトレーニングとファインチューニングのためのフレームワークを詳細に紹介しています。ドメイン固有データで学習させることで、開発者は小型で効率的なリランカーを作成でき、専門的な検索タスクにおいてははるかに大きな汎用モデルを上回る性能を発揮します。

リランカーモデル と 埋め込みモデル の違い

Cross Encoder のリランカーモデルは、テキストのペア（例: クエリとドキュメント）を同時に共有ニューラルネットワークに通し、単一のスコアを出力して関連性を評価します。これは、テキストを個別にベクトル化し、距離指標で類似度を計算する bi‑encoder（埋め込みモデル）とは異なります。

Cross Encoder はすべてのペアを処理しなければならないため計算コストが高くなりますが、一般に精度が高くなります。そのため、業界標準の高性能検索は「取得 → リランク」の二段階パイプラインとなっており、まず高速な埋め込みモデルで候補を取得し、上位 k 件を Cross Encoder で再評価します。

コアトレーニングコンポーネント

リランカーモデルのトレーニングは、CrossEncoderTrainer を通じて統合される 5 つの主要コンポーネントで構成されます。

1. データセット

CrossEncoderTrainer は datasets.Dataset または datasets.DatasetDict オブジェクトをサポートします。データは Hugging Face Hub から取得するか、ローカルの CSV、JSON、Parquet、Arrow、SQL ファイルから読み込めます。選択したロス関数と互換性を保つため、データセットにはラベル用に "label"、"labels"、"score"、または "scores" のいずれかの列が必要です。残りの列はロス関数が要求する入力数と一致している必要があります。

2. ハードネガティブマイニング

モデル性能はネガティブサンプルの質に大きく依存します。完全に無関係な "soft negative" に対し、"hard negative" は見た目は関連しそうだが実際は違うパッセージです。Sentence Transformers は mine_hard_negatives 関数を提供しており、これらの難しい例を抽出することでモデルにより精緻な判別を学習させます。

3. ロス関数

ロス関数は利用可能なデータに基づいて最適化プロセスを導きます。LambdaLoss や ListNetLoss といった学習順位付けロスもありますが、ラベル付きペアに対しては BinaryCrossEntropyLoss がシンプルでありながら高い効果を示します。

4. トレーニング引数

カスタマイズは CrossEncoderTrainingArguments で行い、学習率、バッチサイズ、ウォームアップ比率、精度設定（FP16/BF16）などを調整できます。バッチ内ネガティブを利用するロスの場合は、batch_sampler=BatchSamplers.NO_DUPLICATES の設定が特に推奨されます。

5. 評価指標

単なるロス値だけでなく性能を追跡するため、Sentence Transformers は以下の組み込み評価器を提供します。

  • CrossEncoderClassificationEvaluator: 二値または多クラスラベル用。
  • CrossEncoderCorrelationEvaluator: 類似度スコア（例: STSb データセット）用。
  • CrossEncoderRerankingEvaluator: クエリ、ポジティブ、ネガティブを用いたリランク性能評価用。
  • CrossEncoderNanoBEIREvaluator: 英語リランク向けの軽量評価器。

マルチデータセットトレーニング

CrossEncoderTrainer は、形式が異なる、あるいはロス関数が異なる複数のデータセットを同時にトレーニングできるように設計されています。データセットの辞書と、必要に応じてロス関数の辞書を渡すだけです。サンプリング戦略は次の 2 種類があります。

  • ROUND_ROBIN: 各データセットから均等にサンプリングし、いずれかが尽きるまで繰り返す。
  • PROPORTIONAL: 各データセットのサイズに比例してサンプリングし、すべてのサンプルが使用されるようにする。

パフォーマンス評価と結果

実践例として、ModernBERT-base をベースにしたリランカーを GooAQ データセットの 99k クエリ‑回答ペアでファインチューニングしました。BinaryCrossEntropyLoss とハードネガティブマイニングを組み合わせた結果、得られたモデル（tomaarsen/reranker-ModernBERT-base-gooaq-bce）は、サイズが最大 4 倍の 13 種類の一般的なオープンソースリランカーを上回る性能を示しました。

GooAQ NDCG@10 結果（上位 30 リランク）

Model Parameters Realistic NDCG@10 Evaluation NDCG@10
Retriever only (No reranking) - 59.12 59.12
BAAI/bge-reranker-large 560M 73.20 77.46
mixedbread-ai/mxbai-rerank-large-v2 1.54B 75.40 80.04
ModernBERT-base-gooaq-bce 150M 77.14 83.51
ModernBERT-large-gooaq-bce 396M 79.42 85.81

テクニカルトレーニングのコツ

  • 過学習防止: Cross Encoder は過学習しやすいです。load_best_model_at_end と metric_for_best_model を設定した評価器を使用し、ベストモデルを保存しましょう。
  • ネガティブのバランス: ハードネガティブだけに依存すると、簡単なタスクでの性能が低下します。ランダムネガティブとハードネガティブを混合することでこの問題を緩和できます。
  • 効率性: ドメイン固有データで小型リランカーをファインチューニングすれば、検索精度の向上と同時に、巨大な汎用モデルを使用した場合に比べて推論レイテンシを大幅に削減できます。

Sources