训练 Vela Embedding 和 Reranker
使用 Vela Embedding 高效找到相关文档,再用 Vela Reranker 改善少量候选的排序。两者均从共享的 Vela Encoder适配而来,并支持选择编码器深度和输出维度。
直接使用已发布模型,请参阅嵌入和重排序。
Embedding:双编码器
Embedding 分别将查询和文档编码为归一化向量。预先计算文档向量,再用余弦相似度或点积与查询比较。
使用查询与正例文档,以及有效的负例训练。如果应用还需要比较请求、对相关文本分组或识别重复提问,可以加入语义相似度或释义样本。覆盖索引实际使用的语言和领域。
Reranker:交叉编码器
Reranker 联合读取查询和候选文档,输出相关性分数。它用于检索系统已返回的候选。
训练时提供完整候选列表和相关性标注。从检索系统挖掘难负例,有助于区分看似相关但不正确的结果。未标注相关性的文档应与已确认负例分开。分数用于排序,不代表文档正确的概率。
选择深度、维度和上下文
| 设置 | 可选值 | 权衡 |
|---|---|---|
| 编码器深度 | 3、6、11、22 层 | 更少层数减少编码器计算 |
| 维度 | 64、128、256、512、768 | 更小的嵌入向量降低索引存储和比较成本 |
| 输入上限 | 最多 32,768 tokens | 更长输入需要更多内存和时间 |
训练会监督全部 20 种深度与维度组合。Reranker 为每种组合提供已训练的评分 head;降低维度不会跳过编码器层。评测时选择实际准备部署的组合。
Embedding 的查询和索引文档必须使用相同模型 revision、深度和维度。改变这些设置后重建索引。Reranker 的预算包含查询、文档和特殊 tokens 的总长度。
准备训练
在仓库 checkout 中创建独立环境,安装适合加速器的 PyTorch 构建及训练依赖。ROCm 使用 PyTorch 的 cuda 设备名称。
选择一种起点:
- **新任务:**下载 Vela Encoder,从完整 checkpoint 初始化任务。
- **继续任务:**下载 Vela Embedding 或 Reranker,设置
initialization: "continued_task",保留编码器和 head,重新初始化优化器。 - **恢复中断:**使用
--resume恢复优化器和训练进度。
分别准备训练集、开发集、采样计划和训练配置。训练参考提供 JSON 格式、损失函数、可选教师监督和配置字段。
以下命令假定文件已准备在 /data/retrieval。将 VELA_TRAIN_CONFIG_SHA256 设置为配置文件的 SHA-256:
export PYTHONPATH="$PWD"
python -m src.training.model_embeddings.mmbert_32k.newbase_training \
--config /data/retrieval/task.json \
--config-sha256 "${VELA_TRAIN_CONFIG_SHA256:?Set the configuration SHA-256}" \
--output /data/retrieval/run --device cuda
配置中选择 task: "embedding" 或 task: "reranker",并设置数据、输入预算、训练步数和评测间隔。先用小预算检查首轮开发集结果,再延长训练。
评测 checkpoint
使用与基线相同的开发集评测已保存的 checkpoint:
python -m src.training.model_embeddings.mmbert_32k.newbase_scoring \
--model /data/retrieval/run/step-100 --task embedding \
--known-dev /data/retrieval/validation --split validation \
--output /data/retrieval/evaluation --device cpu --token-budget 32768
把 checkpoint 路径替换成实际保存的步骤;Reranker 使用 --task reranker。先根据开发集选择 checkpoint,再在独立测试集上做最终对比。
Embedding 需要覆盖检索、相似度和多语言迁移。Reranker 使用固定候选列表,对比 nDCG 等排序指标。分别报告短输入、长输入,以及各部署深度与维度的延迟和内存。
完整 MMTEB 分数需要运行所选基准的全部任务,并使用一致评测协议。任务子集可以诊断短板,但不能得出总体分数或榜单排名。
部署结果
通过本地模型绑定选择 checkpoint 和引擎,再通过路由预览检查深度、维度和输入预算。ONNX 部署需要从同一组训练权重导出的图。
早期 mmBERT 流程
原始基座、Embedding 和 Reranker 配方保留在工作流 README。历史配置用于早期 mmBERT 流程;新的 Vela 任务使用上面的 Vela 起始模型。