fairseq LayerDrop 指南：用结构化 Dropout 训练可裁剪 Transformer（Reducing Transformer Depth on Demand）

fairseq LayerDrop 指南：用结构化 Dropout 训练可裁剪 Transformer（Reducing Transformer Depth on Demand）

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

LayerDrop（Structured Dropout，Fan et al., 2019）是一种在训练阶段随机丢弃整个 Transformer 层的正则化与模型裁剪方案：以 0.1～0.2 的低丢弃率训练出的模型，在推理阶段可以按需移除任意比例的层（例如 50%）而不需要重新训练。本文以decoding/IAD/fairseq中内置的 fairseq 实现为主线，完整讲解 LayerDrop 的训练参数、推理裁剪命令、预训练模型评测、论文复现步骤、调参建议与常见问题，并结合 LayerDropModuleList 与 prune_state_dict 的源码，说明其底层工作原理。读完本文，你将掌握"训练一个大而鲁棒的模型 → 按需裁剪成任意深度的小模型"这一完整实战链路。

LayerDrop 是什么：从"随机丢弃层"到"按需裁剪深度"

传统 Dropout 随机丢弃的是神经元（标量级），而 LayerDrop 随机丢弃的是整层 Transformer 模块（结构化级别）。论文《Reducing Transformer Depth on Demand with Structured Dropout》(arXiv:1909.11556) 的核心思想是：让网络在训练阶段"习惯"缺失部分层仍然能正常工作，从而在推理阶段将任意连续子集之外的层直接移除，模型依然保持可用性。

在 fairseq 中，这一机制由 LayerDropModuleList 实现。它是torch.nn.ModuleList的子类，其核心是重写了__iter__方法：

def __iter__(self): dropout_probs = torch.empty(len(self)).uniform_() for i, m in enumerate(super().__iter__()): if not self.training or (dropout_probs[i] > self.p): yield m

关键行为可以从源码中归纳出两点：

1. 每次迭代都会重新随机：每次遍历 LayerDropModuleList 时，都会用均匀分布生成一组与层数等长的随机数，概率小于p的层被跳过。因此同一个 batch 的前向与后向、甚至不同 batch 之间，"哪些层被丢弃"都在变化。
2. 评估/推理时默认保留全部层：当self.training为 False 时，条件not self.training恒为真，迭代会完整遍历所有层。这意味着推理时模型默认按完整深度运行，而"裁剪"是通过加载裁剪后的 checkpoint 或传参实现的（见下文）。

在 Transformer 模型 与 TransformerDecoder 中，当encoder_layerdrop > 0.0或decoder_layerdrop > 0.0时，层的容器从普通ModuleList切换为LayerDropModuleList(p=...)，丢弃概率分别作用于编码器与解码器。

训练带 LayerDrop 的模型：命令行参数与推荐值

在 fairseq 训练命令中追加以下两个参数即可开启 LayerDrop：

--encoder-layerdrop 0.2 --decoder-layerdrop 0.2

对应参数在 transformer.py 的参数定义 中：

使用要点：

• 推荐值 0.2：官方实验表明 0.2 在大多数场景下效果良好，0.1～0.2 是常用区间。
• 编码器与解码器可以分别设置不同的 LayerDrop 值。
• Decoder-only 的语言模型只需--decoder-layerdrop（如 transformer_lm.py 中仅定义了 decoder 侧参数）。
• Encoder-only 的 RoBERTa 只需--encoder-layerdrop，在 roberta/model.py 中生效。

从配置文件看默认值

fairseq 的 Hydra 配置也内置了对应字段，例如 transformer_lm_wiki103.yaml 中的：

decoder_layerdrop: 0 decoder_layers_to_keep: null

可见默认情况下 LayerDrop 关闭（概率为 0），需要显式开启。在 Transformer 语言模型的 Dataclass 定义 中，decoder_layerdrop为float、decoder_layers_to_keep为Optional[str]，与命令行参数一一对应。

官方提供的预训练模型

官方发布了以下已按 LayerDrop 0.2 训练的模型（参见 layerdrop README）：

推理阶段裁剪模型：按需保留指定层

训练完成后，通过--encoder-layers-to-keep/--decoder-layers-to-keep指定要保留的层编号（从 0 开始，逗号分隔）即可裁剪模型：

--encoder-layers-to-keep 0,2,4,6,8,10,12,14 --decoder-layers-to-keep 0,2,4,6,8,10,12,14

