FEATURED · 精选文章

LayerDrop 结构化 Dropout 实战:在 unilm/edgelm(fairseq)中训练可按需剪枝的 Transformer 深度

发布时间 / 2026/9/13 6:12:34
来源 / 创域科博编辑部
栏目 / 资讯中心
LayerDrop 结构化 Dropout 实战:在 unilm/edgelm(fairseq)中训练可按需剪枝的 Transformer 深度 LayerDrop 结构化 Dropout 实战在 unilm/edgelmfairseq中训练可按需剪枝的 Transformer 深度【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm本文基于 edgelm/examples/layerdrop/README.md系统讲解如何在 unilm 仓库内置的 fairseqedgelm中训练 LayerDrop 模型并在推理阶段按需裁剪 Transformer 层数实现在同一模型上弹性调节计算量与深度。读完本文你将掌握 LayerDrop 的训练参数配置、推理时层剪枝方法、预训练模型评估流程以及背后的源码实现原理。LayerDrop 是什么结构化 Dropout 与按需减深LayerDropStructured Dropout出自 Fan et al., 2019 论文Reducing Transformer Depth on Demand with Structured Dropout是一种训练时正则化技术训练过程中Transformer 的每一层以概率p被随机丢弃模型在每次前向传播时只经过一个随机的子网络而在推理时可以任意保留下层集合例如只保留一半层从而按需减少网络深度。这一设计带来了两个直接收益正则化层级别的随机丢弃是一种很强的结构化正则能抑制过拟合效果优于或接近普通的 dropout 微调弹性部署经过 LayerDrop 训练的模型对删层具有鲁棒性推理时可以不加微调直接裁剪层数使同一个 checkpoint 服务于不同算力/时延预算的部署环境。该技术在 unilm 仓库中由 edgelmfairseq 分支完整实现核心代码与配置、测试均可直接查阅。源码级实现LayerDropModuleList 与编解码器集成1. 核心容器LayerDropModuleList实现位于 edgelm/fairseq/modules/layer_drop.py其核心逻辑非常简洁class LayerDropModuleList(nn.ModuleList): def __init__(self, p, modulesNone): super().__init__(modules) self.p p 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关键行为docstring 中亦有明确说明每次迭代重新采样每遍历一次容器就为每一层重新抽取均匀随机数因此每次前向传播经过的子层集合都不同训练/评估双模式训练self.trainingTrue时以概率p丢弃层评估self.trainingFalse时总是遍历全部层即默认行为与普通 Transformer 完全一致LayerDrop 只在训练期生效概率语义p是每层被丢弃的概率代码中dropout_probs[i] self.p才保留该层p越大丢弃越多。2. 编码器与解码器的集成编码器在 edgelm/fairseq/models/transformer/transformer_encoder.py 中读取cfg.encoder.layerdrop当layerdrop 0.0时用LayerDropModuleList(pself.encoder_layerdrop)承载全部编码器层否则退化为普通nn.ModuleList解码器在 edgelm/fairseq/models/transformer/transformer_decoder.py 中做同样的判断读取cfg.decoder.layerdrop。因此训练时的层丢弃发生在for layer in self.layers:遍历阶段编码器与解码器各自独立采样。3. 配置字段layerdrop与layers_to_keep作为标准配置字段定义在 edgelm/fairseq/models/transformer/transformer_config.py# args for Reducing Transformer Depth on Demand with Structured Dropout (Fan et al., 2019) layerdrop: float field(default0, metadata{help: LayerDrop probability}) layers_to_keep: Optional[List[int]] field( defaultNone, metadata{help: which layers to *keep* when pruning} )layerdropLayerDrop 丢弃概率默认0即关闭layers_to_keep剪枝时保留哪些层从 0 开始编号的层索引列表。对应的命令行参数为--encoder-layerdrop / --decoder-layerdrop与--encoder-layers-to-keep / --decoder-layers-to-keep编码器、解码器可分别设置不同取值。预训练模型原文档提供了已用 LayerDrop 0.2 训练好的预训练模型可通过 fairseq 官方模型库获取对应 checkpoint 压缩包模型说明layerdrop_wmt_en_de_12_6WMT16 en-de 上训练的 Transformer LayerDrop 0.212 层编码器、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 上微调评估预训练模型机器翻译BLEU 评估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其中scripts/compound_split_bleu.sh为文档所述评估脚本对生成的文本做 compound splitting 后计算 BLEU4。RoBERTa LayerDropMNLI / QNLI 分类评估MNLI 评估示例使用sentence_classification_head做三分类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训练带 LayerDrop 的模型在训练命令中追加以下两个标志即可开启 LayerDrop--encoder-layerdrop 0.2 --decoder-layerdrop 0.2实践要点来自原文档推荐默认值 0.2在论文实验中表现良好纯解码器语言模型如 GPT 类只需--decoder-layerdrop纯编码器模型如 RoBERTa只需--encoder-layerdrop编码器与解码器的 LayerDrop 概率可以分别设置不同的值。从配置层面看layerdrop默认值为0对应完全关闭因此显式传入非零值才会启用见 edgelm/fairseq/models/transformer/transformer_config.py。edgelm 内置的 Transformer LM 配置如 edgelm/fairseq/config/model/transformer_lm/transformer_lm_wiki103.yaml中decoder_layerdrop: 0、decoder_layers_to_keep: null即为默认关闭状态可直接在此基础上改值。推理时按需剪枝Pruning对已用 LayerDrop 训练好的模型在加载时追加以下标志并给出想保留的层索引列表逗号分隔、从 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同时模型参数量会明显下降。例如 16 层 Transformer 语言模型完整 16 层num. model params: 246933504剪枝到 8 层保留 0,2,4,6,8,10,12,14 共 8 层num. model params: 146163712剪枝的底层机制剪枝并非简单跳过层而是真正从 checkpoint 中剔除对应层的参数。其核心实现是 edgelm/fairseq/checkpoint_utils.py 中的prune_state_dict从模型配置中读取encoder_layers_to_keep/decoder_layers_to_keep对要保留的层索引排序并构造原层号 → 新层号的映射例如保留0,2,4,...时原第 2 层重映射为新的第 1 层用正则^{encoder|decoder}.*\.layers\.(\d)逐 key 扫描 state_dict只搬运保留层的权重并重写层编号embedding 等无层号的支撑参数原样保留剪枝完成后自动清空配置中的*_layers_to_keep字段。同时在 edgelm/fairseq/models/transformer/transformer_legacy.py 的build_model中模型层数会被同步改写为保留层的数量if args.encoder_layers_to_keep: args.encoder_layers len(args.encoder_layers_to_keep.split(,)) if args.decoder_layers_to_keep: args.decoder_layers len(args.decoder_layers_to_keep.split(,))即剪枝后构建出的模型结构本身就是变浅的而不是空转被丢弃的层。剪枝后继续训练 / 仅评估继续训练直接带上--encoder-layers-to-keep/--decoder-layers-to-keep标志即可训练脚本会以剪枝后的浅层模型为起点继续优化仅做评估不训练可能需要传递覆盖参数。以语言模型为例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}--model-overrides会覆盖训练时保存的参数、更新模型配置从而让评估加载剪枝后的浅层模型而非完整模型。复现论文结果WMT16 en-de 翻译遵循 edgelm/examples/scaling_nmt 中的设置RoBERTa 预训练遵循 edgelm/examples/roberta 中的设置Wikitext-103 语言模型遵循 edgelm/examples/language_model 中的设置。以上三个示例目录在本仓库 edgelm 分支中均有对应实现可直接对照使用。调参 Tips目标是高质量大模型LayerDrop 应取较小值如 0.1 或 0.2。过大的 LayerDrop 会带来过强的正则化反而达不到最优性能。由于 LayerDrop 本身具有正则效果可考虑将普通 dropout 略微降低例如减少 0.1以获得最佳效果。目标是剪枝变小若打算激进剪枝如去掉一半以上网络LayerDrop 应取较大值如 0.5若只剪少量层可取较小值如 0.2。原论文实验使用的主要是低值0.1、0.2。推理剪枝的层选择保留的层最好在网络中均匀分布。例如要裁掉 50% 的网络隔层保留每隔一层保留一层是好的选择——这正对应0,2,4,6,...这种等间隔索引列表的推荐用法。FAQ 与常见问题1. LayerDrop 与权重共享结合的效果如何原论文附录在 Wikitext-103 语言模型上做了 LayerDrop Weight Sharing 实验将层按 2 层一组共享权重即第 1、2 层共享第 3、4 层共享依此类推假设网络共 6 层则 1↔2、3↔4、5↔6 分别共享权重。2. LayerDrop 在自己的场景中没有帮助训练期 LayerDrop 主要是正则化手段在模型已经过拟合时最有效。如果模型处于欠拟合状态LayerDrop 可能带来过强正则。建议先用小值如 0.1 或 0.2并同步降低普通 dropout 的量例如减少 0.1。3. 能否先训练普通模型无 LayerDrop、再在微调时开启 LayerDrop如 BERT 场景原论文实验显示效果不佳。RoBERTa 这类模型在预训练阶段训练时间很长仅在下游任务如 MNLI上微调几个 epoch 的 LayerDrop不足以获得成功剪枝所需的鲁棒性。测试验证LayerDrop 的端到端用例edgelm 仓库自带的测试 edgelm/tests/test_binaries.pytest_transformer_layerdrop完整覆盖了训练 剪枝评估链路用 3 层编解码器的transformer_iwslt_de_en结构训练传入--encoder-layerdrop 0.01 --decoder-layerdrop 0.01正常执行generate_main(data_dir)完成一次完整模型解码再通过--model-overrides {encoder_layers_to_keep:0,2,decoder_layers_to_keep:1}执行解码验证剪枝后模型可正确加载与推理。这既是 LayerDrop 功能的回归测试也是最小可复现示例——你可以仿照该测试的参数组织方式快速验证自己的 LayerDrop 训练与剪枝配置是否正确。引用如果该技术对你有所帮助请引用原论文article{fan2019reducing, title{Reducing Transformer Depth on Demand with Structured Dropout}, author{Fan, Angela and Grave, Edouard and Joulin, Armand}, journal{arXiv preprint arXiv:1909.11556}, year{2019} }【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED — 相关阅读

相关资讯

LATEST — 最新资讯

最新发布

TODAY — 本日精选

新闻

WEEKLY — 本周精选

新闻

MONTHLY — 本月精选

新闻