NeurIPS 2020 之预训练语言模型压缩

近期，大规模和超大规模预训练模型近期频繁出现在各大新闻媒体，从聊天到创作，一步一步刷新了我们的认知。但是我们想拿来就用的梦想却被手头的设备要求扼杀在了萌芽中。模型压缩是解决这个问题非常自然的一个想法，它的核心的思路包括：参数量化、权重共享、知识蒸馏，尤其是在BERT为基础的模型压缩工作近些年也层出不穷，如果说以前的工作依然还是从传统的模型压缩到NLP的“拿来就用主义”，那么当前在预训练语言模型的压缩方向有什么新的探索方向和简单高效的方法值得使用呢？本文从NeurIPS 2020中精选了一些工作，主要是以下2个标准做了一些调研：（1）大组出品，具有很强的适用性和可操作性；（2）领域前沿，打开新的视角。文章也同步发布在AI Box知乎专栏（知乎搜索 AI Box专栏），也欢迎大家在知乎专栏的文章下方评论留言，交流探讨！

作者简介：刘沛羽，中国人民大学高瓴人工智能学院博士一年级在读，导师为赵鑫教授，研究方向是语言模型压缩。

01

剪枝系列

剪枝是最直接的压缩策略，目的是识别和去掉对模型表达能力无用或者重要性低的权重或者组件。通常包括结构化剪枝和非结构化剪枝两大类。下面先是一个结构化剪枝的论文。

《Accelerating Training of Transformer-Based Language Models with Progressive Layer Dropping》

https://proceedings.neurips.cc/paper/2020/file/a1140a3d0df1c81e24ae954d935e8926-Paper.pdf

工作简介：本文是来自微软发表自NeurlPS 2020的工作。

问题：由于BERT的预训练过程FLOPS是和Transformer 层数深度呈线性关系的。本文从计算机视觉领域证明有用的随机深度的方法（Stochastic Depth）入手，发现应用在Transformer layer时候对模型的训练稳定性和效果有负面效果，本文则是从这个角度出发进行改进。

解决方法：本文从多个角度分析了随机剪掉层数带来的影响，首先从LN层放置的位置对比了PreLN和PostLN的区别，结论是PreLN可以避免梯度弥散，使训练更稳定；另外从剪枝层数的时间和位置上进行了定性的分析，最好的方法是避免训练前期频繁的剪枝，和避免较低层的频繁剪枝。基于这些分析，本文核心贡献是：（1）实现了一个Switchable-Transformer(ST)模块来控制剪枝，本质上类似一个门控机制，通过一个从伯努利分布中的采样值来控制是否剪枝；下图是模型的结构：

单层Transformer结构下展示变体结构

（2）设计了一种控制剪枝时间和位置的函数（Progressive Layer Dropping Schedule ），从易到难的实现剪枝。

文章的实现方法还是会给实际的操作带来很多启发，另外就是本文的写作逻辑感强，学习很多。

《Deep Transformers with Latent Depth》

https://proceedings.neurips.cc/paper/2020/file/1325cdae3b6f0f91a1b629307bf2d498-Paper.pdf

工作简介：本篇工作来自Facebook AI和University of Edinburgh合作的文章，重点关注如何通过较大或者可变的模型层数来提升模型容量。本文是通过学习一个是否选择或跳过的概率分布来解决这个问题。

问题：本文的应用场景是在机器翻译领域，当需要对不同的语言学习一个统一的机器翻译模型的时候，一个比较大的挑战就是如何使模型学习到有用的信息(positive transfer)在不同语言之间有效迁移，同时还能避免负向的信息（negative transfer）迁移。

解决方法： 具体的实现方法非常巧妙，本文是将是否选择一个层这个问题看作是一个从概率分布中抽样的问题。在训练中，作者同学学习一个Gumbel-softmax 分布来近似这个离散分布问题，在最后测试的时候，最终学习到的离散部分可以作为标准来对层进行剪枝而得到最终的压缩模型，当然也可以使用这个“选择概率”作为一个层权重（soft weights）使用。下面为论文的模型图部分：

模型结构图

非结构化剪枝：和结构化剪枝最大的区别在于，非结构化剪枝策略侧重于数值权重部分，这样以来原始的模型会变得十分稀疏。接下来介绍一篇结合“彩票假设”使用非结构化剪枝的论文。

《The Lottery Ticket Hypothesis for Pre-trained BERT Networks》

https://proceedings.neurips.cc/paper/2020/file/b6af2c9703f203a2794be03d443af2e3-Paper.pdf

工作简介：本文是来自Tianlong Chen与MIT CSAIL的Jonathan Frankle 和 Michael Carbin合作的关于“彩票假设”（Lottery Ticket Hypothesis）在BERT中的应用的工作。

背景： 在本文之前Jonathan Frankle 和 Michael Carbin在2019年ICLR中提出了彩票假设，并且在图像分类问题上（MNIST和CIFAR10）实现了寻找这种赢得彩票的子网络的方法，结果子网络仅为原网络大小的10%-20%。

问题：这篇工作拿到了当年的ICLR best award，既有一个新奇的想法又通过实验证明了它的存在，就像是一个灯塔，必然会引起各大研究小组的注意，于是，直接在自己的研究领域尝试彩票模型是否存在，成为了一个非常直接的motivation。本文便是将这种方法应用在了预训练语言模型上，当然除此之外本篇文章还有其他主要的贡献。

解决方法：一个随机初始化的稠密网络会包含一个子网络，这个子网络可以在不超过原网络的迭代轮数的条件下达到近似原网络的模型性能。以下为原作中的表述。

Lottery Ticket Hypothesis原文

