NeurIPS 2021 Transformer部署难？北大&华为诺亚提出Vision Transformer的后训练量化方法

详细信息如下：

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

• 项目链接：未开源

导言：

最近，Transformer在各种计算机视觉应用中取得了不错的性能。与主流卷积神经网络相比，视觉Transformer通常具有复杂的结构，用于提取功能强大的特征表示，但是在移动设备上部署比较困难。

在本文中，作者提出了一种有效的后训练量化算法，以减少视觉Transformer 的存储和计算成本。量化任务可以看做是分别为权重和输入找到最佳低比特的量化间隔。为了保持注意力机制的功能，作者在传统的量化目标中引入了排名损失（ranking loss），目的是在量化后保持自注意结果的相对顺序。

此外，作者深入分析了不同层次的量化损失与特征多样性之间的关系，并利用每个注意图和输出特征的核范数探索了一种混合精度量化方案。该方法的有效性在多个基准模型和数据集上得到验证，其性能优于SOTA的后训练量化算法。基于在ImageNet数据集上使用的DeiT-B模型，本文的量化模型可以获得81.29%的top-1精度，量化约为8比特。

      01      

Motivation

随着自然语言处理（NLP）任务的应用，基于Transformer的模型在各种计算机视觉（CV）任务中显示出强大的能力，如图像分类、目标检测和图像超分辨率。这些模型通常具有数亿个参数，例如，ViT-L模型中有307M个参数和64G FLOPs，这在推理过程中既占用显存又占用计算。这给模型在资源有限的设备上运行和部署非常困难。

在各种压缩方法（如剪枝和权重分解）中，量化方法能够通过使用较低的位宽来压缩神经网络，而不改变模型结构，这对于精心设计的网络结构（如Transformer）特别有用。通过将浮点数操作调整为整数或位操作，量化权重和输入可以加快推理速度。在NLP中已经有一些基于Transformer的模型的训练感知量化方法。然而，这些方法不是为计算机视觉任务设计的，通常需要额外的微调或者训练。此外，在某些场景下，可能无法获得完整的训练数据来优化量化模型。

后训练量化是一种有效的模型压缩技术，它可以直接量化神经网络模型，而无需进行微调。大多数现有的后训练量化方法是为卷积神经网络（CNN）或递归神经网络（RNN）设计的。这些方法没有考虑到视觉Transformer的特性（例如，CNN中不存在注意机制），不完全适合量化视觉Transformer。然而，视觉Transformer在各种各样的计算机视觉任务中表现出更强的性能。因此，作者希望为视觉Transformer结构探索后训练量化，以减少显存和计算成本。

在本文中，作者研究了具有混合精度的视觉Transformer模型的后训练量化方法，以获得更高的压缩比和加速比。Transformer中的量化过程可以被描述为寻找最佳量化间隔的优化问题，目标是在视觉Transformer中最大化全精度和量化输出之间的相似性。为了更好地保留注意机制的功能，作者深入分析了注意层和传统层（如MLP）之间的差异。

然后，引入排名损失（Rank Loss）来保持注意值的相对顺序。此外，作者提出根据特征多样性，即由注意图和输出特征计算的核范数，来确定每一层的位宽。并交替地搜索所有层中权重和输入的量化区间，以获得最佳量化结果。此外，作者还引入了偏置校正来减小累积量化误差。在几个基准数据集上的实验结果表明，本文的算法比现有的后训练量化方法获得了更好的性能。

      02      

方法

2.1 Preliminaries

标准Transformer接收token嵌入序列作为输入，因此视觉Transformer通常将图像转换为一系列的patch。H和W是原始图像的高度和宽度，（P，P）是每个图像块的分辨率，是Transformer的有效序列长度。通常，视觉Transformer在其所有层中使用恒定的宽度，线性投影层将每个patch的维度投影到d。因此，第一个Transformer层的输入为：

标准Transformer层包括两个主要模块：多头自注意(MSA)模块和多层感知机(MLP)模块。对于第l个Transformer层，假设其输入为，通过key和query的点积计算出的注意力分数，可以表示为：

然后对归一化分数应用Softmax函数得到，多头注意模块的输出为：

MLP模块包含两个的线性层，参数为，其中为MLP中间层的神经元数量。将MLP的输入表示为，则输出计算为：

基于上面的计算，第l层Transformer的前向传播可以表示为：

其中LN表示层归一化。

视觉Transformer的最大计算开销在于MSA和MLP模块中的大矩阵乘法。按照CNN的量化方法，作者对矩阵乘法中涉及的所有权重和输入进行量化。对于权重量化，作者对所有Transformer层的权重以及线性嵌入的参数进行量化。

除了这些权重，作者还量化了所有线性层和矩阵乘法运算的输入，但是不量化softmax操作和LN层，因为这些操作中包含的参数可以忽略不计，并且量化它们可能会带来显著的精度降低。

2.2 Optimization for Post-Training Quantization

对于后训练的量化，需要将浮点数限制为一组有限的值。量化间隔的选择对于量化至关重要，一个常用的选择是使用统一的量化函数，其中数据范围平均分割：

