重训练量化（改进QAT）
原文链接：https://www.yuque.com/yahei/hey-yahei/quantization-retrain_improved_qat
欢迎引用&转载，但烦请注明出处~
Quantize Aware Training(QAT)通过在训练过程中融入量化和反量化过程，来实现量化模型的精度恢复，但考虑一下量化过程
w
q
=
q
(
w
)
=
α
⋅
C
l
a
m
p
(
R
o
u
n
d
(
w
α
)
)
w_q = q(w) = \alpha \cdot Clamp(Round(\frac{w}{\alpha}))
wq=q(w)=α⋅Clamp(Round(αw))
∂
L
∂
w
=
∂
L
∂
w
q
⋅
∂
w
q
∂
w
≈
S
T
E
∂
L
∂
w
q
⋅
1
=
∂
L
∂
w
q
\frac{\partial L}{\partial w} = \frac{\partial L}{\partial w_q} \cdot \frac{\partial w_q}{\partial w} \mathop{\approx} \limits_{STE} \frac{\partial L}{\partial w_q} \cdot 1 = \frac{\partial L}{\partial w_q}
∂w∂L=∂wq∂L⋅∂w∂wqSTE≈∂wq∂L⋅1=∂wq∂L
显然取整的
R
o
u
n
d
(
⋅
)
Round(\cdot)
Round(⋅)求导的时候处处为零，是无法训练的。因此QAT需要引入Straight-Through Estimator(STE)来得到一个近似的梯度，也即令
∂
R
o
u
n
d
(
x
)
∂
x
≈
1
\frac{\partial Round(x)}{\partial x} \approx 1
∂x∂Round(x)≈1，相当于用
y
=
x
y=x
y=x去近似
y
=
R
o
u
n
d
(
x
)
y=Round(x)
y=Round(x)或
y
=
q
(
x
)
y=q(x)
y=q(x)来进行求导。通常，如果比特数比较高的时候，这里带来的误差并不明显；但在低精度场景下，STE会导致训练不稳定。所以一直有研究者试图通过避免或优化STE来改进QAT。
Alpha-Blending(AB)
论文：《Learning low-precision neural networks without Straight-Through Estimator(STE) (IJCAI2019)》
AB算法将全精度权重跟量化权重加权混合得到新的权重用于训练，反向传播时，量化权重的链路不采用STE，故导数为0，因此只会沿着全精度权重的链路反向传播。
w
′
=
(
1
−
α
)
w
+
α
w
q
w' = (1-\alpha)w + \alpha w_q
w′=(1−α)w+αwq
∂
L
∂
w
=
∂
L
∂
w
′
∂
w
′
∂
w
=
∂
L
∂
w
′
(
(
1
−
α
)
+
α
∂
w
q
∂
w
)
=
(
1
−
α
)
∂
L
∂
w
′
\frac{\partial L}{\partial w} = \frac{\partial L}{\partial w'} \frac{\partial w'}{\partial w} = \frac{\partial L}{\partial w'} \left( (1-\alpha) + \alpha \frac{\partial w_q}{\partial w} \right) = (1-\alpha) \frac{\partial L}{\partial w'}
∂w∂L=∂w′∂L∂w∂w′=∂w′∂L((1−α)+α∂w∂wq)=(1−α)∂w′∂L
随着训练的进行，
α
\alpha
α逐渐从0增长到1，增长的方式可以有很多种，比如
α
=
1
−
e
−
λ
⋅
s
t
e
p
\alpha = 1 - e^{-\lambda \cdot step}
α=1−e−λ⋅step
或
α
=
{
0
s
t
e
p
≤
T
0
1
−
(
T
1
−
s
t
e
p
T
1
−
T
0
)
T
0
<
s
t
e
p
≤
T
1
\alpha = \begin{cases} 0 & step \leq T_0 \\ 1 - (\frac{T_1 - step}{T_1 - T_0}) & T_0 < step \leq T_1 \end{cases}
α={01−(T1−T0T1−step)step≤T0T0<step≤T1
并且学习率也随着
α
\alpha
α的增长而衰减
ϵ
=
ϵ
b
a
s
e
⋅
(
1
−
α
)
\epsilon = \epsilon_{base} \cdot (1-\alpha)
ϵ=ϵbase⋅(1−α)
（以下实验中，均采用最小化L2误差的交替优化方案）
（ImageNet实验）
QuantNoise
论文：《Training with Quantization Noise for Extreme Model Compression (2020)》
来自facebook的工作，算是借鉴了Dropout和增量式量化的思路，在训练过程中随机量化一部分权重而保持其他权重为全精度，以保证训练过程中总有一部分权重是不需要使用STE获得近似梯度。
（Dropout示意，训练时随机丢弃部分权重）
（增量式量化示意，训练时从绝对值较大的权重开始，逐步扩大量化范围）
（QuantNoise示意，训练时每次随机量化一部分权重）
PyTorch实现上非常简单，原本QAT的代码为
noise = quantize(w) - w
w_q = w + noise.detach()
QuantNoise只需要乘上一个随机生成的掩膜
noise = quantize(w) - w
w_q = w + noise.detach() * mask
论文里做了很多NLP的实验，可惜我对NLP了解的不多，做不出评判。CV也就跑了EfficientNet-B3，诶就不能跑下MobileNet吗，跑EfficientNet不太好跟其他论文直接比较。
不过看起来EfficientNet int4掉点有点厉害啊（这里的int4应该是同时将权重和激活都量化成int4了）。简单地复现过论文，发现并不能跑出这样的结果，在同等条件下QAT总是比QuantNoise好而且收敛更快。论文其实隐瞒了很多实验细节，倒是指明了QuantNoise激活值是采用PyTorch的HistogramObserver来量化（HistogramObserver实现细节可以参见《后训练量化 - 数据收集方法 - 直方图 - 如何在直方图上计算L2误差？ | Hey~YaHei!》），但对QAT没做更多说明。我怀疑实验中QAT是采用了指数平滑平均的方式量化激活值，也做了几个简单的实验，发现QAT采用指数平滑平均量化激活值，QuantNoise采用HistogramObserver量化激活值时实验结果跟论文实验结果“惊人”地相似……于是跑到github项目上（只开源了一小部分代码）提了issue不过没有得到答复，感觉QuantNoise的实际效果存疑。
全精度辅助模块
论文：《Training Quantized Neural Networks with a Full-precision Auxiliary Module (CVPR2020)》
AB算法混合了量化权重和全精度权重来训练，那么为什么不干脆分成两个分支，共享一套权重，一个分支走量化权重、另一个分支走全精度呢？这样在反传的时候，就既有来自全精度分支的梯度、也有来自量化分支的梯度。
这种想法比较朴素，训练代价（主要是训练时间）相比于原始模型就增加了一倍。而论文则是采用辅助模块来做量化补偿以实现这种想法：
（设计思路）
（如果觉得上一张图比较绕，我们可以试试换种画法）
蓝色部分的block采用量化的模型，每个block的输出（加和之前的特征图）经过Adaptor来做一定的量化补偿后加到全精度分支里，最后分别求loss，两者加和后反传和更新梯度，此时权重梯度既能来源于量化分支又能来源于“全精度”分支。论文里Adaptor采用的是一个1x1卷积，更复杂的设计会带来一些细微的提升，但1x1卷积基本就够用（后边有实验数据）。
这个思路跟辅助任务、知识蒸馏有点像，论文里还特地比较了一下
（辅助任务）
（知识蒸馏）
辅助任务可能存在多个分类器，挑选辅助分类器插入的位置会比较麻烦，而且跟引入全精度分支的思路相比还是差了那么点意思；知识蒸馏两分支则完全独立，训练时用全精度网络来监督量化网络的训练，但训练代价会比较大，再有就是知识蒸馏的是否需要加权，加多大的权重也是需要微调的。
实验细节
量化均采用QIL的算法，保持BN层统计量的更新；对于RetinaNet，各个head之间不共享，论文给出的解释是，在量化模型中，不同尺度特征图语义信息不能被很好的编码，不共享head虽然会增加参数量，但却不会增加计算代价。
（有无辅助模块的比较实验，ImageNet）
（与辅助任务、知识蒸馏方法的比较实验，2bit，ImageNet）
（Adaptor设计比较，2bit，ImageNet）
（3x3卷积稍微好一点点，但不明显）
（与FQN方法比较，4bit，COCO）
（共享head，4bit，COCO）
（Backbone only指的是只量化Bakbone）