这是一篇北大与华为诺亚方舟实验室联合在2021年联合发表的论文,已经被2022年的cvpr收录。
github地址:https://github.com/hahnyuan/PTQ4ViT
ViT量化的难点:
如何找到合适的量化间隔?
如何在量化时保持self-attention map的相对位置?
是否可以混合精度地量化?
第一个和第三的问题更具一般性,作者提出的trick理论上来说可以泛化到其它网络中。
标准的transformer层有两个模块:Multi-head self attention(MSA)和Multi-layer perceptron(MLP):
其中:
作者采用传统的做法,会将所有与矩阵乘相关的weight和input都量化,不量化softmax和normalization。
文章创新点:
1. Ranking aware:作者将寻找最佳的量化间隔的问题建模成一个优化问题,不同于普通的量化优化目标是量化前后的相似度指标,优化的目标更加复杂,还加入了一个自定义的ranking loss。
ranking loss是用来保持量化后attention map的相对顺序:
相似度是用Pearson correlation coefficient来衡量:
在搜索空间(已经通过线性划分)中寻找最佳的过程是一个迭代过程:
固定input的量化间隔,调整weight的量化间隔
→固定weight的量化间隔,调整input的量化间隔
→直到目标函数收敛或迭代结束
2. Bias correction:量化之后还要对bias进行相应的较准,这一步可以用较准集的数据来计算
3. Nuclear Norm Based Mixed-Precision Quantization:还可以利用attention map和MLP的输入的核范数来进行混合精度量化。
根据所有候选量化bit数可能引起的total second-order perturbation来选择最小的量化bit数:
其中,tr()是SVD分解的表示:
把这三个方法结合在一起用,就是文章提出的后量化方法:
实验结果: