【TPrune】TPrune: Efficient Transformer Pruning for Mobile Devices论文翻译理解

1. 背景

transformer促进了神经机器翻译的任务的性能达到了前所未有的高度。之前的研究工作中，为了使transformer在资源受限的平台上运行更友好，压缩模型方法可以分为三类：模型剪枝，迁移学习，高效transformer变型。模型剪枝以它简单的操作和很好的压缩率，在视觉任务的CNN模型上已经获得了极大的成功。尽管如此，之前的transformer压缩没有展现出对现成移动设备上的transformer组件的详尽模型分析和评估。在本文中，我们以行级别分析并裁剪了transformer模型，在真实移动平台也实现了我们的裁剪方法。我们探索了transformer所有组件的性质包括稀疏特性，它导致模型被裁剪。我们把模型分析和裁剪的整个过程叫做TPrune，我们在TPrune中首次提出块级别结构稀疏学习，来分析transformer的性质。基于BSSL得到的性质，我们应用结构霍耶平方得到了最终的裁剪模型。和当前最好的transformer裁剪模型比较，TPrune可以实现更高的压缩率，性能降低更少。实验结果显示，与传统transformer相比，我们的模型实现了1.16-1.92倍的加速，BLEU分数只下降了0-8%。

transformer凭借着多层自注意力机制+前馈神经网络在很多神经机器翻译任务上取得最佳效果。最近，基于自注意力机制的预训练模型的性能随着数据集的变大和模型尺寸的变大进一步得到提高。

transformer模型的变大带来内存消耗和计算成本，使得模型难以部署到移动设备中。但是将transformer放到远程服务器运行会带来时延、隐私安全问题，所以压缩模型大小是减少transformer运行时资源需求的有效方式。

模型剪枝是微调之前的预训练模型，使得参数和激活函数尽可能变成0。有四种裁剪力度。

迁移学习使用一个预训练好的大模型的知识，去指导一个小模型的训练。 Sanh 将老师模型的每两层抽出一层形成学生模型，学生模型基于老师和学生的输出概率分布的KL散度（学生模拟老师的概率分布）去训练。 Sun 使用特征图的均方误差（特征图像素和标签像素的误差平方和取平均值）和学生模型和教师模型之间自注意力head分布情况的KL散度作为损失来迁移transformer模型的知识。

有效transformer变量使用更加有效的操作数代替高成本的transformer模块。 Zhang 引入平均层和门控层，代替之前自注意力模块的位置，通过累计平均操作去总结历史注意力。 Wu 引入长短范围注意力，为局部头部建模和远距离关系指定不同的头部，以拓宽transformer注意力结构的功能。

三种方法中，裁剪最直观，不需要在模型结构设计上付出额外的人力，也不需要重新开始训练模型。

模型剪枝和其他方法是互不影响的，可以叠加使用。在本文中，我们将自己的方法和其他transformer裁剪模型比较，其他裁剪模型往往是基于启发式算法的，而我们详尽的分析了transformer组件，找到最佳的裁剪方案，实现令人满意的稀疏度-精确度的平衡。另外，我们量化地在移动设备上测试了TPrune的实时推理速度，并报告了第一手结果。

本文的贡献有以下几点：
1）我们使用块级别结构稀疏学习（用到了组L1正则化），发现transformer的冗余。
2）使用BSSL后，我们观察了transformer模型，分析模型的稀疏模式，获得了每一个组件的裁剪规律。
3）受CNN启发，经裁剪规律指导，我们将SHS正则化应用到transformer上。
4）我们在移动设备上评估了TPrune模型的性能，稀疏度，计算成本，也比较了TPrune和其他模型剪枝、剪枝正则化领域的最佳模型。

2. 文献关键词
神经机器翻译，模型剪枝，BSSL，SHS

3. 相关工作
1）transformer结构

2）特征
transformer的权重尺度比CNN大得多，所以裁剪起来比CNN困难很多。

为了理解深度神经网络不同类型层的性质，我们部署了lenet, inception-BN, vgg16在nexus5手机上。针对每一个模耗。在运行期间，全连接层和卷积层会消耗更多资源，全连接层消耗内存，卷积层消耗时间。**在transformer的block中，基本组件是MHA和FFN，而二者由全连接层组成，**因此，减少模型参数对transformer部署在资源受限平台非常重要。

3）相关裁剪模型

层级裁剪： Fan 提出层失活，它在推测期间从原始transformer网络中选择任意深度的子网络。训练期间每一层都要学习一个失活率，测试期间分数最高的层才会被采用。

