Sentence Transformers 6.0 MultiVectorEncoder: 訓練とファインチューニングガイド

あらまし

Hugging Face は Sentence Transformers v6.0 に MultiVectorEncoder モデルタイプを追加し、ColBERTスタイルのラテインタラクションリトリーバー（例: multi-vector-encoder/mLateOn-medical）を、単一の RTX 3090 で 14.5 時間でファインチューニングするためのステップバイステップの訓練スクリプトをリリースしました。これにより、医療リトリーブベンチマークで 50 以上の密な、スパースな、語彙的な、マルチベクトルベースラインと比較して 0.9139 NDCG@10 の最高スコアを達成しました。

---

マルチベクトル（ラテインタラクション）モデルとは？

マルチベクトルモデルは、テキスト全体を単一のベクトルに圧縮するのではなく、各トークンごとに1つの埋め込みを保持します。

• リトリーブには MaxSim 演算子を使用：各クエリトークンが最も一致するドキュメントトークンを見つけて、スコアを合計します。
• トークンレベルのマッチングにより、単一ベクトルモデルが平均化してしまう細かい関連性信号を保持でき、インデックスが大きくなる代わりに、より強力なリトリーブが可能になります。
• 詳細なアーキテクチャ、エンコーディング、スコアリング、インデキシングについては、『Sentence Transformers によるマルチベクトル（ラテインタラクション）埋め込みモデル』という補足記事をご覧ください。

"密な埋め込みモデルは、テキスト全体を1つのベクトルに圧縮し、類似度は2つの要約間の1つのドット積で表されます。マルチベクトルモデルは……各トークンごとに小さなベクトルを保持し、MaxSim演算子でクエリをドキュメントとスコアリングします。" – ブログ抜粋

---

なぜマルチベクトルモデルをファインチューニングするのか？

ファインチューニングにより、モデルは特定のドメイン（例：医療、法務、コード）の語彙、クエリスタイル、関連性の概念に適応します。主な観察点：

• ドメイン信号 はトークン単位で捉えられるため、わずかなドメイン内データでも大きな向上が得られます。
• 多くの公開チェックポイントはドキュメントを 180～512 トークンで切り捨てていますが、長文（MIRIAD医療セットの平均 941 トークン）では、切り捨てにより最大 0.24 NDCG@10 の損失が生じます。
• 教師なし チェックポイント（対照学習前だがまだ教師あり学習前）から開始すると、完全に教師ありのチェックポイントよりも一貫して優れたドメイン適応性能を示します。
• 全体のファインチューニングパイプラインは1台のコンシューマーGPUで実行可能であり、多くのチームがドメイン固有のリトリーバーにアクセスできるようになります。

---

訓練コンポーネントの概要

コンポーネント	役割
モデル	MultiVectorEncoder インスタンス – 事前学習済みチェックポイントまたはベーストランスフォーマー上で構築された新規モデル
データセット	datasets.Dataset または DatasetDict – クエリ・ドキュメントペア（または損失に必要な他のフォーマット）
損失	バッチ内ネガティブ（CachedMultiVectorMultipleNegativesRankingLoss）または知識蒸留損失
訓練引数	MultiVectorEncoderTrainingArguments – バッチサイズ、学習率、プロンプトなど
評価者	NDCG@10、acc@1 などに使用する MultiVectorInformationRetrievalEvaluator
トレーナー	上記を統合する MultiVectorEncoderTrainer

以下の各セクションは独立して読むことができます。コードスニペットは pip install -U "sentence-transformers[train]" で実行可能です。

---

モデル選択と設定

既存のチェックポイントのファインチューニング

from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192},  # 全長ドキュメントを許可
)
# 組み込みの長さ制限（例：180～512トークン）を解除
model[0].query_length = None
model[0].document_length = None
# オプション：インデックスを小さくするために句読点トークンをスキップ
import string
model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)

• チェックポイントには既にクエリ/ドキュメントマーカートークン、プロジェクションヘッド、スコアリングスキップリストが含まれています。
• 制限を解除することで、1,400トークンの医療パラグラフを完全に読み込めるようになります。

