Sentence Transformers v6.0 Multi-Vector Encoder リリース

まとめ

Sentence Transformers 6.0 は、ColBERTスタイルのラテインタラクション(マルチベクトル)検索を実装する新しいモデルタイプ MultiVectorEncoder を導入しました。テキストおよびマルチモーダルデータに対して、より強力な意味検索と視覚ドキュメント検索を実現しつつ、より大きなトークンレベルのインデックスを必要とします。


マルチベクトルモデルとは?

マルチベクトルモデルは、全文を1つのベクトルに圧縮するのではなく、各トークンごとに1つの埋め込みを保持します。エンコード後、クエリはドキュメントに対して MaxSim 演算子でスコアリングされ、各クエリトークンについてドキュメント内のすべてのトークンから最も高いコサイン類似度を合計します。これにより、正確なトークンマッチ(例:製品コード)や意味的な言い換えを保持でき、トークン埋め込みが文脈に応じて調整されるためです。

"‘Where do penguins live?’ を ‘Penguins inhabit Antarctica.’ と比較すると、クエリトークン liveinhabit と 0.94 の類似度で最良のマッチを見つけます。" – Hugging Face blog

MaxSim 演算子

クエリトークン集合 $Q$ とドキュメントトークン集合 $D$ に対して:

$$ \text{MaxSim}(Q, D) = \sum_{q_i \in Q} \max_{d_j \in D} ; q_i \cdot d_j $$

埋め込みが L2 正規化されているため、各ドット積は $[-1, 1]$ のコサイン類似度であり、合計スコアは $[-|Q|, |Q|]$ の範囲にあります。

バランスの取れ方

  • 品質の向上 – 1つの正確な語や複数の独立した制約に依存するクエリの検索精度が向上します。
  • インデックスコスト – 1トークンあたり1つのベクトルであるため、ストレージが膨張します(例:128次元モデルで4,874のパッセージは311 MB、384次元の密ベクトルモデルでは7.5 MB)。PLAIDや階層的トークンプーリングなどの圧縮技術でこのオーバーヘッドを削減できます。

インストール

pip install -U sentence-transformers
# 視覚ドキュメント検索のためには image オプションを追加
pip install -U "sentence-transformers[image]"

Sentence Transformers v6.0 は transformers v5.x、torch ≥ 2.2、huggingface-hub v1.x を必要とします。


マルチベクトルモデルの読み込み

from sentence_transformers import MultiVectorEncoder
model = MultiVectorEncoder("lightonai/LateOn")

multi-vector および sentence-transformers というタグが付いたチェックポイントは、PyLate、Stanford-NLP ColBERT、colpali-engine からのものかどうかに関わらず、直接読み込み可能です。ColPali視覚モデルの場合は、小さなリポジトリ設定が必要です(サポートされているモデル を参照)。

モデル構成の確認

model = MultiVectorEncoder("colbert-ir/colbertv2.0")
print(model)
print(model.prompts)

通常の出力では Transformer、トークンレベルの Dense 投影(例:128次元)、スキップリスト用の MultiVectorMask、および Normalize レイヤーが表示されます。クエリおよびドキュメントの長さ制限(例:32トークン、180トークン)も表示されます。


クエリとドキュメントのエンコード

マルチベクトルモデルは 非対称 です。専用のメソッドを使用してください:

queries = ["What is the capital of France?"]
documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
q_emb = model.encode_query(queries)          # shape (n_query_tokens, dim)
d_emb = model.encode_document(documents)    # list of (n_doc_tokens, dim)

各ドキュメントは、トークン数に等しい最初の次元を持つ行列を返すため、パディングなしにテンソルを1つの長方形バッチに積み重ねることはできません。


MaxSimによるスコアリング

scores = model.similarity(q_emb, d_emb)  # すべてのペアに対する MaxSim マトリクス
print(scores)

結果はコサイン類似度の合計からなるテンソルです。合計スコアはクエリ長に比例するため、異なるクエリ長制限を持つモデル間で直接比較できません。上限のあるメトリクスを得るには、MeanMaxSim に切り替えます:

model = MultiVectorEncoder("lightonai/LateOn", similarity_fn_name="meanmaxsim")
print(model.similarity(q_emb, d_emb))  # 値は [0, 1] の範囲

意味検索(網羅的)

小規模なコーパスでは、一度だけコーパス全体をエンコードし、クエリ時に網羅的な MaxSim を実行できます:

from datasets import load_dataset
from sentence_transformers import MultiVectorEncoder