头级裁剪： Voita 使用随机松弛将门控应用到MHA的每一个head上进行裁剪。它修剪了transformer几乎一半的head，但是BLEU分数下降还没有0.25。 Michel 基于提出的代理分数，迭代地mask了MHA的head，来评估每一个注意力head的重要性。

行级裁剪： Kenton 使用正则化器和近似梯度下降来自动调整大小，来修剪transformer模型权重。然而，自动调整大小严重损失模型精度，在阿拉伯语-英语翻译任务中，去掉20%的参数，BLEU分数从27.9降到了20.9。

元素级裁剪： Cheong 实现了迭代范围裁剪，mask掉未达到threshold的参数值，迭代地重训练模型。虽然一般来说，元素剪枝比其他剪枝方法能导致更高的稀疏性，但它只有在特定的硬件设置下才能实现真正的加速。

BERT是一种自监督的方法，具有双向transformer结构，为许多下游语言任务训练了一个预训练模型。迁移学习、裁剪、量化等压缩方法都曾被用于BERT。

4）相关正则化器
L0损失函数表示非0权重的个数。L0是非凸且不可微的，因此，基于梯度的优化方法不能用于L0正则化器上。 Han 以最小的范围，迭代地裁剪了固定百分比的权重，这样一种启发式算法无法达到最优的压缩率。L1作为替代方案出现了， Liu 将L1正则化器应用到权重上实现行级稀疏。但是L1有尺度变化性（权重绝对值的和），梯度会随着原始权重的大小改变。

结构化裁剪将元素级扩展到了组级。通过从原始模型权重中移除权重组，内存消耗减少，实现实时加速。 Wen 将权重矩阵的每行、每列分别作为一组，应用L2正则化器，同时保持每组的L1损失最小。 **Kenton 将L1应用到权重组的L无穷 ？？？**然而尺度变化的问题仍然存在。门控方法被用于transformer中MHA组件的每一个head上。门控值会从参数具体分布中独立得到，然而这种方法会引入门控估计的额外开销，使得可扩展性较差。

5）SHS
之前，SHS应用于CNN，效果很好。H具有尺度不变性，R(αS)=R(S)

霍耶正则化器的极小值接近于L0，这满足权重裁剪的目的。为了使上述公式的范围0-根号N能够匹配L0的范围0-N，针对元素级稀疏，霍耶正则化器调整为霍耶平方正则化器。除了尺度不变性，它还是可微的。它有和L0一样的范围和极小值结构。

在结构化裁剪的场景下，再次调整为结构化霍耶平方正则化器

||w^(g)||2是每组权重的L2损失，可以化简为

4. BSSL

1）transformer组件分析
MHA：dmodel设为512。n为head数，本文设为8。

对于多头注意力机制，QKV都是一样的，因为从a得到QKV，然后从Q中得到八个不同的head：Q1-8

注意力函数计算公式：

2）为了更好地裁剪模型，我们需要先回答关于ransformer组件的几个问题：
a）我们应该按行裁剪、按列裁剪还是行列都裁剪？
b）我们要同等看待encode, decoder吗？
c）浅层和深层的裁剪策略一样吗？
d）MHA中的应该以相同的稀疏度裁剪吗？

3）BSSL
在BSSL中，我们把一个权重矩阵分为几个大小相同的子块。我们把L1损失函数应用到每一个子块的行组、列祖上，权重衰减率是一样的。大小为(r, c)的子块的行级和列级BSSL正则化器公式：

衰减率控制在训练数据上的损失和BSSL正则化惩罚之间的平衡。

BSSL可以有不同的裁剪力度，当子块的大小等于原始权重矩阵时，粒度是行级；当子块的大小为(1, 1)时，粒度为元素级。BSSL可以被用来分析transformer的最佳裁剪粒度：

把L1损失应用到每一个子块的行级、列级组上，进行BSSL分析。在训练期间，每一个子块的行列都应该尽可能变成0。我们观察经过BSSL分析的模型的稀疏特性，来获得关于transformer权重冗余的更好理解。

4）BSSL分析
在 WMT English-German 数据集上训练transformer，应用BSSL正则化器，权重衰减率设为0.001，学习率设为0.001。在本文中，训练后低于0.0001的权重被看作是0。

观察得到：
a）所有的权重矩阵都可以被裁剪到一定程度；

观察得到：

