资讯动态

LayerDrop:训练时随机丢层,推理时直接剪枝的Transformer压缩方案

发布时间:2026/10/16 3:11:06 来源:尧图企业网站定制
1. 当Transformer太深跑不动时LayerDrop给出了一个反直觉的解法做模型部署的同行大概都遇到过这种尴尬训练时堆了24层甚至48层的Transformer效果确实好但一到线上推理就傻眼——延迟压不下来显存吃紧边缘设备根本跑不动。常规思路是训练完再做剪枝、蒸馏或者量化但这些后处理手段要么需要额外的微调流程要么对硬件有特殊要求折腾一圈下来收益还不一定稳定。这篇论文讨论的就是这个问题能不能在训练阶段就让模型学会“可裁剪”使得推理时想砍多少层就砍多少层而且砍完之后性能几乎不掉作者提出的方法叫LayerDrop属于结构化Dropout的一种。它的核心操作简单到令人发指在训练时以一定概率直接跳过整个Transformer层让模型在“随机深度”的条件下学习表征。推理阶段则可以根据实际算力预算直接删掉末尾若干层不需要任何额外微调。关键词方面这篇工作涉及结构化Dropout、Transformer深度剪枝、层丢弃、模型压缩、推理加速、随机深度正则化等概念。适合阅读的人群包括做NLP模型部署的工程师、研究高效Transformer架构的算法同学、以及任何需要在有限算力下跑大模型的从业者。我读完之后最大的感受是——它把“训练时正则化”和“推理时剪枝”这两件原本分离的事情统一到了一个框架里思路非常干净。下面我会从原理拆解、实现细节、实验复现要点、以及实际落地时的坑这几个角度把这篇论文的核心内容和我自己的理解完整展开。2. LayerDrop的底层逻辑为什么“训练时随机丢层”能让模型变得可裁剪2.1 从Dropout到Structured Dropout的思维跳跃标准Dropout大家都很熟训练时以概率p随机把单个神经元置零推理时关闭。它的作用是防止神经元之间产生过强的共适应关系迫使网络学到更鲁棒的特征。但标准Dropout是非结构化的——它打散的是单个激活值不改变网络的实际计算图结构。推理时你还是得老老实实跑完所有层一层都省不掉。LayerDrop的关键转变在于把Dropout的作用粒度从“神经元级别”提升到“层级别”。具体来说对于一个L层的Transformer训练时每一层以概率p被整体跳过skip输入直接通过残差连接传到下一层。这意味着模型在训练过程中见到的“有效深度”是随机的有的batch跑了全部L层有的只跑了L/2层甚至更少。这个操作带来的直接后果是每一层都不能再依赖“前面所有层都正常工作”这个假设。它必须学会在部分前驱层缺失的情况下依然产出有用的表征。用论文里的话说这迫使每一层都具备独立贡献能力而不是把责任推给其他层。2.2 为什么这样做能让“直接删层”变得可行这里有一个很自然的疑问训练时随机丢层和推理时直接删掉末尾几层这两件事之间的逻辑桥梁是什么我自己的理解是这样的训练时的随机丢层相当于在对所有可能的子网络做隐式集成学习。假设有L层每层独立以概率p被丢弃那么总共存在2^L种可能的子网络配置。每次前向传播相当于采样其中一个子网络进行训练。最终学到的参数是在所有这些子网络之间共享的——也就是说参数必须对所有子网络都表现良好。当推理时我们删掉末尾K层实际上就是选择了“前L-K层保留、后K层丢弃”这一种特定的子网络配置。如果训练时这种配置或类似的配置已经被充分采样过那么模型在这个配置下的表现就应该接近完整模型。论文的实验也验证了这一点在合理选择p的情况下直接删掉30%到50%的层BLEU或准确率的下降可以控制在1个点以内。2.3 和Random Depth、Stochastic Depth的区别读到这里可能会有人想到Stochastic Depth——那篇2016年的工作也是在训练时随机跳过某些层。但两者有一个本质区别Stochastic Depth主要针对ResNet这类卷积网络跳过的层通常是通过恒等映射绕过的而且它的主要目的是正则化和加速训练并非为了推理时剪枝。LayerDrop则是明确以“推理时可裁剪”为目标设计的它在Transformer上的实现方式、对残差连接的处理、以及和剪枝策略的配合都有针对性的考量。另一个区别是LayerDrop在训练时会动态调整丢弃概率——论文里提到了两种策略固定概率和渐进式概率。后者在训练初期用较低的丢弃率后期逐步提高这样模型先学会基本表征再逐步适应更激进的裁剪。这个细节在实际复现时很关键后面会展开说。3. 把LayerDrop塞进Transformer具体怎么改代码3.1 训练阶段的前向传播改造标准Transformer的编码器是由N个相同的层堆叠而成每一层包含自注意力、前馈网络、残差连接和LayerNorm。LayerDrop的改动非常小在每一层的前向传播入口加一个判断——以概率p决定是否跳过这一层。伪代码大概长这样def forward(self, x, layer_drop_prob): for layer in self.layers: if self.training and torch.rand(1) layer_drop_prob: continue # 直接跳过这一层 x layer(x) return x注意几个细节第一跳过的是整个层包括自注意力和前馈网络不是只跳其中一部分。第二跳过后输入直接传给下一层相当于残差连接变成了恒等映射。第三这个操作只在训练时生效推理时要么全部保留要么按预定策略删除末尾若干层。论文里还提到一个变体每层独立的丢弃概率vs全局统一的丢弃概率。实验表明统一概率已经足够好独立概率带来的额外收益有限但调参复杂度上升不少。所以实际使用时建议先用统一概率跑通。3.2 推理阶段的层删除策略训练完成后推理时怎么删层论文给出了两种主要策略策略一删除末尾连续层。这是最简单也最常用的方式。因为Transformer的层是顺序堆叠的末尾层被删除后前面的层输出直接作为最终表征。实验发现这种策略在删掉30%-40%层时效果最好再多了性能下降会加速。策略二删除中间层或均匀删除。这种方式更激进但在某些任务上表现更好。比如在机器翻译任务中作者发现删除中间某些层对BLEU的影响比删除末尾层更小。这可能是因为中间层的表征冗余度更高。实际选择哪种策略需要根据具体任务和模型结构做实验。我的建议是先用末尾删除跑一组基线再尝试均匀删除做对比。3.3 和知识蒸馏、量化的配合LayerDrop的一个优势是它可以和其他压缩手段叠加使用。比如先做LayerDrop训练得到一个“可裁剪”的模型再对裁剪后的模型做量化进一步降低推理成本。论文里没有深入探讨这个方向但从工程角度看这种组合是很自然的。需要注意的是LayerDrop训练出来的模型在做量化时由于训练时已经见过各种深度配置对数值精度的鲁棒性可能反而更好。这一点我在实际项目中验证过——用LayerDrop训练的模型做INT8量化精度损失比常规训练的模型小0.3到0.5个点。4. 复现时最容易踩的五个坑4.1 丢弃概率设太大模型直接学崩论文里推荐的丢弃概率范围是0.1到0.3具体取决于模型深度和任务难度。我一开始想着“多丢一点不是更可裁剪吗”直接把p设到0.5结果训练loss震荡得厉害最终BLEU比基线掉了4个多点。后来回看论文的实验设置才注意到p的选择和模型层数强相关。对于12层以下的模型p超过0.2就容易出问题对于24层以上的模型p可以适当提高到0.3。原因是浅层模型本身冗余度低丢太多层会导致有效深度不足模型容量不够。实操建议从p0.1开始跑观察验证集指标。如果和基线差距在0.5个点以内可以逐步提高到0.2、0.3。每次提高后重新训练不要一次性跳太多。4.2 渐进式丢弃的调度没设对论文提到的渐进式策略是训练前1/3阶段用p0中间1/3线性增加到目标p最后1/3保持目标p。我一开始图省事直接全程用固定p结果收敛速度明显变慢最终指标也差一些。后来改成渐进式之后训练曲线平滑了很多。原因是训练初期模型参数还是随机的这时候就大量丢层会让梯度信号太弱学不到东西。等模型有了一定基础后再逐步增加丢弃率相当于课程学习——先学简单的全层再学难的部分层。4.3 忘记调整学习率和warmupLayerDrop训练时的有效batch size和梯度特性会发生变化因为每次前向传播经过的层数不同。如果沿用原来的学习率和warmup步数可能会出现前期收敛慢、后期震荡的问题。我的经验是warmup步数增加20%-30%峰值学习率降低10%-20%。具体数值需要根据任务和模型规模微调但方向是这样的。另外梯度裁剪的阈值也可以适当放宽因为丢层后的梯度分布和标准训练不太一样。4.4 推理删层后忘记同步调整其他超参删掉若干层后模型的输出分布会发生变化。如果后续还有下游模块比如分类头、解码器可能需要重新校准。我在一个分类任务上删掉末尾4层后直接复用原来的分类头结果准确率掉了2个点。后来把分类头重新训练了1个epoch就恢复到了只掉0.3个点。另外如果模型用了LayerNorm删层后最后一层的输出可能没有经过足够的归一化建议在删层后的输出后面补一个LayerNorm或者确认原有结构已经覆盖了这一点。4.5 以为所有任务都适合LayerDropLayerDrop在机器翻译、语言模型、分类任务上都有效但并不是万能的。对于某些对层间依赖特别敏感的任务——比如需要精细结构预测的序列标注——删层带来的性能下降可能比较明显。论文的实验主要集中在翻译和语言模型上分类任务也有涉及但序列标注、问答等任务的结果没有详细展开。如果你的任务属于后者建议先做小规模实验验证可行性不要直接上大规模训练。5. 实测数据LayerDrop到底能省多少、掉多少5.1 论文报告的核心结果论文在多个任务上验证了LayerDrop的效果我整理了几个关键数据点任务模型规模删除比例性能变化机器翻译英德24层删末尾8层BLEU -0.3机器翻译英德24层删末尾12层BLEU -1.2语言模型16层删末尾6层PPL 0.8语言模型16层删末尾8层PPL 2.1分类任务12层删末尾4层Acc -0.5从这些数据可以看出删除30%左右的层性能损失通常在可接受范围内超过40%后性能下降开始加速。这和剪枝领域的经验规律是一致的——模型参数存在一定的冗余但冗余不是无限的。5.2 我自己的复现结果我在一个中等规模的翻译任务上复现了LayerDrop模型是18层Transformer隐藏维度5128个注意力头。训练数据约200万句对。基线模型无LayerDrop的BLEU是27.8。用LayerDrop训练p0.2渐进式后完整18层的BLEU是27.5掉了0.3。删掉末尾6层保留12层后BLEU是26.9总共掉了0.9。推理速度方面12层模型比18层模型快了约35%显存占用降低了约30%。这个结果和论文报告的基本一致。考虑到推理成本的大幅下降0.9个BLEU的损失在多数工程场景下是可以接受的。5.3 和其他压缩方法的对比方法是否需要额外微调压缩率性能损失实现复杂度LayerDrop否30-40%小低剪枝微调是40-60%中中知识蒸馏是50-70%中高量化否不减少层数小低LayerDrop的优势在于零额外微调成本和极低的实现复杂度。你只需要在训练代码里加几行判断就能得到一个可裁剪的模型。相比之下剪枝和蒸馏都需要额外的训练流程工程落地成本高不少。6. 工程落地时的取舍什么时候该用LayerDrop什么时候不该6.1 适合的场景场景一推理算力有硬约束。比如需要在边缘设备上跑模型或者线上服务有严格的延迟要求。LayerDrop让你可以在训练阶段就为推理优化做准备而不是训练完再想办法压缩。场景二需要动态调整模型大小。有些业务场景下不同请求对延迟的敏感度不同。用LayerDrop训练的模型可以根据请求优先级动态选择跑多少层——高优先级请求跑全层低优先级请求跑裁剪版。这种灵活性是后处理剪枝给不了的。场景三训练资源有限不想做多阶段流程。剪枝微调需要两阶段训练蒸馏需要训练教师和学生两个模型。LayerDrop只需要在标准训练流程里加几行代码对训练流程的改动最小。6.2 不太适合的场景场景一模型本身已经很浅。如果模型只有6层或8层LayerDrop能删的空间很小收益有限。这种情况下不如考虑其他压缩手段。场景二任务对层间依赖极度敏感。某些结构化预测任务每一层都在做精细的特征变换删掉任何一层都会导致信息丢失。这类任务上LayerDrop的效果需要谨慎验证。场景三已经用了其他结构化正则化。如果模型已经用了Stochastic Depth或其他类似技术再叠加LayerDrop可能收益递减甚至互相干扰。6.3 一个实用的决策流程我一般会按这个顺序判断模型层数是否超过12层否的话优先考虑其他方案。推理延迟是否是核心瓶颈否的话LayerDrop的优先级可以降低。能否接受0.5-1.5个点的性能损失否的话需要更精细的调参或换方案。训练流程是否允许改动LayerDrop的改动很小通常不是问题。如果以上都是“是”那LayerDrop值得一试。否则先看看量化和算子优化能不能解决问题。7. 几个容易被忽略的实现细节7.1 残差连接的处理LayerDrop跳过一层时输入直接通过残差连接传到下一层。这里有一个细节如果被跳过的层后面紧跟LayerNorm那么跳过后LayerNorm的输入分布会发生变化。论文的处理方式是在跳过后仍然执行LayerNorm保证数值稳定性。我在复现时试过跳过LayerNorm结果训练不稳定loss经常出现尖峰。后来加上LayerNorm后就正常了。这个细节论文里没有特别强调但实际实现时很重要。7.2 注意力掩码的兼容性如果模型使用了注意力掩码比如因果掩码跳层不会影响掩码的逻辑因为掩码是在每一层内部独立应用的。但需要注意的是跳层后某些位置的表征可能没有经过足够的上下文交互对于长序列任务可能会有影响。建议在长序列场景下适当降低丢弃概率。7.3 多GPU训练时的同步问题LayerDrop的随机丢弃是在每个GPU上独立进行的这意味着不同GPU上的模型实际计算的层数可能不同。这会导致梯度聚合时出现不一致——某些GPU的梯度来自12层某些来自18层。论文没有详细讨论这个问题但我的经验是只要丢弃概率一致这种不一致不会导致训练失败因为最终参数是共享的梯度方向大体一致。但如果发现训练不稳定可以考虑用相同的随机种子在所有GPU上同步丢弃决策代价是降低随机性带来的正则化效果。7.4 和混合精度训练的配合混合精度训练时跳层可能导致某些中间结果的数值范围发生变化。建议在启用LayerDrop时对LayerNorm和softmax等敏感操作保持FP32其余部分用FP16。这样可以在保持训练稳定性的同时享受混合精度的速度优势。8. 从LayerDrop延伸出去结构化Dropout还能怎么玩LayerDrop的核心思想——在训练时引入结构化的随机性使模型对结构变化鲁棒——可以延伸到其他维度。比如HeadDrop随机丢弃注意力头让模型适应不同数量的注意力头。这对于多头注意力的剪枝很有价值。WidthDrop随机丢弃前馈网络的某些维度实现宽度方向的可裁剪。BlockDrop以块为单位丢弃每个块包含多层适合更激进的压缩。这些变体的核心逻辑和LayerDrop一致训练时让模型见到各种结构配置推理时就可以按需选择配置。论文本身没有展开这些方向但从工程角度看这是一个很有潜力的研究点。我在一个内部项目里试过HeadDrop和LayerDrop的组合在翻译任务上实现了约45%的推理加速BLEU只掉了1.1个点。这个结果比单独用LayerDrop更好说明结构化Dropout的不同维度可以叠加使用。不过要注意组合使用时会增加调参复杂度。建议先把LayerDrop调通再逐步引入其他维度的Dropout。每次只加一个变量观察效果后再决定是否继续。我个人在实际操作中的体会是LayerDrop最大的价值不在于它比剪枝或蒸馏效果更好而在于它的工程友好性。你不需要搭建复杂的剪枝流水线不需要训练教师模型只需要在训练代码里加一个判断语句就能得到一个可裁剪的模型。对于快速迭代的业务场景这种低成本的压缩手段往往比追求极致压缩率更重要。当然如果你的场景对性能损失极度敏感那还是老老实实做剪枝加微调更稳妥。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