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

前言

• Extremely Low Bit Neural Network: Squeeze the Last Bit Out with ADMM - 原文 ；AAAI 18 Oral。
• 从结果看比之前的一些量化工作都要好上不少，干货较多。

量化(Quantization)

为什么要压缩网络

• 网络的表达能力过强，压缩网络可以起到一定正则的作用
• 网络存在很多冗余参数和计算量，压缩网络可以极大减少储存空间以及运算量

什么是量化

• 对权重的参数进行取整，降低存储精度。低精度数不仅减少存储空间，而且能够规避浮点运算带来的巨大计算量。经过特别设计的芯片还能对低精度数的运算（可用速度更快的位运算）进行更进一步的优化。
• 有几种方法降低精度，常见的方法如直接对权值作聚类分析（更进一步可以对每一个cluster进行Huffman编码，继续压缩存储空间）或是直接添加整数约束/改变激活函数等

• 也可以对梯度和输出量化；对于不同的方法，可能需要专门设计训练过程来计算梯度。

ADMM(Alternative Direction Multiplier Method)

• 因大规模网络的应用，近年来颇受关注的一种优化方法。虽然ADMM在理论分析中收敛慢，但是在许多实际的大规模应用场景中表现稳定，且较为适合分布式计算。
• 更进一步的对ADMM的讨论可见另一篇文章

论文

• 论文实现了将32比特权重压缩至3比特的低比特权重量化。之前的若是要将权重量化到3或更低的比特上的话，在大数据集上效果并不好。
• 论文只介绍了思路，并没有源码

实现算法

• 将离散权重网络训练定义成离散约束问题，以三值网络为例，约束为： $$\underset{W}{\min}\ f(W)\ \ s.t.W\in \mathcal{C}=\{1,0,-1\}^d$$
• 引进 scale 参数扩展可行域，将约束条件扩展至 $W\in \mathcal{C}={a,0,-a}$ 。由于约束空间变大（如下图），收敛速度将大大加快。
  • 注意这里之后再 scale 回去严格二值的话，实际上是会损失精度，这样量化出来的BWN只能算作变种（参考XNOR Net）
  • 个人有点怀疑这样部署到端上后的效果。有些端不支持高位数的scale参数的存储和计算。

• 使用ADMM算法一般约束下的形式，将问题转换为：

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

  其中 $I_C$ 为指示函数，$G$ 符合约束则为0，否则为无穷大。

  其拉格朗日函数形式为：

  $$L_\rho(W,G,\lambda)=f(W)+I_C(G)+\frac{\rho}{2}||W-G+\lambda||^2-\frac{\rho}{2}||\lambda||^2$$

  迭代过程：

  $$\begin{align} W^{k+1}&:=\arg\min_W L_\rho(W, G^k, \lambda^k)\\ G^{k+1}&:=\arg\min_G L_\rho(W^{k+1}, G, \lambda^k)\\ \lambda^{k+1}&:=\lambda^k+W^{k+1}-G^{k+1} \end{align}$$

• 对第一个对 $W$ 更新的过程，采用 Extra-gradient 算法进行梯度下降能够避免收敛过慢的问题：

  $$\begin{align} W^{(p)}&:=W-\beta_p\partial_WL(W)\\ W^{(c)}&:=W-\beta_c\partial_WL(W^{(p)}) \end{align}$$
  • Extra-gradient method 即外梯度算法，由 Korpelevich 提出（可见原文引文[23]）。这篇论文距今已经40年了，既难找又难懂，感觉又是个新坑……涉及变分不等式的部分可以和ADMM一起又单写一篇文章了。
  • （学习完凸优化后补充）其实也不是很复杂，如果不关心数学图景上的原理的话，可以看做用下一步的梯度来取代这一步的梯度做梯度下降。单步计算量和存储量都会更高，但是会获得相对更正确的下降方向。
• 对第二个对 $G$ 更新的过程，我们可以将其看做在非凸集上整数约束的优化问题。令 $W_i, G_i, \lambda_i, \mathcal{C}_i$ 分别为第 $i$ 层权重、辅助变量、拉格朗日乘子和约束集合。根据ADMM算法，迭代时需要计算在 $\mathcal{C}_i$ 上 $W^{k+1}_i+\lambda^k_i$的欧式投影。令 $V_i=W^{k+1}_i+\lambda^k_i$ ，其欧式投影可以表示为： $$\min_{G_i,\alpha_i} ||V_i-G_i||^2$$ $$s.t. G_i\in\{0,\pm\alpha_i, \pm 2\alpha, \cdots, \pm 2^N\alpha\}^{d_i}$$ 从约束中分离 scale 参数： $$\min_{Q_i,\alpha_i} ||V_i-\alpha_i\cdot Q_i||^2$$ $$s.t. Q_i\in\{0,\pm, \pm 2, \cdots, \pm 2^N\}^{d_i}$$ 算法交替优化 $Q_i$ 和 $\alpha_i$ 中的其中一个而固定另外一个：
  • 在 $Q_i$ 固定时： $$\alpha_i=\frac{V^T_i Q_i}{Q^T_i Q_I}$$
  • 在 $\alpha_i$ 固定时： $$Q_i=\Pi_{\{0,\pm 1,\pm 2,\cdots,\pm 2^N\}}(\frac{V_i}{\alpha_i})$$ 对 $Q_i$ 的更新，$\Pi$ 投影只是简单地取离约束集合最近的点（参见ADMM笔记4.3.5）。
• 这种迭代算法能保证很快收敛到局部极小值。
• 最后迭代停止时，直接用 $G$ 取代 $W$ 即可。

论文结果

• 算法在AlexNet、VGG、ResNet、GoogleNet上均做了测试，2~3bit的量化相比state-of-art 的一些二值网络与三值网络高出有1~4%的Top5正确率，只比全精度的baseline低1%多一些。