论文标题： MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices

论文来源： ACL 2020

论文链接： https://arxiv.org/abs/2004.02984

目前很多预训练模型的压缩方法都是和下游任务相关的，一般做法都是把 pre-train 的模型在下游任务上 fine-tune 之后，进行缩小操作（蒸馏，裁剪等等）。据作者说目前还没有与任务无关的轻量级预训练模型，他们要填补这个领域的空白。实验结果上看，这个空白填补的还是很成功的。

Insights

要做一个预训练的小模型，用蒸馏？蒸馏完的模型要么浅，要么窄，但是这两种模型都有自己的缺点，浅的表示能力不足，窄的不好训练（毕竟还是要考虑后续 fine-tune 的）。

MobileBERT 采用的和 BERT-large 一样深的层数，在每一层中的 transformer 中加入了 bottleneck 机制使得每一层 transformer 变得更窄，但是并没有丢失其与 self-attention 和 Feed-Forward 之间的平衡。

作者先训练了一个带有 bottleneck 机制的 BERT-large（IB-BERT），然后把 IB-BERT 中的知识迁移到 Mobile-BERT 中，在这个迁移的策略中，作者做了很多尝试。最后 Mobile-BERT 比 BERT-base 小 4.3 倍，快 5.5 倍，而且在 GLUE 上的效果没怎么降，还是很牛的。

可以看到 Table1 中，MobileBERT 就是作者之前说的那种窄的模型，会很难训练。所以作者采用的方法就是先训练一个 IB-BERT-large，然后把 IB-BERT-large 学到的知识迁移到 MobileBERT 中去。

采用了 bottleneck 机制的 IB-BERT-large 也存在问题，bottleneck 机制会打破原有的 MHA（Multi Head Attention）和 FFN（Feed Forward Network）之间的平衡，原始 bert 中的两部分的功能不同，参数比大概为 1：2。采用了 bottleneck 机制会导致 MHA 中的参数更多，所以作者在这里采用了一个堆叠 FFN 的方法，增加 FFN 的参数，Table1 中也能看出。

为了让模型更快，作者发现最耗时间的是 Layer-Norm 和 gelu，将这两个部分进行替换。把需要均值和方差的 Layer-Norm 替换为 NoNorm 的线性操作，把 gelu 替换为 ReLU，word-embedding 降为 128，然后用一个 3 核卷积操作提高到 512。

3.1 Feature Map Transfer (FMT）

由于在 BERT 中的每一层 transformer 仅获取前一层的输出作为输入，layer-wise 的知识转移中最重要的是每层都应尽可能靠近 teacher。特别是两模型每层 feature-map 之间的均方误差：

3.2 Attention Transfer (AT)

注意机制极大地提高了 NLP 的性能，并且成为 transformer 中至关重要的组成部分。作者使用从经过优化的 teacher 那里得到 self_attention map，帮助训练 MobileBERT。作者计算了 MobileBERT 和 IB-BERT 之间的自注意力之间的 KL 散度：

3.3 Pre-training Distillation (PD)

除了 layer-wise 的知识转移外，作者还在对 MobileBERT 进行预训练时使用知识蒸馏损失。作者使用原始的线性组合掩码语言模型（MLM），下一句预测（NSP）和新的 MLM 知识蒸馏（KD）是我们的预训练蒸馏损失：

鉴于上述 target，在训练中可以有各种组合策略。作者讨论了三种策略。

辅助知识迁移

在这个策略中将中间知识转移视为知识蒸馏的辅助任务。使用 single 损失，这是来自所有层的知识转移损失以及预训练蒸馏损失。

联合知识迁移

IB-BERT 的中级知识（即attention map和feature map）可能不是 MobileBERT 学生的最佳解决方案。因此，作者建议将这两个 Loss 分开。首先在 MobileBERT 上训练所有 layer-wise 知识转移损失，然后通过预训练蒸馏进一步训练它。

渐进知识转移

作者也担心如果 MobileBERT 无法完美模仿 IB-BERT，下层可能会影响更高的层次的知识转移。因此，作者建议逐步培训知识转移的每一层。渐进式知识转移分为 L 个阶段，其中 L 是层数。

表格还是很清晰的，就不多做解释了。

本文介绍了 MobileBERT，它是 BERT 的与任务无关的紧凑型变体。 流行的 NLP 基准测试 GLUE 的经验结果表明，MobileBERT 与 BERT-base 相当，但体积更小，速度更快。

• 稿件确系个人 原创作品 ，来稿需注明作者个人信息（姓名+学校/工作单位+学历/职位+研究方向）

• PaperWeekly 默认每篇文章都是首发，均会添加“原创”标志 返回搜狐，查看更多