LayerDrop 结构化 Dropout 实战指南：基于 fairseq 的 Transformer 深度裁剪与按需部署

LayerDrop 结构化 Dropout 实战指南基于 fairseq 的 Transformer 深度裁剪与按需部署【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm本文基于本仓库中 kosmos-2/fairseq/examples/layerdrop/README.md 展开结合 fairseq 框架源码位于kosmos-2/fairseq/fairseq讲解 LayerDrop 的原理、训练、裁剪、评估与调参实践。读完本文你将掌握如何用--encoder-layerdrop/--decoder-layerdrop训练对深度裁剪鲁棒的 Transformer如何用--encoder-layers-to-keep/--decoder-layers-to-keep在推理时按需裁剪层数以及如何复现论文实验并解决常见问题。LayerDropReducing Transformer Depth on Demand with Structured DropoutFan et al., 2019论文 arXiv:1909.11556是一种结构化 Dropout正则化与剪枝技术在训练阶段按概率随机丢弃整层 Transformer Block使模型学会在任意子深度下都能正常工作在推理阶段则可以将网络按需裁剪到任意子集深度几乎不需要重新训练。它解决的核心问题是在资源受限场景移动端、低延迟服务下按需改变模型深度让一个模型、多个深度成为可能。1. 核心原理训练时随机丢层推理时按需裁剪LayerDrop 的思路非常直观可以用一句话概括训练时随机丢弃整层 Transformer Block推理时按需保留任意层子集。由于模型在训练期间见过各种深度组合它不会对完整深度产生过度依赖从而对层裁剪具备天然鲁棒性。从本仓库源码看训练时的随机丢层由 LayerDropModuleList 实现其核心逻辑如下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其关键行为每次迭代重新采样每次遍历LayerDropModuleList都会重新生成一组均匀分布的随机数dropout_probs因此每个 mini-batch甚至每次前向传播被丢弃的层集合都是不同的而非固定裁剪某几层训练时随机丢层在self.training模式下当dropout_probs[i] p时该层被跳过yield 不会产出该层即每层以概率p被丢弃评估时保留全部层非训练模式eval()下条件恒为真即推理时遍历所有层——这与推理时再按需裁剪的设计一致裁剪动作在推理阶段由用户显式指定而不是在 eval 时自动随机丢层。在 Transformer 编码器中当encoder_layerdrop 0.0时self.layers被实例化为LayerDropModuleList见 transformer_encoder.py解码器同样在decoder_layerdrop 0.0时使用LayerDropModuleList见 transformer_decoder.py。当 LayerDrop 概率为 0 时则退化为普通nn.ModuleList行为与标准 Transformer 完全一致。2. 训练一个带 LayerDrop 的模型2.1 最小启动命令要训练带 LayerDrop 的模型只需在训练命令中追加以下参数推荐值为 0.2这是论文实验中表现良好的取值--encoder-layerdrop 0.2 --decoder-layerdrop 0.2适用对象区分Decoder-only 语言模型如 GPT 类、Wikitext-103 LM只需--decoder-layerdropEncoder-only 模型如 RoBERTa、BERT 类只需--encoder-layerdropEncoder-Decoder 模型如 WMT16 en-de 翻译模型编码器与解码器的 LayerDrop 值可以分别设置。2.2 参数在源码中的定义与默认值从本仓库源码看layerdrop与layers_to_keep定义于 transformer_config.py 的EncDecBaseConfig中作为 Encoder/Decoder 共用的基类配置# 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} )layerdrop默认值为0.0即默认关闭取值范围为[0, 1)的浮点数表示每层在训练时被丢弃的概率layers_to_keep默认None不裁剪指定裁剪时保留的层索引列表。在模型侧的默认配置文件中例如 transformer_lm_wiki103.yaml可以看到decoder_layerdrop: 0 decoder_layers_to_keep: null同时wav2vec2_base.yaml 这类语音模型配置中也会用到encoder_layerdrop: 0.05说明 LayerDrop 在语音预训练模型中同样作为常规正则化手段被启用本仓库的 hubert.py 中也有encoder_layerdrop字段。2.3 训练中发生了什么在训练循环中LayerDropModuleList.__iter__每次被调用都会重新采样丢弃集合。这意味着不同的 step甚至同一 batch 内的不同前向调用可能执行不同深度的子网络模型在训练中见过从完整深度到很浅深度的大量子网络变体从而获得深度鲁棒性由于丢弃是结构化的整层 Block 被跳过其正则化强度远高于普通的 dropout 逐元素随机置零属于强正则化手段因此文档建议在启用 LayerDrop 时适当降低标准 dropout例如减少 0.1。3. 推理时按需裁剪Pruning3.1 裁剪命令对已用 LayerDrop 训练好的模型在推理或继续训练时通过以下参数指定要保留的层索引逗号分隔列表--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 语言模型完整参数量为num. model params: 246933504裁剪到 8 层保留偶数层 0,2,4,...,14后为num. model params: 146163712参数从约 2.47 亿降到约 1.46 亿接近减半。3.2 裁剪的底层实现checkpoint 加载时的 state_dict 重映射裁剪并非简单地在推理时跳过某些层而是在加载 checkpoint 时真正重写模型结构。其实现位于 checkpoint_utils.py核心流程从model_cfg读取encoder_layers_to_keep与decoder_layers_to_keep若两者都为空则直接返回原始state_dict不做裁剪否则打印日志Pruning model to specified layer configuration - this works best if the model was trained with LayerDrop对每个要裁剪的模块构建mapping_dict如保留层0,2,4,...重映射为0,1,2,...与正则表达式匹配encoder.layers.数字/decoder.layers.数字遍历state_dict中所有键将保留层的参数重命名为紧凑的新编号例如原第 14 层 → 新第 7 层而 embedding 等非层参数原样保留裁剪完成后将*_layers_to_keep字段重置为None避免后续重复裁剪。该函数被从 checkpoint 加载模型的流程自动调用因此只要在加载时带上--encoder-layers-to-keep/--decoder-layers-to-keep参数就会得到物理上更小的模型而不是一个仍然带着全部参数、只是推理时跳过某些层的假裁剪。3.3 用裁剪后的模型继续训练如果你希望用裁剪后的模型继续训练只需在训练命令中同样加上--encoder-layers-to-keep/--decoder-layers-to-keep参数即可——裁剪在 checkpoint 加载阶段完成后续训练基于裁剪后的小模型进行。3.4 纯评估场景的 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 覆盖训练参数更新模型配置使被裁剪的模型而非完整模型被执行。类似的训练参数中decoder_layers也会在裁剪配置存在时被同步调整——例如 transformer_lm.py 中有if args.decoder_layers_to_keep: args.decoder_layers len(args.decoder_layers_to_keep.split(,))确保层数与裁剪后的实际结构一致。4. 预训练模型与效果评估4.1 预训练模型清单原文档提供了一批已发布的 LayerDrop 预训练模型对应 examples/layerdrop/README.md模型描述layerdrop_wmt_en_de_12_6Transformer LayerDrop 0.2在 WMT16 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 上微调说明原文档中的模型权重托管在 fairseq 官方外部下载地址本仓库仅包含使用它们的文档与代码。若要获取权重请参照文档中的下载链接自行获取。4.2 机器翻译模型评估# Example for Machine Translation 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关键参数说明--beam 8beam search 束宽--lenpen 0.4长度惩罚系数length penalty用于平衡生成长度--remove-bpe生成后移除 BPE 标记还原为原始词--gen-subset test在测试集上生成scripts/compound_split_bleu.sh对复合词进行拆分后计算 BLEU-4示例输出 BLEU4 30.17。4.3 RoBERTa LayerDrop 的 MNLI/QNLI 评估RoBERTa 在 fairseq 中通过 models/roberta/model.py 提供 Python 接口。MNLI 评估示例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 评估示例注意其 label 通过label_dictionary反解字符串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评估脚本要点RobertaModel.from_pretrained三个参数分别为模型目录、checkpoint 文件名与数据目录encode(sent1, sent2)负责 tokenize 与拼接含分隔符predict(sentence_classification_head, tokens)走的是 RoBERTa 的分类头MNLI 的dev_matched.tsv中第 9、10 列为两个前提-假设句子最后一列为标签QNLI 的dev.tsv中第 2、3、4 列分别为 question、sentence、label。5. 复现论文实验论文中的三类实验分别对应本仓库 fairseq 的示例配置WMT16 en-de 机器翻译参照 examples/scaling_nmt/README.md 的数据与配置RoBERTa 预训练/微调参照 examples/roberta 目录下的 README 与脚本Wikitext-103 语言模型参照 examples/language_model 目录下的 README 与脚本。复现时的最小改动就是在上述配置的模型参数上叠加--encoder-layerdrop 0.2或按模型类型只加对应端的参数。6. 调参建议Tips原文档给出了三条经过实验验证的调参经验追求大模型最佳性能LayerDrop 取值宜小如 0.1 或 0.2。过大的 LayerDrop 会引入过多正则化可能导致模型无法达到最佳性能。由于 LayerDrop 本身是一种强正则化为达到最优效果可适当降低标准 dropout例如减少 0.1。追求裁剪后变小的部署LayerDrop 取值宜大如 0.5以便支持激进裁剪例如裁掉一半以上网络若只需裁掉少数几层取较小值如 0.2即可。论文实验以 0.1、0.2 的小值为参考基准。推理时裁剪的层分布裁剪时应尽量让保留的层均匀分布在整个网络中。例如要裁掉 50% 的网络保留隔一取一每两层取一层的效果较好——这可以避免保留层集中在网络某一端造成的信息瓶颈。原文档中 16→8 层裁剪正是采用保留偶数层0,2,4,...,14的方式。7. FAQ常见问题与经验总结原文档针对社区高频问题给出了明确答复Q1论文附录中的层共享实验是如何做的在 Wikitext-103 语言模型上将 LayerDrop 与 Weight Sharing 结合按每 2 层为一组共享权重即网络层 1~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 应在预训练阶段就启用让模型从训练初期就学会适应深度变化。8. 总结LayerDrop 是本仓库 fairseq 框架中一项低成本、高回报的结构化正则化技术训练期一个--encoder-layerdrop 0.2参数即可让模型获得深度鲁棒性推理期通过--encoder-layers-to-keep/--decoder-layers-to-keep或--model-overrides即可在加载 checkpoint 时物理裁剪模型checkpoint_utils.py 中的 state_dict 重映射是其核心实现一个模型、多种深度的按需部署。其实现关键路径为训练随机丢层modules/layer_drop.py 的LayerDropModuleList.__iter__编码器/解码器接入transformer_encoder.py 与 transformer_decoder.py配置定义transformer_config.pycheckpoint 裁剪checkpoint_utils.py。无论是需要大模型的最佳精度还是需要部署时深度减半的小模型LayerDrop 都提供了一条统一、可复现的技术路径值得在 Transformer 类模型的训练与部署流程中作为默认选项考虑。引用格式Fan, Angela, Edouard Grave, and Armand Joulin. Reducing Transformer Depth on Demand with Structured Dropout. arXiv preprint arXiv:1909.11556 (2019).【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考