一、论文基本信息论文题目DynaBERT: Dynamic BERT with Adaptive Width and Depth作者Lu Hou、Zhiqi Huang、Lifeng Shang、Xin Jiang、Xiao Chen、Qun Liu发表会议NeurIPS 2020 Spotlight方法简称DynaBERT官方代码论文中给出了 Huawei Noah’s Ark Lab 的开源实现。论文摘要明确说明DynaBERT 可以通过选择不同width和depth来灵活调整模型大小和推理延迟。(arXiv)这篇论文的核心目标是不要把 BERT 压缩成一个固定大小的小模型而是训练一个可以按资源预算切换宽度和深度的动态 BERT。也就是说DynaBERT 训练完成后可以直接抽取不同子网络例如宽一些、深一些的子网络精度更高速度更慢。窄一些、浅一些的子网络速度更快精度略低。这样一个模型就能适配多种设备和多种资源条件。二、论文要解决的问题BERT、RoBERTa 这类预训练语言模型效果强但计算和存储开销都很大。论文指出边缘设备部署 BERT 面临两个现实问题第一不同设备硬件性能差异很大第二同一设备在不同运行状态下可用资源也会变化。传统压缩方法通常会得到一个固定模型。例如蒸馏得到 6 层 BERT或者剪枝得到某个固定宽度的模型。这样做的问题是一个固定模型只能对应一个固定计算预算。如果设备资源充足固定小模型无法利用更多计算换取更高精度。如果设备资源紧张固定模型又可能仍然太大。所以 DynaBERT 要解决的问题是能不能训练一个 BERT让它在推理时根据资源预算选择不同宽度和不同深度这个问题比单纯剪 head、剪 layer 更进一步。它不是只做一个方向的压缩而是同时考虑宽度方向attention heads 和 FFN 中间层神经元数量。深度方向Transformer layer 数量。论文明确指出DynaBERT 通过选择不同宽度和深度的子网络来调整模型大小和延迟并且训练完成后每个子网络不需要额外 fine-tuning。三、核心思想DynaBERT 的核心思想可以概括为先把 BERT 训练成宽度自适应模型再进一步训练成宽度和深度都自适应的动态模型。它的训练过程分成两个阶段。第一阶段训练DynaBERTW。这个模型只支持adaptive width也就是不同宽度子网络。宽度变化主要通过减少 attention heads 和 FFN 中间层神经元实现。第二阶段训练完整DynaBERT。在 DynaBERTW 的基础上再加入adaptive depth也就是不同层数子网络。最终得到一个既能变窄、又能变浅的模型。论文图 1 和方法部分明确说明训练流程包括两个阶段先训练 width-adaptive DynaBERTW再训练 width-and-depth-adaptive DynaBERT并通过知识蒸馏把大模型知识传给小子网络。这里最重要的是DynaBERT 不是训练多个独立小模型而是训练一个共享权重的动态 BERT。同一个模型内部包含很多子网络部署时只需要选一个满足资源约束的子网络运行。四、DynaBERT 的“宽度”是什么在 CNN 中宽度通常指通道数。但 BERT 是 Transformer 结构不能简单把宽度理解为卷积通道。论文把 BERT 的宽度拆成两个部分第一Multi-Head Attention 中的 attention heads 数量。第二Feed-Forward Network 中间层神经元数量。在一个 Transformer layer 中MHA 的多个 heads 是并行计算的FFN 中间层的神经元也可以看作并行计算单元。因此DynaBERT 通过改变这两部分数量来控制宽度。论文明确说明Transformer layer 的宽度可以通过改变 MHA 中 attention heads 数量和 FFN intermediate layer 中神经元数量来调整。例如原始 BERT-base 每层有 12 个 headsFFN 中间层维度为 3072。如果宽度倍率是 0.5那么子网络可以只保留约一半 heads 和一半 FFN 中间神经元。所以 DynaBERT 的 width-adaptive 本质上是剪 attention heads。剪 FFN intermediate neurons。这属于结构化压缩因为它删除的是完整结构单元而不是零散权重。五、Network Rewiring为什么要重新排列 heads 和 neurons如果直接取前一半 heads、前一半 FFN neurons可能会出现问题前面的 heads 不一定比后面的 heads 重要。所以 DynaBERT 在训练宽度自适应模型前先做Network Rewiring。它的做法是先估计每个 attention head 和 FFN neuron 的重要性。然后把更重要的 heads 和 neurons 排到更靠前的位置。这样所有小宽度子网络都会优先使用最重要的结构单元。论文明确说Network Rewiring 的目的是让更重要的 attention heads 和 neurons 被更多子网络共享做法是在 fine-tuned BERT 中计算它们的重要性并按重要性从高到低重新排列。这一步非常关键。因为 DynaBERT 的子网络通常是取“前若干 heads / 前若干 neurons”。如果不 rewiring小子网络可能拿到的是随机顺序下的结构单元如果做了 rewiring小子网络优先拿到高重要性结构单元。可以这样理解Network Rewiring 是为了让小模型先继承大模型中最有价值的部分。六、重要性怎么计算DynaBERT 采用类似 Taylor pruning 的思想。对于一个 attention head论文估计的是如果移除这个 head训练损失会变化多少。如果移除后损失变化大说明这个 head 重要。如果移除后损失变化小说明这个 head 冗余。对于 FFN 中间层神经元也采用类似思路看与该神经元相关的权重被移除后对训练损失的影响有多大。论文说明它使用一阶 Taylor 展开来估计 attention heads 和 neurons 的重要性并且在开发集上计算这些重要性分数。所以 DynaBERT 的 rewiring 不是随机重排而是基于 loss sensitivity 的结构重要性排序。这点和前面 head pruning 论文有明显联系Attention head 的重要性越高它越应该被更多子网络共享。七、第一阶段训练宽度自适应 DynaBERTW完成 rewiring 后DynaBERT 开始训练DynaBERTW。DynaBERTW 只支持宽度变化不改变层数。它的训练方式是知识蒸馏。具体来说Teacherrewired 后的完整 BERT。Student不同宽度的 DynaBERTW 子网络。蒸馏内容包括三部分输出 logits。embedding 表示。每一层 hidden states。论文明确说明在训练 DynaBERTW 时会把 fixed teacher model 的 logits、embedding 和 hidden states 蒸馏给不同 width multiplier 的学生子网络。这一步的作用是让不同宽度的子网络都模仿完整 BERT 的行为。如果只用真实标签训练小宽度子网络很容易性能不足蒸馏能给小模型更丰富的监督信号。八、为什么需要 DynaBERTW 作为中间老师论文没有直接从完整 BERT 一步训练到 width-depth adaptive DynaBERT而是先训练 DynaBERTW。原因是完整 BERT 和很小的子网络之间能力差距太大。如果直接让小而浅的子网络模仿完整 BERT训练难度会很高。所以 DynaBERTW 扮演了一个teacher assistant。它先学会宽度自适应然后再作为第二阶段的老师指导同时变窄和变浅的 DynaBERT。论文明确指出DynaBERTW 可以缩小 student 和 teacher 之间的模型规模差距消融实验也显示使用 DynaBERTW 作为 teacher assistant 在 SST-2、CoLA、MRPC 上都带来更好平均表现。所以这个两阶段设计不是多余步骤而是为了降低训练难度。可以简单理解为先学会变窄再学会又变窄又变浅。九、第二阶段训练宽度和深度都自适应的 DynaBERT第二阶段开始训练完整 DynaBERT。此时模型不仅支持不同宽度还支持不同深度。深度变化通过删除部分 Transformer layers 实现。论文采用类似 LayerDrop 中的Every Other策略也就是均匀删层让保留下来的层在整个网络深度上分布更平衡。论文明确说明当 depth multiplier 小于 1 时使用 Every Other 策略均匀丢层并将学生保留层的 hidden states 与 teacher 对应层匹配。第二阶段仍然使用知识蒸馏。Teacher最大深度的 DynaBERTW。Student不同宽度、不同深度的 DynaBERT 子网络。蒸馏目标仍然包括logits。embedding。hidden states。这样训练完成后DynaBERT 就可以直接抽取很多子网络例如1.0 width 1.0 depth0.75 width 1.0 depth0.5 width 0.75 depth0.25 width 0.5 depth不同子网络共享同一套权重只是运行时选择不同结构规模。十、DynaBERT 和 LayerDrop、US-Net、OFA 的关系DynaBERT 和 LayerDrop 都可以调整深度但 LayerDrop 主要解决depth adaptive即删 layer。DynaBERT 不只删 layer还会同时改变 width即删 attention heads 和 FFN neurons。DynaBERT 和 US-Net 都支持宽度变化但 US-Net 主要用于 CNN / MobileNet 这类视觉模型DynaBERT 面向 BERT 的 Transformer 结构因此宽度定义完全不同。DynaBERT 和 OFA 都属于“一次训练多种子网络部署”的思想。但 OFA 主要用于 CNN 架构搜索支持 kernel size、depth、width、resolution 等DynaBERT 则专门针对 BERT把自适应维度设计成attention heads。FFN neurons。Transformer layers。因此DynaBERT 的价值在于把动态宽度和动态深度思想真正落到 BERT 结构上。十一、实验设置论文主要在GLUE benchmark和SQuAD v1.1上评估。GLUE 包括多个自然语言理解任务如 MNLI、QQP、QNLI、SST-2、CoLA、STS-B、MRPC、RTE 等。论文在 GLUE 开发集上报告结果STS-B 使用 Spearman correlationCoLA 使用 Matthews correlation其他任务使用 accuracy在 QQP 和 MRPC 的 test set 上报告 F1。SQuAD v1.1 是抽取式问答任务包含约 100k 个问答样本评价指标是 EM 和 F1。效率指标方面论文不只看参数量还比较#parameters。FLOPs。Nvidia K40 GPU latency。Kirin 810 A76 ARM CPU latency。论文明确说明实验在不同效率约束下评估 DynaBERT 和 DynaRoBERTa包括参数量、FLOPs、GPU 延迟和 ARM CPU 延迟。这点很重要因为 DynaBERT 不是只追求理论压缩率而是考虑真实推理延迟。十二、实验结果解读12.1 GLUE很多任务可以在变窄或变浅后保持精度论文在 GLUE 开发集上展示了不同 width multiplier 和 depth multiplier 的结果。一个很明显的现象是最大尺寸子网络不一定总是最好某些任务在较小宽度或较小深度下也能达到相近甚至更好的结果。论文对此解释为原始 BERT / RoBERTa 中存在冗余。例如表 1 中BERTBASE 在 CoLA 上是 58.1而 DynaBERT 某些配置可以达到 60.8BERTBASE 在 SST-2 上是 92.9DynaBERT 多个配置都在 93 左右。这个结果说明BERT 的宽度和深度都存在冗余不是所有 heads、FFN neurons 和 layers 都必需。尤其在某些分类任务上适当缩小模型不一定会导致精度下降甚至可能带来一定正则化效果。12.2 宽度压缩比深度压缩更稳论文有一个很重要的观察在相同压缩倍率下压缩宽度通常比压缩深度更稳。论文在 GLUE 结果分析中明确指出使用某个 width multiplier 通常比使用相同的 depth multiplier 有更高准确率这说明相比深度方向宽度方向对压缩更鲁棒。这很好理解。减少宽度是减少 attention heads 和 FFN neurons但仍然保留完整层级处理流程。减少深度是直接删 Transformer layers会缩短表示变换链路对语义建模影响更大。所以 DynaBERT 的结论是BERT 更容易变窄变浅相对更敏感。这对实际剪枝很有启发如果要先压缩优先考虑 heads 和 FFN neurons可能比直接删层更安全。12.3 SQuAD小子网络仍然有竞争力在 SQuAD v1.1 上BERTBASE 的结果是81.5 EM / 88.7 F1。DynaBERT 最大子网络达到82.6 / 89.7。即使在较小配置下例如 0.5 width 1.0 depth也能达到81.9 / 89.20.5 width 0.75 depth 也有81.7 / 89.0。这说明在问答任务中DynaBERT 的较小子网络仍然可以保持接近 BERTBASE 的效果。而当宽度和深度都压得很小例如 0.25 width 0.5 depth性能下降到76.6 / 85.0说明极限压缩下仍然会损失较明显精度。这个结果比较合理DynaBERT 提供的是一条精度—效率曲线而不是保证所有小模型都无损。12.4 与 DistilBERT、TinyBERT、LayerDrop 对比论文将 DynaBERT 与 DistilBERT、TinyBERT、LayerDrop 等方法比较。需要注意的是DistilBERT 和 TinyBERT 的每个大小通常对应一个独立模型而 DynaBERT 的不同点来自同一个模型中的不同子网络。论文也特别指出TinyBERT 和 DistilBERT 的不同结果来自不同模型而 DynaBERT / DynaRoBERTa 的不同结果来自一个模型内部的不同子网络。这点是 DynaBERT 的核心优势不是为每个资源预算训练一个模型而是一个模型覆盖多个预算。论文在 SQuAD 上也指出在相同参数量和 FLOPs 下DynaBERT 子网络明显优于 TinyBERT 和 DistilBERT。十三、消融实验说明了什么论文做了比较完整的消融。13.1 Network Rewiring 很重要在训练 DynaBERTW 时vanilla DynaBERTW 平均准确率为 79.0加入 Network Rewiring 后提升到 81.4再加入蒸馏和数据增强后提升到 82.9。这说明把重要 heads 和 neurons 放在前面让它们被更多子网络共享是有效的。这一步对小宽度子网络尤其重要。13.2 知识蒸馏和数据增强很重要同样在 DynaBERTW 消融中加入 distillation 和 data augmentation 后平均性能进一步提升 1.5 个点左右。这说明小子网络需要从大模型中学习更丰富的中间表示而不只是学习 hard label。数据增强也能给蒸馏提供更多任务相关样本使子网络学得更稳。13.3 两阶段训练比直接训练更好论文还比较了是否使用 DynaBERTW 作为 teacher assistant。结果显示使用 DynaBERTW 的 DynaBERT 在 SST-2、CoLA、MRPC 上平均都更好。这说明先训练宽度自适应模型再训练宽度和深度都自适应的模型比一步到位更稳定。十四、方法优点14.1 一个模型支持多种资源预算DynaBERT 最大优点是只训练一个模型就能得到很多不同大小、不同延迟的子网络。这对部署非常有用。不同设备可以选择不同子网络同一设备资源状态变化时也可以切换到更小或更大的子网络。14.2 同时支持宽度和深度很多方法只支持删层或者只剪 heads。DynaBERT 同时支持减少 attention heads。减少 FFN neurons。减少 Transformer layers。因此它的结构选择空间比单一深度压缩更丰富。14.3 Network Rewiring 让小子网络更强通过重要性排序DynaBERT 让高重要性 heads 和 neurons 被所有小子网络优先使用。这比随机截取宽度更合理。14.4 蒸馏目标比较完整DynaBERT 不只蒸馏最终 logits还蒸馏 embedding 和 hidden states。这让小子网络不仅学最终预测也学习中间表示。14.5 考虑真实延迟论文同时报告参数量、FLOPs、GPU latency 和 ARM CPU latency。这使它比只看参数量或 FLOPs 的压缩方法更接近实际部署需求。十五、方法局限15.1 训练流程比较复杂DynaBERT 不是简单剪枝。它需要fine-tuned teacher。importance estimation。network rewiring。DynaBERTW 训练。DynaBERT 训练。知识蒸馏。数据增强。可能还需要 final fine-tuning。这使得复现和迁移到新任务的成本较高。15.2 主要是任务级压缩不是预训练阶段通用压缩论文中强调如果没有特别说明文中的 BERT 指的是task-specific BERT也就是已经针对下游任务 fine-tuned 的模型。这意味着 DynaBERT 通常需要针对每个下游任务单独训练动态模型。如果要得到一个通用的预训练动态 BERT再适配所有任务原论文并没有完全解决。15.3 子网络选择不是输入自适应DynaBERT 的“dynamic”主要指根据资源预算选择不同 width/depth 子网络。它不是根据每个输入句子的难度动态选择子网络。也就是说它是resource-adaptive。不是input-adaptive。如果希望简单样本走小模型、困难样本走大模型还需要额外的输入难度判断或早退机制。15.4 子网络是前缀式结构宽度压缩时DynaBERT 保留前若干 heads 和前若干 FFN neurons。因此必须依赖 rewiring把重要单元排到前面。这种方式简单但也有局限它不能为每个子网络自由选择任意 head 组合或 neuron 组合。15.5 只压缩部分结构DynaBERT 压缩 MHA heads、FFN intermediate neurons 和 Transformer layers但 embedding hidden size 不变。论文也说明embedding dimension 上的 neurons 因为通过 skip connections 跨层连接不能灵活缩放。所以它不是对 BERT 所有维度都完全弹性。十六、整体评价DynaBERT 的核心价值在于它把 BERT 压缩从“一个预算训练一个小模型”推进到“一个模型支持多个宽度和深度”。它不是单纯 head pruning也不是单纯 layer pruning而是把二者结合起来同时考虑 FFN 中间层神经元压缩。这篇论文真正重要的地方有三点第一它定义了 BERT 的宽度弹性attention heads FFN intermediate neurons。第二它通过 Network Rewiring 让重要结构单元被更多小子网络共享。第三它用两阶段蒸馏训练让一个模型能够支持多个 width-depth 子网络。从实验结果看DynaBERT 的很多子网络在 GLUE 和 SQuAD 上都能保持较强性能在相同参数量、FLOPs 或延迟约束下也常常优于固定大小的 BERT 压缩方法。因此DynaBERT 的定位可以概括为面向部署预算变化的动态 BERT。它的目标不是一次性得到最小模型而是提供一组可以按需选择的高效子网络。十七、一句话总结《DynaBERT: Dynamic BERT with Adaptive Width and Depth》提出一种可同时调节宽度和深度的动态 BERT宽度通过减少 attention heads 和 FFN 中间层神经元实现深度通过减少 Transformer layers 实现方法先训练宽度自适应 DynaBERTW再训练宽度和深度都自适应的 DynaBERT并通过 Network Rewiring 与知识蒸馏让重要结构单元被更多子网络共享从而实现一个模型覆盖多种参数量、FLOPs 和延迟预算。
