之前在做一个 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遍历查询的所有 tokenj遍历文档的所有 token。这个计算方式允许查询中的每个词面或语义单元在文档中找到自己最匹配的位置而不是强迫整个句子变成一个向量。1.4 多向量模型适合哪些场景多向量嵌入模型并不是要完全替代单向量模型而是在特定场景下优势更明显RAG 检索中的精排环节。先用单向量模型做粗排召回再用多向量模型对候选文档做精排能明显提升效果。长文档检索。文档中哪怕只有一小段与查询高度相关多向量模型也能通过局部 token 匹配抓住。FAQ 匹配与客服问答。问题和标准答案之间往往存在大量词面重合和句式差异多向量模型对这种场景更友好。法律、医疗、科研文献检索。这些领域的文本通常很长且查询往往针对某一句话或某一个术语展开。当然多向量模型的代价也很直接存储量更大每个文档需要保存不止一个向量检索时计算量更大不能像单向量那样直接用余弦相似度建索引。这些需要在工程上做权衡后面会详细展开。2. 环境准备与关键依赖2.1 安装依赖本文的代码基于 Python 3.10 及以上版本推荐使用 PyTorch 2.x并安装sentence-transformers、datasets、transformers这几个核心库。示例环境是 Linux CUDAWindows 和 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, paddingTrue, truncationTrue, max_length64, return_tensorspt, ) doc_enc tokenizer( docs, paddingTrue, truncationTrue, max_length256, return_tensorspt, ) 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_size8, shuffleTrue, collate_fncollate_fn, )这里有一个容易忽略的细节查询和文档的最大长度要分开设置。查询通常短文档可以长一些。这样既能保留文档中的关键信息又能控制显存占用。5.3 自定义带 Mask 的 MaxSim 函数由于 batch 内文本长度不同填充padding会产生大量无效的[PAD]token。如果不处理这些 tokenMaxSim 计算时它们也会参与点积从而干扰分数。正确做法是计算相似度时用 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) # 文档侧 maskpadding 位置设为 -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,) # 查询侧 maskpadding 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, temperature1.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, devicequery_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(temperature1.0) optimizer AdamW(backbone.parameters(), lr2e-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_idsq_ids, attention_maskq_mask) doc_outputs backbone(input_idsd_ids, attention_maskd_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(fepoch {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, paddingTrue, truncationTrue, max_length64, return_tensorspt, ).to(device) with torch.no_grad(): out backbone(input_idsenc[input_ids], attention_maskenc[attention_mask]) return out.last_hidden_state[0], enc[attention_mask][0] def encode_doc(text): enc tokenizer( text, paddingTrue, truncationTrue, max_length256, return_tensorspt, ).to(device) with torch.no_grad(): out backbone(input_idsenc[input_ids], attention_maskenc[attention_mask]) return out.last_hidden_state[0], enc[attention_mask][0] def recall_at_k(query, docs, labels, k1): 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)), keylambda i: scores[i], reverseTrue)[: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_k5): 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(keylambda x: x[1], reverseTrue) return results[:top_k]这里逐条编码文档在文档量