ベーストランスフォーマーからの新規モデル構築

from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder(
    "answerdotai/ModernBERT-base",
    model_kwargs={"torch_dtype": "float32"},
)

• パイプラインはランダムなトークンレベルプロジェクション（128次元）と通常の ColBERTモジュールを追加します。
• たとえ強力な密なバックボーン（例：Alibaba-NLP/gte-modernbert-base）でも、25,000ペアで訓練するとチェックポイント性能に近づきます。

どのスタートポイントを選ぶべきか？

25,000のMIRIADペアでの実証比較：

チェックポイント	ゼロショット NDCG@10	ファインチューニング後	Δ
lightonai/mLateOn-unsupervised	0.9087	0.9398	+0.0311
lightonai/mLateOn	0.9277	0.9319	+0.0042
lightonai/LateOn-unsupervised	0.9026	0.9206	+0.0180
lightonai/LateOn	0.9185	0.9105	–0.0080
lightonai/GTE-ModernColBERT-v1	0.9198	0.9007	–0.0191
新規ヘッド（gte-modernbert-base）	–	0.9177	–
結論： 教師なしで学習されたチェックポイントが最も適応しやすく、完全に教師ありのチェックポイントはしばしば性能が低下する。

---

データセットの準備

トレーナーは任意の datasets.Dataset（Hubまたはローカル）を受け入れます。必要なフォーマットは損失に依存します：

• ラベル列（label または score）：損失が教師信号を必要とする場合。
• 入力列 は順序付き。最初の列がクエリとして扱われ、以降の列がドキュメント（router_mapping で上書き可能）。

Hubから読み込む（MIRIADの例）

from datasets import load_dataset
train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train")
print(train_dataset)
# Dataset({features: ['question', 'passage_text'], num_rows: 4_467_542})

• 各行は (question, passage_text) ペアを提供。パラグラフの平均長は 941 トークン。

ローカルCSV/JSONの例

from datasets import load_dataset
dataset = load_dataset("csv", data_files="my_file.csv")
# または JSON
# dataset = load_dataset("json", data_files="my_file.json")

• カスタム前処理が必要な場合、Dataset.from_dict を使用します。

---

損失関数の選択

質問-回答ペアには、バッチ内ネガティブ と GradCache を使用することが推奨されます：

from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss
loss = CachedMultiVectorMultipleNegativesRankingLoss(
    model=model,
    mini_batch_size=16,  # メモリ制御用のチャンクサイズ；ブログでは有効バッチサイズ128
)

• mini_batch_size はメモリを制御しますが、有効な対照バッチサイズは128のままです。
• スケールパラメータ：MaxSimスコアではデフォルトの scale=1.0 を維持してください。密なモデルのデフォルト（scale=20.0）を使用すると勾配が飽和します。
• 知識蒸留には MultiVectorDistillKLDivLoss を使用（ブログ例では使用せず）。

---

訓練引数

最良の結果をもたらしたキーパラメータ：

from sentence_transformers import MultiVectorEncoderTrainingArguments
from sentence_transformers.base.sampler import BatchSamplers
args = MultiVectorEncoderTrainingArguments(
    output_dir="models/mLateOn-medical",
    num_train_epochs=1,
    per_device_train_batch_size=128,   # GradCacheによる有効バッチ
    per_device_eval_batch_size=16,
    learning_rate=1e-4,
    warmup_steps=0.05,
    prompts={"question": "[Q] ", "passage_text": "[D] "},
    fp16=False,
    bf16=True,
    batch_sampler=BatchSamplers.NO_DUPLICATES,
    eval_strategy="steps",
    eval_steps=0.1,
    save_strategy="steps",
    save_steps=0.05,
    logging_steps=0.01,
    run_name="mLateOn-medical",
)

• プロンプト は明示的に指定する必要があります。モデルは保存されたマーカーを自動的に適用しません。
• max_length を未設定のままにすることで、訓練時に完全なドキュメント長を認識できます。
• 5e⁻⁶ → 2e⁻⁴ のスイープの結果、より高い学習率（1e-4）が最良の性能を示しました。

