用 Sentence Transformers 跑 ColBERT

很多人对 Embedding 的认知还停留在把一段话压成一个向量（Single Vector），但这种做法在处理长文本或复杂查询时，信息丢失太严重。ColBERT 这种 Late Interaction（延迟交互）机制的核心就是不搞这种“暴力压缩”，而是给每个 Token 都保留一个向量，检索时再通过 MaxSim 做匹配。

我最近在实操 Sentence Transformers 的多向量支持，发现它现在对这种模式的兼容度很高。如果你想在自己的项目里部署一套精度比传统 RAG 高得多的检索系统，可以参考这个逻辑。

一、环境准备与模型加载
首先得确保 sentence-transformers 版本足够新。加载模型时，不能用普通的 SentenceTransformer 类，得找那些原生支持多向量输出的模型（比如 ColBERT 家族）。

pip install -U sentence-transformers

from sentence_transformers import SentenceTransformer

# 加载一个支持多向量的模型，例如 ColBERT 风格的模型
model = SentenceTransformer('colbert-ir/colbertv2.0')

# 对文本进行编码，此时输出的不是 [batch_size, dim] 
# 而是 [batch_size, seq_len, dim]，每个 token 都有自己的向量
embeddings = model.encode(["这是我的测试句子", "多向量检索精度更高"], output_value='token_embeddings')

二、核心计算逻辑：MaxSim
多向量模型最关键的不是编码，而是怎么算分数。它不是简单的余弦相似度，而是计算查询端每个 token 向量与文档端所有 token 向量的最大相似度之和。

import torch

def maxsim(query_emb, doc_emb):
    # query_emb: [1, q_len, dim]
    # doc_emb: [1, d_len, dim]
    # 计算所有 token 之间的点积相似度
    scores = torch.matmul(query_emb, doc_emb.transpose(-1, -2)) # [1, q_len, d_len]
    # 对文档维度取最大值，然后对查询维度求和
    max_scores = torch.max(scores, dim=-1).values # [1, q_len]
    return torch.sum(max_scores)

# 实操演示
q_vec = model.encode("检索精度", output_value='token_embeddings')
d_vec = model.encode("多向量检索的精度确实比单向量高", output_value='token_embeddings')
score = maxsim(q_vec, d_vec)
print(f"相似度得分: {score.item()}")

三、避坑指南与实战建议
这种架构虽然精度暴力提升，但存储压力极大。因为每个文档的存储空间变成了原来的 seq_len 倍。在实际部署时，建议关注以下几点：

• 量化压缩： 必须使用量化技术（如 ColBERT 的 1-bit 量化），否则索引库会撑爆内存。
• 两阶段检索： 不要直接用多向量做全库扫描。建议先用 BM25 或轻量级单向量模型粗筛出 Top 100，再用 Late Interaction 做精排。
• 内存对齐： 在 GPU 上跑 matmul 时，注意 Tensor 的维度对齐，否则性能损耗很严重。