BERT-Of-Theseus: Compressing BERT by Progressive Module Replacing

论文信息:论文代码

概述

随着深度学习的发展,许多大型的神经网络被提出并且获得了非常好的效果。尤其在NLP领域,预训练加微调几乎已经成为了一种新的范式。超大的参数量使得这些模型所向披靡,但同样也使得它们需要消耗大量计算资源。

有一系列模型压缩的方法被提出以解决这个问题,其中最著名的就是知识蒸馏。知识蒸馏构建一个大的教师网络和一个较小的学生网络,目标是让学生网络学会模拟教师网络并替换教师网络。但这种方式最终的性能取决于知识蒸馏损失函数的设计,也就是如何教学生网络学会教师网络的行为。

本论文受到哲学领域著名的“Ship of Theseus”实验的启发,提出了一种逐步使用小模块替换BERT中模块的模型压缩方式,取名为Theseus压缩方法。在这种方法中,被替换模块和替换模块被称为前驱模块(predecessor)和后继模块(successor),对应于知识蒸馏方法中的教师网络和学生网络。

Theseus压缩方法的主要过程如上图所示。首先为每个前驱模块指定一个后继模块,然后在训练阶段以一定概率使用后继模块替换对应的前驱模块。当训练收敛后,再把所有后继模块串起来作为新模型用于预测。这样之后,前驱模块就被压缩成了后继模块。

Theseus压缩方式和知识蒸馏的思想是类似的,但拥有三个优点:1)不需要手动设计额外的知识蒸馏损失函数;2)不使用Transformer特定的特征进行压缩,是一种通用的做法;3)允许前驱模块和后继模块充分交互(具体可以参考方法部分)。

本文论文的创新点可以总结如下:

  1. 提出了Theseus压缩方法,只使用一个损失函数和一个超参数就能够进行模型压缩。
  2. 使用Theseus压缩的BERT模型在保证98%性能的基础上比原模型快1.94倍,优于知识蒸馏方法。

方法

Theseus压缩方法的基本思想和知识蒸馏方法类似,但是只需要一个任务特定的损失函数,而且这种方式是一种通用的模型压缩方式。

模块替换

假设前驱模型(P)和后继模型(S)都含有(n)个模块,即(P={prd_1,dots,prd_n})(S={scc_1,dots,scc_n}),其中(scc_i)是用来替换(prd_i)的。假设第(i)个模块输入为(y_i),前驱模型的前向过程可以表示为:

[y_{i+1} = prd_i(y_i) ]

压缩时,对于第(i+1)个模块,(r_{i+1})是一个独立的从伯努利分布采样的变量,即以概率(p)取值为1,以概率(1-p)取值为0:

[r_{i+1} sim ext{Bernoulli(p)} ]

那么第(i+1)个模块的输出变成

[y_{i+1} = r_{i+1} odot scc_i(y_i) + (1-r_{i+1}) odot prd_i(y_i) ]

其中(odot)表示逐元素乘操作,(r_{i+1} in {0,1})。通过这种方式,前驱模块和后继模块在训练时一起运作。并且由于引入了类似Dropout的随机性,也相当于为训练过程添加了正则化。

训练的损失函数就是任务特定的损失函数,比如分类问题中的交叉熵损失函数:

[L=-sum_{jin |X|}sum_{cin C}[mathbb{1}[z_j=c]cdot log P(z_j=c|x_j)] ]

其中(x_j)是训练集的第(j)个样本,(z_j)(x_j)的标签。(c)(C)表示一个类标签和标签集合。在反向传播时,所有前驱模块的权重将会被冻结,只有后继模块参数会被更新。对于前驱模型的嵌入层和输出层除了进行权重冻结外,在训练阶段还会直接当作后继模块。通过这种方式,可以在前驱模块和后继模块之间计算梯度,从而可以进行更深层次的交互。

后继模块微调和推断

当训练收敛时,所有的后继模块将会串起来组成后继模型(S):

[egin{aligned} S = {scc_1,dots,scc_n} \ y_{i+1} = scc_i(y_i) end{aligned} ]

因为(scc_i)(prd_i)小得多,因此整个前驱模型被压缩成了更小的后继模型。然后后继模型将会再次使用损失函数进行微调优化,最后使用微调的后继模型进行推断。

渐进替换策略

尽管设置一个固定的替换概率(p)已经可以压缩模型,但本论文还探索了基于课程学习的渐进式替换策略。在这种策略中,替换概率随着时间增加:

[p_d = min (1, heta(t))=min (1, kt+b) ]

其中(t)是训练步数,(k > 0)是系数,(b)是基本的替换概率。非渐进式替换和渐进式替换的替换概率如下图(a),(b)所示:

通过渐进式替换策略,之前独立的训练和微调过程可以统一起来,形成一个端到端,从易到难的学习过程。在初始阶段,前驱模块作用仍然较大,使得模型损失能够平滑下降。之后,替换概率增大,逐渐过渡到了后继模块微调阶段。

在压缩的初始阶段,( heta(t)<1),替换模块数量是(ncdot p_d)(n)个后继模块的平均学习率为:

[lr^{'} = (n cdot p_d / n)cdot lr = (kt+b) cdot lr ]

其中(lr)是初始学习率,(lr^{'})是考虑所有后继模块的等效学习率。因此,当应用渐进替换策略时,同时对学习率使用warm-up机制有助于模型训练。

实验

具体的实验细节请参照原文,总之就是模型参数大幅压缩,但是性能却不怎么损失。这里我就放几张实验效果图。

原文地址:https://www.cnblogs.com/weilonghu/p/12584456.html