5）分析总结
前人发现：
a）从图中可以看出，一些head（subblock)可以完全移除，说明只有一部分head很重要；
b）self-MHA的列级稀疏度为53.5%，但是 encoder-decoder MHA的列级稀疏度只有22%。说明模型更倾向于裁剪encoder部分的head，或者说是encoder-decoder MHA head比 self-MHA head更重要。
c）在6层的self-MHA中，**浅层（1，2层）的head比深层（5，6层）的留下更多。**层1的列级稀疏度为13.3%，层2为23.3%，层5为36.7%，层6为93%；
d）在6层的encoder-decoder MHA中，**浅层（1，2层）的head比深层（5，6层）的留下更少。**层1的列级稀疏度为68%，层2为63%，层5为10%，层6为16.7%；

我们发现：
a）尽管我们应用的是block-wise惩罚，最终实现的也只是line-wise稀疏度，这就是为什么我们提出的想法是line-wise的原因；
b）W^Q, W^K 总能裁剪到相同的列级稀疏度，W^V但不太一样，从注意力计算公式可以看出，Q和K的转置相乘，但是V和乘积的softmax结果有关；
c）FFN组件在encoder中可以被压缩，但是无法在decoder中被压缩；
d）所有权重将按行或列进行修剪，因为所有子块都希望利用整个嵌入维度，因此只能修剪内部维度。换句话说，修剪后的维度始终是每个组件（MHA和FFN）的内部维度。

5. 模型剪枝
我们在**权重组的L2范数上应用L1（其实就是BSSL）**来生成结构化稀疏性。然而，SSL也有缺点，关于组的规模。应用SSL正则化在一组权重Wg上，这组权重的L2范数是ng，单个权重wi的梯度就是Wi/ng（把L2范数对Wi求导，导数是Wi/ng）。梯度依赖于权重本身大小和所属组的L2范数。SSL是尺度变化的，我们采用尺度不变的SHS去裁剪模型。这里的group指的是一行/一列权重。

encoder的都可以被裁剪；decoder只有self-MHA, encoder-decoder MHA的W^Q, W^K可以被裁剪（列级）。

之前的工作，训练包括两个环节：裁剪和微调。裁剪过程中，模型根据正则化器训练（SSL, SHS），旨在尽力去除冗余权重；微调过程中，不使用正则化器训练上一步得到的模型，针对0权重的梯度都被mask掉。

6. 实验结果

1）实验过程
使用tensorflow去实现TPrune，tensorflow是一个被广泛使用的深度学习框架，用于在不同的硬件平台上训练和测试。针对transformer的训练和裁剪，使用tensor2tensor，专为 机器学习研究 设计的tensorflow库。它融合了许多数据集、预训练末、预训练超参来加速模型开发。针对transformer在移动平台上的实时运行，我们使用 TFLite 模型基准工具，将模型结构转换成TFLite格式，把深度学习网络部署在IOS, Android平台上的嵌入式软件上。

使用T2T提供的 WMT English-German 数据集去验证TPrune的有效性。这个数据集包括3万2千个token，共享源-目标词汇表。剪枝过程中，batch大小16Ktoken，step800K。BLEU分数也作为一项性能评价指标，测试数据集为 newstest2014 。

使用T2T提供的预训练transformer模型，最初在 newstest2014 上测试的BLEU分数为29.2。

为了加速训练，模型使用4个 NVIDIA TITAN RTX GPUs 在Linux系统下并行训练。每个GPU的batch为4K。在四个不同的手机上测试模型，都基于移动CPU。

2）实验结果

图7展示了分别使用SSL, SHS正则化器裁剪的模型的稀疏度和BLEU分数，权重衰减率分别为0.00001和0.0001。因为衰减率比较小，经过400Kstep的训练后，模型依然能保持较好的性能。
经过400Kstep训练，SSL, SHS的BLEU分数差不多，但是SHS的稀疏度更高，为8.16%，而SSL只有2.2%。这是因为当衰减率比较小时，大多数权重都不是0，SHS能更有效的去掉接近0的权重，而SSL在所有权重上都施加了相同的惩罚（都是除以L2），导致更少的权重被移除。当我们需要一个高性能、低稀疏度的模型时，SHS比SSL表现更好。

图8展示了分别使用SSL, SHS正则化器裁剪的模型的稀疏度和BLEU分数，权重衰减率变大了，分别为0.00005和0.001。因为衰减率变大，经过400Kstep的训练后，稀疏度都升到了45%左右。SSL的BLEU为21.97，SHS为21.66。

这时，SSL的表现和SHS居然差不多了。更大的衰减率对于SSL, SHS有不同的表现：稀疏度增加，SSL的BLEU会持续降低；稀疏度增加，SHS的BLEU先迅速下降，然后逐渐恢复到和SSL差不多的水平。这种稀疏度和性能之间不同的平衡让我们可以灵活的设计transformer的裁剪方式。

