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

每个文档返回一个矩阵,其第一维等于其标记数量,因此在不进行填充的情况下无法将张量堆叠成单一矩形批次。


使用 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️⃣ 使用密集模型检索 top-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 标记向量占用 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 上 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 维)在 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

所有 Hub 上标记为 multi-vector 的模型均兼容;视觉模型可能需要 revision,直到其仓库配置 PR 合并为止。


致谢

实现基于 ColBERT(Khattab & Zaharia, 2020)、LightOn 的 PyLate 和 fast-plaid、ColPali 研究团队以及 Clavié、Chaffin 和 Adams 的标记聚合工作。额外感谢 MTEB 基准贡献者及所有检查点作者。


进一步资源

  • 文档 – 使用方法、预训练模型、自定义模型创建、效率优化和 API 参考(原始帖子中的链接)。
  • 示例脚本 – 语义搜索、检索-重排、标记聚合、热图和 NanoBEIR 评估。
  • 训练指南 – LateOn/mLateOn 的概述、损失函数和配方脚本。
  • 配套博客文章 – 密集训练、重排器、稀疏编码器、多模态嵌入和 Matryoshka 嵌入。

Sources

相关