Sentence Transformers を用いた埋め込みモデルのトレーニングとファインチューニング

Hugging Face は、Sentence Transformers ライブラリを使用した埋め込みモデルのトレーニングとファインチューニングのための効率的なワークフローを詳述しています。このフレームワークにより、開発者は汎用埋め込みモデルを、検索拡張生成(RAG)やセマンティック検索、パラフレーズマイニングなどのタスクで必要とされる特定の類似性概念に合わせて適応させることができます。

タスク固有の類似性のためのファインチューニングの必要性

ファインチューニングは重要です。なぜなら、アプリケーションごとに「類似性」の定義が異なるからです。例えば、異なる企業(例:Apple と NVIDIA)に関する2つのニュース見出しは、ニュース分類モデルでは両方が「テクノロジー」カテゴリに属するため類似と見なされるかもしれませんが、セマンティックテキスト類似性や検索モデルでは異なる出来事を記述しているため非類似と見なす必要があります。

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

  1. Dataset: トレーニングおよび評価データで、通常は datasets.Dataset または datasets.DatasetDict インスタンスとしてロードされます。
  2. Loss Function: 利用可能なデータと対象タスクに基づき、モデルの性能を定量化し最適化を導く関数です。
  3. Training Arguments: SentenceTransformersTrainingArguments を通じて指定できるオプションパラメータで、トレーニングの効率、トラッキング、デバッグを制御します。
  4. Evaluator: トレーニング前・中・後に具体的な指標でモデル性能を評価するためのオプションツールです。
  5. Trainer: SentenceTransformerTrainer は、モデル、データセット、ロス関数、その他のコンポーネントを統合するトレーナーです。

データセットの要件とフォーマット

データセットは Hugging Face Hub(しばしば sentence-transformers タグが付く)から取得するか、CSV、JSON、Parquet、Arrow、SQL 形式でローカルにロードできます。選択したロス関数と互換性を保つため、データセットは特定のフォーマット規則に従う必要があります:

  • Labels: ロス関数がラベルを必要とする場合、データセットは label または score という名前の列を含んでいる必要があります。
  • Inputs: ラベル以外のすべての列は入力として扱われます。これらの列の数と順序はロス関数の要件に合わせる必要があります(例:(anchor, positive, negative) のトリプレット形式)。

ロス関数とトレーニング引数

ロス関数はトレーニング対象のモデルと共に初期化されます。ロスの選択は利用可能なデータに依存します(例:浮動小数点の類似度スコアを持つペアには CoSENTLoss)。

SentenceTransformersTrainingArguments を使用してトレーニング性能を調整できます。これには num_train_epochsper_device_train_batch_sizewarmup_ratio、および fp16bf16 といったハードウェア固有の設定が含まれます。「バッチ内ネガティブ」を利用するロスの場合、batch_sampler=BatchSamplers.NO_DUPLICATES 引数の使用が推奨されます。

モデル評価

トレーナーは評価ロスを提供できますが、専用の評価器はタスク固有の指標を提供します。利用可能な評価器は以下の通りです:

  • EmbeddingSimilarityEvaluator: 類似度スコアを持つペア向け(例:STSb ベンチマークを使用)。
  • TripletEvaluator: (anchor, positive, negative) のペア向け(例:AllNLI データセットを使用)。
  • InformationRetrievalEvaluator: クエリ、コーパス、関連文書向け。
  • BinaryClassificationEvaluator: クラスラベルを持つペア向け。

複数の評価器は単一の SequentialEvaluator に統合でき、トレーニング中にさまざまな指標を同時に追跡できます。

高度なトレーニングワークフロー

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

高性能モデルはしばしば複数のデータセットを同時にトレーニングする必要があります。SentenceTransformerTrainer はデータセットの辞書と対応するロス関数の辞書を受け取ることでこれをサポートし、同一トレーニング実行内でデータセットごとに異なるロスを適用できます。

複数データセットからのサンプリングは MultiDatasetBatchSamplers を用いて、以下の2つの戦略で処理できます:

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

SentenceTransformerTrainer への移行

Sentence Transformers v3.0 のリリースに伴い、従来の SentenceTransformer.fit メソッドは内部で SentenceTransformerTrainer を使用するようになりました。レガシーコードは引き続き機能しますが、Hugging Face はマルチGPUトレーニングや改善されたロスロギングなどの高度な機能を活用するために新しいトレーナーアプローチの採用を推奨しています。

パフォーマンス例

提供された例では、MultipleNegativesRankingLoss を用いて AllNLI のトリプレットで microsoft/mpnet-base モデルをファインチューニングした結果、パフォーマンスが大幅に向上しました。ベースモデルは開発セットで 68.32% のスコアだったのに対し、ファインチューニング後のモデルは開発セットで 90.04%、テストセットで 91.5% を達成し、コサイン類似度によるトリプレット精度で測定されました。

Sources