企业官网建设流程全解析

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

【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm

SimLM 是 unilm 仓库中一套面向稠密段落检索（Dense Passage Retrieval, DPR）的检索导向预训练框架，其核心思想是用「表征瓶颈（Representation Bottleneck）」压缩输入信息，并将语言建模目标替换为基于词元替换（Replaced Language Modeling）的自监督目标。本文以 SimLM README 为骨架，完整继承其可用模型、依赖、数据下载、复现命令与预期结果，并结合 源码目录 深入解析预训练模型、双塔编码器、交叉编码器重排器与知识蒸馏的训练细节，帮助读者在 MS-MARCO 段落排序任务上从预训练、微调、检索到重排完整复现 SimLM 的结果。

一、核心思想与整体流程

SimLM 的检索导向预训练可以概括为一句：把输入信息压缩到一个表征瓶颈中，并用被替换的语言建模（Replaced LM, RLM）目标进行自监督学习。预训练完成后，仓库提供一条四阶段的监督微调流水线，用于训练高质量的稠密检索器：

四阶段流水线在脚本中对应如下关系：

1. 用 BM25 难负例训练双塔（biencoder）检索器，对应 train_biencoder_marco.sh；
2. 用双塔对训练集检索，产出可被挖掘的难负例，复用 search_marco.sh；
3. 用交叉编码器（cross-encoder）教师打分，做知识蒸馏再训练双塔，对应 train_kd_biencoder.sh 与 gen_kd_teacher_scores.sh；
4. 训练交叉编码器重排器，对双塔 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 上加载的模型：

所有模型都可以通过 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 给出的核心依赖为：

python>=3.7 transformers==4.15 datasets==2.0.0 torch>=1.7 deepspeed==0.6.0 pytrec_eval

安装命令为：

pip install -r requirements.txt

requirements.txt 中实际还包含ir_datasets==0.5.0、pyserini==0.15.0、tqdm、numpy等包，用于数据加载、BM25 基线与评估。

DeepSpeed 配置由 ds_config.json 提供，其要点是：启用 fp16（enabled: "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/qrels；
• data/msmarco_distillation/：知识蒸馏训练数据，含kd_train.jsonl、kd_dev.jsonl；
• data/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

预期结果：

源码层面的实现链路如下：

• 编码段落：encode_marco.sh 调用 encode_main.py，以torch.multiprocessing.spawn按 GPU 数分片并行编码。每条段落按title + contents拼接（_psg_transform_func），p_max_len=144，编码结果按encode_shard_size（默认 200 万）切分保存为shard_{gpu}_{idx}张量。
• 近邻检索：search_marco.sh 调用 search_main.py。它对每个 GPU 分片上的 query 集合，逐 shard 加载段落向量，做torch.mm(query, passage.T)并取search_topk（train 用 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

预期结果：

从 rerank_marco.sh 看，重排入口 rerank_main.py 以rerank_depth=200对双塔 top-200 候选做交叉编码器打分，rerank_max_length=192，输出到{OUTPUT_DIR}/rerank.{split}.msmarco.txt。

5.3 用 BM25 难负例训练双塔检索器

GPU 需求：4 张 V100（32GB）。

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

预期结果：

关键超参可从 train_biencoder_marco.sh 中提取：model_name_or_path=intfloat/simlm-base-msmarco、t=0.02、train_n_passages=16、per_device_train_batch_size=16、q_max_len=32、p_max_len=144、num_train_epochs=3、learning_rate=2e-5、use_scaled_loss=True、share_encoder=True。

5.4 用知识蒸馏训练双塔

GPU 需求：4 张 V100（32GB）。

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_hn=False、kd_cont_loss_weight=0.2、train_n_passages=24、num_train_epochs=6、learning_rate=3e-5，并按metric_for_best_model=mrr加载最优模型。蒸馏损失的具体实现见下文 7.2。

5.5 训练交叉编码器重排器

GPU 需求：8 张 V100（32GB）。

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-discriminator，rerank_max_length=192、rerank_use_rdrop=True、train_n_passages=64、rerank_forward_factor=4、learning_rate=3e-5，脚本注释提示 electra-large 下learning_rate > 1e-5会导致训练不稳定。