设置该参数后，日志会打印类似| Pruning model to specified layer configuration的信息（见 checkpoint_utils.py 中的日志），且模型参数量显著下降。原文档给出了一个 16 层 Transformer 语言模型的对比：

# 16 层完整模型 num. model params: 246933504 # 裁剪到 8 层后 num. model params: 146163712

裁剪的底层原理：prune_state_dict

裁剪并非运行时跳过层，而是真正从 checkpoint 中删除对应层的权重并重新映射键名。这一过程由 prune_state_dict 完成，它会在加载 checkpoint 时被自动调用，无需手动触发。其流程可以概括为：

1. 读取model_cfg中的encoder_layers_to_keep/decoder_layers_to_keep；若两者都为空则直接返回原 state_dict。
2. 对每个要裁剪的部分构造一次"裁剪通道"：将保留层编号解析为有序列表，并建立"原层号 → 新层号"的映射（例如保留0,2,4,6时，原层 0→新层 0、原层 2→新层 1、原层 4→新层 2……），同时用正则^{layer}.*\.layers\.(\d+)匹配对应编码器/解码器的层权重键名。
3. 遍历原 state_dict 的所有键：没有\.layers\.<数字>\.结构的键（如 embedding、layernorm 等支撑层）原样保留；匹配到的层权重则通过映射改写层号后放入新 state_dict。
4. 裁剪完成后，将model_cfg中的*_layers_to_keep置为 None，避免后续重复裁剪。

值得注意的细节是：在 Transformer.build_model 中，一旦设置了--encoder-layers-to-keep，args.encoder_layers会被更新为保留层数量，因此裁剪后的模型架构（层数）与裁剪后的权重是严格对齐的。

仅做评测时的 model-overrides

如果你只想用脚本做纯评测（不训练），直接传裁剪参数可能不生效，需要借助--model-overrides把训练时的参数覆盖为裁剪配置。语言模型评测示例：

fairseq-eval-lm /path/to/wikitext-103 \ --path /path/to/model/checkpoint.pt \ --model-overrides "{'decoder_layers_to_keep':'0,2,4,6,8,10,12,14'}"

该命令会用 override 更新模型参数，使裁剪后的模型（而非完整模型）被加载执行。

测试用例印证

仓库测试 test_binaries.py 中的 test_transformer_layerdrop 完整演示了"训练 → 裁剪 → 生成"链路：先用--encoder-layerdrop 0.01 --decoder-layerdrop 0.01训练一个 3 层小 Transformer，再通过--model-overrides "{'encoder_layers_to_keep':'0,2','decoder_layers_to_keep':'1'}"将其裁剪为编码器保留 2 层、解码器保留 1 层后执行generate_main。这既验证了训练参数与裁剪参数的正确性，也说明fairseq 允许"裁剪后继续训练"——直接加上保留层参数即可从裁剪模型恢复训练。

评测预训练模型

机器翻译（WMT16 en-de）

fairseq-generate /path/to/bped/wmt/data --path nmt_checkpoint.pt \ --beam 8 --lenpen 0.4 \ --batch-size 64 \ --remove-bpe \ --gen-subset test > wmt16_gen.txt bash scripts/compound_split_bleu.sh wmt16_gen.txt # prints BLEU4 = 30.17

RoBERTa + LayerDrop（MNLI / QNLI）

MNLI 微调模型评测（使用 RobertaModel 加载）：

from fairseq.models.roberta import RobertaModel roberta_layerdrop = RobertaModel.from_pretrained( '/path/to/MNLI/model', checkpoint_file='mnli_checkpoint.pt', data_name_or_path='/path/to/MNLI/data/MNLI-bin' ) label_map = {0: 'contradiction', 2: 'neutral', 1: 'entailment'} ncorrect, nsamples = 0, 0 roberta_layerdrop.cuda() roberta_layerdrop.eval() with open('/path/to/MNLI/data/dev_matched.tsv') as fin: fin.readline() for index, line in enumerate(fin): tokens = line.strip().split('\t') sent1, sent2, target = tokens[8], tokens[9], tokens[-1] tokens = roberta_layerdrop.encode(sent1, sent2) prediction = roberta_layerdrop.predict('sentence_classification_head', tokens).argmax().item() prediction_label = label_map[prediction] ncorrect += int(prediction_label == target) nsamples += 1 print('| Accuracy: ', float(ncorrect)/float(nsamples)) # prints | Accuracy: 0.9026999490575649

QNLI 微调模型评测：

