Extremely Low Bit Neural Network: Squeeze the Last Bit Out with ADMM

Abstract

在DeepLearning日益发展的同时，模型压缩的关注度也越来越大。继BWN和TWN之后， 这篇论文在超低比特量化领域的一篇新思想文章，发表在AAAI2018上，作者是阿里的。
该文的主要思想是将超低bit quantization建模成一个离散约束优化问题(discretely constrained optimization problem)。借助ADMM(Alternating Direction Method of Mutipliers)思想，将连续参数从网络的离散约束中分离出来，并将原来的难题转化为几个子问题。该文针对这些子问题采用了extragradient以及iterative quantization算法（比传统算法更快收敛）。实验在图像分类以及目标检测(SSD)上都比最新的方法要好。

解决了什么？为什么写这篇文章？

• 之前的工作将pretrained weights量化到4-12bit效果很好，但是当只用1bit或者2bit表示weights时，只是在小datasets(MNIST、CIFAR10)上表现很好，在大datasets通常会产生很大的loss。
• 该文提出了一个独特的strategy，将超低bit weights作为一个离散约束的非凸优化问题(MIP)。借助ADMM算法，本文主要思想是使用离散空间中的辅助变量将离散约束中的连续变量解耦。不像之前量化方法修改连续参数中的特定梯度，本文同时优化连续和离散空间，使用增广拉格朗日函数连接了两个空间的解。

ADMM算法

主要参考这篇博客http://mullover.me/2016/01/19/admm-for-distributed-statistical-learning/
若优化问题可以表示为

minf(x)+g(z)s.t.Ax+Bz=c(1)$$minf(x)+g(z)s.t.Ax+Bz=c(1)$$

其中x∈Rs,z∈Rn,A∈Rp×n,B∈Rp×mandc∈Rp.$x\in {\mathbb{R}}^s,z\in {\mathbb{R}}^n,A\in {\mathbb{R}}^{p\times n},B\in {\mathbb{R}}^{p\times m}andc\in {\mathbb{R}}^p.$ x$x$与z$z$是优化变量，f(x)+g(z)$f(x)+g(z)$是目标函数，Ax+Bz=c$Ax+Bz=c$是等式约束条件。
等式(1)$(1)$的增广拉格朗日函数(augmented Lagrangian)可以表示为：

Lρ(x,z,y)=f(x)+g(z)+yT(Ax+Bz−c)+(ρ/2)∥Ax+Bz−c∥22(2)$$L_{\rho }(x,z,y)=f(x)+g(z)+y^T(Ax+Bz-c)+(\rho /2)\Vert Ax+Bz-c{\Vert }_2^2(2)$$

其中y$y$是拉格朗日乘子, ρ>0$\rho >0$是惩罚参数。增广就是由于加入了二次惩罚项。
则ADMM由三步迭代组成：

xk+1:=argminxLρ(x,zk,yk)$$x^{k+1}:=arg\underset{x}{min}{L}_{\rho }(x,z^k,y^k)$$

zk+1:=argminzLρ(xk+1,z,yk)$$z^{k+1}:=arg\underset{z}{min}{L}_{\rho }(x^{k+1},z,y^k)$$

yk+1:=yk+ρ(Axk+1+Bzk+1−c)$$y^{k+1}:=y^k+\rho (A{x}^{k+1}+B{z}^{k+1}-c)$$

可以看出，每次迭代分为三步:
1.求解与x$x$相关的最小化问题，更新变量x$x$
2.求解与x$x$相关的最小化问题，更新变量x$x$
3.更新ρ$\rho$

Objective function

记作f(W)$f(W)$为一个NN的loss fuction, W=$W=${W1,W2,...,WL$W_1,W_2,...,W_L$}, Wi$W_i$代表网络的第i$i$层的权重，Wi$W_i$是一个di$d_i$维的向量Rdi${\mathbb{R}}^{d_i}$, W∈Rd,d=∑idi$W\in {\mathbb{R}}^d,d=\sum _i{d}_i$.本文中由于是训练超低bit进行量化，特别的，网络中的weights被严格限制为0或者是2的幂次方，这样float-point操作就可以被更快的bit移位操作所代替。假设training一个三元网络{−1,0,+1$-1,0,+1$},这样就training这样的网络就可以在数学公式上作为MIP：

minWf(W)s.t.W∈C={−1,0,+1}d$$\underset{W}{min}f(W)s.t.W\in C=\{-1,0,+1{\}}^d$$

