Noise injection and clamping estimation for neural network quantization

Noise injection and clamping estimation for neural network quantization

文章鏈接

代碼鏈接

Introduction

主要工作：提出了NICE(noise injection and clamping estimation)量化方式，主要分爲以下兩個步驟：

1：Noise injection during training that emulates the quantization noise introduced at inference time

2：Statistics-based initialization of parameter and activation clamping, for faster model convergence. In
addition, activation clamp is learned during train time

這種方法的好處在於不需要在訓練的時候，修改網絡，並且它將網絡中的所有參數量化爲fixed point(integer)。

Related work

Expressiveness based methods ：如使用更大的網絡或者加一個線性的尺度因子來彌補精度的損失。

Keeping full-precision copy of quantized weights：使用兩份權值，一份是量化後的，一份是full precision的。前向時使用quantized weight，更新時使用full precision weight，如使用STE(straight-through estimator)來估算梯度，使用隨機(stochastic)或者確定(deterministic)的方式量化參數。

Distillation：使用distillation。用student network來學習teacher network的輸出，還挺有用的。

Model parametrization 、Optimization techniques 、Generalization bounds

Method

文章並不量化網絡的第一層和最後一層，因爲這兩層對網絡的性能影響很大。

這種方法的優勢在於它在反向更新之後能迅速影響前向，相對於那些直接量化權重的策略而言，它們在反向時幾乎不會更新來量化的權重。

爲了實現一個類似drop-out的效果，作者使用了Bernoulli分佈的一個mask$M$，來量化部分權值並向其他部分添加噪聲。經驗取$M \sim Ber(0.05)$。並且，在前向時，使用的量化權重爲：$\hat{w}=(1-M)\mathcal{Q}_{\Delta}(w)+M(w-e)$，其中$e \sim Uni(-\Delta/2, \Delta/2)$，$\Delta$表示quantization bin的大小。

Gradual Quantization

文章是先量化一部分參數，然後讓剩下的參數adapt to the changes。對於梯度的量化，首先將網絡的層分爲$N$個相同大小的block$\{B_1,\ldots,B_N\}$，然後在第$i$ 階段(at the i-th stage)，將噪聲從$B_i$塊注入，它前面的blocks $\{B_1,\ldots,B_{i-1}\}$都量化掉，後面的blocks ${B_{i+1},\ldots,B_{N}}$都保持full precision，進行一次梯度更新，當完成$N$個過程後，量化所有的層並使用STE方式訓練。

這樣訓練就可以使得後面的層可以根據前面量化後的層的輸入值的變化進行相應的調整，網絡不會一下子變化特別大。之後的使用STE訓練是爲了使網絡收斂，對於一個用來量化的已經訓好的網絡，作者發現最好的block size爲帶activation的一個single layer，但是使用上面的方式並沒有使效果變得更好。

Clamping and Quantization

爲了量化網絡的權重，將權值限定在$[-c_w,c_w]$：
$w_c=Clamp(w,-c_w,c_w)=max(-c_w,min(x, c_w)) \tag{1}$式中，對每一層定義一個$c_w$，並將其初始化爲$c_w=\operatorname{mean}(w)+\beta\times \operatorname{std}(w)$，其中$w$是該層的權重，$\beta$是一個超參數，對於一個給定的$c_w$，量化限定後的權重到$B_w$ bits：
$\hat{w}=\left[w_{c} \frac{2^{B_{w}-1}-1}{c_{w}}\right] \frac{c_{w}}{2^{B_{w}-1}-1} \tag{2}$式中，$\left[\cdot\right]$表示rounding操作。同樣的，激活值的量化也是一同樣的方式。如ReLU被替換成clamped ReLU
$a_c=\operatorname{Clamp}(a,0,c_a) \tag{3}$其中，$a$表示某一層的線性部分的輸出，$c_a$是限定的範圍，$a_c$表示在量化之前截斷的一個非負的值，$c_a$是一個集合，每一層都有一個$c_a$，並通過反向傳播學習更新。同樣的，量化$a_c$到$B_a$ bits的方式也是
$\hat{a}=\left[a_{c} \frac{2^{B_{a}-1}-1}{c_{a}}\right] \frac{c_{a}}{2^{B_{a}}-1} \tag{4}$由於$\operatorname{Round}$函數是不可微的，因此使用STE進行反向傳播，對於$c_a$的更新，可以計算$\hat{a}$對$c_a$的偏導爲：
$\frac{\partial \hat{a}}{\partial a_{c}}= \left\{ \begin{array}{ll} {1,} & {a_{c} \in\left[0, c_{a}\right]} \\ {0,} & {\text { otherwise }} \end{array} \right. \tag{5}$對於每一層的biases的量化就更爲複雜，因此他們的尺度取決於activations和weights尺度。對於每一層，將截取偏差的值並初始化爲：
$c_{b}=(\underbrace{\frac{c_{a}}{2^{B_{a}}-1}}_{\text {Activation scale }} \cdot \underbrace{\frac{c_{w}}{2^{B_{w}-1}-1}}_{\text {Weight scale }}) \cdot(\underbrace{2^{B_{b}-1}-1}_{\text {Maximal bias value }}) \tag{6}$式中，$B_b$表示bias的bitwidth，bias的截斷和量化方式和weights相同。

Experiment Results

使用預訓練好的FP-32模型。然後用NICE進行量化，用PyTorch實現，在ImageNet和CIFAR-10上測試分類準確率，在the MSR joint denoising and demosaicing dataset上測試迴歸。

ImageNet

量化Resnet-18/34/50，使用NICE來fine-tune網絡，訓了120個epoch，學習率爲$10^{-4}$，momentum爲0.9，weight decay設爲$4\times 10^{-5}$，結果如下：

Ablation study

爲了探究NICE每一部分的重要性，作者在ImageNet上使用了Resnet-18。從結果可以看出，對於high bitwidth，即5.5這種，noise addition和gradual training比clamp learning貢獻更大。對於low bitwidth，結果則相反。

文章還對結果做了進一步的闡述，並且說明了在硬件上的加速，不是關注的重點，就不闡述了。