北京大学提出 PTQ4ViT | 双均匀量化+Hessian引导度量,推进Transformer模型落地

2023-11-01 20:07:59 浏览数 (1)

量化是一种非常有效的神经网络压缩方法,已在卷积神经网络(CNN)上取得了巨大成功。最近,视觉Transformer在计算机视觉中展现出巨大潜力。然而,以前的PTQ方法在视觉Transformer上表现不佳,即使在8位量化中也会导致1%以上的准确率下降。因此,作者分析了在视觉Transformer上进行量化的问题。作者观察到在Softmax和GELU函数之后的激活值分布与高斯分布非常不同。作者还观察到,常见的量化度量,如均方误差和余弦距离,无法准确确定最佳缩放因子。 在本文中,作者提出了双均匀量化方法,以减小这些激活值的量化误差。作者还提出使用Hessian引导度量来评估不同缩放因子,从而以较小的代价提高校准的准确性。为了实现视觉Transformer的快速量化,作者开发了一个高效的框架,PTQ4ViT。 实验证明,在ImageNet分类任务上,量化的视觉Transformer实现了近乎无损的预测准确性(在8位量化下准确度下降不到0.5%)。

1、简介

自注意力模块是Transformer的基本构建块,用于捕捉全局信息。受到Transformer在自然语言处理(NLP)任务上的成功启发,研究人员将自注意力模块引入了计算机视觉。他们用自注意力模块替代了卷积神经网络(CNNs)中的卷积层,将这些网络称为视觉Transformer。视觉Transformer在许多计算机视觉任务上与CNNs相媲美,有巨大的潜力用于各种应用。

然而,无论是CNN还是视觉Transformer,都需要大量计算资源和能源。神经网络规模越来越大,限制了它们在各种硬件设备上的部署,如移动电话和物联网设备,并增加了碳排放。因此,需要对这些神经网络进行压缩。量化是压缩神经网络的最有效方式之一。浮点值被量化为低位宽的整数,从而减少了内存消耗和计算成本。

有两种类型的量化方法,即量化感知训练(QAT)和训练后量化(PTQ)。虽然QAT可以生成准确度下降较小的量化网络,但网络的训练需要训练数据集、长时间的优化和超参数的调整。因此,在没有训练数据集或需要快速部署时,QAT不切实际。而PTQ通过在训练后使用未标定的校准图像来量化网络,从而实现了快速量化和部署。

尽管PTQ在CNN上取得了巨大成功,但将其直接应用于视觉Transformer会导致准确度下降超过1%,即使采用8位量化。因此,作者分析了视觉Transformer上的量化问题。作者收集了视觉Transformer中激活值的分布,并观察到存在一些特殊的分布情况。

  1. 在Softmax之后的值在[0,1]范围内具有非常不平衡的分布,其中大多数值非常接近于零。尽管较大值的数量非常少,但它们代表两个Patch之间的高关注度,这在注意力机制中至关重要。这需要一个较大的缩放因子,以使量化范围覆盖大值。然而,较大的缩放因子将小值量化为零,导致较大的量化误差。
  2. 在GELU函数之后的值具有不对称分布,其中正值具有较大的分布范围,而负值具有非常小的分布范围。很难通过均匀量化来很好地量化正值和负值。因此,作者提出了双均匀量化,分别量化两个范围内的值。为了在硬件设备上实现高效处理,作者设计了数据格式并限制了两个范围的缩放因子。

第二个问题是用于确定最佳缩放因子的度量在视觉Transformer上不准确。在先前的PTQ方法中有各种度量,包括均方误差、余弦距离以及量化前后层输出之间的皮尔逊相关系数。然而,作者观察到它们不准确,无法准确评估不同的缩放因子候选者,因为只使用了局部信息。因此,作者提出使用Hessian引导度量来确定量化参数,这更加准确。所提出的方法如图1所示。

