技术背景
深度学习最近在object recognition, speech recognition, machine translation, games, image generation等各个领域突飞猛进,都取得了state of the art的效果。越来越多的学术界以及工业界的研究人员将深度学习应用到传统的领域, 推动这些领域的变革,以及整个AI领域的发展。相对于传统的浅层学习, 深度学习模型复杂, 参数多,训练和Inference的计算量大。以AlexNet为例:
- 网络深度达到8层
- 参数规模超过 60M
- 一张图片的Inference需要约1G FLOPS
Object recognition中的其他重要的模型:VGG, GoogleNet , ResNet等都比AlexNet复杂, 参数规模有增减, 但计算量都增加了不少。比如100层的ResNet一张图片的Inference计算量超过了10G FLOPS。
这些问题给深度学习在服务器端、移动端以及其他嵌入式设备的应用带来挑战,主要集中在两个层面
- 参数规模大:影响模型的部署和更新,同时会降低Cache命中率,增加Inference时间
- Inference的latency过高: 计算量多大, AlexNet上一张图片的Inference在CPU上需要上百毫秒
这些都限制深度学习的发展。针对上述问题,一方面可以通过硬件来加速:Nvidia不停的发布加速深度学习训练、Inference的GPU,Google也推出了TPU等,以及很对团队研究通过使用FPGA来加速。另一方面, 可以通过
结合硬件结构, 设计针对深度学习模型的量化压缩、Inference加速算法。主要有:
- parameter pruning & sharing
- Low-rank factorization
- Transfered/compact convolutional filters
- Knowledge distillation
在阿里巴巴购物搜索引擎中, 深度学习模型得到了大量的应用, 在业务上取得了非常不错的效果。为了能在线上部署复杂的深度学习模型,降低latency, 我们做了大量的尝试:一方面优化GPU调用以及设计FPGA硬件来加速Inference,同时也在设计压缩、量化算法来加速Inference。这里我们主要介绍我们在模型量化压缩及Inference加速算法上的工作。该工作的论文Alternating Multi-bit Quantization for Recurrent Neural Networks已经被深度学习顶级会议ICLR 2018接收
方案
深度学习模型的量化压缩算法是目前的一个研究热点,每年在NIPS, ICLR等计算机顶级会议上都有很多相关的文章, 也获得了学术界和工业界的认可,比如Song Han的Deep Compression算法获得了2016年ICLR的“Best Paper Award"。在阿里巴巴购物搜索中, 我们训练了大量的深度学习模型中, 运行在阿里巴巴的服务器集群上, 使用CPU, GPU计算。由于GPU成本较高以及FPGA研发周期长, 而服务器集群有大量的CPU,因此我们设计了一套基于CPU的深度学习模型量化压缩与Inference加速算法AQN: Alternating multi-bit Quantzation Network, 并在搜索的相关业务上上线, 取得了非常不错的效果。
通用矩阵乘法(GEMM)是深度神经网络(DNN)的核心(https://petewarden.com/2015/04/20/why-gemm-is-at-the-heart-of-deep-learning/), 也是深度学习模型中参数最多, Inference计算最耗时的地方。大多数的深度学习框架中的卷积计算(CNN)也是使用GEMM来实现的, 以及循环神经网络(RNN)中包含了大量的GEMM计算。因此优化GEMM是最关键的一步。为了简单起见, 我们将GEMM简化为 $C = MatMul(A, B)$ ,其中$MatMul$是标准矩阵乘法。通常而言, $C$为深度学习网络中某个节点的输出, $A$为节点的输入, $B$为节点的参数, 模型的一部分。
为了优化GEMM,基于综合比较,我们选择了基于binary的量化算法, 并在此基础上提出了新的量化算法(AQN)。Courbariaux M等首先提出了BinaryConnect, 将GEMM的每个参数以及每个输入都量化成1个bit表示的$\{-1, 1\}$的算法,大幅压缩模型大小,但模型精度降低了很多。Mohammad Rastegari等在此基础上提出了XNOR Net,模型大小不怎么变化, 但相对BinaryConnect模型精度大幅提升。 为了介绍AQN, 我们首先介绍XNOR Net。对于矩阵 $A \in \Re^{m \times n} $ 的binary 量化过程为:
可以表示为 :
$$ A \approx \alpha \times A_b $$
其中:
$$ \alpha \in \Re \ and \ \alpha \geq 0 $$
$$ A_b \in \{-1, +1\}^{m \times n} $$
由于$A_b$中所有的element都是$\{+1, -1\}$, 因此存储该矩阵的时候可以用一个bit来表示矩阵中的一个element,当该bit为0时元素值为-1, 该bit为1时表示元素值为1。相比较于全精度(float), 参数的压缩比例约为32倍。参数$\alpha$以及$A_b$的推导过程为:
对于矩阵量化的误差为:
$$ error = \left\lVert A - \alpha A_b \right\rVert ^ 2 $$
最小化$error$可以得到:
$$ \alpha = \sum abs(A_{ij}) / (m*n) $$
$$ A_b = sign(A) $$
我们定义二值化的op为:
$$ \alpha, A_b=binarize(A) $$
将binary network应用到GEMM中。
$$ \alpha_A, A_b = binarize(A) $$
$$ \alpha_B, B_b = binarize(B) $$
$$ GEMM(A,B) \approx \alpha_A \alpha_B B\_GEMM(A_b, B_b) $$
$B\_GEMM$表示二值矩阵的乘法。因此问题转化为二值矩阵的乘法。 针对二值矩阵的乘法高效计算可以采用XNOR+POPCNT算法(POPCNT统计值为1的bit数)。首先假设$A_b$中的每个元素已经通过bit存储 ,那么矩阵乘法可以表示为
$$ C_{ij}=K-2POPCNT(XNOR(A_{i,}, B_{,j})) $$
在$binarize$ op中, 参数$\alpha = \sum(A_{ij})/ (m * n)$, 是整个矩阵元素绝对值的平均值。一个优化的方法可以考虑根据原始矩阵$A$的每行计算一个系数$\alpha$, 或者每行中每$X$个连续的元素计算一个$\alpha$。 当$X=1$时, 退化成没有量化。选择适当的$X$不会显著增加计算量,但能提升模型的性能。为了计算方便, 一般$X$是64的倍数。
使用1个bit来表示一个参数还是会引入了较大的误差。为了降低这些误差, 可以引入multi-bit XNOR。通过每次增加一个bit,量化上次的残差,可以降低最终的误差。假如 矩阵A通过$x$个bit量化, 那么
$$ \alpha_{A1}, A_{b1} = binarize(A) $$
$$ \alpha_{A2}, A_{b2} = binarize(A-\alpha_{A1}A_{b1}) $$
$$ \cdot\cdot\cdot $$
因此multi-bit量化可以定义为:
$$ \left[\alpha_{A1},\cdots,\alpha_{Ax}\right],\left[A_{b1},\cdots,A_{bx}\right] = multi-bit(A, x) $$
基于multi-bit的MatMul可以表示为:($A$用$x$个bit量化, $B$用$y$个bit量化)
$$ GEMM(A,B) \approx \sum_{i=1}^{x}\sum_{j=1}^{y}\alpha_{Ai}\alpha_{Bi}B\_GEMM(A_{bi}, B_{bj}) $$
上面的多个bit方案, 只能保证每增加一个bit的时候, 对当前的残差的量化最优的,但是不能保证整个过程是一个最优或是次优的量化方案。通过传统的数值方法不能得到最优解, 因为目标函数不是关于$\{\alpha_i\},\{A_{bi}\}$的凸函数。因此我们设计了一种交替迭代的量化方案:AQN: Alternating multi-bit Quantzation Network。考虑这样一个事实,对于$ error = \left\| A - \sum_{i=1}^{n} \alpha_i A_{b_i} \right\|^2$, 当给定$\{W_{b_i}\}$, $\{\alpha_i\}$是有解析解的. 令
$$ A_b^{\ast}=\left[vec(A_{b_1})^T, vec(A_{b_2})^T,...,vec(A_{b_n})^T\right] $$
其中vec为将矩阵展开成向量操作。矩阵$A_b^{\ast} $的shape为$mn \times M_A$。那么可以得出
$$ \left[\alpha_{A1}, ..., \alpha_{An}\right] = (({A^{\ast}_b}^T {A^{\ast}_b)}^{-1}{A_b^{\ast}}^T A)^T $$
。 同理, 当$\{\alpha_i\}$给定的时候, $\{A_{b_i}\}$有全局最优解。因此我们采用交替迭代的方法, 迭代求解的算法为:
- inputs: matrix ${A}$, iteration times $n$
- get $\{\alpha_i, A_{b_i}\}$ with multi-bit XNOR method
REPEAT
- $\left[\alpha_{A1}, ..., \alpha_{An}\right] = (({A^{\ast}_b}^T {A^{\ast}_b)}^{-1}{A_b^{\ast}}^T A)^T$
- $\{A_{bi}\} = \mathop{\arg\min}_{\{A_{bi}\}} \left\| {A - \sum_{i}^{m}{\alpha_{i} A_{bi}}}\right\|^2 $}
- n = n-1
- UNTIL{n>0}
- return $\{\alpha_i\}, \{A_{b_i}\}$
实现
我们在X86-64指令集上基于Caffe框架实现了AQN算法,包含训练和Inference部分, 并且在spell-check这样一个task上上线。基于Caffe框架的一个重要原因是该框架简单, 代码量不大。由于Caffe是一个基于C++的DL框架, 通过简单的封装可以部署在线上。
目前有几个开源的BNN算法lib, 但他们对intel intrinsic的应用并不好,参考这些开源代码,我们实现了AQN的training和inference模块。training和inference是分开, 是因为training和inference在我们的场景中是不一样的。 training部分我们采用单机GPU训练, inference主要是在搜索的QP上,基于Intel broadwell 架构的CPU。
在training的过程中, 主要是对Caffe中的InnerProductLayer的input和$W$分别做基于AQN的approximation,然后再调用GEMM计算矩阵乘法。bp的时候分别采用量化后的结果计算$W_q$以及向下传的residual。其中一个重要的细节就是clip。把$W$ clip到一个$\left[-x, x\right]$区间相当于增加了正则, 通常$x=1$。
在approximation之后再做GEMM,会增加training的计算量, 但由于training主要是离线的, 因此这些overhead还是可以接受的。 我们在GPU上实现了AQN的training部分, 和原始的InnerProduct相比, training的时间没有增加多少, 下面训练过程中的forward和backward流程
目前有一些基于binary量化的$B\_GEMM$实现,参考这些实现,基于avx2指令集, 我们在caffe中实现的AQN的Inference代码,
- 初始化的时候, 将$W$通过AQN 量化得到$\{\alpha_{i}\}$以及$\{W_{bi}\}$
- forward的时候, 对input量化, 得到对应的量化结果, 利用$B\_GEMM$(XOR+POPCNT)计算二值矩阵的乘法
实验
我们在搜索的query纠错、以及一些开源的数据集合上对比了由于量化带来的模型精度影响, 以及对GEMM速度的提升。query纠错是一个基于LSMT的深度学习模型。LSTM是其中计算最主要的计算单元,其中包含了大量的GEMM算子。我们对比了使用多bit的AQN后对模型精度的影响以及LSTM单元的加速情况:
在上面的实验中, 在$input$和$W$各使用2个bit量化的时候, 模型的accuracy有部分下降, 但LSTM单元的加速可以提升3倍,这可以确保模型能在线部署。同时我们在一个已经train好的LSTM模型上使用AQN进行量化, 我们分析了量化后的模型和原始模型之间数值上的RMSE, 并和目前最新的方法做了对比
对上面的对比可以看出, 经过AQN量化后的模型, 在数值上和原模型最接近, 大幅降低了数值上的误差。为了判断量化最终对模型性能的影响, 我们在开源的数据集Peen Tree Bank (PTB)上,在training的过程中, 结合AQN进行训练, 并最终和其他目前最新的方法对比了在测试数据集上的模型性能
在上面的实现对比中, AQN对模型性能影响最小,效果比其他对比的方法好。在有些时候甚至超过了全精度的模型, 因为量化的过程相当于引入了模型的正则, 降低了噪音数据对模型的影响。
对了验证AQN的inference的计算性能,我们和Intel的科学计算库MKL对比了一下矩阵乘法的latency。统计了在基于Intel的broadwell架构的CPU上, 采用单线程计算了在不同场景下的latency
上面是对M比较小的时候的结果, 因为在实际的场景中, 需要inference的数据只有一条, 比如:用户输入的query等。在AQN实现的过程中, 我们发现目前的Intel CPU指令集对POPCNT支持的不是很好, 但后下一代的CPU体系架构中, 相关的指令集得到了优化。相信在后面的CPU体系中, 基于 binary量化的相关算法, 在Inference上的加速能更上一个台阶。同时,很多研究将基于binary量化的模型应用到GPU,FPGA上, 也取得了非常不错的效果。
我们还考察了基于ternary的算法。ternary的一个想法是, 矩阵
$$ A \approx \alpha A_t $$
其中
$$ A_t \in \{-1, 0, 1\}^{m \times n} $$
基于XOR和POPCNT的ternary, binary矩阵运算需要对ternary部分引入mask,这样可以在保持模型性能的同时, 降低对POPCNT的使用。
参考文献
- Pete Warden. Why GEMM is at the heart of deep learning
Yu Cheng, Duo Wang, Pan Zhou and Tao Zhang. A Survey of - Model Compression and Acceleration for Deep Neural Networks
- Mohammad Rastegari, Vicente Ordonez, Joseph Redmon and Ali Farhadi. XNOR-Net: ImageNet Classification Using Binary Convolutional Neural Networks
- Zefan Li, Bingbing Ni, Wenjun Zhang, Xiaokang Yang and Wen Gao. Performance Guaranteed Network Acceleration via High-Order Residual Quantization
- Matthieu Courbariaux, Itay Hubara, Daniel Soudry, Ran El-Yaniv and Yoshua Bengio. Binarized Neural Networks: Training Deep Neural Networks with Weights and Activations Constrained to +1 or -1
- Fengfu Li, Bo Zhang, Bin Liu. Ternary Weight Networks
Abhisek Kundu, Kunal Banerjee, Naveen Mellempudi, Dheevatsa Mudigere, Dipankar Das, Bharat Kaul, Pradeep Dubey. Ternary Residual Networks - Hande Alemdar, Vincent Leroy, Adrien Prost-Boucle, Frédéric Pétrot. Ternary Neural Networks for Resource-Efficient AI Applications
- Naveen Mellempudi, Abhisek Kundu, Dheevatsa Mudigere, Dipankar Das, Bharat Kaul, Pradeep Dubey. Ternary Neural Networks with Fine-Grained Quantization
- Yiwen Guo, Anbang Yao, Hao Zhao, Yurong Chen. Network Sketching: Exploiting Binary Structure in Deep CNNs
- Haojin Yang, Martin Fritzsche, Christian Bartz, Christoph Meinel. BMXNet: An Open-Source Binary Neural Network Implementation Based on MXNet
- https://en.wikipedia.org/wiki/AVX-512
- Eriko Nurvitadhi, David Sheffield, Jaewoong Sim, Asit Mishra, Ganesh Venkatesh and Debbie Marr Accelerator Architecture Lab, Intel Corporation. Accelerating Binarized Neural Networks:
Comparison of FPGA, CPU, GPU, and ASIC - Mitsuru Ambai, Takuya Matsumoto, Takayoshi Yamashita, Hironobu Fujiyoshi. Ternary Weight Decomposition and Binary Activation Encoding for Fast and Compact Neural Networks