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

オープンソースとオープンサイエンスを通じて、人工知能の進歩と民主化を目指しています。

Hugging Face Blog · May 28, 2024

Training and Finetuning Embedding Models with Sentence Transformers

この記事は中国語版もご覧いただけます简体中文。

Sentence Transformers は、検索拡張生成、セマンティック検索、意味的テキスト類似度、パラフレーズマイニングなど、幅広い用途向けの埋め込みモデルを利用・学習するためのPythonライブラリです。このブログ記事では、特定のタスクでのパフォーマンスを向上させるためにSentence Transformerモデルをファインチューニングする方法を紹介します。また、この方法を使って新しいSentence Transformerモデルをゼロから学習することもできます。

Sentence Transformersのファインチューニングには、データセット、損失関数、学習引数、評価器、トレーナーなど、いくつかの要素が含まれます。これらの各要素について詳しく説明し、効果的なモデルを学習するための使用例を提供します。

なぜファインチューニングするのか？
学習の構成要素
データセット
損失関数
学習引数
評価器
- STSbを使ったEmbeddingSimilarityEvaluator
- AllNLIを使ったTripletEvaluator
トレーナー
- コールバック
複数データセットの学習
非推奨
追加リソース
- 学習例
- ドキュメント

なぜファインチューニングするのか？

Sentence Transformerモデルのファインチューニングにより、特定のタスクでのパフォーマンスを大幅に向上させることができます。これは、各タスクに固有の類似性の概念が必要だからです。ニュース記事の見出しの例を考えてみましょう：

「Apple launches the new iPad」
「NVIDIA is gearing up for the next GPU generation」

ユースケースによっては、これらのテキストに対して類似または非類似の埋め込みを望む場合があります。たとえば、ニュース記事の分類モデルでは、両方ともTechnologyカテゴリに属するため、これらのテキストを類似とみなすかもしれません。一方、意味的テキスト類似度や検索モデルでは、意味が異なるため非類似とみなすべきです。

学習の構成要素

Sentence Transformerモデルの学習には、以下の構成要素が含まれます：

データセット：学習と評価に使用するデータ。
損失関数：モデルのパフォーマンスを定量化し、最適化プロセスを導く関数。
学習引数（オプション）：学習のパフォーマンスやトラッキング/デバッグに影響するパラメータ。
評価器（オプション）：学習の前・中・後にモデルを評価するためのツール。
トレーナー：モデル、データセット、損失関数、その他の構成要素をまとめて学習を実行。

それでは、これらの各構成要素について詳しく見ていきましょう。

データセット

SentenceTransformerTrainerは、学習と評価にdatasets.Datasetまたはdatasets.DatasetDictのインスタンスを使用します。Hugging Face Datasets Hubからデータを読み込むことも、CSV、JSON、Parquet、Arrow、SQLなどのさまざまな形式のローカルデータを使用することもできます。

注：Sentence Transformersでそのまま動作するHugging Faceデータセットの多くにはsentence-transformersタグが付いているため、https://huggingface.co/datasets?other=sentence-transformersを閲覧することで簡単に見つけることができます。タスクに役立つ可能性のある学習用データセットを見つけるために、これらのデータセットを閲覧することを強くおすすめします。

Hugging Face Hubのデータ

Hugging Face Hubのデータセットからデータを読み込むには、load_dataset関数を使用します：

from datasets import load_dataset

train_dataset = load_dataset("sentence-transformers/all-nli", "pair-class", split="train")
eval_dataset = load_dataset("sentence-transformers/all-nli", "pair-class", split="dev")

print(train_dataset)
"""
Dataset({
    features: ['premise', 'hypothesis', 'label'],
    num_rows: 942069
})
"""

sentence-transformers/all-nliのような一部のデータセットには、異なるデータ形式の複数のサブセットがあります。データセット名に加えてサブセット名を指定する必要があります。

ローカルデータ（CSV、JSON、Parquet、Arrow、SQL）

一般的なファイル形式のローカルデータがある場合は、load_datasetを使って簡単に読み込むことができます：

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")
# or
dataset = load_dataset("json", data_files="my_file.json")

前処理が必要なローカルデータ

ローカルデータに前処理が必要な場合は、datasets.Dataset.from_dictを使って辞書形式のリストでデータセットを初期化できます：

from datasets import Dataset

anchors = []
positives = []
# Open a file, perform preprocessing, filtering, cleaning, etc.
# and append to the lists

dataset = Dataset.from_dict({
    "anchor": anchors,
    "positive": positives,
})