---

リトリーブ用の評価者

最も有用な評価者は、ホールドアウトクエリセットとディストラクターを含むコアスから構築された MultiVectorInformationRetrievalEvaluator です：

from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator
# コアス、クエリ、関連ドキュメントマッピングの構築（ブログの完全ループを参照）
evaluator = MultiVectorInformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant_docs,
    name="miriad-dev",
    batch_size=16,
)

• 飽和を避けるために、ハードディストラクターのパラグラフを含めるべきです。ブログでは、10,000のゴールドセットに約19万のランダムな訓練パラグラフを追加しました。

---

完全な訓練スクリプト

以下のスクリプトは mLateOn-medical モデルを再現します：

import logging, string, traceback
from datasets import load_dataset
from sentence_transformers import (
    MultiVectorEncoder,
    MultiVectorEncoderModelCardData,
    MultiVectorEncoderTrainer,
    MultiVectorEncoderTrainingArguments,
)
from sentence_transformers.base.sampler import BatchSamplers
from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator

logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO)

def main():
    # 1️⃣ 教師なしチェックポイントを読み込み、制限を解除
    model = MultiVectorEncoder(
        "lightonai/mLateOn-unsupervised",
        model_kwargs={"torch_dtype": "float32"},
        processor_kwargs={"model_max_length": 8192},
        model_card_data=MultiVectorEncoderModelCardData(
            language="en",
            license="apache-2.0",
            model_name="mLateOn ファインチューニング済み MIRIAD 医療リトリーブ",
        ),
    )
    model[0].query_length = None
    model[0].document_length = None
    model[2].skiplist_words = list(string.punctuation)
    model[2].resolve_with_tokenizer(model.tokenizer)

    # 2️⃣ 医療QAペア100万件を読み込み
    train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train").select(range(1_000_000))

    # 3️⃣ 損失の定義（GradCacheバッチ内ネガティブ）
    loss = CachedMultiVectorMultipleNegativesRankingLoss(model=model, mini_batch_size=16)

    # 4️⃣ 軽量の開発評価者（500クエリ）
    eval_split = load_dataset("tomaarsen/miriad-4.4M-split", split="eval")
    corpus, queries, relevant_docs, passage_to_id = {}, {}, {}, {}
    for idx, row in enumerate(eval_split):
        if row["passage_text"] not in passage_to_id:
            pid = f"p{len(passage_to_id)}"
            passage_to_id[row["passage_text"]] = pid
            corpus[pid] = row["passage_text"]
        if idx < 500:
            qid = f"q{idx}"
            queries[qid] = row["question"]
            relevant_docs[qid] = {passage_to_id[row["passage_text"]]}  # 関連ドキュメント
    dev_evaluator = MultiVectorInformationRetrievalEvaluator(
        queries=queries, corpus=corpus, relevant_docs=relevant_docs, name="miriad-dev", batch_size=16
    )

    # 5️⃣ 訓練引数（前のセクションを参照）
    args = MultiVectorEncoderTrainingArguments(
        output_dir="models/mLateOn-medical",
        num_train_epochs=1,
        per_device_train_batch_size=128,
        per_device_eval_batch_size=16,
        learning_rate=1e-4,
        warmup_steps=0.05,
        prompts={"question": "[Q] ", "passage_text": "[D] "},
        fp16=False,
        bf16=True,
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.05,
        logging_steps=0.01,
        run_name="mLateOn-medical",
    )

    # 6️⃣ トレーナーと訓練
    trainer = MultiVectorEncoderTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        loss=loss,
        evaluator=dev_evaluator,
    )
    trainer.train()

    # 7️⃣ 保存とHubへのプッシュ（オプション）
    model.save_pretrained("models/mLateOn-medical/final")
    try:
        model.push_to_hub("mLateOn-medical")
    except Exception:
        logging.error("モデルのアップロードに失敗しました: \n" + traceback.format_exc())

if __name__ == "__main__":
    main()