作者开发了一个用于视觉Transformer的PTQ框架,使用双均匀量化,名为PTQ4ViT。实验表明,量化的视觉Transformer(ViT、DeiT和Swin)在ImageNet分类任务上实现了接近无损的预测准确性(8位量化下不到0.5%的下降)。作者的贡献如下:

  • 发现PTQ在视觉Transformer上的问题在于post-Softmax和post-GELU激活的特殊分布以及不准确的度量。
  • 提出双均匀量化来处理特殊分布,可以在包括CPU和GPU在内的现有硬件设备上高效处理。
  • 提出使用Hessian引导度量来确定最佳缩放因子,替代不准确的度量。
  • 量化网络实现了接近无损的预测准确性,使PTQ在视觉Transformer上可行。

2 背景和相关工作

2.1 视觉Transformer

在过去几年里,卷积神经网络(CNN)在计算机视觉领域取得了巨大成功。卷积层是CNN的基本组成部分,用于利用局部信息提取特征。最近,视觉Transformer挑战了CNN在计算机视觉中的地位,它使用自注意力模块来利用全局信息。DETR是第一项工作,将对象检测Head替换为Transformer,直接回归边界框并取得了与基于CNN的Head相媲美的结果。ViT是第一个替换了所有卷积层的架构,在图像分类任务上取得了更好的结果。在ViT之后,提出了各种视觉Transformer架构以提升性能。视觉Transformer已成功应用于下游任务,在计算机视觉任务中具有巨大潜力。

Transformer的输入是一系列向量。图像被划分成多个Patch,并使用线性投影层将每个Patch投影为向量。这些向量形成视觉Transformer的输入序列。作者将这些向量表示为

X in R^{N×D}

,其中

N

是Patch的数量,

D

是隐藏大小,即线性投影后向量的大小。

视觉Transformer包含若干块。如图1所示,每个块由多头自注意力模块(MSA)和多层感知器(MLP)组成。MSA生成不同Patch之间的关注度以提取具有全局信息的特征。典型的MLP包含两个全连接层(FC),并在第一层之后使用GELU激活函数。

输入序列首先被馈送到MSA的每个自注意力Head中。在每个Head中,序列被线性投影到3个矩阵中,query

Q = XW^Q

,key

K = XW^K

,和value

V = XW^V

。然后,矩阵乘法

QK^T

计算Patch之间的关注度分数。使用Softmax函数将这些分数归一化为关注概率

P

。Head的输出是矩阵乘法

P V

。这个过程可以表示为公式:

其中

d

是Head的隐藏大小。多个Head的输出被连接在一起,作为MSA的输出。视觉Transformer具有大量的内存、计算和能耗消耗,这阻碍了它们在现实世界应用中的部署。研究人员提出了许多用于压缩视觉Transformer的方法,如Patch修剪、知识蒸馏和量化。

2.2 量化

网络量化是压缩神经网络的最有效方法之一。权重值和激活值从浮点数转化为低位宽的整数,大大降低了内存消耗、数据移动和能耗。均匀对称量化是最常用的方法,它将浮点值

x

投影到具有缩放因子

k

位整数值

x_q

其中

round

将一个值投影到整数,并clamp将输出限制在k位整数可以表示的范围内。作者提出了双均匀量化,分别量化两个范围内的值。PWLQ也使用了多个量化范围。然而,他们的方法面向CNN,不适用于ViT。他们使用额外的位来表示使用哪个范围,比作者的方法多占用12.5%的存储空间。此外,他们使用FP32计算来对齐这两个范围,这不高效。作者的方法使用位移操作,避免了格式转换和额外的FP32乘法和FP32加法。

有两种类型的量化方法,即QAT和PTQ。QAT方法将量化与网络训练结合在一起。它优化量化参数,以最小化标记训练数据集上的任务损失。QAT可以用于量化Transformer。Q-BERT使用Hessian谱来评估不同张量对混合精度的敏感性,实现了3位权重和8位激活量化。尽管QAT实现了较低的位宽,但它需要训练数据集、长时间的量化和超参数调整。PTQ方法使用少量未标记图像来量化网络,比QAT快得多,不需要任何标记的数据集。

PTQ方法应确定每层的激活和权重的缩放因子∆。Choukroun等人提出最小化量化前后张量之间的均方误差(MSE)。EasyQuant使用余弦距离来提高CNN上的量化性能。最近,刘等人首次提出了一种PTQ方法来量化视觉Transformer。Pearson相关系数和排名损失用作确定缩放因子的度量标准。然而,这些度量标准不准确,因为它们只使用了局部信息来评估不同的缩放因子候选者。

