Sentence Transformers v6.0 多向量編碼器發行版

簡要重點

Sentence Transformers 6.0 引入了第四種模型類型 MultiVectorEncoder,實現類 ColBERT 的晚期互動(多向量)檢索,適用於文字與多模態資料,提供更強大的語意搜尋與視覺文件檢索能力,但需要更大的詞元層級索引。


什麼是多向量模型?

多向量模型會保留每個詞元的嵌入向量,而非將整個段落壓縮成單一向量。編碼後,查詢會使用 MaxSim 運算子與文件進行評分,該運算子會對每個查詢詞元,在所有文件詞元中找出最高的餘弦相似度並求和。此方法能保留精確的詞元匹配(例如產品代碼),並捕捉語意上的同義表達,因為詞元嵌入是上下文相關的。

"將『企鵝住在哪裡?』這個查詢與『企鵝棲息於南極』進行比對,查詢詞元 live 可以在 inhabit 上找到最佳匹配,相似度達 0.94。" – Hugging Face 博客

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|]$ 範圍內。

取捨關係

  • 品質提升 – 詞元層級匹配能改善依賴單一精確詞或多個獨立條件的查詢檢索表現。
  • 索引成本 – 每個詞元一個向量會大幅增加儲存空間(例如,4,874 個段落使用 128 維模型需 311 MB 原始資料,而 384 維密集模型僅需 7.5 MB)。可透過 PLAID 或層次化詞元聚合等壓縮技術降低此開銷。

安裝

pip install -U sentence-transformers
# 若需視覺文件檢索,請安裝影像額外套件
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-vectorsentence-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)

每個文件返回的矩陣第一維等於其詞元數量,因此在未補齊的情況下無法直接堆疊成單一矩形批次。


使用 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 秒,每次查詢評分約需 120 毫秒。


檢索並重新排序模式

結合快速的密集雙編碼器與多向量重新排序器,避免建立完整的晚期互動索引:

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)

原始 608 k 個詞元向量佔用 311 MB(float32),但透過 PLAID 的質心加殘差方案可壓縮至約 92 MB。


視覺文件檢索

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))

單一頁面可能產生數百個詞元向量(例如 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.pytext_similarity_map.py 指令碼,可列印詞元層級貢獻並標示來源段落。


使用詞元聚合降低索引大小

HierarchicalTokenPooling 會以餘弦距離的 Ward 連結方式聚類文件的詞元向量,並以質心取代每個群集,可達約 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 上約 0.4 % 的 NDCG 下降。更高因子的效益遞減,建議在實際資料上評估後再決定。


加快推論速度

  • 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 維)在 13 個資料集中的 9 個上超越對應的密集模型 lightonai/DenseOn(0.6764),平均 NDCG@10 達 0.6868


從 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

所有標記為 multi-vector 的模型皆相容;視覺模型可能需要 revision,直到其倉儲配置 PR 合併為止。


致謝

本實作建基於 ColBERT(Khattab & Zaharia, 2020)、LightOn 的 PyLate 與 fast‑plaid、ColPali 研究團隊,以及 Clavié、Chaffin 與 Adams 的詞元聚合工作。額外感謝 MTEB 基準貢獻者與所有檢查點作者。


進一步資源

  • 文件 – 使用方式、預訓練模型、自訂模型建立、效率與 API 參考(連結見原始貼文)。
  • 範例指令碼 – 語意搜尋、檢索並重新排序、詞元聚合、熱力圖與 NanoBEIR 評估。
  • 訓練指南 – 概述、損失函數與 LateOn/mLateOn 的食譜指令碼。
  • 搭配部落格文章 – 密集訓練、重新排序器、稀疏編碼器、多模態嵌入與 Matryoshka 嵌入。

Sources

相關