Reducing Transformer Depth on Demand with Structured Dropout 解读

一、论文基本信息论文题目Reducing Transformer Depth on Demand with Structured Dropout作者Angela Fan、Edouard Grave、Armand Joulin发表ICLR 2020方法名称LayerDrop实现fairseq 中提供了 LayerDrop 的训练与预训练模型说明包括 WMT En-De、RoBERTa Base、RoBERTa Large 等模型。(GitHub)这篇论文的核心问题是能不能只训练一个深层 Transformer然后在推理阶段按需要直接裁剪成不同深度的浅层模型而不需要重新训练或微调论文提出的答案是LayerDrop训练时随机丢弃整个 Transformer layer让模型在训练过程中见过大量“缺层”的子网络推理时就可以直接删除部分层得到任意深度的子网络并且性能下降较小。论文摘要明确指出LayerDrop 是一种 structured dropout既有正则化作用又能让模型在推理阶段高效剪枝一个大模型可以抽取出不同深度的子网络而且不需要额外 fine-tuning。(arXiv)二、论文要解决的问题Transformer 模型通常层数多、参数大、推理成本高。机器翻译、语言建模、摘要、问答、BERT/RoBERTa 这类模型都可能包含上亿参数每次推理都要完整执行所有层。论文指出Transformer 每层都会带来大量参数访问和计算因此训练和推理阶段都存在明显的内存与延迟压力。传统压缩方法通常有两个问题。第一如果要得到不同大小的模型通常需要分别训练或微调。例如想要 12 层、8 层、6 层、3 层模型往往要分别训练或蒸馏多个模型。第二普通深层模型直接删层通常效果很差。因为模型训练时从未见过缺少某些层的情况推理时突然删层会破坏表示流。所以这篇论文要解决的问题是如何训练一个深层 Transformer使它天然具备“可删层”的鲁棒性换句话说它想让一个模型同时具备两种能力完整深度时性能强。删掉部分层后仍然能正常工作。三、核心思想LayerDrop 的核心思想很简单训练时随机丢弃整个 Transformer layer推理时根据目标深度直接删除部分 layers。普通 dropout 是随机丢神经元或激活值。LayerDrop 不丢单个神经元而是丢一个完整的结构单元也就是 Transformer 的整层。训练时每一层以一定概率被跳过。因为 Transformer 有 residual connection所以跳过一层后信息仍然可以通过残差路径继续向后传播。这样训练出来的模型会经常遇到各种“缺层子网络”。久而久之它就会学会在不同深度下工作。论文中把这种做法称为structured dropout因为它不是随机删除零散权重而是按结构化单位删除一组权重作者也明确说明这种 structured dropout 可以作用于 heads、FFN matrices 或 layers但最终主要聚焦在 layer 级别因为删层能直接带来深度减少和推理加速。这里最关键的地方是LayerDrop 不是训练后才想办法剪枝而是在训练阶段就让模型适应剪枝。四、为什么重点剪 Layer而不是剪 Head论文也讨论了不同结构化 dropout 单元包括 attention head、FFN 子层、Transformer layer 等。作者发现只 drop attention heads 效果不如 drop layer 或 drop sub-layer而且删 head 不一定带来明显运行时间优势因为 attention heads 通常是并行计算的。相比之下删除完整 Transformer layer 可以直接减少深度降低顺序计算链路因此更适合做推理加速。这点很重要。Head pruning 属于结构化剪枝但它主要减少注意力内部的并行分支LayerDrop 关注的是depth pruning也就是减少 Transformer 的层数。因此这篇论文真正做的是Transformer layer-level structured pruning。不是 attention head pruning也不是 FFN neuron pruning。五、训练阶段LayerDrop 怎么做训练阶段LayerDrop 对每个 Transformer layer 进行随机保留或随机丢弃。如果某一层被保留就正常执行 self-attention、FFN、residual、LayerNorm 等操作。如果某一层被丢弃就直接跳过该层让输入通过残差路径继续向后传递。这样每个 mini-batch 实际看到的网络深度都可能不同。例如一个 12 层 encoder如果 LayerDrop rate 为 0.2那么训练时平均有约 80% 的层被激活但具体哪些层被激活是随机的。论文实验中多数任务使用LayerDrop rate 0.2作者也建议如果目标是非常小的推理模型可以使用更大的 LayerDrop rate比如 0.5。这带来两个效果。第一正则化效果。模型不能过度依赖某一层因为该层训练时可能被随机丢掉。第二剪枝鲁棒性。模型在训练时已经见过很多浅层子网络因此推理时删层不会突然破坏模型。六、推理阶段如何选择保留哪些层训练好带 LayerDrop 的模型后推理阶段要从深层模型中选出一个浅层子网络。论文讨论了三种策略。6.1 Every Other最简单的方法是每隔一层保留一层也就是均匀地从整个网络中抽层。例如 12 层模型想剪到 6 层可以保留第 1、3、5、7、9、11 层或者类似的均匀间隔策略。论文发现这个非常简单的策略效果 surprisingly well甚至在很多任务上很难被复杂策略明显超过。它的直觉是均匀保留不同深度位置的层可以让浅层、中层、深层信息都保留一部分。6.2 Search on Valid第二种方法是在验证集上搜索不同层组合找出效果最好的子网络。这种方法更精细但计算成本较高而且可能过拟合验证集。论文也认为它相比 Every Other 只能带来很小提升。6.3 Data Driven Pruning第三种方法是学习每层的 drop rate然后根据 learned scores 选择保留哪些层。但作者同样发现这种更复杂的方法只带来 marginal gains。实际中简单的 Every Other 已经很强。所以这篇论文一个很实用的结论是训练时用 LayerDrop推理时直接均匀删层往往就能得到很好的浅层子网络。七、实验任务与数据集论文覆盖了很多 NLP 任务而不是只在一个数据集上验证。任务数据集模型 / 设置指标机器翻译WMT English-GermanTransformer BigWMT16 训练newstest2013 验证newstest2014 测试BLEU语言建模WikiText-10316 层 / 40 层 Transformer with Adaptive InputsPerplexity摘要CNN-DailyMailTransformer baseROUGE长文本问答ELI5Transformer Big multitaskROUGE语言理解MNLI、MRPC、QNLI、SST-2RoBERTa / RoBERTa LayerDropAccuracy论文在实验设置中明确列出了这些任务WMT14 En-De 机器翻译、WikiText-103 语言建模、CNN-DailyMail 摘要、ELI5 长答案问答以及多个自然语言理解任务。八、实验结果解读8.1 机器翻译LayerDrop 稳定深层 Transformer并提升 BLEU在 WMT En-De 上普通 Transformer Big 结果为 29.3 BLEUDynamicConv 为 29.7 BLEU而 12 层 encoder 6 层 decoder 的 Transformer LayerDrop 达到30.2 BLEU。论文还指出普通 Transformer 在 12 encoder layers 设置下会发散而 LayerDrop 可以稳定训练更深的模型。这个结果说明LayerDrop 不只是为了剪枝它本身也是一种有效正则化手段。更深模型原本难训练LayerDrop 通过随机跳层降低过拟合和训练不稳定性使模型能够从更大深度中获益。8.2 语言建模WikiText-103 上降低 PPL在 WikiText-103 上Adaptive Inputs baseline 的 test PPL 是 18.7加入 LayerDrop 后 16 层模型达到18.3进一步训练 40 层 Transformer LayerDropPPL 降到17.7。论文解释说WikiText-103 虽然很大但对于非常深的 Transformer 仍然容易出现过拟合和训练不稳定。LayerDrop 随机减少每次 forward 的活跃层数能减轻内存压力并增强稳定性。这个结果强调LayerDrop 可以帮助训练更深的 Transformer。8.3 摘要与长答案问答ROUGE 有提升在 CNN-DailyMail 摘要上Transformer baseline 的 ROUGE-1 / ROUGE-2 / ROUGE-L 是 40.1 / 17.6 / 36.8加入 LayerDrop 后6 encoder 8 decoder 设置达到41.1 / 18.1 / 37.5。在 ELI5 长文本问答上Transformer Multitask baseline 为 28.9 / 5.4 / 23.1加入 LayerDrop 后提升到29.4 / 5.5 / 23.4。这说明LayerDrop 的正则化效果不仅适用于翻译和语言建模也适用于生成式摘要和长答案生成。8.4 RoBERTa / BERT-like 模型剪到 6 层、3 层仍有竞争力论文在 RoBERTa 上也做了实验。RoBERTa Large LayerDrop 在 MNLI、MRPC、QNLI、SST-2 上与普通 RoBERTa Large 相比表现相当或略有提升。例如 Books Wiki 数据下RoBERTa Large 在 MNLI-m / MRPC / QNLI / SST-2 上为 89.0 / 90.2 / 93.9 / 95.3加入 LayerDrop 后为89.2 / 90.2 / 94.2 / 95.4。更重要的是论文把 RoBERTa LayerDrop 剪到 6 层和 3 层与从头训练的小模型、DistilBERT 等进行比较。6 层 RoBERTa LayerDrop 在 MNLI 和 SST-2 上分别达到 82.9 和 92.5使用更多预训练数据后6 层模型达到 84.1 和 93.2。3 层 RoBERTa LayerDrop more data 在 MNLI 和 SST-2 上达到 82.2 和 92.0。这个结果说明LayerDrop 训练出的深模型可以直接抽取出浅层 BERT-like 模型而且效果优于不少从头训练的小模型。8.5 剪枝实验LayerDrop 让删层更稳定论文专门比较了三类模型普通 Transformer 训练后直接删层。每个目标深度单独从头训练小模型。先训练一个 LayerDrop 深模型再按目标深度删层。结果显示LayerDrop 模型在不同剪枝比例下都明显优于“普通模型直接删层”并且通常优于“从头训练同等深度的小模型”。论文强调LayerDrop 模型只训练一次然后可以抽取不同深度子网络不需要为每个深度重新训练。这是这篇论文最核心的实验结论LayerDrop 让一个深层模型变成一个可按需裁剪深度的模型族。九、LayerDrop 与普通 Dropout、Stochastic Depth、模型蒸馏的区别9.1 与普通 Dropout 的区别普通 dropout 丢的是神经元或激活值主要用于正则化。LayerDrop 丢的是完整 Transformer layer除了正则化还能让模型在推理时直接删层。所以普通 dropout 主要改善训练。LayerDrop 同时服务于训练正则化和推理剪枝。9.2 与 Stochastic Depth 的区别Stochastic Depth 最早用于 ResNet训练时随机跳过 residual blocks主要是为了训练很深的 CNN。LayerDrop 的思想与其相似但目标不同。论文明确说之前的 Stochastic Depth 主要关注加速和稳定训练而这篇论文重点关注 Transformer 推理阶段的结构化剪枝。9.3 与知识蒸馏的区别DistilBERT 这类方法需要 teacher-student 训练过程。LayerDrop 不需要额外 teacher也不需要针对每个目标深度单独蒸馏。它直接从一个训练好的大模型中抽取浅层子模型。所以蒸馏是“大模型教小模型”。LayerDrop 是“大模型训练时就包含大量浅层子模型”。十、方法优点10.1 一个模型支持多个深度LayerDrop 最大优势是只训练一个深层模型就能在推理时抽取不同深度的子网络。这对部署很有用。不同设备、不同延迟预算下可以选择不同层数而不需要保存和训练多个模型。10.2 不需要额外 fine-tuning论文多次强调生成任务上的 pruned networks 不需要额外 fine-tuning。附录也显示对语言模型剪到 8 层后再 fine-tune 只带来很小提升。这说明LayerDrop 的训练过程已经让子网络具备较强鲁棒性。10.3 实现简单LayerDrop 的实现非常简单只需要在训练时对每一层加随机保留 / 跳过机制。论文也在摘要和结论中强调LayerDrop 像 dropout 一样容易实现。(arXiv)10.4 能稳定训练更深模型论文在 WMT、WikiText-103、RoBERTa 等任务中都展示了 LayerDrop 的正则化效果。它可以减少每次 forward 中活跃层数降低内存使用并提高深层模型训练稳定性。10.5 属于真正结构化剪枝LayerDrop 推理时删除的是完整 Transformer layer。这会直接减少顺序深度、参数访问和计算路径因此比只 mask 少量权重更有部署意义。十一、方法局限11.1 剪枝粒度较粗LayerDrop 删除的是整层 Transformer。这种粒度非常粗。如果某一层内部只有部分 attention heads 或 FFN neurons 冗余LayerDrop 无法精细处理。它适合降低深度但不适合做细粒度结构压缩。11.2 Every Other 简单有效但不一定最优论文发现 Every Other 策略很强但它本质上仍然是启发式方法。对于某些任务或模型最优保留层组合可能不是均匀间隔。虽然 Search on Valid 和 Data Driven Pruning 可以略微改进但收益有限且成本更高。11.3 大幅剪枝仍会损失性能LayerDrop 能让模型更鲁棒但不是说删多少层都无损。当剪掉 50%、75% 层时性能仍然会下降只是下降幅度比普通删层或从头训练更可控。11.4 需要在训练时提前使用 LayerDrop普通已经训练好的 Transformer如果训练时没有 LayerDrop直接删层效果通常很差。所以 LayerDrop 更像是一种training-time pruning-aware regularization而不是对任意已有模型都能直接生效的后处理剪枝方法。11.5 主要适合残差堆叠式 TransformerLayerDrop 依赖 Transformer 层之间的残差结构。如果模型结构不是标准残差堆叠直接跳过整层可能不稳定。对于现代 LLM、ViT、MoE 或多模态模型LayerDrop 思想可以迁移但具体跳哪些层、如何保持表示稳定需要重新设计。十二、整体评价这篇论文的核心价值在于它把“删层”从训练后被动剪枝变成训练时主动适配。普通深层 Transformer 直接删层通常会损坏性能因为模型从未学过缺层推理。LayerDrop 在训练阶段随机丢掉完整层相当于让模型提前适应各种深度子网络。这样推理时就可以根据速度、内存、延迟需求直接抽取浅层模型。它的贡献可以概括为三点第一提出 LayerDrop用 structured dropout 训练可删层 Transformer。第二证明 LayerDrop 既能正则化深层 Transformer又能稳定训练更深模型。第三证明一个大模型可以按需抽取不同深度的子网络不需要为每个深度重新训练或蒸馏。这篇论文和 head pruning 的思路不同。Head pruning 问的是哪些 attention heads 可以删LayerDrop 问的是一个 Transformer 能不能在训练时就学会缺少若干层后仍然工作所以它的重点不是评估某一层的重要性而是构造一种训练机制让整个模型天然适合做depth pruning。十三、一句话总结《Reducing Transformer Depth on Demand with Structured Dropout》提出 LayerDrop在训练时随机跳过完整 Transformer layers使模型对缺层结构具有鲁棒性推理时可以直接按需求删除部分层得到不同深度的子网络而不需要重新训练或额外微调。它的核心贡献是把 Transformer 的深度剪枝前移到训练阶段使一个深层模型能够按需裁剪成多个浅层模型并在机器翻译、语言建模、摘要、问答和 RoBERTa 下游任务中保持较强性能。