论文笔记: Extremely Low Bit Neural Network: Squeeze the Last Bit Out with ADMM

读了AAAI18年发表的《Extremely Low Bit Neural Network: Squeeze the Last Bit Out with ADMM》这篇论文,收获很大。这篇论文是阿里Rong Jin团队(现在是达摩院机器智能研究团队)的论文,金榕是CMU的博士,曾任密歇根州立大学终身教授,15年到阿里。第一作者是Cong Leng。

文章把神经网络的量化问题看成是一个优化问题,使用了ADMM方法把不连续的量化和连续的参数优化分开来解,避免了采用STE方法中因为前向和反向使用的weight不是不同所导致的不稳定,采用ADMM方法量化后的效果很好。

方法

N-bit的对称量化的问题formulate为:
$$\min_{W} f(W,X)\quad s.t.\quad W \in C= \{ -2^{N-1}-1,…,-1,0,1,…,2^{N-1} \} $$
在训练的时候,需要给每一层的$W_i$乘上一个scaling factor $\alpha_i$,C就变成$C=C_1\times C_2 \times … \times C_L$
定义一个指示函数$I_C(W)$,当$W \in C$等于0,否则为正无穷。问题可以转化为:
$$\min_{W} f(W)+I_C(W)$$
引入一个辅助变量 G,问题转化为:
$$\min_{W,G} f(W)+I_C(G)\quad s.t. W=G$$
这个问题就是ADMM需要的形式了,引入拉格朗日乘子$\mu \in R^d$,其增广拉格朗日方程为:
$$L_\rho (W,G,\mu)=f(W)+I_C(G)+\rho/2||W-G||^2+\mu^T(W-G)$$
对朗格朗日乘子进行变量变换$\lambda=(1/\rho)\mu$,得到
$$L_\rho (W,G,\mu)=f(W)+I_C(G)+\rho/2||W-G+\lambda||^2-\rho/2||\lambda||^2$$
就可以轮流迭代的解下面三个方程来得到量化后的权重G:
$$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}$$

第一个问题:排除不变的变量,等价于优化:
$$\arg\min_{W} f(W)+\rho/2||W-G+\lambda||^2$$
相当于一个神经网络训练,加上一个特殊的正则化因子

第二个问题:排除不变的变量,等价于优化:
$$\arg\min_{G} I_C(G)+\rho/2||W-G+\lambda||^2$$
进一步转化为:
$$\arg\min_{G}||W+\lambda-G||^2 s.t. G\in C$$
就是一个最小化量化到$W+\lambda$的误差的问题,论文对每一层单独解,使用5次交替优化$\alpha$和$Q$

结果

  • 在resnet-50上,binary 比BWN高了5个点,到了68.7%的top1;ternary比TWN高了7个点,到了72.5%的top1
  • 在VOC2007上,Darkent+SSD ternary到了60.9,低了3个点;VGG16+SSD ternary到了76.2mAP,低了2个点
    • 1x1卷积超低bit量化对网络的性能有较大影响,如果使用8bit,darknet在ternary下降2个点

评价

这篇论文第一次把ADMM方法引入进超低bit的神经网络量化中来,取得了非常不错的效果。也第一次把object detection做了超低bit的量化,并且分析到了1x1卷积的作用,整个文章非常solid,很有实用性。但文章并没有分析为什么ADMM就能很好地用在这种非凸优化的量化任务上来;为什么STE的方法就不稳定了,如果提升稳定性的STE能比得过ADMM吗?对activation的量化也没有做,后续应该有相应的工作值得调研。

Your browser is out-of-date!

Update your browser to view this website correctly. Update my browser now

×