辞書の各キーは、生成されるデータセットの列になります。

データセットの形式

データセットの形式が選択した損失関数と一致していることを確認することが重要です。これには2つの点を確認する必要があります：

損失関数がLabelを必要とする場合（Loss Overviewテーブルに記載）、データセットに"label"または"score"という名前の列が必要です。
"label"または"score"以外のすべての列はInputsとみなされます（Loss Overviewテーブルに記載）。これらの列の数は、選択した損失関数で有効な入力の数と一致する必要があります。列の名前は重要ではなく、順序のみが重要です。

たとえば、損失関数が(anchor, positive, negative)のトリプレットを受け付ける場合、データセットの1列目、2列目、3列目はそれぞれanchor、positive、negativeに対応します。つまり、1列目と2列目には近い埋め込みが望ましいテキストを、1列目と3列目には遠い埋め込みが望ましいテキストを含める必要があります。そのため、損失関数によってはデータセットの列の順序が重要になります。

["text1", "text2", "label"]という列を持つデータセットで、"label"列に浮動小数点の類似度スコアが含まれている場合を考えてみましょう。このデータセットは、以下の理由でCoSENTLoss、AnglELoss、CosineSimilarityLossで使用できます：

データセットにこれらの損失関数で必要な"label"列があります。
データセットに2つの非ラベル列があり、これらの損失関数で必要な入力の数と一致します。

データセットの列が正しく順序付けされていない場合は、Dataset.select_columnsを使って並べ替えてください。また、sample_id、metadata、source、typeなどの余分な列は、Dataset.remove_columnsを使って削除してください。削除しない場合、これらの列も入力として扱われます。

損失関数

損失関数は、与えられたデータバッチでのモデルのパフォーマンスを測定し、最適化プロセスを導きます。損失関数の選択は、利用可能なデータとターゲットタスクによって異なります。包括的なオプション一覧については、Loss Overviewを参照してください。

ほとんどの損失関数は、学習対象のSentenceTransformerのmodelだけで初期化できます：

from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from sentence_transformers.losses import CoSENTLoss

# Load a model to train/finetune
model = SentenceTransformer("FacebookAI/xlm-roberta-base")

# Initialize the CoSENTLoss
# This loss requires pairs of text and a floating point similarity score as a label
loss = CoSENTLoss(model)

# Load an example training dataset that works with our loss function:
train_dataset = load_dataset("sentence-transformers/all-nli", "pair-score", split="train")
"""
Dataset({
    features: ['sentence1', 'sentence2', 'label'],
    num_rows: 942069
})
"""

学習引数

SentenceTransformersTrainingArgumentsクラスでは、学習のパフォーマンスやトラッキング/デバッグに影響するパラメータを指定できます。オプションですが、これらの引数を試すことで学習効率の向上や学習プロセスの洞察を得ることができます。

Sentence Transformersのドキュメントでは、最も有用な学習引数の一部を概説しています。Training Overview > Training Argumentsで確認することをおすすめします。

SentenceTransformersTrainingArgumentsの初期化例を以下に示します：

from sentence_transformers.training_args import SentenceTransformerTrainingArguments

args = SentenceTransformerTrainingArguments(
    # Required parameter:
    output_dir="models/mpnet-base-all-nli-triplet",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    warmup_ratio=0.1,
    fp16=True,  # Set to False if your GPU can't handle FP16
    bf16=False,  # Set to True if your GPU supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # Losses using "in-batch negatives" benefit from no duplicates
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=100,
    save_strategy="steps",
    save_steps=100,
    save_total_limit=2,
    logging_steps=100,
    run_name="mpnet-base-all-nli-triplet",  # Used in W&B if `wandb` is installed
)

eval_strategyはtransformersバージョン4.41.0で導入されました。それ以前のバージョンでは代わりにevaluation_strategyを使用してください。

評価器

SentenceTransformerTrainerにeval_datasetを指定して学習中の評価損失を取得できますが、学習中に具体的なメトリクスを取得すると便利な場合もあります。そのため、学習の前・中・後に有用なメトリクスでモデルのパフォーマンスを評価する評価器を使用できます。eval_datasetと評価器の両方、どちらか一方、またはどちらも使用しないことも可能です。評価はeval_strategyとeval_stepsの学習引数に基づいて行われます。

Sentence Transformersに付属する実装済みの評価器は以下の通りです：

また、SequentialEvaluatorを使って複数の評価器を1つにまとめて、SentenceTransformerTrainerに渡すこともできます。