corpus = list(dict.fromkeys(load_dataset("sentence-transformers/natural-questions", split="train[:5000]")["answer"]))
model = MultiVectorEncoder("lightonai/LateOn")
corpus_emb = model.encode_document(corpus, convert_to_tensor=True)

query = "when did richmond last play in a preliminary final"
q_emb = model.encode_query([query], convert_to_tensor=True)
score_vec = model.similarity(q_emb, corpus_emb)[0]
top_scores, top_idx = score_vec.topk(3)
for s, i in zip(top_scores.tolist(), top_idx.tolist()):
    print(f"{s:.4f}  {corpus[i][:100]}")

RTX 3090 では、4,874のパッセージのコーパスは約20秒でエンコードされ、各クエリのスコアリングは約120msです。


リトリーブ&リランクパターン

高速な密ベクトルバイエンコーダーとマルチベクトルリランカーを組み合わせて、完全なラテインタラクションインデックスを構築する必要を回避します:

from sentence_transformers import SentenceTransformer, MultiVectorEncoder, util

retriever = SentenceTransformer("jinaai/jina-embeddings-v5-text-nano-retrieval")
reranker  = MultiVectorEncoder("perplexity-ai/pplx-embed-v1-late-0.6b", trust_remote_code=True)

# 1️⃣ 密モデルでトップ-kを取得
corpus_emb = retriever.encode_document(corpus, convert_to_tensor=True)
hits = util.semantic_search(retriever.encode_query([query], convert_to_tensor=True), corpus_emb, top_k=50)[0]
candidates = [corpus[h['corpus_id']] for h in hits]

# 2️⃣ MaxSimで候補を再スコアリング
q_emb = reranker.encode_query([query])
d_emb = reranker.encode_document(candidates)
rerank_scores = reranker.similarity(q_emb, d_emb)[0]
print(rerank_scores.argsort(descending=True)[:3])

選択された候補のみがマルチベクトルとしてエンコードされるため、メモリを大幅に削減しつつ、ラテインタラクションの品質を維持できます。


インデックスのオプション

複数のベクトルデータベースが、MaxSimをサポートするネイティブなマルチベクトルフィールドを提供しています:

  • Qdrant (v1.10+)
  • Weaviate (v1.29+)
  • Vespa (長文対応 ColBERT)
  • LanceDB (v0.15.0+)
  • VectorChord (Postgres拡張)
  • Milvus (v2.6.4、配列型構造)
  • fast-plaid (Rust実装、近似だが高速)

各システムは encode_document から返されるトークンレベルのテンソルのリストをインジェストします。fast-plaid(正確なインデックス)の例:

from fast_plaid import search
fast_plaid = search.FastPlaid(index="nlp-index", device="cuda")
fast_plaid.create(documents_embeddings=document_emb)
results = fast_plaid.search(queries_embeddings=q_emb.unsqueeze(0), top_k=3)

608kトークンベクトルはfloat32で311MBを占めますが、PLAIDの重心+残差スキームで約92MBに圧縮できます。


視覚ドキュメント検索

ColPaliスタイルのモデルは、ページ画像を画像パッチのトークン列として扱います。同じAPIが利用可能です:

model = MultiVectorEncoder("vidore/colqwen2.5-v0.2")
queries = ["What is the variable on the y-axis?", "Total outlay is maximum in which year?"]
images = [".../doc1.jpg", ".../doc2.jpg", ".../doc3.jpg", ".../doc4.jpg"]
q_emb = model.encode_query(queries)
d_emb = model.encode_document(images)
print(model.similarity(q_emb, d_emb))

1ページあたり数百のトークンベクトル(例:755 × 128)が得られるため、トークンプーリングが価値あるものになります(次のセクションを参照)。


音声および動画検索

オムニモーダルモデル vidore/colqwen-omni-v0.1 は、トランスクリプションステップなしにテキスト、画像、音声、動画をサポートします。

model = MultiVectorEncoder("vidore/colqwen-omni-v0.1", model_kwargs={"dtype": torch.bfloat16})
# 音声の例
audio = [...]  # ロウ波形のリスト(16 kHz モノ)
q_emb = model.encode_query(["medicine for car nausea"])
d_emb = model.encode_document(audio, batch_size=2)
print(model.similarity(q_emb, d_emb)[0].topk(3))

ゼロショット音声検索が可能であるのは、モデルが画像-テキストペアで訓練されており、クロスモーダルトークンの対応を学習しているためです。


解釈可能性

MaxSimのトークンごとの分解により、正確な寄与度の割り当てが可能になります:

  • 画像用ヒートマップ – ページのパッチにトークンレベルのスコアを重ねて表示。
  • テキスト類似度マップ – 各クエリトークンに対応する最良のマッチングドキュメントトークンをリスト表示。