由于weights被严格限制为0或者2的幂次方，所以C={−2N,...−2−1,−20,0,+20,+21,...,+2N}$C=\{-2^N,...-2^{-1},-2^0,0,+2^0,+2^1,...,+2^N\}$, N是bits的数量。在这里引入一个scaling factor参数α$\alpha$来限制C$C$,C={...,−2α,−α,0+α,+2α,...}$C=\{...,-2\alpha ,-\alpha ,0+\alpha ,+2\alpha ,...\}$.α>0$\alpha >0$.这里应该注意的是不同层的α$\alpha$是不同的。换句话说，对于一个L$L$层的网络来说，实际是引入了L个不同的scaling factors {α1,...,αL}$\{{\alpha }_1,...,{\alpha }_L\}$。这样低比特量化网络的目标函数可以表示为:

minWf(W)s.t.W∈C=C1×C2×...×CL(3)$$\underset{W}{min}f(W)s.t.W\in C=C_1\times C_2\times ...\times C_L(3)$$

Ci={0,±αi,±2αi,...,±2Nαi},αi>0.$C_i=\{0,\pm {\alpha }_i,\pm 2{\alpha }_i,...,\pm 2^N{\alpha }_i\},{\alpha }_i>0.$其实αi${\alpha }_i$在每层中是不需要花更多的计算量的，由于可以在最后乘以这个因子就可以。αi${\alpha }_i$可以帮助扩展constraint space.如下图，

在二维的空间中，上图是一个三元网络，被限制为{−1,0,+1}$\{-1,0,+1\}$，可能的解是九个离散点，当增加α$\alpha$因子之后，约束空间被扩展为四条线。这种约束空间的large expansion可以是优化变得更容易。

Decouple with ADMM

公式(3)其实是一个NP难问题，由于weights被限制在一个离散的空间。这里引入一个辅助变量，受离散空间限制，让它等于原始变量。首先，定义一个indicator function IC,W∈C$I_C,W\in C$, 目标函数(3)变成

minWf(W)+IC(W)(4)$$\underset{W}{min}f(W)+I_C(W)(4)$$

其中如果W∈C$W\in C$, IC(W)=0$I_C(W)=0$, 否则IC(W)=+∞$I_C(W)=+\mathrm{\infty }$
引入辅助变量G$G$，重写公式(4),

minW,Gf(W)+IC(G),s.t.W=G(5)$$\underset{W,G}{min}f(W)+I_C(G),s.t.W=G(5)$$

,公式(5)的增广拉格朗日函数为：

Lρ(W,G,μ)=f(W)+IC(G)+(ρ/2)∥W−G∥2+<μ,W−G>(6)$$L_{\rho }(W,G,\mu )=f(W)+I_C(G)+(\rho /2)\Vert W-G{\Vert }^2+<\mu ,W-G>(6)$$

λ=(1/ρ)μ$\lambda =(1/\rho )\mu$,公式(6)可以被转化为：

Lρ(W,G,μ)=f(W)+IC(G)+(ρ/2)∥W−G+λ∥2−(ρ/2)∥λ∥2(7)$$L_{\rho }(W,G,\mu )=f(W)+I_C(G)+(\rho /2)\Vert W-G+\lambda {\Vert }^2-(\rho /2)\Vert \lambda {\Vert }^2(7)$$

根据ADMM，这个问题的求解迭代步骤:

Wk+1:=argminWLρ(W,Gk,λk)$$W^{k+1}:=arg\underset{W}{min}{L}_{\rho }(W,G^k,{\lambda }^k)$$

Gk+1:=argminWLρ(Wk+1,G,λk)$$G^{k+1}:=arg\underset{W}{min}{L}_{\rho }(W^{k+1},G,{\lambda }^k)$$

λk+1:=λk+Wk+1−Gk+1$${\lambda }_{k+1}:={\lambda }^k+W^{k+1}-G^{k+1}$$

这三部分别是proximal step, projection step以及dual update.跟之前的工作不同的是，这里是同时在连续空间以及离散空间优化问题，这样在learning过程中两个空间的解通过ADMM算法就联系在一起了。

Algorithm subroutines