roberta = RobertaModel.from_pretrained( '/path/to/QNLI/model', checkpoint_file='qnli_checkpoint.pt', data_name_or_path='/path/to/QNLI/data/QNLI-bin' ) label_fn = lambda label: roberta.task.label_dictionary.string( [label + roberta.task.target_dictionary.nspecial] ) ncorrect, nsamples = 0, 0 roberta.cuda() roberta.eval() with open('/path/to/QNLI/data/dev.tsv') as fin: fin.readline() for index, line in enumerate(fin): tokens = line.strip().split('\t') sent1, sent2, target = tokens[1], tokens[2], tokens[3] tokens = roberta.encode(sent1, sent2) prediction = roberta.predict('sentence_classification_head', tokens).argmax().item() prediction_label = label_fn(prediction) ncorrect += int(prediction_label == target) nsamples += 1 print('| Accuracy: ', float(ncorrect)/float(nsamples)) # prints | Accuracy: 0.9480139117700896

注意：MNLI 数据集的dev_matched.tsv中目标标签位于最后一列（tokens[-1]），而 QNLI 的dev.tsv目标标签位于第 4 列（tokens[3]），两份数据的列布局不同，评测脚本需按数据集分别处理。

复现论文结果

要复现论文中的实验结果，官方给出的路线是沿用 fairseq 中对应的实验配置：

1. WMT16 en-de 机器翻译：参考 fairseq 的examples/scaling_nmt设置（在本仓库中对应 decoding/IAD/fairseq/examples/scaling_nmt）。
2. RoBERTa 预训练：参考 fairseq 的examples/roberta设置（在本仓库中对应 decoding/IAD/fairseq/examples/roberta）。
3. Wikitext-103 语言模型：参考 fairseq 的examples/language_model设置，并结合 transformer_lm_wiki103.yaml 这类配置文件调整超参。

调参 Tips

1. 追求大模型的最佳性能：LayerDrop 应取较小值（0.1 或 0.2）。过大的 LayerDrop 意味着过强的正则化，模型可能无法达到最优性能。由于 LayerDrop 本身增加了正则效果，可以适度降低标准 dropout（例如减少 0.1）来平衡。
2. 追求裁剪出更小的模型：若打算激进裁剪（比如去掉一半以上网络），LayerDrop 可取较大值（如 0.5）；若只需裁剪少量层，0.2 这类较小值即可。论文实验主要采用 0.1 和 0.2 的低值。
3. 裁剪时均匀分布保留层：推理裁剪时，最好让保留的层在网络中均匀散布。例如要移除 50% 的网络，隔层保留（keep every other layer）是良好实践。这与 prune_state_dict 的映射逻辑直接相关——层号映射按保留列表顺序重排，保留层分布越均匀，裁剪后的深度结构越平滑。

常见问题（FAQ）

Q1：权重共享 + LayerDrop 的实验是怎么做的？论文附录（OpenReview 链接见原文档）在 Wikitext-103 语言模型上做了 LayerDrop 与权重共享结合的实验：每 2 层共享一组权重，即相邻层成对共享。例如网络共 6 层时，层 1 与 2 共享、层 3 与 4 共享、层 5 与 6 共享。

Q2：LayerDrop 在我的场景里没有效果？训练阶段 LayerDrop 起正则化作用，它最适用于已经过拟合的网络。若网络处于欠拟合状态，LayerDrop 可能带来过多正则化。建议使用较小值（0.1 或 0.2），并同步减少标准 dropout（如减少 0.1）。

Q3：能否不用 LayerDrop 训练、只在微调时加 LayerDrop（比如 BERT）？根据官方实验，这种做法效果不佳。RoBERTa 这类模型在预训练阶段已训练很久，仅在下游任务（如 MNLI）上微调几个 epoch 的 LayerDrop，不足以获得成功裁剪所需的鲁棒性。换言之，LayerDrop 的鲁棒性主要来自预训练阶段的长期训练，而非下游微调。

小结

LayerDrop 提供了一条"大模型训练、按需深度部署"的实用路径：训练期以 0.1～0.2 的概率结构化丢弃整层，让模型对层缺失鲁棒；推理期通过--encoder-layers-to-keep/--decoder-layers-to-keep指定保留层，配合 prune_state_dict 完成 checkpoint 级裁剪与参数映射，即可在几乎不损失可用性的前提下显著缩小模型体积。无论是机器翻译、RoBERTa 式编码器还是 decoder-only 语言模型，fairseq 都提供了统一且经过测试验证的参数入口（test_transformer_layerdrop 可作为最小可运行参考），值得在需要"弹性深度部署"的场景中直接采用。

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