評価データがない場合でも、一般的なベンチマークでのモデルのパフォーマンスを追跡したい場合は、Hugging Faceのデータを使ってこれらの評価器を使用できます：

STSbを使ったEmbeddingSimilarityEvaluator

STS Benchmark（別名STSb）は、「A man is feeding a mouse to a snake."のような短いテキストの意味的テキスト類似度を測定するための一般的なベンチマークデータセットです。

Hugging Faceのsentence-transformers/stsbデータセットを自由に閲覧してください。

from datasets import load_dataset
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator, SimilarityFunction

# Load the STSB dataset
eval_dataset = load_dataset("sentence-transformers/stsb", split="validation")

# Initialize the evaluator
dev_evaluator = EmbeddingSimilarityEvaluator(
    sentences1=eval_dataset["sentence1"],
    sentences2=eval_dataset["sentence2"],
    scores=eval_dataset["score"],
    main_similarity=SimilarityFunction.COSINE,
    name="sts-dev",
)
# Run evaluation manually:
# print(dev_evaluator(model))

# Later, you can provide this evaluator to the trainer to get results during training

AllNLIを使ったTripletEvaluator

AllNLIはSNLIとMultiNLIのデータセットを結合したもので、どちらも自然言語推論用のデータセットです。このタスクは伝統的に、2つのテキストが含意、矛盾、またはどちらでもないかを判定するためのものです。含意文と矛盾文は埋め込みモデルの学習に有用な(anchor, positive, negative)トリプレット形式となるため、埋め込みモデルの学習にも採用されています。

このスニペットでは、アンカーテキストと含意テキストの類似度が、アンカーテキストと矛盾テキストの類似度よりも高いとモデルが判断する頻度を評価するために使用されます。テキストの例は「An older man is drinking orange juice at a restaurant.」です。

Hugging Faceのsentence-transformers/all-nliデータセットを自由に閲覧してください。

from datasets import load_dataset
from sentence_transformers.evaluation import TripletEvaluator, SimilarityFunction

# Load triplets from the AllNLI dataset
max_samples = 1000
eval_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split=f"dev[:{max_samples}]")

# Initialize the evaluator
dev_evaluator = TripletEvaluator(
    anchors=eval_dataset["anchor"],
    positives=eval_dataset["positive"],
    negatives=eval_dataset["negative"],
    main_distance_function=SimilarityFunction.COSINE,
    name=f"all-nli-{max_samples}-dev",
)
# Run evaluation manually:
# print(dev_evaluator(model))

# Later, you can provide this evaluator to the trainer to get results during training

トレーナー

SentenceTransformerTrainerは、モデル、データセット、損失関数、その他の構成要素を学習のためにまとめます：

from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    SentenceTransformerModelCardData,
)
from sentence_transformers.losses import MultipleNegativesRankingLoss
from sentence_transformers.training_args import BatchSamplers
from sentence_transformers.evaluation import TripletEvaluator

# 1. Load a model to finetune with 2. (Optional) model card data
model = SentenceTransformer(
    "microsoft/mpnet-base",
    model_card_data=SentenceTransformerModelCardData(
        language="en",
        license="apache-2.0",
        model_name="MPNet base trained on AllNLI triplets",
    )
)

# 3. Load a dataset to finetune on
dataset = load_dataset("sentence-transformers/all-nli", "triplet")
train_dataset = dataset["train"].select(range(100_000))
eval_dataset = dataset["dev"]
test_dataset = dataset["test"]

# 4. Define a loss function
loss = MultipleNegativesRankingLoss(model)

# 5. (Optional) Specify training arguments
args = SentenceTransformerTrainingArguments(
    # Required parameter:
    output_dir="models/mpnet-base-all-nli-triplet",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    warmup_ratio=0.1,
    fp16=True,  # Set to False if GPU can't handle FP16
    bf16=False,  # Set to True if GPU supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss benefits from no duplicates
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=100,
    save_strategy="steps",
    save_steps=100,
    save_total_limit=2,
    logging_steps=100,
    run_name="mpnet-base-all-nli-triplet",  # Used in W&B if `wandb` is installed
)

# 6. (Optional) Create an evaluator & evaluate the base model
dev_evaluator = TripletEvaluator(
    anchors=eval_dataset["anchor"],
    positives=eval_dataset["positive"],
    negatives=eval_dataset["negative"],
    name="all-nli-dev",
)
dev_evaluator(model)

# 7. Create a trainer & train
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=loss,
    evaluator=dev_evaluator,
)
trainer.train()