• 1.Proximal step
  对于这个步骤，在连续space优化（？？为什么是连续空间，W和G都是离散空间的啊~想不通），需要最小化下面这个公式
  
  Lρ(W,Gk,λk)=f(W)+(ρ/2)∥W−Gk+λk∥2$$L_{\rho }(W,G^k,{\lambda }^k)=f(W)+(\rho /2)\Vert W-G^k+{\lambda }^k{\Vert }^2$$
  
  利用standard梯度下降方法可以得到W$W$的梯度：
  
  ∂WL=∂Wf+ρ(W−Gk+λk)$${\mathrm{\partial }}_{W}L={\mathrm{\partial }}_{W}f+\rho (W-G^k+{\lambda }^k)$$
  
  然而该文发现用venilla梯度下降方法收敛很慢，因为第二个平方项在整个lost中所占比重很大，因此SGD迅速将优化器拉回到当前量化的权重以至于第二项消失(Since the second quadratic term occupies a large proportion of the whole lost, SGD will quickly pull the optimizer to the currently quantized weights so that the second term vanishes, and stack in that point没太懂?)。这就导致了loss的次优解，由于网络的损失函数没有得到充分的优化。
  为了克服这个困难，本文采用了extragradient方法。extragradient方法的一次迭代由两个简单的步骤组成，prediction和correction:
  
  W(p):=W−βp∂WL(W)$$W^{(p)}:=W-{\beta }_p{\mathrm{\partial }}_{W}L(W)$$
  
  
  W(c):=W−βc∂WL(W(p))$$W^{(c)}:=W-{\beta }_c{\mathrm{\partial }}_{W}L(W^{(p)})$$
  
  其中βp${\beta }_p$和βc${\beta }_c$都是学习率。extragradient方法的显著特点是使用了额外的梯度步骤，这可以当做在优化过程中的一个guide。这种额外的迭代考虑了曲率信息，导致比stardard梯度下降方法更好的收敛。对于prediction step，算法将快速move到接近Gk−λk$G^k-{\lambda }^k$值的一个点进而导致了平方项消失；对于correction step，算法最小化loss函数f(W)$f(W)$。这两个步骤就避免陷入局部最小。在实践中，这种方法极大地加速了收敛。
  2.projection step
  对于辅助变量G$G$来说，所有Gi$G_i$都被解耦,因此每层的辅助变量Gi$G_i$都可以被独立的被优化。寻找一种Euclidean映射可以将(Wk+1i+λki)$(W_i^{k+1}+{\lambda }_i^k)$映射到离散的Ci$C_i$上。这里记作(Wk+1i+λki)$(W_i^{k+1}+{\lambda }_i^k)$为Vi$V_i$.这种映射可以被表示为：
  
  minGi,αi∥Vi−Gi∥2$$\underset{G_i,{\alpha }_i}{min}\Vert V_i-G_i{\Vert }^2$$
  
  
  s.t.Gi∈{0,±αi,±2αi,...,±2Nαi}di$$s.t.G_i\in \{0,\pm {\alpha }_i,\pm 2{\alpha }_i,...,\pm 2^N{\alpha }_i{\}}^{d_i}$$
  
  将scaling因子拿掉，变为：
  
  minGi,αi∥Vi−αi⋅Qi∥2$$\underset{G_i,{\alpha }_i}{min}\Vert V_i-{\alpha }_i\cdot Q_i{\Vert }^2$$
  
  
  s.t.Qi∈{0,±1,±2...,±2N}di$$s.t.Q_i\in \{0,\pm 1,\pm 2...,\pm 2^N{\}}^{d_i}$$
  
  本文提出了一种量化迭代方法，该方法交替优化αi${\alpha }_i$和Qi$Q_i$，一个固定另一个变。当Qi$Q_i$固定的时候，这个问题变成了一个单变量优化，αi=VTiQiQTiQi${\alpha }_i=\frac{V_i^T{Q}_i}{Q_i^T{Q}_i}$。当αi${\alpha }_i$固定，Qi$Q_i$实际是Viαi$\frac{V_i}{{\alpha }_i}$在{0,±1,±2,...,±2N}$\{0,\pm 1,\pm 2,...,\pm 2_N\}$上的投影，Qi=∏{0,±1,±2,...,±2N}(Viαi)$Q_i=\prod _{\{0,\pm 1,\pm 2,...,\pm 2_N\}}(\frac{V_i}{{\alpha }_i})$。

实验

1.Image Classification

可以看出无论是在哪个网络，本文超低比特量化都最突出。
需要着重看的是，为什么GoogleNet会降低的比较多，本文的猜测+实验是由于1×1$1\times 1$的卷积核用这样强的一个调节器来优化会导致网络的欠拟合。在网络的不同部分用不同的bit来量化会更为有效。

2.Object Detection