3 方法

首先介绍视觉Transformer的基本PTQ方法。然后,将分析使用基本PTQ时的量化问题,并提出解决这些问题的方法。最后,作者将介绍作者的PTQ框架,PTQ4ViT。

3.1 视觉Transformer的PTQ方法

在视觉Transformer中,全连接层和QKT以及PV的计算中使用了矩阵乘法。在本文中,作者将其表示为

O = AB

,作者将专注于其量化。A和B使用具有缩放因子

ΔA

ΔB

的对称均匀量化分别量化为

k

位。根据公式,作者有

A_q = Ψ_k(A,ΔA)

B_q = Ψk(B, ΔB)

。在PTQ中,使用量化前后输出之间的距离作为确定缩放因子的度量标准,公式如下:

其中

hat O

是矩阵乘法在量化后的输出,

hat O = ΔAΔBA_qB_q

作者使用余弦距离作为度量标准来计算距离。作者通过线性划分

[alpha frac{A_{max}}{2^{k−1}}, beta frac{A_{max}}{2^{k−1}}]

[alpha frac{B_{max}}{2^{k−1}}, betafrac{B_{max}}{2^{k−1}}]

来创建

∆_A

∆_B

的搜索空间,分别划分为

n

个候选值。

A_{max}

B_{max}

A

B

的最大绝对值。

α

β

是两个控制搜索范围的参数。作者交替搜索最佳缩放因子

∆^∗_A

∆^∗_B

在搜索空间中。首先,固定

∆_B

,搜索最佳的

∆_A

以最小化距离

(O,hat O)

。其次,固定

∆_A

,搜索最佳的

∆_B

以最小化距离

(O,hat O)

∆A

∆B

交替优化若干轮。

A

B

的值是通过未标记的校准图像收集的。作者逐层搜索激活或权重的最佳缩放因子。然而,在作者的实验中,基本PTQ导致了量化视觉Transformer准确度下降超过1%。

3.2 双均匀量化

CNN中的激活值通常被认为是高斯分布的。因此,大多数PTQ量化方法都是基于这一假设来确定缩放因子的。然而,作者观察到Softmax后的值和GELU后的值的分布非常特殊,如图2所示。

具体来说,

  1. Softmax后的激活分布非常不均衡,其中大多数值非常接近零,只有少数值接近一。
  2. 经过GELU函数后的值具有高度不对称的分布,其中无界的正值很大,而负值具有非常小的分布范围。如图3所示,作者展示了使用不同缩放因子的均匀量化的量化点。

对于Softmax后的值,大的值意味着两个Patch之间存在高度的相关性,这在自注意机制中非常重要。较大的缩放因子可以降低这些大值的量化误差,导致较小的值被量化为零。而较小的缩放因子会使大值被量化为小值,显著降低了两个Patch之间的关注强度。

对于经过GELU函数后的值,使用对称均匀量化很难很好地量化正值和负值。非均匀量化可以用来解决这个问题。它可以根据分布设置量化点,确保整体量化误差很小。然而,大多数硬件设备无法有效处理非均匀量化的值。只有在专门设计的硬件上才能实现加速。

作者提出了双均匀量化,可以在包括CPU和GPU在内的现有硬件设备上高效处理。如图4所示,双均匀量化具有两个量化范围,R1和R2,分别由两个缩放因子ΔR1和ΔR2控制。k位的双均匀量化可以表示为:

对于Softmax后的值,

R1 = [0, 2^{k−1}∆^s_{R1})

范围内的值可以使用小的

∆^s_{R1}

来很好地量化。为了避免校准数据集的影响,作者保持

∆^s_{R2}

固定为

1/2^{k−1}

。因此,

R2 = [0, 1]

可以覆盖整个范围,大值可以在

R2

中很好地量化。对于经过GELU后的激活值,负值位于

R1 = [−2 ^{k−1}∆^g_{R1}, 0]

,正值位于

R2=[0, 2^{k−1}∆^g_{R2}]

。作者还保持

∆^g_{R1}

固定,使

R1

只覆盖负数的整个范围。由于分别使用不同的量化参数来处理正值和负值,可以有效减小量化误差。在校准网络时,作者搜索最佳的

