fairseq LayerDrop 指南用结构化 Dropout 训练可裁剪 TransformerReducing Transformer Depth on Demand【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilmLayerDropStructured DropoutFan et al., 2019是一种在训练阶段随机丢弃整个 Transformer 层的正则化与模型裁剪方案以 0.10.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关键行为可以从源码中归纳出两点每次迭代都会重新随机每次遍历 LayerDropModuleList 时都会用均匀分布生成一组与层数等长的随机数概率小于p的层被跳过。因此同一个 batch 的前向与后向、甚至不同 batch 之间哪些层被丢弃都在变化。评估/推理时默认保留全部层当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 的参数定义 中参数类型默认值说明--encoder-layerdropfloat0编码器每层被丢弃的概率--decoder-layerdropfloat0解码器每层被丢弃的概率--encoder-layers-to-keepstrNone裁剪时保留的编码器层编号逗号分隔--decoder-layers-to-keepstrNone裁剪时保留的解码器层编号逗号分隔使用要点推荐值 0.2官方实验表明 0.2 在大多数场景下效果良好0.10.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模型描述layerdrop_wmt_en_de_12_6Transformer LayerDrop 0.2WMT16 en-de 训练12 层编码器 6 层解码器roberta_layerdrop.baseRoBERTa Base LayerDrop 0.2roberta_layerdrop.largeRoBERTa Large LayerDrop 0.2roberta_layerdrop.large.mnliroberta_layerdrop.large在 MNLI 上微调roberta_layerdrop.large.qnliroberta_layerdrop.large在 QNLI 上微调推理阶段裁剪模型按需保留指定层训练完成后通过--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 时被自动调用无需手动触发。其流程可以概括为读取model_cfg中的encoder_layers_to_keep/decoder_layers_to_keep若两者都为空则直接返回原 state_dict。对每个要裁剪的部分构造一次裁剪通道将保留层编号解析为有序列表并建立原层号 → 新层号的映射例如保留0,2,4,6时原层 0→新层 0、原层 2→新层 1、原层 4→新层 2……同时用正则^{layer}.*\.layers\.(\d)匹配对应编码器/解码器的层权重键名。遍历原 state_dict 的所有键没有\.layers\.数字\.结构的键如 embedding、layernorm 等支撑层原样保留匹配到的层权重则通过映射改写层号后放入新 state_dict。裁剪完成后将model_cfg中的*_layers_to_keep置为 None避免后续重复裁剪。值得注意的细节是在 Transformer.build_model 中一旦设置了--encoder-layers-to-keepargs.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-defairseq-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.17RoBERTa LayerDropMNLI / QNLIMNLI 微调模型评测使用 RobertaModel 加载from fairseq.models.roberta import RobertaModel roberta_layerdrop RobertaModel.from_pretrained( /path/to/MNLI/model, checkpoint_filemnli_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.9026999490575649QNLI 微调模型评测roberta RobertaModel.from_pretrained( /path/to/QNLI/model, checkpoint_fileqnli_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 中对应的实验配置WMT16 en-de 机器翻译参考 fairseq 的examples/scaling_nmt设置在本仓库中对应 decoding/IAD/fairseq/examples/scaling_nmt。RoBERTa 预训练参考 fairseq 的examples/roberta设置在本仓库中对应 decoding/IAD/fairseq/examples/roberta。Wikitext-103 语言模型参考 fairseq 的examples/language_model设置并结合 transformer_lm_wiki103.yaml 这类配置文件调整超参。调参 Tips追求大模型的最佳性能LayerDrop 应取较小值0.1 或 0.2。过大的 LayerDrop 意味着过强的正则化模型可能无法达到最优性能。由于 LayerDrop 本身增加了正则效果可以适度降低标准 dropout例如减少 0.1来平衡。追求裁剪出更小的模型若打算激进裁剪比如去掉一半以上网络LayerDrop 可取较大值如 0.5若只需裁剪少量层0.2 这类较小值即可。论文实验主要采用 0.1 和 0.2 的低值。裁剪时均匀分布保留层推理裁剪时最好让保留的层在网络中均匀散布。例如要移除 50% 的网络隔层保留keep every other layer是良好实践。这与 prune_state_dict 的映射逻辑直接相关——层号映射按保留列表顺序重排保留层分布越均匀裁剪后的深度结构越平滑。常见问题FAQQ1权重共享 LayerDrop 的实验是怎么做的论文附录OpenReview 链接见原文档在 Wikitext-103 语言模型上做了 LayerDrop 与权重共享结合的实验每 2 层共享一组权重即相邻层成对共享。例如网络共 6 层时层 1 与 2 共享、层 3 与 4 共享、层 5 与 6 共享。Q2LayerDrop 在我的场景里没有效果训练阶段 LayerDrop 起正则化作用它最适用于已经过拟合的网络。若网络处于欠拟合状态LayerDrop 可能带来过多正则化。建议使用较小值0.1 或 0.2并同步减少标准 dropout如减少 0.1。Q3能否不用 LayerDrop 训练、只在微调时加 LayerDrop比如 BERT根据官方实验这种做法效果不佳。RoBERTa 这类模型在预训练阶段已训练很久仅在下游任务如 MNLI上微调几个 epoch 的 LayerDrop不足以获得成功裁剪所需的鲁棒性。换言之LayerDrop 的鲁棒性主要来自预训练阶段的长期训练而非下游微调。小结LayerDrop 提供了一条大模型训练、按需深度部署的实用路径训练期以 0.10.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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考