本文采用的是非结构化剪枝的策略，这一点和Frankel他们的方法差不多，具体来说就是定义一个mask矩阵用来对网络权重进行剪枝，迭代的进行这么一个流程：首先将原网络训练到  step，根据权重的大小剪枝剩下权重的10%，然后将现在的权重赋值为  step时候的权重。现有的模型如果换成更大规模的预训练语言模型，比如GPT-3，这一策略后续还可以从剪枝的目标（不用权重的幅值做标准？），剪枝的顺序（每次剪枝某一定比例？）继续探索。

本文后面对子网络学习到的信息迁移能力进行了研究，发现从某一个特定任务下寻找到的子网络，通用性是和产生子网络的任务数据规模相关的，尤其是MLM任务下得到的子网络基本可以迁移到其他所有的子任务。这里放出来一张图作者的图。其中纵轴列表示源任务，在源任务寻找到70%稀疏度下的子网络，再在横轴目标任务上训练，深色的方格表示这个目标任务的训练依然可以“赢得彩票”。

模型的迁移性分析

扩展部分： 既然提到了彩票系列，本文这里大概梳理一下Lottery Ticket的发展历程，笔者调研的过程中发现在假设提出后出现了非常多的相关文章而且都发在了非常不错的会议上，这里只是选择其中几篇进行梳理。

• ICLR 2019，《The Lottery Ticket Hypothesis Finding Sparse, Trainable Neural Networks》 本文作者中的Frankle和Carbin是彩票假设最初提出的论文作者，在2019年论文中提出了彩票假设。
• NeurlPS 2019，Uber工作《Deconstructing Lottery Tickets: Zeros, Signs, and the Supermask》在针对彩票模型的研究中获得了意外衍生品：“超级掩膜”（Super Mask）
• 2020 EMNLP《When BERT Plays the Lottery, All Tickets Are Winning》这篇在彩票假设应用在BERT上较Tianlong Chen发的更早一些，也是应用彩票假设在BERT模型中，但是Tianlong Chen这篇工作中强调是使用的方法是非结构化剪枝策略，并且结论的落脚点也有一些不同。
• 田渊栋老师相关有4篇工作，从彩票假设的图像领域直接扩展到了强化学习、NLP领域，这里就不细说了，知乎中有一篇帖子介绍的比较详细，这里贴出来供大家参考：https://zhuanlan.zhihu.com/p/93988943

本文这里不是主要介绍彩票假设的工作，这里只做抛砖引玉，可以方便大家继续跟进相关研究。

总结： 剪枝方法思想简单，在NeurlPS中依然有比较高的曝光率，传统的方法结合新的思想往往有种让人耳目一新的感觉，“原来还可以这么用”，这一点在彩票假设的论文中和第二篇facebook的工作中体会尤为明显。

02

知识蒸馏系列

该系列的研究工作始于Hinton在2015年发表的工作《Distilling the Knowledge in a Neural Network》，后续有很多研究沿用知识蒸馏的思想，主要的工作有围绕蒸馏损失函数的设计以及蒸馏的位置不同，这里列举NeurlPS 2020上一篇蒸馏的工作。

《DynaBERT: Dynamic BERT with Adaptive Width and Depth》

https://proceedings.neurips.cc/paper/2020/file/6f5216f8d89b086c18298e043bfe48ed-Paper.pdf

本篇工作来自于2020年华为实验室刘群老师组在NeurlPS会议发表的工作，其实本文之前已经在很多博客和知乎上有讨论。

问题：本文的核心问题在于和以前的方法只能生成固定大小的小模型或者只能是深度不同的小模型不同，本文是通过同时自适应的调整小模型的宽度和深度，可以得到不同大小的小模型。本文压缩的重点依然是着重推理的速度。

方法：本文在训练的中的贡献在于2个点，（1）本文的自适应的宽度和深度是通过一个可学习的参数  和  来实现的，但是为了保证每次裁剪的都是不重要的部分，本文需要实现一个方法来判断结构的重要性，即文中提到的Network rewiring的方法。在所有的Pruning思想的压缩文章中都有这个部分，本文的想法比较直接，就是如果去掉待裁剪的结构，整体的loss是不是会产生很大的变化来判断这个结构是不是重要的。这个部分承担了DynaBERT核心的压缩的工作（2）在蒸馏部分，则是为了弥补压缩带来的效果损失。本文没有同时对宽度和深度进行压缩，而是2-step的方式，首先使用最大深度条件下，压缩得到一个宽度最优的模型  ，然后以此作为初始化，来压缩得到深度最优的DynaBERT，这里有个细节很重要，就是如果在压缩深度的时候只考虑固定的宽度，作者提到可能会出现遗忘以前学习的深度有关的知识的情况。

To avoid catastrophic forgetting of learned elasticity in the width direction, we still train over different widths in each iteration

模型结构：

DynaBERT模型结构图

总结来说，本文的压缩的思想整体比较直接，有借鉴意义的是其中自适应调整宽度和深度的部分为了学习合适离散值，本文依然使用可学习的 和 作为系数来表示drop掉的部分，同时为了保证每次drop掉的都是不重要的部分，还引入了一个重要性计算的模块Network revwiring。另外本文的实验部分比较充分，其中举了一个定性分析（Looking into DynaBERT）的例子的比较有意思，从注意力权重的可视化图中发现，压缩后的DynaBERT可以比原始的BERT模型中注意力部分学习到更多的句法和语义模式。

注意力权重可视化

总结： 蒸馏依然是当前模型压缩的学术界和工业界研究的热点，其最大的优势在于思想实现非常简单，对于研究来说可以从损失函数设计以及损失函数的位置就有很多的尝试策略，探索方向比较多。

点击下方“阅读原文”前往知乎专栏

↓↓↓