之前在做一个 RAG 检索项目时，我发现用常见的单向量嵌入模型（Sentence Transformers 系列）做召回，总会在一些“关键词重合度高、语义又有点相关”的查询上表现不够稳定，尤其在长文档和细粒度匹配场景下，一个向量把整段文本压扁，信息丢失很明显。后来把目光转向多向量嵌入模型，也就是 ColBERT 这类“每个 token 一个向量”的思路，并基于 Sentence Transformers 完成了训练和微调。这篇文章就把整个过程的原理、数据准备、训练代码、推理落地方案和踩坑点完整整理出来，既适合刚接触 Embedding 模型的新手，也适合正在做 RAG / 语义检索落地的开发者直接参考。
1. 从单向量到多向量：为什么需要多向量嵌入模型
1.1 单向量嵌入模型的基本流程
Sentence Transformers 是目前最常用的句子嵌入工具库。它的核心思路是：把 BERT 这类预训练语言模型当作编码器，对输入句子做编码，再通过池化层（Mean Pooling 或 CLS Pooling），把整句话压缩成一个固定维度的向量。比如BAAI/bge-base-zh-v1.5、all-MiniLM-L6-v2这类模型，输入任意长度的文本，最终输出一个 768 维或 384 维的向量。
from sentence_transformers import SentenceTransformer model = SentenceTransformer("BAAI/bge-base-zh-v1.5") embedding = model.encode("深度学习是机器学习的一个重要分支") print(embedding.shape)两个句子的相似度，就变成两个向量的余弦相似度或点积：
from sentence_transformers import util emb1 = model.encode("如何学习深度学习") emb2 = model.encode("深度学习入门指南") score = util.cos_sim(emb1, emb2) print(score.item())这种方式实现简单、检索速度快，也是目前 RAG 系统里最主流的向量化方案。绝大多数向量数据库（Milvus、FAISS、Qdrant 等）都天然支持这种固定维度向量的索引和检索。
1.2 单向量模型的短板
单向量模型最大的问题在于“信息压缩”。无论输入是一句话还是一整篇长文档，最终都只能用一个向量来表示。这里会带来几个实际痛点：
- 长文本压缩丢失细节。一篇 2000 字的文档，被压成 768 维向量后，很多局部信息会被平均掉，导致召回不准。
- 对“局部精确匹配”不敏感。例如查询是“数据库索引优化”，文档里恰好有一段专门讲“如何优化 MySQL 索引”，但由于整篇文档还有大量其他内容，均值池化会把这段关键信息稀释掉，最终相似度反而不高。
- 同义改写与词面重合的权衡很难做。单向量模型擅长语义相似，但对“关键词重合度高但语义不同”“词面不完全重合但语义相同”这两类情况，往往只能顾一头。
这些问题在 FAQ 匹配、长文档检索、法律/医学文本检索等场景中会特别明显。
1.3 多向量嵌入模型（ColBERT 思路）怎么工作
多向量嵌入模型并不把一个句子压缩成一个向量，而是让每个 token 都对应一个向量。以 ColBERT 为代表的这类模型，采用“晚期交互（Late Interaction）”机制：
- 查询文本经过编码器，得到一组 token 向量：
(query_len, hidden_dim)。 - 文档文本经过同一个编码器，得到一组 token 向量：
(doc_len, hidden_dim)。 - 计算相似度时，查询的每个 token 向量与文档的所有 token 向量做点积，取最大值，再把所有 token 的最大值求和。
用公式表达就是：
score(q, d) = Σ_i max_j sim(q_i, d_j)其中i遍历查询的所有 token，j遍历文档的所有 token。这个计算方式允许查询中的每个词面或语义单元，在文档中找到自己最匹配的位置，而不是强迫整个句子变成一个向量。
1.4 多向量模型适合哪些场景
多向量嵌入模型并不是要完全替代单向量模型，而是在特定场景下优势更明显：
- RAG 检索中的精排环节。先用单向量模型做粗排召回，再用多向量模型对候选文档做精排，能明显提升效果。
- 长文档检索。文档中哪怕只有一小段与查询高度相关，多向量模型也能通过局部 token 匹配抓住。
- FAQ 匹配与客服问答。问题和标准答案之间往往存在大量词面重合和句式差异，多向量模型对这种场景更友好。
- 法律、医疗、科研文献检索。这些领域的文本通常很长，且查询往往针对某一句话或某一个术语展开。
当然，多向量模型的代价也很直接：存储量更大，每个文档需要保存不止一个向量；检索时计算量更大，不能像单向量那样直接用余弦相似度建索引。这些需要在工程上做权衡，后面会详细展开。
2. 环境准备与关键依赖
2.1 安装依赖
本文的代码基于 Python 3.10 及以上版本，推荐使用 PyTorch 2.x，并安装sentence-transformers、datasets、transformers这几个核心库。示例环境是 Linux + CUDA，Windows 和 macOS 也基本兼容，只是是否使用 GPU 会有差异。
pip install sentence-transformers pip install datasets pip install transformers pip install torch版本需要根据你的项目实际情况调整，本文示例以常见环境为例，重点演示训练与微调思路。如果你的网络环境不方便直接从 HuggingFace 下载模型，可以提前把模型下载到本地目录，再用本地路径加载。
2.2 硬件环境说明
微调一个 Embedding 模型，和微调一个大语言模型（例如 Qwen、Llama 这种十亿参数以上的模型）差别很大。Sentence Transformers 的底座通常是一个 BERT 规模的模型，参数量在 1 亿到 3 亿左右，一张 16GB 显存的显卡就能跑小 batch 的训练。
如果你只是跑通代码、验证思路，数据量很小的情况下 CPU 也可以完成训练，只是速度慢一些。建议至少准备一张 8GB 以上显存的 GPU，这样在 batch size 和序列长度上会从容很多。这里也顺带回应一个常见疑问：为什么很多人觉得“人工智能训练需要钱多和 GPU 多”？那通常是指大语言模型的全参微调、预训练，而本文这种 embedding 模型微调属于轻量任务，不必有太大压力。
2.3 与 LLM 大模型微调（LoRA）的区别
最近“LoRA 微调”“大模型微调”这些词很热，很多人容易混淆。这里做一个简单区分：
- LoRA 微调针对的是大语言模型的生成能力，目标是让模型回答风格、领域知识发生改变，数据是“指令 + 回答”的形式。
- Embedding 模型微调针对的是语义表征能力，目标是让模型输出的向量能更好地区分“相关”和“不相关”的文本对，数据是“查询 + 文档”的形式。
两者的训练目标和数据形态完全不同。微调一个 embedding 模型通常不需要 LoRA，因为底座模型本身不大；但如果你的显存很紧张，其实也可以考虑对底层的 Transformer 加 LoRA 做参数高效微调，不过这不是本文的重点。先把常规的全参微调流程跑通，再根据资源情况决定是否引入 LoRA 会更稳妥。
3. 多向量嵌入模型的核心概念拆解
3.1 token 级嵌入是什么
在 BERT 类模型中，输入文本会被分词器拆成若干个 token，每个 token 经过编码器后都会得到一个向量。比如输入“深度学习”可能会被拆成['深', '度', '学', '习']四个 token（中文按字或按子词切分），编码器输出就是(4, hidden_dim)的矩阵。
单向量模型在编码后会接一个池化层，把这个矩阵压成一个向量。多向量模型则跳过池化层，直接保留每个 token 的向量。这样做的好处是：查询中的“深度”可以单独去文档中寻找“深度”相关的语义单元，而不需要整个句子对齐。
3.2 Late Interaction 与 MaxSim 计算
Late Interaction 是 ColBERT 论文中提出的交互方式，意思是查询和文档先独立编码，最后在计算相似度时才做交互。这样做的好处是文档向量可以离线预计算，在线只需要对查询做编码，再与文档库中预先存好的向量做交互计算。
MaxSim 是 Late Interaction 的核心步骤：
对查询中的每个 token i： 计算 q_i 与文档所有 token 的相似度 取最大值 max_j sim(q_i, d_j) 最后把所有 token 的最大值加和这个过程可以理解为：查询中每个“语义碎片”都在文档里找自己的最佳对应点，然后把这些最佳匹配累加起来。相比于单向量余弦相似度，这种计算能保留“局部匹配”的信息。
3.3 查询与文档的编码差异
在实际使用中，查询和文档的输入长度策略往往不同。查询通常比较短，一般截断到 32 或 64 个 token；文档可能很长，ColBERT 类模型通常支持 512 个 token 的输入。这样做有两个原因：
- 查询短、文档长更符合检索场景的真实分布。
- 限制序列长度可以显著降低显存占用和计算量，尤其是多向量模型需要保留所有 token 的向量，序列越长开销越大。
在训练时，应该让查询和文档分别走不同的最大长度设置，而不是统一用一个长度。
3.4 手动实现 MaxSim
为了后面训练和推理方便，我们先写一个最朴素的 MaxSim 计算函数。假设查询向量是(q_len, dim)，文档向量是(d_len, dim)：
import torch def maxsim_score(query_emb, doc_emb): # query_emb: (q_len, dim) # doc_emb: (d_len, dim) sim = torch.matmul(query_emb, doc_emb.T) # (q_len, d_len) max_sim = sim.max(dim=-1).values # 每个查询 token 的最大相似度 return max_sim.sum() # 累加这个函数看起来简单，但它是多向量检索的核心。后面的训练损失和推理打分都会围绕它展开。
4. 数据准备：训练多向量模型需要什么样的数据
4.1 训练数据格式
Embedding 模型的训练数据通常有两种基本形式：
- 二元组（pair）：
(query, positive_doc)，表示查询和文档相关。 - 三元组（triplet）：
(query, positive_doc, negative_doc)，表示正例相关、负例不相关。
多向量模型的微调主要采用对比学习思路。一个 batch 内包含若干(query, positive_doc)对，模型需要学会把相关对的分数拉高，把不相关对的分数压低。
以 HuggingFacedatasets库为例，可以这样组织数据：
from datasets import Dataset data = { "query": [ "什么是黑洞？", "怎么做酸面包？", "Python 中如何捕获异常？", ], "positive_doc": [ "黑洞是一种引力极强的天体，连光都无法逃脱。", "酸面包需要面粉、水和天然酵母经过长时间发酵。", "Python 中可以使用 try-except 语句捕获异常。", ], } dataset = Dataset.from_dict(data) print(dataset)如果数据量不大，也可以直接用一个 JSON Lines 文件或者 CSV 文件存储，读取后转成Dataset即可。
4.2 硬负样本为什么重要
训练时，如果只把 batch 内其他样本的 positive 文档当作负样本，这些负样本往往和查询完全不相关，模型很容易学会“区分完全无关的文本”这种简单任务，泛化能力有限。
更好的做法是引入硬负样本（hard negatives）。硬负样本指的是：表面上与查询有一点相关，但实际答案并不正确的文档。它们能逼着模型学习更细致的语义差异。生成硬负样本的常见方法：
- 用 BM25 或单向量模型先做召回，取排名靠前但不是正确答案的文档。
- 从用户点击日志中找“曝光但未点击”的文档。
- 人工构造容易混淆的干扰项。
硬负样本并不是越多越好，每个查询搭配 1 到 3 个优质硬负样本通常在效果和成本上比较平衡。
4.3 数据清洗与数量建议
训练数据质量直接决定微调效果。下面几点是实践中验证过的经验：
- 去掉空文本、超长重复文本。
- 查询和文档不要出现“完全一样”或者“只有标点差异”的样本，否则模型学不到有用的东西。
- 如果做领域微调，建议先收集几千到上万条高质量的查询-文档对；数据量太小（比如只有几百条）时，可以先从公开中文语料中筛选相近领域的数据，再叠加你的人工标注数据。
- 训练集和验证集要按查询维度划分，避免同一个查询的多个相关文档同时出现在训练集和验证集里，造成评估虚高。
5. 完整微调实战：用 Sentence Transformers 训练多向量模型
5.1 加载预训练模型
多向量模型可以基于已经开源的多向量底座继续微调，也可以从一个普通 BERT 类模型开始训练成多向量模型。前者收敛更快，也更推荐。这里以常见可用的多向量模型answerai/answerai-colbert-small-v1为例：
from sentence_transformers import SentenceTransformer model_name = "answerai/answerai-colbert-small-v1" model = SentenceTransformer(model_name)在较新的 Sentence Transformers 版本中，加载这类 ColBERT 风格模型后，encode返回的是 token 级嵌入对象，而不是普通的定长向量。如果你的版本提示输出类型不一致，可以在使用前打印一下输出结构，以实际版本为准。
为了训练时灵活控制输入和损失计算，我们也可以直接用transformers加载同一个模型，得到底层的AutoModel和分词器：
from transformers import AutoTokenizer, AutoModel model_name = "answerai/answerai-colbert-small-v1" tokenizer = AutoTokenizer.from_pretrained(model_name) backbone = AutoModel.from_pretrained(model_name)需要说明的是，不同多向量模型的底层结构可能有差异，例如answerai/answerai-colbert-small-v1是基于 BERT 结构加线性投影，而其他模型可能是别的结构。因此实际代码中请先打印模型的配置，确认输出维度，再继续后续步骤。
5.2 准备数据集与 DataLoader
我们把上一节的Dataset转成 PyTorch 的DataLoader。关键在于写一个collate_fn，分别对查询和文档做分词、填充、截断。
from torch.utils.data import DataLoader def collate_fn(batch): queries = [item["query"] for item in batch] docs = [item["positive_doc"] for item in batch] query_enc = tokenizer( queries, padding=True, truncation=True, max_length=64, return_tensors="pt", ) doc_enc = tokenizer( docs, padding=True, truncation=True, max_length=256, return_tensors="pt", ) return { "query_input_ids": query_enc["input_ids"], "query_attention_mask": query_enc["attention_mask"], "doc_input_ids": doc_enc["input_ids"], "doc_attention_mask": doc_enc["attention_mask"], } dataloader = DataLoader( dataset, batch_size=8, shuffle=True, collate_fn=collate_fn, )这里有一个容易忽略的细节：查询和文档的最大长度要分开设置。查询通常短，文档可以长一些。这样既能保留文档中的关键信息，又能控制显存占用。
5.3 自定义带 Mask 的 MaxSim 函数
由于 batch 内文本长度不同，填充（padding）会产生大量无效的[PAD]token。如果不处理这些 token，MaxSim 计算时它们也会参与点积，从而干扰分数。正确做法是：计算相似度时，用 attention_mask 把 padding 位置变成负无穷大，使它们在取最大值时被自动忽略。
def maxsim_with_mask(query_emb, query_mask, doc_emb, doc_mask): # query_emb: (q_len, dim) # query_mask: (q_len,) # doc_emb: (doc_len, dim) # doc_mask: (doc_len,) sim = torch.matmul(query_emb, doc_emb.T) # (q_len, doc_len) # 文档侧 mask：padding 位置设为 -inf doc_mask = doc_mask.unsqueeze(0).float() # (1, doc_len) sim = sim * doc_mask + (1 - doc_mask) * (-1e9) max_sim = sim.max(dim=-1).values # (q_len,) # 查询侧 mask：padding token 不参与累加 query_mask = query_mask.float() max_sim = max_sim * query_mask return max_sim.sum()这个函数是后面训练损失的基础。小 batch 情况下，逐个 pair 循环计算是可以接受的；生产环境再考虑向量化优化。
5.4 自定义 Late Interaction 对比损失
有了带 mask 的 MaxSim，就可以构造对比损失了。核心思路是：一个 batch 中有 B 个(query, positive_doc)对，把每个查询和 batch 内所有文档都算一遍分数，得到(B, B)的分数矩阵。对角线上的分数是正例，其他位置是 batch 内负样本。然后用交叉熵损失训练。
import torch import torch.nn as nn import torch.nn.functional as F class LateInteractionContrastiveLoss(nn.Module): def __init__(self, temperature=1.0): super().__init__() self.temperature = temperature def forward(self, query_emb, query_mask, doc_emb, doc_mask): # query_emb: (B, q_len, dim) # query_mask: (B, q_len) # doc_emb: (B, d_len, dim) # doc_mask: (B, d_len) batch_size = query_emb.size(0) scores = [] for i in range(batch_size): q_emb = query_emb[i] # (q_len, dim) q_mask = query_mask[i] score_i = [] for j in range(batch_size): d_emb_j = doc_emb[j] # (d_len, dim) d_mask_j = doc_mask[j] score = maxsim_with_mask(q_emb, q_mask, d_emb_j, d_mask_j) score_i.append(score) scores.append(torch.stack(score_i)) scores = torch.stack(scores) # (B, B) scores = scores / self.temperature labels = torch.arange(batch_size, device=query_emb.device) return F.cross_entropy(scores, labels)上面这个写法是教学版，双重循环在 batch 很小时没问题，但训练速度慢。下面给一个向量化版本，思路是扩展维度后一次算出所有 pair 的交互矩阵：
def compute_all_scores(query_emb, query_mask, doc_emb, doc_mask): # query_emb: (B, q_len, dim) # doc_emb: (B, d_len, dim) B = query_emb.size(0) # 扩展为 (B, 1, q_len, dim) 与 (1, B, d_len, dim) q = query_emb.unsqueeze(1) # (B, 1, q_len, dim) d = doc_emb.unsqueeze(0) # (1, B, d_len, dim) sim_all = torch.matmul(q, d.transpose(-1, -2)) # (B, B, q_len, d_len) # 文档 mask 扩展为 (1, B, 1, d_len) d_mask = doc_mask.unsqueeze(0).unsqueeze(2).float() # (1, B, 1, d_len) sim_all = sim_all * d_mask + (1 - d_mask) * (-1e9) max_sim = sim_all.max(dim=-1).values # (B, B, q_len) q_mask = query_mask.unsqueeze(0).unsqueeze(-1).float() # (1, B, q_len, 1) max_sim = max_sim * q_mask scores = max_sim.sum(dim=-1) # (B, B) return scores这个向量化版本在 batch 为 8、查询长度 32、文档长度 256 时内存可以接受，但 batch 和长度变大后会快速增长。实际训练时建议先跑一个小 batch 验证内存，再逐步调大。
5.5 完整训练循环
把上面的模块组合起来，就是一个完整的微调脚本：
import torch from torch.optim import AdamW device = torch.device("cuda" if torch.cuda.is_available() else "cpu") backbone = backbone.to(device) criterion = LateInteractionContrastiveLoss(temperature=1.0) optimizer = AdamW(backbone.parameters(), lr=2e-5) num_epochs = 3 gradient_accumulation_steps = 2 global_step = 0 for epoch in range(num_epochs): for step, batch in enumerate(dataloader): q_ids = batch["query_input_ids"].to(device) q_mask = batch["query_attention_mask"].to(device) d_ids = batch["doc_input_ids"].to(device) d_mask = batch["doc_attention_mask"].to(device) # 编码查询和文档，保留 token 级输出 query_outputs = backbone(input_ids=q_ids, attention_mask=q_mask) doc_outputs = backbone(input_ids=d_ids, attention_mask=d_mask) # 注意：这里需要根据底座模型确认最后一层向量位置 q_emb = query_outputs.last_hidden_state d_emb = doc_outputs.last_hidden_state loss = criterion(q_emb, q_mask, d_emb, d_mask) loss = loss / gradient_accumulation_steps loss.backward() if (step + 1) % gradient_accumulation_steps == 0: optimizer.step() optimizer.zero_grad() if step % 50 == 0: print(f"epoch {epoch}, step {step}, loss {loss.item():.4f}") global_step += 1这里有一个值得注意的点：last_hidden_state是否适合直接作为多向量模型的 token 编码，取决于你选用的底座模型。有些 ColBERT 风格模型在AutoModel之外还带有一个额外的线性投影层，此时需要调用对应模型的完整 forward 逻辑。建议在训练前先用一小批数据打印输出的 shape，并核对每个 token 的输出维度是否符合预期。
5.6 训练效果验证
训练过程中不能只看 loss 降低，还要关注实际的检索效果。一个简单的验证方式是：准备一个小的评估集，包含若干查询和文档库，计算每个查询的召回是否命中正确答案。
def encode_query(text): enc = tokenizer( text, padding=True, truncation=True, max_length=64, return_tensors="pt", ).to(device) with torch.no_grad(): out = backbone(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]) return out.last_hidden_state[0], enc["attention_mask"][0] def encode_doc(text): enc = tokenizer( text, padding=True, truncation=True, max_length=256, return_tensors="pt", ).to(device) with torch.no_grad(): out = backbone(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]) return out.last_hidden_state[0], enc["attention_mask"][0] def recall_at_k(query, docs, labels, k=1): q_emb, q_mask = encode_query(query) scores = [] for doc in docs: d_emb, d_mask = encode_doc(doc) scores.append(maxsim_with_mask(q_emb, q_mask, d_emb, d_mask).item()) top_k = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k] return any(labels[i] == 1 for i in top_k)这个函数在实际项目中会写成批量评估，但思路完全一致：对每个查询，把所有候选文档打一次分，看正确答案是否进入 top-k。
5.7 保存与加载模型
训练结束后，把模型和分词器保存到本地目录：
backbone.save_pretrained("models/my_colbert_finetuned") tokenizer.save_pretrained("models/my_colbert_finetuned")之后加载就可以直接使用：
from transformers import AutoModel, AutoTokenizer model = AutoModel.from_pretrained("models/my_colbert_finetuned") tokenizer = AutoTokenizer.from_pretrained("models/my_colbert_finetuned")如果你希望继续使用 Sentence Transformers 的高层 API 来管理模型，可以把保存目录作为一个新模型路径交给SentenceTransformer加载。需要注意，这样加载后模型是否能输出正确的多向量结构，取决于原模型是不是 Sentence Transformers 支持的架构，不同版本行为可能有差异，建议先做一次小规模推理验证。
6. 推理与检索落地
6.1 单条查询的相似度计算
训练完成后，推理阶段的核心逻辑与验证时一致。对一条查询和一批文档，先分别编码，再用 MaxSim 打分：
def search(query, docs, top_k=5): q_emb, q_mask = encode_query(query) results = [] for idx, doc in enumerate(docs): d_emb, d_mask = encode_doc(doc) score = maxsim_with_mask(q_emb, q_mask, d_emb, d_mask).item() results.append((idx, score)) results.sort(key=lambda x: x[1], reverse=True) return results[:top_k]这里逐条编码文档在文档量
