对于动辄上百 M 大小的神经网络来说,模型压缩能够减少它们的内存占用、通信带宽和计算复杂度等,以便更好地进行应用部署。最近,来自 Facebook AI 的研究者提出了一种新的模型量化压缩技术 Quant-Noise,让神经网络在不降低性能的前提下,内存占用减少 90% 以上。
我们都知道,对于神经网络来说,参数量越大、层数越多,就代表着输出的结果越精细。当然,这也意味着许多性能优越的神经网络体积会非常庞大。比如当前的标准机器翻译架构 Transformer,一层就可能包含数百万个参数。即使是一些优化过性能和参数效率的模型(比如 EfficientNet),也仍然需要几十到几百 MB。这就使得它们的应用范围限制在机器人或者虚拟助手等领域。
所以我们在应用部署之前,必须面对一个问题:如何对模型进行压缩?
剪枝和蒸馏是模型压缩中常用的两种方法,通过减少网络权重的数量来删减参数。还有一种方法就是「量化」,不同的是,它是通过减少每个权重的比特数来压缩原始网络。
标量量化(scalar quantization)等流行的后处理量化方法是让训练网络的浮点权重以一个低精度表征去表示,比如说定宽整数。这些后处理量化方法的好处在于压缩效率很高,并且能够加速支持硬件上的推理。但缺点在于,这些近似值造成的误差会在前向传播的计算过程中不断累积,最终导致性能显著下降。
现在,来自Facebook 的研究者提出了一种新的模型量化压缩技术 Quant-Noise,可对模型进行极致压缩,同时在实际应用部署时保持高性能。
- 论文地址:https://arxiv.org/abs/2004.07320
- 项目地址:https://github.com/pytorch/fairseq/tree/master/examples/quant_noise
在这项研究中,研究者提出了一种仅量化权重子集而非整个网络的压缩方案。在每次前向传播时仅量化网络的随机部分,对大多数权重使用无偏梯度进行更新。
该方法可以在训练过程中采用更简单的量化方案,这对于具有可训练参数的量化模块来说是非常有用的,比如乘积量化(Product Quantizer,PQ)算法。
总体来说,该方法的亮点在于:
- 将名为「Quant-Noise」的量化噪声应用到权重的随机子集上,来学习更适用于 int4、 int8 和 PQ 算法等各种量化方法的网络;
- 在 PQ 算法上添加 Quant-Noise,实现准确率和模型大小综合结果上的 SOTA 水平。使用该方法进行压缩后,在自然语言处理方面,RoBERTa 在压缩到 14MB 的前提下,在 MNLI 上实现了 82.5% 的准确率;在计算机视觉方面,EfficientNet-B3 在压缩到 3.3 MB 的前提下,在 ImageNet 上实现了 80.0% 的 top-1 准确率。
- 使用 Quant-Noise 训练的网络,通过结合 PQ 算法和 int8 来量化网络的权重和 activation,获得在固定精度计算条件下的极致压缩效果,实现了 ImageNet 数据集上的 79.8% top-1 准确率和 WikiText-103 数据集上的 21.1 困惑度。
此外,使用 Quant-Noise 压缩后的模型性能几乎与原始模型一致,同时将内存占用量减少至原来的十分之一甚至二十分之一。这大大超过了 PyTorch 和 Tensorflow 中所用的 int8 的 4 倍压缩。目前,相关代码也已经开源。
原理:量化神经网络
在本章中,研究者介绍了量化原理以及几种标准量化方法,并详细解释了如何将标量和乘积量化相结合。需要注意的是,为了定义清楚,研究者主要专注于固定实矩阵(fixed real matrix)W∈R ^ n×p 的示例,并假设此矩阵被拆分为 m×q 个块(b_kl):
这些块的属性由量化方法来确定,而码本(codebook)是包含 K 个向量的集合,即 C = {c [1], ..., c [K]}。此量化方法使用如下方法压缩矩阵 W:向每个块 b_kl 分配一个指向码本 C 中「码字 c」的索引,同时存储码本 C 以及结果索引(作为索引矩阵 I 的条目 I_kl), 而不是使用实际权重。在推理过程中,该方法重建原始矩阵 W 的近似值
,使得 b_kl = c [I_kl]。研究者将标量量化(如 int8,即每个块 b_kl 由一个权重组成)与向量量化(将多个权重共同量化)区分开来。
上图显示了研究者在训练过程中如何将量化噪声应用于权重子集,从而改善量化模型的性能(完整视频请参照链接)。
定点标量量化
定点(Fixed-point)标量量化方法用低精度定点表示代替了浮点表示。通过在能够兼容的硬件上使用定点运算,实现减少模型内存占用并加快推理速度的效果。
Quant-Noise 的具体实现方法
深度网络训练过程中不会接触到 quantization drift 引起的噪声,从而导致性能欠佳。如何使网络对量化具有一定的鲁棒性?解决方案就是在训练期间引入量化噪声。
量化感知训练(Quantization Aware Training,QAT)通过对前向传播过程中的权重进行量化来实现这一点。这种转换是不可微的,且使用直通估计器(straight through estimator,STE)来逼近梯度。STE 在梯度中引入了一个偏置,该偏置取决于权重的量化级别,因此也取决于压缩率。
因此,研究者提出了一种简单的修改方法 Quant-Noise,通过随机改善 QAT 来控制该偏置。其思路是:量化权重中的随机部分,而不是像 QAT 那样量化整个网络,使未经量化的权重使用无偏梯度。该方法的一般形式可以模拟训练期间的量化和剪枝效果。
使用 Quant-Noise 训练网络
仍以实矩阵 W 为例。在训练时,Quant-Noise 方法的运行方式如下:
首先,计算与目标量化方法相关的块 b_kl;然后在每个前向传播过程中,随机选择这些块的子集,并添加一些失真;在反向传播过程中,使用 STE 逼近失真的权重,进而计算所有权重的梯度值。换句话说,给定索引元组 J ⊂ {(k, l)},其中 1≤k≤m,1≤l≤q,同时将失真或噪声函数 ϕ 作用在块上,于是研究者定义算子 ψ (· | J),对于每个块 b_kl 均应用以下转换:
Quant-Noise 的效果怎么样?
Quant-Noise 对不同量化方法的影响
如下表 1 所示,研究者首先展示了在语言建模和图像分类任务上,Quant-Noise 对 int4、int8 和 iPQ 三种不同量化方案的影响。结果表明,Quant-Noise 对 iPQ 等高性能量化方案表现出非常好的效果,而 QAT 一般导致性能下降,即使与量化作为后处理步骤的情况相比也会这样。
表 1:在 QAT 和 Quant-Noise 两种设置下,int4、int8 和 iPQ 三种量化方案在语言建模和图像分类任务上的量化结果对比。对于语言建模任务,研究者在 Wikitext-103 基准上训练 Transformer,得出了测试困惑度(PPL);对于图像分类任务,研究者在 ImageNet-1k 基准上训练了 EfficientNet-B3,得出了验证 Top-1 准确率。
研究者分析得出,Quant-Noise 能够与剪枝(pruning)和权重共享(weight sharing)方法兼容并相互补充。如下表 2 所示,研究者分别展示了 WikiText-103 上语言建模、MNLI 上预训练句子表征和 ImageNet-1k 上目标分类的结果,这些结果在不同的任务和基准上保持了高度一致:Quant-Noise 在强大的 iPQ 基线方法上实现了性能的大幅提升。
表 2:几种不同的压缩方法(权重共享、剪枝和 Quant-Noise)在语言建模、句子表征和图像分类任务上对未量化模型和量化模型性能的影响。
与 SOTA 方法的对比
研究者在相同任务上将 Quant-Noise 与当前 SOTA 方法进行对比。他们在有竞争力的模型上应用了最佳量化设置,在结合权重共享和剪枝方法时将这些模型的内存占用缩小到原来的 1/94-1/20,从而为高性能模型提供了极致压缩。
如下图 2 所示,研究者得出了在语言建模、MNLI 和 ImageNet 上性能与模型大小的权衡情况。
图 2:模型性能随模型大小的变化情况。
控制变量测试
如下表 3 所示,研究者分别展示了在未利用 Quant-Noise 训练、利用 Quant-Noise 微调以及利用 Quant-Noise 训练三种不同的设置下,Adaptive Input 架构的困惑度和 RoBERTa 的准确率变化情况。可以看出,直接利用 Quant-Noise 训练可以实现最低的困惑度和最高的准确率。
表 3:未利用 Quant-Noise 训练、利用 Quant-Noise 微调和利用 Quant-Noise 训练三种不同设置下的困惑度和准确率对比。
如下图 3 所示,不同 Quant-Noise 值对语言建模模型 Transformer 性能的影响。对于 iPQ,高噪声值对模型性能带来较大影响;而对于 int8 量化及其噪声函数而言,较高的噪声值对模型性能造成轻微的负面影响,但不如前者那么严重。
图 3:量化参数的影响。