3.3、用于Softmax量化的Log Int Softmax
多头自注意力(MSA)是基于Transformer的架构中最重要的组件之一,但由于Token数量的二次复杂性,即图像分辨率除以Patch size,它被认为是资源最密集的组件。随着模型性能被证明受益于更高的分辨率和更小的Patch size,当分辨率增加和Patch size减小时,注意力图的存储和计算成为瓶颈,直接影响推理的吞吐量和延迟。因此,更小的注意力映射和更高效的推理成为迫切需要。
1、Attention Map的Log2量化
为了将注意力映射压缩到较小的大小并加快推理速度,将注意力映射量化到较低的位宽。当实验用均匀量化将注意力图从8位量化到4位时,所有的Vision Transformer都表现出严重的性能下降。
例如,在具有4位均匀量化注意力图的ImageNet上,DeiT-T仅导致8.69%的top-1准确率,比8位情况降低了63.05%。
受DynamicViT中稀疏注意力思想的启发,作者探讨了注意力图的分布,如图3所示。观察到以一个相当小的值为中心的分布,而只有少数异常值具有接近1的较大值。对ViT-L中的所有注意力图进行平均,约98.8%的值小于1/16。与只为这么多值分配1个bin的4位均匀量化相比,log2方法可以为它们分配12个bin。
此外,遵循对感知损失进行排名的目的,log2量化可以在全精度注意力图和量化注意力图之间保持很大的阶数一致性。因此,我们避免了注意力图的4位量化中的极端退化,并以减少50%的内存占用实现了与8位均匀量化相同的性能。从两个方面证明了量化与MSA相结合是合适的。
首先,与方程6相比,Softmax的固定输出范围使函数校准Free:
这确保了注意力图的量化不会受到校准数据的波动的影响。
其次,它还介绍了在量化注意力图(AttnQ)和值(VQ)之间将MatMul转换为BitShift的优点,如下所示:
其中。注意到用AttnQ的结果直接右移可能导致严重的截断误差。使用作为量化输出,其标度等于,这确保了左移操作以防止截断误差。
2、Softmax的量化推理
Softmax一般用于Vision Transformer,可表述为:
在实践中,减去输入的最大值,以避免溢出:
因此,Softmax可以写成:
现在,所有的输入 都变成了非正的。将它分解为,其中是一个非负整数,中的一个实数。那么,的指数可以写成:
其中是右移。I-BERT使用一个二阶多项式来近似区间内的指数函数:
图6绘制了i-exp的结果,它与指数函数几乎相同。这两个函数之间的最大差距只有,与的8位量化误差相比,这是相对可以忽略不计的。
在(算法1)的基础上,提出了Log-IntSoftmax(LIS),如算法2所示。
首先,减去输入的最大值,以确保所有的输入都是非正的,并将结果发送到。然后,将原来的Softmax替换为其倒数,以确保整数除法的结果大于1。最后但并非最不重要的是,对反向Softmax的输出执行一个量化,如下:
与。,是量化的输入和scale。
函数可以用整数算法计算,如算法2所示。为了获得整数的,引入了先找到前一个函数,该函数返回输入中最重要的一个位的索引。整个过程可以在整数域中完成。例如,如果 的输入为 ,χ将计算为,舍入结果将为。
Softmax的以下一步是注意力映射附加数与值¥之间的矩阵乘法。考虑到是的量化值,其和为,,矩阵乘法可以写为:
正如看到的,乘法被转换为具有输出scale 的BitShift运算符,其中N在4位量化中等于。基于上述过程,作者实现了对Softmax的完全定点推理,并利用BitShift加速了注意力图和值的计算。
3、Integer-only推理
先前的工作选择不量化Softmax,因为Softmax中计算量的可忽略性和量化可能导致精度显著下降。然而,数据在CPU和GPU/NPU之间移动,进行去量化和再量化,会给硬件设计带来很大困难,这是一个不可忽视的消耗。
将量化与指数函数的多项式近似相结合,提出了Log Int Softmax,一个仅整数、更快、低功耗的Softmax:
其中。整数函数可以很容易地实现,方法是使用BitShift找到值为1的第1个位索引(称之为find first One函数),并在其后面添加位的值。
正常MSA和本文的方法之间的差异如图4所示,每个阶段的数据类型都被标记。在具有左侧所示的未量化Softmax的多头自注意力中,和的矩阵乘法需要在Softmax之前去量化到全精度,然后再量化。
当采用Log IntSoftmax时,如右侧所示,整个数据类型可以是纯整数,并单独计算和并行计算量化尺度。值得注意的是,LIS在注意力图上使用4位表示,这显著减少了内存占用。
4、Softmax的量化的实现
class QIntSoftmax(nn.Module): def __init__(self, log_i_softmax=False, quant=False, calibrate=False, last_calibrate=False, bit_type=BIT_TYPE_DICT['int8'], calibration_mode='layer_wise', observer_str='minmax', quantizer_str='uniform'): super(QIntSoftmax, self).__init__() self.log_i_softmax = log_i_softmax self.quant = quant self.calibrate = calibrate self.last_calibrate = last_calibrate self.bit_type = bit_type self.calibration_mode = calibration_mode self.observer_str = observer_str self.quantizer_str = quantizer_str self.module_type = 'activation' self.observer = build_observer(self.observer_str, self.module_type, self.bit_type, self.calibration_mode) self.quantizer = build_quantizer(self.quantizer_str, self.bit_type, self.observer, self.module_type) @staticmethod def log_round(x): x_log_floor = x.log2().floor() big = x_log_floor extra_mask = (x - 2**big) >= 2**(big - 1) big[extra_mask] = big[extra_mask] + 1 return big @staticmethod def int_softmax(x, scaling_factor): def int_polynomial(x_int, scaling_factor): coef = [0.35815147, 0.96963238, 1.] # ax**2 + bx + c coef[1] /= coef[0] coef[2] /= coef[0] b_int = torch.floor(coef[1] / scaling_factor) c_int = torch.floor(coef[2] / scaling_factor**2) z = x_int + b_int z = x_int * z z = z + c_int scaling_factor = coef[0] * scaling_factor**2 return z, scaling_factor def int_exp(x_int, scaling_factor): x0 = -0.6931 # -ln2 n = 30 # sufficiently large integer x0_int = torch.floor(x0 / scaling_factor) x_int = torch.max(x_int, n * x0_int) q = torch.floor(x_int / x0_int) r = x_int - x0_int * q exp_int, exp_scaling_factor = int_polynomial(r, scaling_factor) exp_int = torch.clamp(torch.floor(exp_int * 2**(n - q)), min=0) scaling_factor = exp_scaling_factor / 2**n return exp_int, scaling_factor x_int = x / scaling_factor x_int_max, _ = x_int.max(dim=-1, keepdim=True) x_int = x_int - x_int_max exp_int, exp_scaling_factor = int_exp(x_int, scaling_factor) exp_int_sum = exp_int.sum(dim=-1, keepdim=True) return exp_int, exp_int_sum def forward(self, x, scale): if self.log_i_softmax and scale is not None: exp_int, exp_int_sum = self.int_softmax(x, scale) softmax_out = torch.round(exp_int_sum / exp_int) rounds = self.log_round(softmax_out) mask = rounds >= 2**self.bit_type.bits qlog = torch.clamp(rounds, 0, 2**self.bit_type.bits - 1) deq_softmax = 2**(-qlog) deq_softmax[mask] = 0 return deq_softmax else: x = x.softmax(dim=-1) if self.calibrate: self.quantizer.observer.update(x) if self.last_calibrate: self.quantizer.update_quantization_params(x) if not self.quant: return x x = self.quantizer(x) return x
4、实验
4.1、Image Classification on ImageNet
为了证明所提出的方法的有效性,作者在ImageNet中使用各种Vision Transformer进行了广泛的实验,即ViT、DeiT、Swin Transformer。
表1中报告了Top-1的准确度结果。很明显,目前所有的方法都无法捕捉到完全量化的Vision Transformer,而FQ-ViT做到了这一点,即使在注意力图上的比特非常低,也能实现几乎无损的量化。同时,FQ-ViT显著超过了ViT的PTQ,其LayerNorm和Softmax没有量化。
例如,在所有模块量化为8位的情况下,FQ-ViT在DeiT-B上实现了81.20%的准确率,并且当注意力图压缩为4位时,它仍然可以实现80.85%的准确率。
4.2、Object Detection on COCO
作者还对目标检测基准COCO进行了实验。选择Swin系列检测器进行实验,结果见表2。可以观察到,所有当前的方法在完全量化的检测器上都具有较差的性能。FQ-ViT显著提高了量化精度,并在Mask R-CNN(Swin-S)上实现了47.2mAP,在Cascade Mask R-CNN上实现了50.8mAP,在权重/激活上实现了8位,在注意力图上实现了4位。
4.3、消融实验
为了研究本文中PTF和LIS的效果,作者在各种策略下对ViTB进行了充分量化,并在表3中报告了结果。作者设计了两个基线。”Baseline#8”索引模型由MinMax完全量化,具有8位权重、激活和注意力图,而“Baseline#4”在注意力图上具有较低的位宽(4位)。
从结果来看有几个观察结果。首先,具有PTF和LIS的模型优于基线,并实现了几乎无损的精度。其次,得益于BitShift运算符,PTF只引入了少量的BitOP,而LIS则减少了这一数量。
4.4、可视化结果
作者将注意力图可视化,以查看均匀量化和LIS之间的差异,如图5所示。当两者都使用8位时,均匀量化集中在高激活区域,而LIS在低激活区域保留更多纹理,这保留了注意力图的更多相对秩。在8位的情况下,这种差异不会产生太大的差异。然而,当量化到较低的位宽时,如6位和4位的情况所示,均匀量化会急剧退化,甚至使所有关注区域失效。相反,LIS仍然表现出类似于8位的可接受性能。
5、参考
[1].FQ-ViT: Post-Training Quantization for Fully Quantized Vision Transformer.