• 実行時間：単一の RTX 3090 で 14.5 時間（ピークVRAM 17.5 GB）。
• データ効率：10万ペア（約75分）で、全100万ペア実行の結果と NDCG@10 で 0.012 以内の差異。

---

インデックスサイズと最適化

マルチベクトルインデックスは、各トークンがベクトルを生成するため、大きくなります（941トークンのパラグラフあたり約878ベクトル）。20万件のパラグラフのfp16ストレージは約45GBです。

トークンプーリング

HierarchicalTokenPooling(pool_factor=4) により、ベクトル数を4分の1に削減し、NDCG@10損失は < 0.0033 に抑えられます。

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
pooling = HierarchicalTokenPooling(pool_factor=4)
embeddings = model.encode_document(passages, token_pooling=pooling)

• ベクトル数の1/4 → 約11GB、NDCG@10 ≈ 0.8991。

量子化とプルーニング（PLAID）

1ビットリジッド量子化と適度なプルーニングを使用：

設定	保持されるベクトル	インデックスサイズ	NDCG@10
1ビット PLAID、すべてのベクトル	100%	3.37GB	0.8984
1ビット PLAID + プルーニング	65%	2.23GB	0.8830
1ビット PLAID + プルーニング	42%	1.45GB	0.8642

• 量子化により13倍以上のサイズ削減が可能で、NDCG損失は0.02未満。マルチベクトルリトリーブのストレージは密モデルと同等になります。

---

評価結果

ファインチューニングされた multi-vector-encoder/mLateOn-medical と、20万件のパラグラフ医療ベンチマーク（1,000件のホールドアウトクエリ）で比較された50以上のベースライン：

モデル	ファミリー	NDCG@10	acc@1
multi-vector-encoder/mLateOn-medical	マルチベクトル（ファインチューニング済み）	0.9139	0.849
lightonai/mLateOn	マルチベクトル（ゼロショット）	0.8520	0.758
lightonai/GTE-ModernColBERT-v1（長さ制限解除）	マルチベクトル（ゼロショット）	0.8502	0.763
Qwen/Qwen3-Embedding-4B	密（ゼロショット）	0.7817	0.669
voyageai/voyage-4-nano	密（ゼロショット）	0.7563	0.638
BM25	語彙的	0.7501	0.641
naver/splade-v3	スパース（ゼロショット）	0.6853	0.574
結論： ドキュメント長が長い場合、ラテインタラクションモデルが優位。ファインチューニングにより、最も強力なゼロショットモデルよりさらに +0.062 NDCG@10 の向上。

---

実用的な教訓

1. ドメインファインチューニングには教師なしチェックポイント（例：lightonai/mLateOn-unsupervised）から始める。
2. ドキュメント長の制限を解除して、データに合わせる。そうでなければ、長文では最大 0.24 NDCG@10 の損失が生じる。
3. GradCache（CachedMultiVectorMultipleNegativesRankingLoss）を使用して、1台のGPUで大きな有効バッチサイズを実現する。
4. 句読点スキップリストを適用して、品質を損なわずにインデックスを約10%小さくする。
5. インデックス圧縮（トークンプーリング、PLAID量子化）に投資して、マルチベクトルストレージを密モデルと同等の規模に抑える。
6. わずかなデータ（10万ペア）でも強力な性能が得られるため、多くの組織にとって実現可能。

---

追加リソース

• 訓練例：MIRIAD医療ファインチューニング、MS MARCOの対照学習および知識蒸留、マルチモーダル ColPali、PEFT LoRAアダプター。
• ドキュメント：インストール、クイックスタート、使用法、カスタムモデル作成、事前学習モデル一覧、訓練と損失の概要、APIリファレンス。
• 補足記事（使用法）：Sentence Transformers によるマルチベクトル（ラテインタラクション）埋め込みモデル – エンコーディング、インデキシング、サービングをカバー。

---

"マルチベクトルインデックスは大きすぎるという反論は、適切に設定されたインデックスでは成立しない。" – オマール・ハッタブ氏による量子化測定への謝辞。