论文(ICLR 2020):https://arxiv.org/abs/2004.09576
暂无源码
LSQ+ 概要
Motivation
- LSQ 只适用于激活函数为 ReLU 的模型,即只适用于 Activation 为非负分布的模型,不适用于激活函数为 Swish\Leaky-ReLU 等的模型。
- 基于梯度的可学习量化方案其实不太稳定,需要更完美的超参数。
对于 Motivation 1,事实是 LSQ 只考虑 Activation 为非负分布,即所有的值均大于等于 0 的,那么理论上就不适用于 Swish\Leaky-ReLU 等非对称分布、存在负值的激活函数。在公式 (1)中, LSQ 设定 Activation 中 n = 0 , p = 2 b − 1 n=0, p=2^b-1 n=0,p=2b−1,这会将 Activation 中的负值全部 c l a m p clamp clamp 到 0 。
而在实验过程中,作者发现量化的有关参数的初始化还会对模型最终的收敛性能产生影响,引申出 Motivation 2。
Method
Proposal:Learnable asymmetric quantization
作者提出了一种通用的非对称激活量化方案,在训练过程中不仅学习量化尺度参数 s s s ,还学习量化偏移 (零点) 参数 β \beta β 来处理非对称分布的Activation,如公式 (2) 所示。
式中 n n n 为下界, p p p 为上界, c l a m p ( ) clamp() clamp() 为钳制函数, ⌊ ⌉ \lfloor~\rceil ⌊ ⌉ 为四舍五入 (近似取整) 函数。 x ˉ \bar{x} xˉ 是离散化的整数, x ^ \hat{x} x^ 是离散化的小数, s s s 和 β \beta β 均是可学习参数。于是用于更新参数 s s s 的梯度和更新参数 β \beta β 的梯度分别为:
在公式 (3) 和 (4) 中,作者均使用了STE。与对称量化相比,作者方案对Activation 的非对称量化在推理期间是不消耗额外成本的,因为额外的偏移项可以在编译时预计算并合并到 Conv 层的 Bias 中:
需要注意的是,作者提出的方案只针对 Activation,而对于 Weight 采用的方案仍和 LSQ 一致,即权重量化器没有量化零点 β \beta β ,默认 Weight 为对称分布的。
为了进一步做对照试验来验证 LSQ+ 的有效性,作者提出了 如 Table 1 所示的配置。其中 Config 3 和 Config 4 其实就是作者所提出的方案 (关键在于 β \beta β 是 tranable )。
Initialization of quantization parameters
作者发现在进行极低比特的量化时,QAT训练后的最终性能对量化参数的初始化高度敏感,特别是在模型有深度可分离卷积的情况下。于是作者进一步定义 s s s 和 β \beta β 的初始化方法。
- 对于权重量化器来说, s s s 的初始化没有采用 LSQ 的方法,而是改为:
- 对于激活量化器来说,以往业界公认的初始化方式为, s s s 和 β \beta β 考虑分布的 min 和 max 值,那么就有边界和对应的 s s s 和 β \beta β 堆导分别为公式 (6) (7) 。但实际考虑分布的 min、max 是不妥当的 (大家公认不考虑离群点) ,于是作者提出通过MSE来进行最优化搜索,即如公式 (8)。而公式 (8) 没有封闭形式的解,因此作者将公式 (3) (4) 嵌入到PyTorch的自动求导过程中,通过梯度下降优化几个 Batch 的数据上来获得 s i n i t , β i n i t {s_{init},β_{init}} sinit,βinit 。
Experiment
所有实验均在ImageNet数据集上进行,都用到了预训练模型。
Table 2 和 Table 3 展示的是各种 Config 对带有 Swish 激活函数的模型表现。可以看到作者自己的方案 Config 3、4 的表现还是较好的。
Table 4 展示的是各种 Config 对仅有ReLU激活函数的模型表现。可以看到作者的表现也是比较好的。
接着为了检验量化参数初始化对模型性能的影响,作者将自己的方法和传统的 Min-Max 初始化 (如公式 7 ) 、LSQ 初始化做了比较。Table 5 展示的是在 EfficientNet-B0 上不同初始化方法间的性能。可以看到作者的表现也可以。
接着作者讨论了一下 Initialization 的意义。如 Fig 1 所示,可以看到相比于 LSQ ,LSQ+的初始化方案会更接近于最终 tuning 的结果,也就是说这是有利于快速收敛的,之前的实验也证明性能更好一些。
最后作者讨论了下设置量化零点为可学习参数的意义。 Fig 2 是作者在 Config 4 下的 W4A4 实验,可以看到每一层的 β \beta β 参数都起到了作用 (大部分 β \beta β 非零),且各个层的 β \beta β 都不同,具有多样性。观察还可以看出大部分 β \beta β 为负值,因为在 Config 4 下需要把 Weight 分布右移。这说明 β \beta β 可学习还是非常合理的。最后作者还补充了 Table 6 来进一步验证 β \beta β 可学习的意义。
个人思考
该论文的创新点也是很纯粹简洁的,就是对 Activation 量化器使用了可学习的量化零点,并制定更合理的初始化。作者做的实验表现都挺好,只是没有给出源码。但本人在复现之后跑了几次 MobileNet 上的实验,发现 MobileNet 上 LSQ+ 的效果比 LSQ还要差,这我没想明白。