∆^s_{R1}

∆^g_{R1}

均匀对称量化使用k位带符号整数数据格式。它由一个符号位和

k - 1

位表示数量的位组成。

为了有效存储双均匀量化的值,作者设计了一种新的数据格式。最高有效位是范围标志,表示使用哪个范围(0表示

R1

,1表示

R2

)。其他

k-1

位组成无符号数字,表示数量。由于相同范围内的值的符号相同,符号位被删除。

不同范围内的数据需要在矩阵乘法中进行乘法和累积。为了在CPU或GPU上有效处理双均匀量化的值,作者将两个范围约束为

∆_{R2} = 2^m∆_{R1}

,其中

m

是无符号整数。假设

a_q

R1

中量化,

b_q

R2

中量化,那么这两个值可以对齐:

作者将

b_q

左移

m

位,这与将该值乘以

2^m

相同。位移操作在CPU或GPU上非常高效。没有这个约束,需要进行乘法来对齐缩放因子,这比位移操作要昂贵得多。

3.3 Hessian引导度量

接下来,作者将分析确定每层缩放因子的度量标准。先前的工作通过贪心策略逐层确定输入和权重的缩放因子。他们使用各种度量标准,如MSE和余弦距离,来衡量原始输出和量化输出之间的距离。内部输出的变化被认为与任务损失呈正相关,因此用于计算距离。

作者在图5中绘制了不同度量标准的性能。作者观察到,与视觉Transformer上的任务损失(交叉熵)相比,MSE、余弦距离和Pearson相关系数不准确。基于它们得出的最佳缩放因子与基于任务损失的最佳缩放因子不一致。例如,在blocks.6.mlp.fc1:activation上,它们表明最佳缩放因子约为

0.4 frac{A_{max}}{2^{k−1}}

,而根据任务损失,最佳缩放因子约为

0.75 frac{A_{max}}{2^{k−1}}

。使用这些度量标准,作者得到次优的缩放因子,导致准确度下降。

最后一层的输出在PTQ中的使用可以更准确。然而,使用它来确定内部层的缩放因子是不切实际的,因为它需要多次执行网络来计算最后一层的输出,这会消耗太多时间。

为了实现高准确性和快速量化,作者建议使用Hessian引导度量来确定缩放因子。在分类任务中,任务损失为

L = CE(hat y, y)

,其中CE为交叉熵,

hat y

为网络的输出,

y

为真实标签。当作者将权重视为变量时,损失的期望是权重的函数

E[L(W)]

。量化对权重产生了小的扰动

ϵ

,权重为

hat W = W ϵ

。作者可以通过Taylor级数展开来分析量化对任务损失的影响。

其中

overline g(W)

是梯度,

overline H(W)

是Hessian矩阵。目标是找到缩放因子以最小化影响:

min∆(E[L(hat W)] − E[L(W)])

。基于[14]中的逐层重构方法,优化可以近似表示为:

其中

O^l

hat O^l

分别是量化前和量化后的第

l

层的输出。如图5所示,Hessian引导度量所指示的最佳缩放因子更接近任务损失(CE)所指示的缩放因子。尽管它仍然与任务损失存在差距,但Hessian引导度量显著提高了性能。例如,在blocks.6.mlp.fc1:activation上,Hessian引导度量所指示的最佳缩放因子对任务损失的影响较小。

3.4 PTQ4ViT框架

为了实现快速的量化和部署,作者开发了一种高效的用于视觉Transformer的PTQ框架,PTQ4ViT。其流程如算法1所述。它支持双均匀量化和Hessian引导度量。

有两个量化阶段:

  1. 第一阶段是收集每一层在量化之前的输出和输出梯度。第l层的输出
O^l

是通过在校准数据集上的前向传播计算的。梯度

frac{∂L} {∂O^l_1} , . . . ,frac{∂L} {∂O^l_a}

是通过反向传播计算的。

  1. 第二阶段是逐层搜索最佳的缩放因子。在搜索空间中使用不同的缩放因子来量化第
l

层的激活值和权值。然后计算层的输出

hat O^l

。作者搜索最佳的缩放因子

∆^∗

,以最小化公式(7)。