其中， 是量化间隔，是量化位宽，是表示权重或输入的张量。表示对张量中超出量化域范围的元素进行剪裁。

Similarity-Aware Quantization for Linear Operation

对于第l层Transformer层的MSA模块和MLP模块中的线性操作，原始输出可计算为。权重和输入的均匀量化相应的反量化操作可以描述为：

其中表示量化层的输出。从量化和反量化的公式可以看出，量化间隔实际上控制量化过程中的clip阈值，这在很大程度上影响原始输出特征图和量化特征图之间的相似性。因此，作者致力于优化权重的量化间隔和输入的量化间隔来提高和的相似度，其中是从校准数据集得到的样本。具体而言，校准数据集比普通训练数据集少得多。在第l个Transformer层中，相似感知量化可被公式化为：

其中是原始和量化输出特征图之间的相似度。在本文中，作者采用皮尔逊相关系数作为相似性的度量：

Ranking-Aware Quantization for Self-Attention

自注意层是Transformer的关键组成部分，因为它可以计算特征的全局相关性，这使得Transformer与卷积神经网络的计算是不同的。对于自注意的计算，作者发现，如上图所示量化后，注意图的相对顺序发生了变化，这可能会导致显著的性能下降。因此，作者引入排名损失（Rank Loss）来解决量化过程中的该问题：

其中，表示基于成对排序的损失函数，而γ表示权衡的超参数。排名损失可以表示为：

其中，是参数为θ的Hinge函数，（，）是矩阵A的大小。给定一对样本，仅当顺序正确且相差margin内时，损失函数才为0。

为了解决上述的优化问题，作者提出了一种Transformer层均匀量化的替代搜索方法。首先，输入的量化间隔是固定的，权重的量化间隔调整来进行优化。然后，固定，并对进行优化，以调整输入的量化间隔。交替优化和，直到目标函数收敛或超过最大迭代。此外，为了快速收敛，和分别根据权重或输入的最大值进行初始化。

Bias Correction

为了进一步减少量化引起的输出偏差误差，作者在每次搜索迭代后引入偏差校正方法。假设输入和权重的量化误差定义为：

如果输出误差的期望值不为零，则输出的平均值将发生变化。这种分布的变化可能导致性能的下降。作者通过以下方式纠正这种变化：

从偏置输出中减去输出的期望误差，可以确保每个输出单元的平均值得以保留。

2.3 Mixed-Precision Quantization for Vision Transformer

不同的Transformer层适用于不同的结构，并且具有不同的灵敏度，将相同数量的位宽分配给所有层是次优的。因此，作者探索了混合精度量化，将更多的比特分配给更敏感的层，以保持性能。考虑到Transformer层的独特结构，作者将MSA或MLP模块中的所有操作分配为相同的位宽。这也将有利于硬件实现，因为权重和输入分配有相同的位宽。

奇异值分解（SVD）是线性代数中一种重要的矩阵分解方法，其公式可写成：

其中，对角线项被称为的奇异值。核范数是奇异值之和，表示矩阵的数据相关性。在本文中，作者使用MSA模块中注意图的核范数和MLP模块中的输出特征来估计Transformer层的灵敏度。核范数可以用来减小混合精度设置的搜索空间，同时对更敏感的层使用更高的位宽，反之亦然。

此外，作者采用Pareto frontier来确定位宽。其主要思想是基于每个候选位宽配置引起的总二阶扰动，根据以下指标对其进行排序：

给定一个目标模型大小，根据值对候选位宽配置进行排序，并选择具有最小的位宽配置。

      03      

实验

3.1. Results and Analysis

Image classification

从上表可以看出，基于分类任务，在多个模型和多个数据集上，本文的后训练量化方法都优于其他后训练量化方法。

Object Detection

为了验证本文方法的泛化性，作者在目标检测上也做了实验，可以看出本文的方法在检测任务上，依旧具有性能上的优越性。

3.2. Ablation study

上表展示了本文提出模块的消融实验，可以看出，每个模块对于性能提升和量化效率的提升都是非常有用的。

      04      

总结

在本文中，作者开发了一种视觉Transformer后训练量化方案，其中每个层的位宽根据注意图的核范数和Transformer层的输出特征而变化。为了解决量化的最优化问题，作者提出通过搜索最佳量化间隔来保持量化后的特征图与原始特征图之间的相似性。

此外，作者还深入分析了注意层和传统层之间的差异，并引入了排名损失来保持注意值的相对顺序，偏差校正用于减少累积量化误差。最后，每个Transformer层的最佳量化间隔使用替代搜索策略进行优化。实验结果表明，与传统的训练后量化方法相比，该方法在网络精度和存储开销方面都有很大的改进。

▊ 作者简介

研究领域：FightingCV公众号运营者，研究方向为多模态内容理解，专注于解决视觉模态和语言模态相结合的任务，促进Vision-Language模型的实地应用。

知乎/公众号：FightingCV

END