リポジトリには heatmap.py および text_similarity_map.py スクリプトが含まれており、トークンごとの寄与を出力し、ソースドキュメントを強調表示します。


インデックスサイズの削減のためのトークンプーリング

HierarchicalTokenPooling は、ドキュメントのトークンベクトルを(コサイン距離に基づくウォード連結で)クラスタリングし、各クラスタをその重心で置き換え、およそ 1 / pool_factor の削減を実現します。

from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
pool = HierarchicalTokenPooling(pool_factor=2)
pooled_emb = model.encode_document(docs, token_pooling=pool)

Natural Questionsコーパスでは、pool_factor=2 でインデックスサイズが半分(311 MB → 156 MB)になり、BEIRでNDCGが約0.4%低下するのみです。より高い因子は限界効果となり、自身のデータで評価してから採用してください。


推論の高速化

  • GPU – fp16 + Flash Attention(attn_implementation="flash_attention_2")で、fp32より約2.4倍のスループットが得られます。
  • CPU – OpenVINO(サポートされている場合)および int8量子化により、<0.5%の精度低下でわずかな高速化が実現できます。
  • attend=False を使用するクエリ拡張マスクを持つモデルは、マスクされたトークンが無視されるため、Flash Attentionを使用できません。代わりにデフォルトの sdpa 実装を使用してください。

評価

MultiVectorNanoBEIREvaluator は、13サブセットの NanoBEIRベンチマークをそのまま実行できます:

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorNanoBEIREvaluator
model = MultiVectorEncoder("lightonai/LateOn")
eval = MultiVectorNanoBEIREvaluator()
print(eval(model))

結果は、LateOn(マルチベクトル、128次元)が平均NDCG@10で 0.6868 を達成し、13データセットのうち9で密ベクトル対応モデル lightonai/DenseOn(0.6764)を上回りました。


PyLate / colpali-engine からの移行

PyLate Sentence Transformers
pylate.models.ColBERT(...) MultiVectorEncoder(...)
model.encode(..., is_query=True) model.encode_query(...)
model.encode(..., is_query=False) model.encode_document(...)
pylate.scores.colbert_scores model.similarity
pylate.indexes.PLAID PyLateをそのまま使用するか、fast-plaid / Qdrant に切り替える
colpali-engine Sentence Transformers
ColQwen2.from_pretrained(...) MultiVectorEncoder(...)
processor.process_queries(...) model.encode_query(...)
processor.process_images(...) model.encode_document(...)
processor.score_multi_vector(...) model.similarity(...)

プレフィックス、クエリ拡張、スキップリストの取り扱いについての詳細は 移行ガイド を参照してください。


サポートされているモデル(抜粋)

テキスト検索(選定)

モデル パラメータ数 次元 NanoBEIR
lightonai/LateOn-regularized 149 M 128 0.6897
lightonai/LateOn 149 M 128 0.6868
LiquidAI/LFM2.5-ColBERT-350M 353 M 128 0.6864
mixedbread-ai/mxbai-edge-colbert-v0-32m 32 M 64 0.6524
colbert-ir/colbertv2.0 110 M 128 0.6053

視覚ドキュメント検索(選定)

モデル パラメータ数 次元 NanoViDoRe
webAI-Official/webAI-ColVec1.1-8b 8.4 B 640 0.6580
tencent/EVIE-Preview-4.5B 4.54 B 128 0.6405
vidore/colqwen2.5-v0.2 3.8 B 128 0.5402
vidore/colpali-v1.3 2.9 B 128 0.4802

Hubに multi-vector タグが付いたすべてのモデルは互換性があります。視覚モデルは、リポジトリ設定のPRがマージされるまで revision が必要な場合があります。


謝辞

本実装は、ColBERT(Khattab & Zaharia, 2020)、LightOnのPyLateおよびfast-plaid、ColPali研究チーム、Clavié、Chaffin、Adamsのトークンプーリング研究に基づいています。MTEBベンチマーク貢献者およびすべてのチェックポイント作者に感謝します。


その他のリソース

  • ドキュメント – 使用法、事前学習モデル、カスタムモデル作成、効率性、APIリファレンス(元の投稿内のリンク)。
  • 例題スクリプト – 意味検索、リトリーブ&リランク、トークンプーリング、ヒートマップ、NanoBEIR評価。
  • トレーニングガイド – 概要、損失関数、LateOn/mLateOn用のレシピスクリプト。
  • 関連ブログ記事 – 密ベクトル学習、リランカー、スパースエンコーダー、マルチモーダル埋め込み、マトリョーシカ埋め込み。

Sources

関連