在第一阶段,作者需要存储

O^l

frac{∂L}{∂O^l}

,这会消耗大量的GPU内存。因此,当它们生成时,作者将这些数据传输到主内存。在第二阶段,当第l层的量化完成后,作者将

O^l

frac{∂L}{∂O^l}

再次传输到GPU内存并销毁它们。为了充分利用GPU的并行性,作者批量计算不同缩放因子下的

hat O^l

和对损失的影响。

4 实验

4.1 在ImageNet分类任务上的结果

作者选择了不同的视觉Transformer架构,包括ViT,DeiT和Swin。结果如表1所示。从这个表格中,作者可以看到基本PTQ在某些视觉Transformer上即使在8位量化时也导致了超过1%的准确度下降。PTQ4ViT在8位量化时的准确度下降不到0.5%。对于6位量化,基本PTQ导致了较高的准确度下降(平均为9.8%),而PTQ4ViT实现了较小的准确度下降(平均为2.1%)。

作者观察到Swin上的准确度下降不如ViT和DeiT明显。在8位量化时,4个SwinTransformer的预测准确度下降不到0.15%。原因可能是Swin在非重叠窗口内局部计算自注意力。它使用更少的块来计算自注意力,减少了post-Softmax值后的不平衡。

作者还观察到较大的视觉Transformer对量化不太敏感。例如,在8位量化时,ViT-S/224/32、ViT-S/224、ViT-B/224和ViT-B/384的准确度下降分别为0.41、0.38、0.29和0.17,而在6位量化时分别为4.08、2.75、2.89和2.65。原因可能是较大的网络具有更多的权重并生成更多的激活,使它们更能抵抗由量化引起的扰动。

表2展示了不同PTQ方法的结果。EasyQuant是一种流行的后训练方法,它交替搜索权重和激活的最佳缩放因子。然而,在8位量化时,准确度下降超过3%。刘等人提出使用Pearson相关系数和排名损失作为确定缩放因子的度量标准,从而提高了Top-1准确度。由于不同层对量化的敏感性不同,他们还使用混合精度技术,在4位量化时获得了良好的结果。

在8位量化和6位量化时,PTQ4ViT优于其他方法,平均提高了1%以上的预测准确度。在4位量化时,PTQ4ViT的性能不佳。尽管偏差校正可以改善PTQ4ViT的性能,但4位量化时的结果低于刘等人的混合精度。这表明混合精度对低位宽的量化很重要。

4.2 消融研究

接下来,作者对提出的双均匀量化和Hessian引导度量的效果进行了消融研究。实验结果如表3所示。作者可以看到,提出的方法提高了量化后视觉Transformer的Top-1准确度。

具体而言,仅使用Hessian引导度量可以在8位量化时略微提高准确度,而在6位量化时显着提高准确度。例如,在ViT-S/224上,8位量化时准确度提高了0.46%,而6位量化时提高了6.96%。同时使用它们可以进一步提高准确度。

基于Hessian引导度量,使用双均匀量化在post-Softmax激活或post-GELU激活上可以提高性能。作者观察到,仅使用双均匀量化而不使用Hessian引导度量会显著降低Top-1准确度。

例如,在8位量化时,同时使用Hessian引导度量和双均匀量化时,ViT-S/224的Top-1准确度达到了81.00%,而不使用Hessian引导度量时,降至79.25%,甚至低于基本PTQ的80.47%的Top-1准确度。这也证明了仅考虑局部信息的度量是不准确的。

5 结论

在本文中,作者分析了用于视觉Transformer的PTQ的问题。作者观察到post-Softmax激活和post-GELU激活都具有特殊的分布。作者还发现常见的量化度量不准确,无法确定最佳的缩放因子。

为了解决这些问题,作者提出了双均匀量化和Hessian引导度量。它们可以降低量化误差并以较小的代价提高预测准确度。为了实现对视觉Transformer的快速量化,作者开发了一个高效的框架,PTQ4ViT。实验证明,作者在ImageNet分类任务上实现了准确度接近零损失,使PTQ对视觉Transformer是可接受的。

6 参考

[1]. PTQ4ViT: Post-Training Quantization for Vision Transformers with Twin Uniform Quantization.

0 人点赞