# (Optional) Evaluate the trained model on the test set, after training completes
test_evaluator = TripletEvaluator(
    anchors=test_dataset["anchor"],
    positives=test_dataset["positive"],
    negatives=test_dataset["negative"],
    name="all-nli-test",
)
test_evaluator(model)

# 8. Save the trained model
model.save_pretrained("models/mpnet-base-all-nli-triplet/final")

# 9. (Optional) Push it to the Hugging Face Hub
model.push_to_hub("mpnet-base-all-nli-triplet")

この例では、microsoft/mpnet-base（まだSentence Transformerモデルではないベースモデル）からファインチューニングしています。これには、all-mpnet-base-v2のような既存のSentence Transformerモデルをファインチューニングする場合よりも多くの学習データが必要です。

このスクリプトを実行すると、tomaarsen/mpnet-base-all-nli-tripletモデルがアップロードされました。コサイン類似度を用いたトリプレット精度、つまりcosine_similarity(anchor, positive) > cosine_similarity(anchor, negative)の割合は、開発セットで90.04%、テストセットで91.5%でした。参考までに、学習前のmicrosoft/mpnet-baseモデルのdevセットでのスコアは68.32%でした。

ベースモデル、言語、ライセンス、評価結果、学習・評価データセット情報、ハイパーパラメータ、学習ログなど、すべての情報は自動生成されたモデルカードに保存されます。手間をかけずに、アップロードしたモデルには、潜在的なユーザーがモデルの適合性を判断するために必要なすべての情報が含まれます。

コールバック

Sentence Transformersのトレーナーは、以下を含むさまざまなtransformers.TrainerCallbackサブクラスをサポートしています：

WandbCallback（wandbがインストールされている場合、W&Bに学習メトリクスをログ）
TensorBoardCallback（tensorboardが利用可能な場合、TensorBoardに学習メトリクスをログ）
CodeCarbonCallback（codecarbonがインストールされている場合、学習中の炭素排出量を追跡）

必要な依存関係がインストールされていれば、これらは自動的に使用され、明示的に指定する必要はありません。

これらのコールバックの詳細や独自のコールバックの作成方法については、Transformers Callbacksのドキュメントを参照してください。

複数データセットの学習

トップパフォーマンスのモデルは、複数のデータセットを同時に使用して学習されることがよくあります。SentenceTransformerTrainerは、データセットを同じ形式に変換することなく複数のデータセットで学習できるようにすることで、このプロセスを簡素化します。各データセットに異なる損失関数を適用することも可能です。複数データセット学習の手順は以下の通りです：

train_datasetとeval_datasetとして、datasets.Datasetインスタンスの辞書（またはdatasets.DatasetDict）を使用します。
（オプション）異なるデータセットで異なる損失関数を使用する場合、データセット名から損失関数への辞書を使用します。

各学習/評価バッチには、1つのデータセットからのサンプルのみが含まれます。複数のデータセットからバッチをサンプリングする順序は、MultiDatasetBatchSamplers列挙型によって決定され、multi_dataset_batch_samplerを介してSentenceTransformersTrainingArgumentsに渡すことができます。有効なオプションは以下の通りです：

MultiDatasetBatchSamplers.ROUND_ROBIN：いずれかのデータセットが使い果たされるまで、各データセットからラウンドロビン方式でサンプリングします。この戦略は各データセットのすべてのサンプルを使用しない場合がありますが、各データセットからの等しいサンプリングを保証します。
MultiDatasetBatchSamplers.PROPORTIONAL（デフォルト）：各データセットのサイズに比例してサンプリングします。この戦略は各データセットのすべてのサンプルが使用されることを保証し、大きなデータセットからより頻繁にサンプリングされます。

マルチタスク学習は非常に効果的であることが証明されています。たとえば、Huang et al. 2024は、MultipleNegativesRankingLoss、CoSENTLoss、およびインバッチネガティブなしでハードネガティブのみを使用するMultipleNegativesRankingLossのバリエーションを採用して、中国語で最先端のパフォーマンスを達成しました。また、MatryoshkaLossを適用して、モデルがMatryoshka Embeddingsを生成できるようにしました。

複数データセット学習の例を以下に示します：

from datasets import load_dataset
from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer
from sentence_transformers.losses import CoSENTLoss, MultipleNegativesRankingLoss, SoftmaxLoss

# 1. Load a model to finetune
model = SentenceTransformer("bert-base-uncased")