表1总结了在batch大小8K，运行400Kstep时，SSL, SHS的稀疏度和BLEU。得出：**稀疏度越低，SHS表现越好；稀疏度越高，SSL表现越好。**通过增加step数目也能提高性能。

自然策略：对encoder和decoder的所有权重矩阵施加行级和列级SHS正则化。

在稀疏度为20%时，三种策略的BLEU分别为 34.36, 31.13, 和 28.98，随着稀疏度加大，性能差距越来越大，证明我们的模型能在保持性能的同时有效去除冗余权重。

图10和11显示了在不同的模型稀疏度下，各个层的稀疏度。有以下规则：

字符串长度：图12展示了在pixel3上运行的不同模型稀疏度和字符串长度下的运行时间。在四个线程下测试。**相同的模型稀疏度，字符串越长，运行时间越长。**例如，当字符串从10增加到80时，运行时间从50ms增加到184ms。注意这里的结果没有严格遵循理论公式，计算成本不是和字符串长度平方成正比的。**除了时间复杂度以外，许多因素都会影响运行时间，例如内存带宽、芯片上缓存大小、内存时延等等。**即使是很小的字符串长度，也会消耗一定的时延，在这种情况下，内存时延是造成较长运行时间的主要原因。

移动设备：图13展示了在不同设备上的模型运行时间和加速情况。在四个线程，字符串长度40下测试。稀疏度为0时，Nexus 5, Pixel 2, Pixel 3, and LG G8 ThinQ的运行时间为190ms, 159ms, 100ms, and 70ms。增加模型稀疏度仍然能减少运行时间。模型稀疏度为45%时，运行时间减少到121ms, 81ms, 57ms, and 40ms。图13b展示了不同设备上不同模型稀疏度下的正则化加速效果。当模型稀疏度从10%到45%，Nexus 5加速1.07 × −1.57×， Pixel 2加速1.2 × −1.96×， Pixel 3 和 LG G8 ThinQ加速1.07 × −1.75×。

多线程：图14展示了不同线程下的运行时间。在字符串40条件下。四个设备的总体趋势差不多。 Pixel 3的加速为1.6×, 1.5×, 1.57×, 1.63×, and 1.19×在1，2，4，8，16线程时。单纯增加线程数不能一直减少运行时间。 当线程从8加到16，在Pixel 2, Pixel 3, 和 LG G8 ThinQ 上运行时间急剧增加，因为大幅度的划分模型带来了额外开销。当线程从4加到8，Nexus 5的运行时间增加了，这是因为Nexus 5和其他设备比起来受限的CPU多线程处理能力。总之，4个线程表现效果最佳。

不同裁剪模型：图15展示了不同模型稀疏度下，两种模型的运行时间。在四个线程，字符串长度40下测试。当模型稀疏度从4%到28%，head-wise比SHS表现好，因为head-wise有更大的裁剪粒度和更简单的裁剪策略；但是28%是head-wise能实现的最高稀疏度了，因为这时已经移除了全部head。然而，SHS可以继续裁剪到45%，加速 1.92×。另外，在相同稀疏度下，head-wise带来的性能下降比SHS严重得多。

稀疏度-准确度权衡
我们以800Kstep，16Kbatch大小训练模型。当模型稀疏度较小时，SHS收敛的很快。在模型稀疏度为9.76%时，仅仅经过200K fine-tuning steps，BLEU分数就收敛到了26.89。当模型稀疏度为45.07%，需要400K steps收敛到24.38。当稀疏度低于30%，fine-tuning step依然能达到原始transformer的性能。

表2总结了不同稀疏度下模型的性能。原始transformer作为baseline。为了对比，我们也使用了另外一个当前最佳裁剪模型head-wise。我们把head-wise从原始框架挪到T2T框架，确保实验使用的模型和数据集一样。所有加速结果都是在 Pixel 3 上测量的。和baseline模型对比，SHS pruned model实现了16%稀疏度，BLEU完全没下降，加速在1.21×。当稀疏度到20.29%，BLEU稍稍降到26.78，或者说是原模型性能的99.48%。当稀疏度到45%，就可以获得24.78的BLEU，1.92×的加速。
对比下，head-wise Transformer在稀疏度为4.29%时没有降低BLEU，但加速仅有1.05×。增加稀疏度会使BLEU快速下降。当稀疏度为25.71%，BLUE降到了只有10.1。更糟糕的是，因为有限的head数，模型无法继续裁剪了。
使用layer-wise pruning，50%的稀疏度只降低5%的BLEU，实现了比TPrune更好的稀疏度-准确度平衡。幸运的是，layer-wise pruning和Line-wise pruning互不影响，所以可以结合起来实现更好的性能。