Sentence Transformers v6.0 が MultiVectorEncoder に対応——ColBERT スタイルモデルのファインチューニング完全ガイド

記事のサマリー（TL;DR）

• Sentence Transformers v6.0 が MultiVectorEncoder を追加。ColBERT スタイルのマルチベクトルモデルをファインチューニング可能になった
• RTX 3090 単体・14.5時間で訓練した医療ドメインモデルが NDCG@10 = 0.9139 を達成。汎用最強の密ベクトルモデル（Qwen3-Embedding-4B / 0.7817）を 0.13 以上上回る
• 45 GB の生インデックスも 1-bit PLAID 量子化で 3.37 GB まで圧縮可能。マルチベクトルの「インデックスが重い」という懸念を実質解消

社内ドキュメント・専門ドメイン検索基盤を構築する日本企業が注目すべき点

RAG（検索拡張生成）を業務に導入する際、汎用の密ベクトルモデルでは語彙・クエリスタイル・関連性の定義が自社ドメインと合わず精度が伸び悩むケースが多い。本記事が示す手法は「コンシューマ GPU 1枚・数時間」でドメイン特化モデルを構築できる点が実用的だ。

医療・法務・金融・社内技術文書など、長文パッセージ（500〜1,400 トークン超）を扱う場合は特に効果が顕著で、汎用モデルがトークン上限で切り捨てている部分を学習時から考慮できる。kintone や Salesforce など業務 SaaS に蓄積されたドキュメントを検索対象とする構成でも、ドメインデータさえ用意すれば同じレシピが適用できる。また BigQuery に統合した社内データから (クエリ, パッセージ) ペアを生成し訓練データとする流れは、データ基盤が整っている組織であれば現実的な選択肢となる。

詳細

マルチベクトルモデルとは何か

密ベクトルモデルはテキスト全体を単一ベクトルに圧縮し、2つのベクトル間のドット積で類似度を算出する。対してマルチベクトルモデル（遅延インタラクション / ColBERT スタイル）はその圧縮を行わず、トークンごとに小さなベクトルを保持する。クエリとドキュメントのスコアリングには MaxSim 演算子 を使用する。各クエリトークンが最も一致するドキュメントトークンを見つけ、スコアを合算する仕組みだ。

トークンレベルのマッチングにより、単一ベクトルが平均化せざるを得なかった細粒度のシグナルを保持できる。その代償はインデックスサイズの増大だが、後述する量子化で実用域に抑えられる。

なぜファインチューニングが必要か

汎用の検索モデルは Web 検索・法律文書・コード検索・学術論文レビューといった用途ごとに異なる語彙・クエリスタイル・関連性の概念を持つ。マルチベクトルモデルはトークン単位でマッチングするため、少量のドメイン内データでも大きく性能が向上する。

また既存の ColBERT チェックポイントの多くはドキュメントを 180〜300 トークンで切り捨て、密ベクトルモデルでも 256〜512 トークンが上限のものが多い。本記事の医療評価データではパッセージが平均 941 トークンあり、この切り捨てだけで最大 0.24 NDCG@10 のロスが発生することが計測されている。LightOn がコード検索で汎用 LateOn を超えるべく LateOn-Code を訓練したケースと同じ構図だ。

訓練コンポーネント

マルチベクトルモデルの訓練には以下のコンポーネントが必要になる。

1. モデル: ファインチューニング対象のチェックポイント、またはゼロから構築するアーキテクチャ
2. データセット: 訓練・評価に使用するデータ
3. 損失関数: モデルの性能を測定し最適化を導く関数
4. 訓練引数（任意）: 訓練速度・追跡・デバッグに影響するパラメータ
5. 評価器（任意）: 訓練前・中・後にモデルを評価するクラス
6. トレーナー: すべてのコンポーネントを統合するクラス

モデル：出発点の選び方が性能を左右する

既存マルチベクトルモデルのファインチューニング

既存の MultiVectorEncoder モデルをファインチューニングする場合、アーキテクチャを意識する必要はなく、チェックポイントを読み込むだけでよい。重要なのは ドキュメント長の設定だ。多くの公開チェックポイントは 180〜512 トークンで切り捨てるため、医療パッセージのような長文では上限を解除する必要がある。

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192},
)
# 長さキャップを解除
model[0].query_length = None
model[0].document_length = None