5.6 在目标语料上预训练 SimLM

GPU 需求：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为初始 encoder，train_file=passages.jsonl.gz，rlm_max_length=144、rlm_encoder_mask_prob=0.3、rlm_decoder_mask_prob=0.5、rlm_generator_model_name=google/electra-base-generator、rlm_freeze_generator=True、rlm_generator_mlm_weight=0.2、all_use_mask_token=True、max_steps=80000、learning_rate=3e-4、warmup_steps=4000、weight_decay=0.0。

六、统一参数体系：Arguments 数据类

所有训练/推理脚本共享 Arguments，它继承自 transformers 的TrainingArguments，并补充了检索任务专属字段。可按用途分组理解：

• 数据与任务：data_dir、task_type（ir/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（默认 256）；do_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、必须可用 CUDA；do_rerank时要求train_n_passages > 1且train_n_passages % rerank_forward_factor == 0；do_kd_biencoder且use_scaled_loss时要求not kd_mask_hn。这些断言解释了脚本中超参组合为什么必须成对出现。

七、关键源码实现解析

7.1 替换语言建模预训练（ReplaceLM）

ReplaceLM 是 SimLM 预训练的核心模块，其结构与前向流程如下：

• 结构：encoder是加载的 BERT 掩码语言模型；decoder是 encoder 最后rlm_decoder_layers（默认 2）层的深拷贝；generator是rlm_generator_model_name指定的 Electra 生成器，若rlm_freeze_generator为真则冻结参数并置于 eval 模式。
• 前向（forward）：
  1. 将输入切分为enc_前缀的 encoder 输入与dec_前缀的 decoder 输入；
  2. _replace_tokens用 generator 对被掩码位置采样出替换词元，得到 encoder/decoder 的input_ids；encoder 侧会回传 generator 的 MLM loss（g_mlm_loss），decoder 侧在no_grad下采样；
  3. encoder 前向得到cls_hidden（取最后一层隐藏态的[CLS]），decoder 输入嵌入取enc侧[CLS]与 decoder 输入嵌入的拼接；
  4. 逐层通过decoder的若干层 transformer，得到隐藏态后用 encoder 的cls头算decoder_mlm_loss；
  5. 总损失为loss = decoder_mlm_loss + encoder_out.loss + g_mlm_loss * rlm_generator_mlm_weight，并记录replace_ratio（被替换词元占比）。

替换词元的采样逻辑在_replace_tokens：generator 输出 logits 后经Categorical采样，仅对labels >= 0的位置（即被掩码位置）用采样词元替换原词元，其余位置保留。

数据侧的掩码由 DataCollatorForReplaceLM 完成：对title + contents做 whole-word mask，encoder 用rlm_encoder_mask_prob，decoder 在此基础上按rlm_decoder_mask_prob - rlm_encoder_mask_prob的概率再叠加一层掩码（因两次独立掩码，需除以1 - rlm_encoder_mask_prob修正）。all_use_mask_token=True时，被选中的位置全部填[MASK]（不做 80:10:10 扰动），这正是 train_rlm.sh 的取值。

7.2 双塔模型（BiencoderModel）

BiencoderModel 的前向分两条训练路径：

• 普通对比学习（do_kd_biencoder=False）：_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_biencoder=True）：在本地 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_pooler=False）后做 L2 归一化。BiencoderModelForInference去掉训练逻辑，仅返回q_reps/p_reps，供编码与检索阶段使用。

训练过程的指标统计在 BiencoderTrainer：用batch_mrr、accuracy(topk=(1,3))维护 mrr、Acc@1、Acc@3 三个 AverageMeter，并在logging_steps节点打印；_save会调用model.save，区分共享塔与双塔的保存目录结构。

7.3 交叉编码器重排器（Reranker）

Reranker 基于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_rdrop=True时（仅训练态）：把输入沿 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_size=256、kd_gen_score_n_neg=1000、rerank_max_length=192。Arguments中kd_gen_score_n_neg默认 30，脚本中提升到 1000，用于为每个 query 生成包含上千难负例的软标签，供后续train_kd_biencoder.sh训练双塔时作为kd_labels使用。

八、常见问题（FAQ）

8.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。引用论文可使用如下 BibTeX：

@inproceedings{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