北大华为鹏城联合首次提出视觉 Transformer 后量化算法！-轻识

点击上方“视学算法”，选择加"星标"或“置顶”

重磅干货，第一时间送达

AI 科技评论报道

编辑 | 陈大鑫

AI 科技评论今天为大家介绍一篇由北京大学、华为诺亚实验室联合鹏城实验室的研究人员最新提出的一种针对视觉Transformer的后训练量化方法，这种方法表现优于CNN中的量化算法。

这篇论文发现量化前后除了数据分布以外，注意力层的排序也存在明显差异，因此研究人员提出使用量化前后分布的相似度损失和注意力层的排序损失来联合优化以得到最优的量化步长。

除此之外，论文还提出了一种基于矩阵核范数的混合精度量化算法。在不进行额外训练的情况下，Deit-B的8bit模型在ImageNet图像分类任务上可以实现81.29%的精度。

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

1

引言

近期，Transformer结构在计算机视觉任务中取得了令人瞩目的成果，比如图像分类、目标检测、图像超分等等。这些模型通常在大数据集上进行训练，并且具有相当多的权重参数。例如，ViT-L模型具有307M参数，需要64G FLOPs的计算量。这使得在一些计算资源受限的设备（比如手机和智能汽车）上使用这些模型带来了很大的挑战。因为，实际应用中十分需要对基于Transformer的这些模型进行压缩和加速。

量化是一种很有效的压缩神经网络的方法，它是通过使用低比特宽度的权重值来代替全精度值而不会改变原本的网络结构，这一点对于一些精心设计的网络（比如Transformer）十分友好。在自然语言处理任务中，已经有一些研究人员针对BERT进行训练感知的量化。然而，这些方法需要额外的训练，增加了计算资源的消耗。此外，在某些场景下，训练数据并不是都可以直接获得了，而在边缘设备上进行重训练的代价是不可忍受的。后训练量化是一种可以在不进行额外训练的情况下进行压缩的算法，然而现有的算法都是针对卷积神经网络而设计（比如忽略了注意力层的影响），因此论文提出了一种针对视觉Transformer的后训练量化算法。

研究人员首先将后训练量化建模为寻找最优的量化步长问题。为了更好地保留注意力层的功能，他们分析了注意力层量化前后特征的不同，引入了注意力层排序的损失函数。并与量化前后特征分布的相似度损失进行了联合优化。此外，论文还提出考虑到不同网络层的特征多样化不同，可以根据注意力特征和输出特征权重矩阵的核范数来决定每个网络层的量化比特。

2

方法

图1 视觉Transformer后训练量化算法框架图

视觉Transformer量化

Vision Transformer的计算量主要集中在多层自注意力（MSA）和多层感知机（MLP）模块。论文中对所有矩阵乘法的输入和权重进行了量化，与BERT量化中的算法类似，它们也没有对softmax操作和layer normalization层的参数进行量化，因为这些层的参数量很少并且量化它们对结果会产生很大的不良影响。

后训练量化优化

对于量化而言，我们需要将浮点数约束为整数，其中量化步长的选择尤为关键。均匀量化函数是一种被广泛使用的选择，这样数据范围会被均匀分割：

，

其中代表量化步长，代表量化比特宽度，代表被量化的数据（比如权重或者网络层的输入）。Clip操作代表超出量化范围的数据会被截断在范围之内。

对于线性操作的相似度感知量化：对于MSA模块和MLP模块当中的线性操作，原本的输出可以被表示为

。那对于权重和输入进行均匀量化所对应的反量化函数可以描述为：

其中代表量化网络层的输出。通过上述两个公式，我们可以看出量化步长实际上也控制了量化过程中的截断阈值，从而影响了量化前后特征的相似度。相似度感知量化算法可以表示为：

其中

代表原始和量化输出特征之间的相似度，输入数据是由一个远小于训练数据集

的校验数据集得到。论文当中采用了比尔森相关系数作为特征相似度的度量：

对于注意力层的排序感知量化：自注意力层是Transfomer结构中十分重要的部分，也是transformer与传统的卷积神经网络不同的地方。对于自注意力层，作者发现量化使得注意力特征的相对顺序变得不同，会带来很大的性能下降，因为，他们在量化过程中引入了一个排序损失：

其中

表示成对的排序损失函数，表示权衡系数。损失函数的公式为：

为了解决上述的优化问题，作者提出了一种简单且有效的迭代搜索方法来寻找最优量化步长。首先，网络层输入的量化步长会被固定，然后搜索最优的权重量化步长。然后，权重的量化步长会被固定，搜索最优的输入量化步长。二者的量化步长交替迭代直到收敛或者达到了最大的迭代次数。为了快速收敛，量化步长根据权重和输入的最大值进行初始化，并且将量化搜索空间

均匀分为了C个选择。

误差纠正：为了进一步减少量化带来的误差，在搜索之后作者采用了偏差纠正的方法。因为如果输出的期望值不为0，那么输出的均值就会改变。这种偏差可能会对之后的网络层带来影响。因此作者对每个网络层的输出进行了偏差纠正。

在实现过程中，误差的期望值可以通过校验数据集来计算，然后在网络层的bias参数中减去它。

混合比特量化：不同的transformer网络层拥有不同的数据分布，因为有不同的量化“敏感度”。本文作者提出了混合精度量化，对于更加“敏感”的网络层分配更多的比特宽度。考虑到transformer中特殊的结构，作者对MSA和MLP模块内的所有操作使用了相同的比特宽度，这样对硬件的实现也会比较友好。

在论文中，作者提出使用MSA模块中注意力层特征和MLP中输出特征矩阵的核范数来作为度量网络层“敏感度”的方法。与HAWQ-V2中的方法类似，作者使用了一种帕累托最优的方式来决定网络层的量化比特。该方法的主要思想是对每个候选比特组合进行排序，具体的计算方式如下所示：

给定一个目标模型大小，会对所有的候选比特组合进行排序并寻找值最小的候选比特组合作为最终的混合比特量化方案。

3

实验

图像分类任务实验

作者首先在图像分类任务上对后训练量化算法进行了验证。从表1可以看出，在VIT和DeiT两种经典的transformer模型上，论文的量化算法均优于之前的卷积神经网络量化算法。在较小的数据集CIFAR-10上，论文的算法都接近或超过了全精度的transformer模型的性能，在ImageNet数据集上，Deit-B模型也取得了 81.29% 的Top-1准确率。

表1图像分类任务上的后训练量化结果