さらに、句読点トークンをドキュメント側スコアリングから除外するスキップリストを追加した。品質のわずかな向上に加え、ドキュメントインデックスを 9.6% 削減 できる。

ベーストランスフォーマーからのゼロ構築

任意のベーストランスフォーマーを MultiVectorEncoder に渡すと、ランダム初期化されたトークンレベルプロジェクションが自動付与される。

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

古典的な ColBERT パイプライン（Transformer → トークンレベル Dense で 128 次元に射影 → MultiVectorMask → Normalize）が構成される。プロジェクションはランダムスタートのため訓練が必須だが、Alibaba-NLP/gte-modernbert-base に新規プロジェクションを付けて 25k ペアだけ訓練した結果、既存チェックポイント出発点から 0.03 以内に到達した。

どの出発点を選ぶべきか

MIRIAD の医療質問・パッセージペア 25k を使い、6種類の出発点を同一レシピで訓練して比較した結果は以下の通り。

教訓: -unsupervised チェックポイント（大規模対比事前訓練済み、教師ありファインチューニング前）がドメイン適応で最も効果的だ。汎用教師あり訓練済みのチェックポイントは学習が進まず、むしろ後退することもある。モデルファミリーが事前教師ありチェックポイントを公開しているなら、そこを出発点にすること。なければ、検索用途で事前訓練された強力なバックボーンへの新規プロジェクションが次善策になる。

データセット

MultiVectorEncoderTrainer は datasets.Dataset または datasets.DatasetDict を使用する。Hugging Face Hub からの読み込みも、CSV・JSON・Parquet などローカルファイルも対応している。

本記事で使用するのは tomaarsen/miriad-4.4M-split：MIRIAD から生成した 440 万件の医療質問と、それぞれ答えを含むソースパッセージ（平均 941 トークン）のペアだ。

データセット形式のポイント

• 先頭列がクエリ、以降の列がドキュメントとして扱われる（列名は不問、順序のみ有効）
• 知識蒸留形式は (query, document_1, ..., document_N, scores) の形式で、scores は N 個の教師スコアのリスト
• resolve_ids を使えば ID と別テキストデータセットを保持する KD データセットもオンザフライで解決可能

損失関数

クエリ・パッセージのペアデータに対する標準的な損失関数は CachedMultiVectorMultipleNegativesRankingLoss（インバッチネガティブ）だ。GradCache バリアントにより、GPU に乗るバッチサイズに縛られず実効的なコントラスティブバッチサイズを自由に設定できる。

from sentence_transformers.multi_vector_encoder.losses import CachedMultiVectorMultipleNegativesRankingLoss

loss = CachedMultiVectorMultipleNegativesRankingLoss(
    model=model,
    mini_batch_size=16,  # メモリ境界となるチャンクサイズ。品質には影響しない
)

重要な落とし穴: マルチベクトルの対比損失は scale=1.0 がデフォルト。密ベクトルの scale=20.0 をそのままコピーすると softmax が飽和し、勾配が消滅する。MaxSim スコアはクエリトークン数分の合算になるため、すでに十分な値域を持つためだ。

訓練引数

実際の訓練で使用した設定例：

args = MultiVectorEncoderTrainingArguments(
    output_dir="models/mLateOn-medical",
    num_train_epochs=1,
    per_device_train_batch_size=128,   # GradCache により実効バッチサイズとなる
    learning_rate=1e-4,                # 5e-6〜2e-4 のスイープで最良
    warmup_steps=0.05,
    prompts={"question": "[Q] ", "passage_text": "[D] "},
    bf16=True,
    batch_sampler=BatchSamplers.NO_DUPLICATES,
)

max_length は意図的に未設定。512 トークンに制限すると速度は約 2 倍になるが、NDCG@10 が 0.015 低下し、データを増やしても回復しない。

learning_rate=1e-4 は通常より高めだが、スイープの結果これが最良だった。

評価器

