ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀，每一条都经过真实交付验证。

首页 / 资讯中心 / 正文

SimLM：面向稠密段落检索的表征瓶颈预训练与四阶段微调实践指南

2026/9/14 5:23:08 拓冰建站 浏览量
SimLM：面向稠密段落检索的表征瓶颈预训练与四阶段微调实践指南 SimLM面向稠密段落检索的表征瓶颈预训练与四阶段微调实践指南【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilmSimLM 是 unilm 仓库中一套面向稠密段落检索Dense Passage Retrieval, DPR的检索导向预训练框架其核心思想是用「表征瓶颈Representation Bottleneck」压缩输入信息并将语言建模目标替换为基于词元替换Replaced Language Modeling的自监督目标。本文以 SimLM README 为骨架完整继承其可用模型、依赖、数据下载、复现命令与预期结果并结合 源码目录 深入解析预训练模型、双塔编码器、交叉编码器重排器与知识蒸馏的训练细节帮助读者在 MS-MARCO 段落排序任务上从预训练、微调、检索到重排完整复现 SimLM 的结果。一、核心思想与整体流程SimLM 的检索导向预训练可以概括为一句把输入信息压缩到一个表征瓶颈中并用被替换的语言建模Replaced LM, RLM目标进行自监督学习。预训练完成后仓库提供一条四阶段的监督微调流水线用于训练高质量的稠密检索器四阶段流水线在脚本中对应如下关系用 BM25 难负例训练双塔biencoder检索器对应 train_biencoder_marco.sh用双塔对训练集检索产出可被挖掘的难负例复用 search_marco.sh用交叉编码器cross-encoder教师打分做知识蒸馏再训练双塔对应 train_kd_biencoder.sh 与 gen_kd_teacher_scores.sh训练交叉编码器重排器对双塔 top-200 结果重排对应 train_reranker_marco.sh 与 rerank_marco.sh。从源码结构看预训练入口是 train_rlm.py核心模型是 ReplaceLM双塔训练入口是 train_biencoder.py核心模型是 BiencoderModel重排训练入口是 train_cross_encoder.py核心模型是 Reranker。三者共享统一的参数类 Arguments这让所有脚本都能在命令行用同一套参数体系进行配置。二、可用模型与加载方式README 提供了四个可在 Huggingface 上加载的模型模型说明intfloat/simlm-base-msmarco在 MS-MARCO 段落语料上预训练的 SimLMintfloat/simlm-base-msmarco-finetuned在 MS-MARCO 上经蒸馏微调的双塔intfloat/simlm-msmarco-rerankerMS-MARCO 上的交叉编码器重排器intfloat/simlm-base-wiki100w在 DPR 版本 Wikipedia 段落语料上预训练的 SimLM所有模型都可以通过 Huggingface transformers API 加载from transformers import AutoModel, AutoTokenizer model AutoModel.from_pretrained(intfloat/simlm-base-msmarco-finetuned) tokenizer AutoTokenizer.from_pretrained(intfloat/simlm-base-msmarco-finetuned)从源码看双塔推理时由 BiencoderModelForInference.build 决定权重加载方式当model_name_or_path是本地目录且存在query_model/、passage_model/子目录时会分别加载查询塔与段落塔否则加载共享权重。若add_pooler为真且目录存在pooler.pt还会加载额外的投影层权重。三、环境依赖与安装README 给出的核心依赖为python3.7 transformers4.15 datasets2.0.0 torch1.7 deepspeed0.6.0 pytrec_eval安装命令为pip install -r requirements.txtrequirements.txt 中实际还包含ir_datasets0.5.0、pyserini0.15.0、tqdm、numpy等包用于数据加载、BM25 基线与评估。DeepSpeed 配置由 ds_config.json 提供其要点是启用 fp16enabled: auto、优化器使用 AdamW、调度器为 WarmupDecayLR、Zero 优化阶段为 stage 2并开启allgather_partitions、reduce_scatter、contiguous_gradients以降低显存并通信重叠。各训练脚本均通过--deepspeed ds_config.json传入该配置。四、下载预处理数据以下脚本会下载 MS-MARCO 段落排序任务所需的三个预处理数据包并解压到data/目录bash scripts/download_msmarco_data.sh从 download_msmarco_data.sh 看脚本分别下载并解压msmarco_bm25_official.zip、msmarco_distillation.zip、msmarco_reranker.zip三个压缩包对应后续三种训练/评测场景的数据目录data/msmarco_bm25_official/BM25 难负例训练与双塔评测数据含passages.jsonl.gz、train.jsonl、dev.jsonl及各 split 的 queries/qrelsdata/msmarco_distillation/知识蒸馏训练数据含kd_train.jsonl、kd_dev.jsonldata/msmarco_reranker/重排器训练与评测数据含train.jsonl、dev.jsonl与各 split 的.msmarco.txt输入。五、复现 MS-MARCO 段落排序结果5.1 评测已微调的双塔检索器export DATA_DIR./data/msmarco_bm25_official/ export OUTPUT_DIR./tmp/ # Encode all the corpus passages bash scripts/encode_marco.sh intfloat/simlm-base-msmarco-finetuned # Perform nearest-neighbor search for queries bash scripts/search_marco.sh intfloat/simlm-base-msmarco-finetuned dev bash scripts/search_marco.sh intfloat/simlm-base-msmarco-finetuned trec_dl2019 bash scripts/search_marco.sh intfloat/simlm-base-msmarco-finetuned trec_dl2020预期结果dev MRR10dev R50dev R1kTREC DL 2019 nDCG10TREC DL 2020 nDCG1041.187.898.771.469.7源码层面的实现链路如下编码段落encode_marco.sh 调用 encode_main.py以torch.multiprocessing.spawn按 GPU 数分片并行编码。每条段落按title contents拼接_psg_transform_funcp_max_len144编码结果按encode_shard_size默认 200 万切分保存为shard_{gpu}_{idx}张量。近邻检索search_marco.sh 调用 search_main.py。它对每个 GPU 分片上的 query 集合逐 shard 加载段落向量做torch.mm(query, passage.T)并取search_topktrain 用 200、其他 1000。检索完成后_compute_and_save_metrics会合并多 worker 的 topk 结果若存在{split}_qrels.txt则调用 trec_eval 与 compute_mrr 计算指标并写metrics_{split}.json。5.2 评测已发布的交叉编码器重排器export DATA_DIR./data/msmarco_reranker/ export OUTPUT_DIR./tmp/ bash scripts/rerank_marco.sh intfloat/simlm-msmarco-reranker $DATA_DIR/dev.msmarco.txt dev bash scripts/rerank_marco.sh intfloat/simlm-msmarco-reranker $DATA_DIR/trec_dl2019.msmarco.txt trec_dl2019 bash scripts/rerank_marco.sh intfloat/simlm-msmarco-reranker $DATA_DIR/trec_dl2020.msmarco.txt trec_dl2020 # Will not compute metrics since test labels are not available bash scripts/rerank_marco.sh intfloat/simlm-msmarco-reranker $DATA_DIR/test.msmarco.txt test预期结果dev MRR10dev R50dev R1kTREC DL 2019 nDCG10TREC DL 2020 nDCG1043.889.298.674.672.7从 rerank_marco.sh 看重排入口 rerank_main.py 以rerank_depth200对双塔 top-200 候选做交叉编码器打分rerank_max_length192输出到{OUTPUT_DIR}/rerank.{split}.msmarco.txt。5.3 用 BM25 难负例训练双塔检索器GPU 需求4 张 V10032GB。export DATA_DIR./data/msmarco_bm25_official/ export OUTPUT_DIR./checkpoint/biencoder/ # Train bi-encoder bash scripts/train_biencoder_marco.sh # Encode corpus passages bash scripts/encode_marco.sh $OUTPUT_DIR # Evaluate on each split bash scripts/search_marco.sh $OUTPUT_DIR dev bash scripts/search_marco.sh $OUTPUT_DIR trec_dl2019 bash scripts/search_marco.sh $OUTPUT_DIR trec_dl2020 bash scripts/search_marco.sh $OUTPUT_DIR test # Predictions for training datasets can be used as mined hard negatives bash scripts/search_marco.sh $OUTPUT_DIR train预期结果dev MRR10dev R50dev R1k38.085.898.3关键超参可从 train_biencoder_marco.sh 中提取model_name_or_pathintfloat/simlm-base-msmarco、t0.02、train_n_passages16、per_device_train_batch_size16、q_max_len32、p_max_len144、num_train_epochs3、learning_rate2e-5、use_scaled_lossTrue、share_encoderTrue。5.4 用知识蒸馏训练双塔GPU 需求4 张 V10032GB。export DATA_DIR./data/msmarco_distillation/ export OUTPUT_DIR./checkpoint/distilled_biencoder/ # Train bi-encoder with knowledge distillation bash scripts/train_kd_biencoder.sh # Encode corpus passages bash scripts/encode_marco.sh $OUTPUT_DIR # Evaluate on each split bash scripts/search_marco.sh $OUTPUT_DIR dev bash scripts/search_marco.sh $OUTPUT_DIR trec_dl2019 bash scripts/search_marco.sh $OUTPUT_DIR trec_dl2020 bash scripts/search_marco.sh $OUTPUT_DIR test结果预期接近intfloat/simlm-base-msmarco-finetuned。从 train_kd_biencoder.sh 看蒸馏训练使用do_kd_biencoder、kd_mask_hnFalse、kd_cont_loss_weight0.2、train_n_passages24、num_train_epochs6、learning_rate3e-5并按metric_for_best_modelmrr加载最优模型。蒸馏损失的具体实现见下文 7.2。5.5 训练交叉编码器重排器GPU 需求8 张 V10032GB。export DATA_DIR./data/msmarco_reranker/ export OUTPUT_DIR./checkpoint/cross_encoder_reranker/ # Train cross-encoder re-ranker bash scripts/train_reranker_marco.sh # Re-rank top-200 outputs by biencoder retrievers bash scripts/rerank_marco.sh $OUTPUT_DIR $DATA_DIR/dev.msmarco.txt bash scripts/rerank_marco.sh $OUTPUT_DIR $DATA_DIR/trec_dl2019.msmarco.txt trec_dl2019 bash scripts/rerank_marco.sh $OUTPUT_DIR $DATA_DIR/trec_dl2020.msmarco.txt trec_dl2020 bash scripts/rerank_marco.sh $OUTPUT_DIR $DATA_DIR/test.msmarco.txt test结果预期接近intfloat/simlm-msmarco-reranker。从 train_reranker_marco.sh 看重排器基于google/electra-base-discriminatorrerank_max_length192、rerank_use_rdropTrue、train_n_passages64、rerank_forward_factor4、learning_rate3e-5脚本注释提示 electra-large 下learning_rate 1e-5会导致训练不稳定。5.6 在目标语料上预训练 SimLMGPU 需求8 张 V100至少 16GB。export DATA_DIR./data/msmarco_bm25_official/ export OUTPUT_DIR./checkpoint/replaced_lm/ bash ./scripts/train_rlm.sh预训练完成后按照上述监督微调流程评测模型质量。从 train_rlm.sh 看预训练以bert-base-uncased为初始 encodertrain_filepassages.jsonl.gzrlm_max_length144、rlm_encoder_mask_prob0.3、rlm_decoder_mask_prob0.5、rlm_generator_model_namegoogle/electra-base-generator、rlm_freeze_generatorTrue、rlm_generator_mlm_weight0.2、all_use_mask_tokenTrue、max_steps80000、learning_rate3e-4、warmup_steps4000、weight_decay0.0。六、统一参数体系Arguments 数据类所有训练/推理脚本共享 Arguments它继承自 transformers 的TrainingArguments并补充了检索任务专属字段。可按用途分组理解数据与任务data_dir、task_typeir/qa、train_file、validation_file、train_n_passages每个 query 的段落数含正负、q_max_len默认 32、p_max_len默认 144。双塔训练share_encoder、use_scaled_loss、loss_scale、t温度默认 0.05、l2_normalize、t_warmup、full_contrastive_loss、add_pooler、out_dimension默认 768。编码与检索do_encode、encode_in_path、encode_save_dir、encode_shard_size默认 200 万、encode_batch_size默认 256do_search、search_split、search_batch_size默认 128、search_topk默认 200、search_out_dir。重排do_rerank、rerank_max_length默认 256、rerank_in_path、rerank_out_path、rerank_split、rerank_batch_size、rerank_depth、rerank_forward_factor、rerank_use_rdrop。蒸馏do_kd_gen_score、kd_gen_score_split、kd_gen_score_batch_size、kd_gen_score_n_neg默认 30、do_kd_biencoder、kd_mask_hn默认 True、kd_cont_loss_weight默认 1.0。替换语言建模预训练rlm_generator_model_name默认google/electra-base-generator、rlm_freeze_generator、rlm_generator_mlm_weight默认 0.2、all_use_mask_token、rlm_num_eval_samples默认 4096、rlm_max_length默认 144、rlm_decoder_layers默认 2、rlm_encoder_mask_prob默认 0.3、rlm_decoder_mask_prob默认 0.5。__post_init__中还包含若干一致性断言必须存在data_dir、必须可用 CUDAdo_rerank时要求train_n_passages 1且train_n_passages % rerank_forward_factor 0do_kd_biencoder且use_scaled_loss时要求not kd_mask_hn。这些断言解释了脚本中超参组合为什么必须成对出现。七、关键源码实现解析7.1 替换语言建模预训练ReplaceLMReplaceLM 是 SimLM 预训练的核心模块其结构与前向流程如下结构encoder是加载的 BERT 掩码语言模型decoder是 encoder 最后rlm_decoder_layers默认 2层的深拷贝generator是rlm_generator_model_name指定的 Electra 生成器若rlm_freeze_generator为真则冻结参数并置于 eval 模式。前向forward将输入切分为enc_前缀的 encoder 输入与dec_前缀的 decoder 输入_replace_tokens用 generator 对被掩码位置采样出替换词元得到 encoder/decoder 的input_idsencoder 侧会回传 generator 的 MLM lossg_mlm_lossdecoder 侧在no_grad下采样encoder 前向得到cls_hidden取最后一层隐藏态的[CLS]decoder 输入嵌入取enc侧[CLS]与 decoder 输入嵌入的拼接逐层通过decoder的若干层 transformer得到隐藏态后用 encoder 的cls头算decoder_mlm_loss总损失为loss decoder_mlm_loss encoder_out.loss g_mlm_loss * rlm_generator_mlm_weight并记录replace_ratio被替换词元占比。替换词元的采样逻辑在_replace_tokensgenerator 输出 logits 后经Categorical采样仅对labels 0的位置即被掩码位置用采样词元替换原词元其余位置保留。数据侧的掩码由 DataCollatorForReplaceLM 完成对title contents做 whole-word maskencoder 用rlm_encoder_mask_probdecoder 在此基础上按rlm_decoder_mask_prob - rlm_encoder_mask_prob的概率再叠加一层掩码因两次独立掩码需除以1 - rlm_encoder_mask_prob修正。all_use_mask_tokenTrue时被选中的位置全部填[MASK]不做 80:10:10 扰动这正是 train_rlm.sh 的取值。7.2 双塔模型BiencoderModelBiencoderModel 的前向分两条训练路径普通对比学习do_kd_biencoderFalse_compute_scores用共享/双塔编码器得到q_reps、p_reps经dist_gather_tensor跨卡聚合为all_q_reps、all_p_reps再调用full_contrastive_scores_and_labels生成全对比分数与标签。若l2_normalize且t_warmup未开启分数乘以1/t。当use_scaled_loss为真时交叉熵损失乘以world_size或loss_scale以补偿跨卡聚合带来的负样本规模。知识蒸馏do_kd_biencoderTrue在本地 query 的train_n_passages个段落组上用group_log_scores与教师软标签kd_labels的 log-softmax 计算 KL 散度kd_loss可选地用kd_mask_hn把组内非正例位置打分为-inf总损失为kd_cont_loss_weight * ce_loss kd_loss。_encode取last_hidden_state[:, 0]即[CLS]表示经pooler默认为恒等add_poolerFalse后做 L2 归一化。BiencoderModelForInference去掉训练逻辑仅返回q_reps/p_reps供编码与检索阶段使用。训练过程的指标统计在 BiencoderTrainer用batch_mrr、accuracy(topk(1,3))维护 mrr、Acc1、Acc3 三个 AverageMeter并在logging_steps节点打印_save会调用model.save区分共享塔与双塔的保存目录结构。7.3 交叉编码器重排器RerankerReranker 基于AutoModelForSequenceClassification前向支持两种机制rerank_forward_factor 1的两阶段前向先用no_grad对整个train_n_passages组前向选出k train_n_passages // rerank_forward_factor个分数最高的样本保证第 0 位正例不被剔除logits[:, 0] inf再对选出的子集做带梯度的前向。这是一种「前向选 top、反向精算」的省显存技巧config.py的断言要求train_n_passages % rerank_forward_factor 0。rerank_use_rdropTrue时仅训练态把输入沿 batch 维复制一倍做两次前向计算双向 KL 的 R-Drop 正则与两项交叉熵outputs.loss rdrop_loss ce_loss否则直接用train_n_passages // rerank_forward_factor组做交叉熵。推理侧 RerankerForInference 仅做eval前向供rerank_main.py调用。7.4 教师打分知识蒸馏数据生成蒸馏需要教师分数。gen_kd_teacher_scores.sh 调用 gen_teacher_scores.py参数do_kd_gen_score、kd_gen_score_batch_size256、kd_gen_score_n_neg1000、rerank_max_length192。Arguments中kd_gen_score_n_neg默认 30脚本中提升到 1000用于为每个 query 生成包含上千难负例的软标签供后续train_kd_biencoder.sh训练双塔时作为kd_labels使用。八、常见问题FAQ8.1 必须使用 DeepSpeed 启动器吗README 强烈建议使用 DeepSpeed 启动训练以获得更快的训练速度与更低的显存占用。若 DeepSpeed 不可用可将 shell 脚本中的 deepspeed 命令行替换为 PyTorch 启动器# Uncomment this line to use pytorch launcher and delete the deepspeed command python -u -m torch.distributed.launch --nproc_per_node 4 src/train_biencoder.py # deepspeed src/train_biencoder.py --deepspeed ds_config.json各训练脚本如 train_biencoder_marco.sh中都保留了注释掉的 pytorch launcher 行可参照切换。8.2 MS-MARCO 段落数据中的 title 字段来自哪里title 数据来自 RocketQA也被 coCondenser 用于训练。从 encode_main.py 看段落编码时确实以title作为text、contents作为text_pair拼接这与数据侧的字段约定一致。九、致谢与引用README 说明部分代码基于 Tevatron。引用论文可使用如下 BibTeXinproceedings{wang-etal-2023-simlm, title {S}im{LM}: Pre-training with Representation Bottleneck for Dense Passage Retrieval, author Wang, Liang and Yang, Nan and Huang, Xiaolong and Jiao, Binxing and Yang, Linjun and Jiang, Daxin and Majumder, Rangan and Wei, Furu, booktitle Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), month jul, year 2023, address Toronto, Canada, publisher Association for Computational Linguistics, pages 2244--2258, }十、总结与延伸阅读SimLM 将「表征瓶颈 替换语言建模」作为预训练目标再用 BM25 难负例、知识蒸馏、交叉编码器重排构成的四阶段流水线把预训练模型打磨为高质量稠密检索器。仓库中所有训练/推理入口都围绕 Arguments 统一参数体系组织配合 ds_config.json 的 ZeRO-2 配置可在 4~8 张 V100 上完整复现 README 给出的 MS-MARCO 结果。若想深入可依次阅读预训练ReplaceLM、DataCollatorForReplaceLM、train_rlm.sh双塔BiencoderModel、BiencoderTrainer、train_biencoder_marco.sh、train_kd_biencoder.sh重排Reranker、train_reranker_marco.sh、rerank_main.py检索与评估encode_main.py、search_main.py、metrics.py。【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考