# 2. Loadseveral Datasets to train with
# (anchor, positive)
all_nli_pair_train = load_dataset("sentence-transformers/all-nli", "pair", split="train[:10000]")
# (premise, hypothesis) + label
all_nli_pair_class_train = load_dataset("sentence-transformers/all-nli", "pair-class", split="train[:10000]")
# (sentence1, sentence2) + score
all_nli_pair_score_train = load_dataset("sentence-transformers/all-nli", "pair-score", split="train[:10000]")
# (anchor, positive, negative)
all_nli_triplet_train = load_dataset("sentence-transformers/all-nli", "triplet", split="train[:10000]")
# (sentence1, sentence2) + score
stsb_pair_score_train = load_dataset("sentence-transformers/stsb", split="train[:10000]")
# (anchor, positive)
quora_pair_train = load_dataset("sentence-transformers/quora-duplicates", "pair", split="train[:10000]")
# (query, answer)
natural_questions_train = load_dataset("sentence-transformers/natural-questions", split="train[:10000]")

# Combine all datasets into a dictionary with dataset names to datasets
train_dataset = {
    "all-nli-pair": all_nli_pair_train,
    "all-nli-pair-class": all_nli_pair_class_train,
    "all-nli-pair-score": all_nli_pair_score_train,
    "all-nli-triplet": all_nli_triplet_train,
    "stsb": stsb_pair_score_train,
    "quora": quora_pair_train,
    "natural-questions": natural_questions_train,
}

# 3. Load several Datasets to evaluate with
# (anchor, positive, negative)
all_nli_triplet_dev = load_dataset("sentence-transformers/all-nli", "triplet", split="dev")
# (sentence1, sentence2, score)
stsb_pair_score_dev = load_dataset("sentence-transformers/stsb", split="validation")
# (anchor, positive)
quora_pair_dev = load_dataset("sentence-transformers/quora-duplicates", "pair", split="train[10000:11000]")
# (query, answer)
natural_questions_dev = load_dataset("sentence-transformers/natural-questions", split="train[10000:11000]")

# Use a dictionary for the evaluation dataset too, or just use one dataset or none at all
eval_dataset = {
    "all-nli-triplet": all_nli_triplet_dev,
    "stsb": stsb_pair_score_dev,
    "quora": quora_pair_dev,
    "natural-questions": natural_questions_dev,
}

# 4. Load several loss functions to train with
# (anchor, positive), (anchor, positive, negative)
mnrl_loss = MultipleNegativesRankingLoss(model)
# (sentence_A, sentence_B) + class
softmax_loss = SoftmaxLoss(model)
# (sentence_A, sentence_B) + score
cosent_loss = CoSENTLoss(model)

# Create a mapping with dataset names to loss functions, so the trainer knows which loss to apply where
# Note: You can also just use one loss if all your training/evaluation datasets use the same loss
losses = {
    "all-nli-pair": mnrl_loss,
    "all-nli-pair-class": softmax_loss,
    "all-nli-pair-score": cosent_loss,
    "all-nli-triplet": mnrl_loss,
    "stsb": cosent_loss,
    "quora": mnrl_loss,
    "natural-questions": mnrl_loss,
}

# 5. Define a simple trainer, although it's recommended to use one with args & evaluators
trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=losses,
)
trainer.train()

# 6. Save the trained model and optionally push it to the Hugging Face Hub
model.save_pretrained("bert-base-all-nli-stsb-quora-nq")
model.push_to_hub("bert-base-all-nli-stsb-quora-nq")

非推奨

Sentence Transformer v3リリース以前は、すべてのモデルがSentenceTransformer.fitメソッドを使って学習されていました。このメソッドを非推奨とするのではなく、v3.0以降はこのメソッドが裏でSentenceTransformerTrainerを使用するようになりました。つまり、既存の学習コードは引き続き動作し、マルチGPU学習や損失ログなどの新機能でアップグレードされることになります。ただし、新しい学習アプローチの方がはるかに強力であるため、新しい学習スクリプトは新しいアプローチで書くことを推奨します。

追加リソース

学習例

以下のページには、説明とコードへのリンク付きの学習例が含まれています。学習ループに慣れるために、これらを閲覧することをおすすめします：

ドキュメント

また、Sentence Transformersについて詳しく学ぶために、以下のページも役立つ場合があります：

インストール
クイックスタート
使い方
事前学習済みモデル
学習の概要（このブログ記事はTraining Overviewドキュメントの要約です）
データセットの概要
損失関数の概要
APIリファレンス

最後に、興味を持っていただける可能性のある高度なページをいくつか紹介します：

0 Comments

No comments yet. Be the first to share your thoughts.