ドメインファインチューニングには MultiVectorInformationRetrievalEvaluator が最も有効だ。評価セットが飽和する場合（ほぼ全クエリが 0.97 超など）は、訓練分割から重複排除したパッセージをディストラクターとして追加し、スコアが分散するようにコーパスを調整する。

トレーナーと完全なスクリプト

以下が multi-vector-encoder/mLateOn-medical を訓練した完全レシピの要点だ。

1. lightonai/mLateOn-unsupervised（対比事前訓練済み・教師あり前）を読み込む
2. ドキュメント長キャップを解除（query_length・document_length を None に）
3. 句読点スキップリストを設定
4. MIRIAD から 100 万件の医療質問・パッセージペアを読み込む
5. CachedMultiVectorMultipleNegativesRankingLoss でインバッチネガティブ訓練
6. MultiVectorInformationRetrievalEvaluator で進捗を監視
7. RTX 3090 単体、ピーク VRAM 17.5 GB、14.5 時間で訓練完了

スケーリング実験では、100k ペア（訓練 75 分）で 100 万ペアとの差が 0.012 NDCG@10 以内に収まる。利益の大半は最初の 1 時間で得られる。

コールバックとマルチデータセット訓練

MultiVectorEncoderTrainer は WandbCallback・TensorBoardCallback・CodeCarbonCallback などの transformers.TrainerCallback サブクラスをサポートしており、report_to 引数で有効化できる。

複数データセットの同時訓練も可能で、データセット名から損失関数へのマッピングを辞書で渡すことで、各データセットに異なる損失関数を適用できる。サンプリング順は ROUND_ROBIN（各データセットから均等）または PROPORTIONAL（データセットサイズ比。デフォルト）から選択する。

評価結果：ドメイン特化の威力

MIRIAD 評価セット（1,000 件の医療質問 × 200,000 パッセージ：10k ゴールドパッセージ + 19 万件のディストラクター）での評価結果（主要モデルのみ抜粋）：

ファインチューニング済みモデルは最強のゼロショットモデルを +0.062 NDCG@10 上回る。1位ヒット率（acc@1）は最強ゼロショット（75.8%）に対してファインチューニング済みは 84.9% と、rank-1 エラーを 3 分の 1 以上削減した。

注目すべき知見：

• アーキテクチャ比較: DenseOn と LateOn は訓練データ・アーキテクチャを揃えてヘッドのみ異なるが、遅延インタラクション版が +0.12。多言語ペア（mDenseOn / mLateOn）も +0.13 で再現
• スケールでは補えない: Qwen3-Embedding-4B は本記事モデルの約 33 倍のパラメータを持つが、0.13 届かない。8B 版は 4B より低スコア
• BM25 の善戦: MIRIAD では質問がパッセージから生成されるため語彙的重複が大きく、トークン数無制限の BM25 が多くのニューラルモデルを上回る。ただしこの傾向は MIRIAD 固有の特性で、他ドメインへの転用は禁物

インデックスの最適化

マルチベクトルの公正な懸念はインデックスサイズだ。本評価では 1 パッセージあたり平均 878 ベクトルを保持し、20 万パッセージのコーパスが fp16 で約 45 GB になる（密ベクトルなら 1 GB 未満）。

HierarchicalTokenPooling によるベクトル数削減：

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling

pooling = HierarchicalTokenPooling(pool_factor=4)
document_embeddings = model.encode_document(passages, token_pooling=pooling)

• pool_factor=4（ベクトル数 1/4 → 約 11.2 GB）でも NDCG@10 = 0.8991
• ベクトル数 1/10 でも 0.8765 を維持

さらに Omar Khattab が fast-plaid で計測した 1-bit 残差量子化との組み合わせ：

生埋め込みの 45 GB から 1-bit PLAID 単体で 3.37 GB（13 分の 1） に圧縮し、NDCG@10 の損失はわずか 0.0155 に留まる。最後の行（1.45 GB）は Qwen3-Embedding-8B の fp16 埋め込み（1.64 GB）より小さく、スコアは 0.0895 高い。

量子化はプーリングやプルーニングと組み合わせ可能だが、まず量子化を優先するのが効率的だ。インデックスはチェックポイントと同じ注意を払う価値がある。