AAAI 2020 | 超低精度量化BERT，UC伯克利提出用二阶信息压缩神经网络

2020 年 1 月 18 日 机器之心

机器之心发布

机器之心编辑部

2020 年 2 月 7 日-2 月 12 日，AAAI 2020 将于美国纽约举办。不久之前，大会官方公布了今年的论文收录信息：收到 8800 篇提交论文，评审了 7737 篇，接收 1591 篇，接收率 20.6%。为向读者们分享更多的优质内容、促进学术交流，在 AAAI 2020 开幕之前，机器之心策划了多期线上分享。

上周四，加州大学伯克利分校 Zhewei Yao 博士分享了他的 AAAI 论文《Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT》，本文对此论文进行了详细解读。该研究介绍了一种使用二阶信息进行模型压缩的新型系统性方法，能够在图像分类、目标检测和自然语言处理等一系列具有挑战性的任务中产生前所未有的小模型。

论文地址： https://arxiv.org/pdf/1909.05840.pdf

研究贡献

该论文对基于 BERT 的模型执行超低精度量化，旨在最小化性能下降幅度，同时保持硬件效率。为此，研究者使用了多项新技术，并提出了新模型 Q-BERT。

该研究的贡献如下：

研究者对二阶信息（即 Hessian 信息）进行大量逐层分析，进而对 BERT 执行混合精度量化。研究发现，与计算机视觉领域中的神经网络相比，BERT 的 Hessian 行为存在极大的不同。因此，该研究提出一种基于 top 特征值均值和方差的敏感度度量指标，以实现更好的混合精度量化。这与仅利用均值的研究 [8] 相反。

研究者提出新的量化机制——组量化（group-wise quantization），该方法能够缓解准确率下降问题，同时不会导致硬件复杂度显著上升。具体而言，组量化机制将每个矩阵分割为不同的组，每个组拥有独立的量化范围和查找表。

研究者调查了 BERT 量化中的瓶颈，即不同因素如何影响 NLP 性能和模型压缩率之间的权衡，这些因素包括量化机制，以及嵌入、自注意力和全连接层等模块。

方法

该研究提出的 BERT 量化方法包括：基于 Hessian 信息的混合精度量化，以及用于组量化机制的技术。

和 [7] 一样，微调后的 BERT_BASE 模型包含三部分：嵌入、基于 Transformer 的编码器层、输出层。BERT_BASE 模型的参数大小为：嵌入 91MB、编码器 325MB、输出 0.01MB。由于输出层的规模极小，该研究并未对这部分执行量化。也就是说，该研究用不同方式对嵌入和编码器参数执行量化。

量化过程

通用神经网络推断通常按照权重和激活的浮点数精度来执行。量化将网络权重限制为一个有限集，如下所示：

其中 Q 是量化算子，z 是实数输入张量（激活或权重），(t_j , t_j+1] 表示一个区间，j 为实数 (j = 0, . . . , 2^k − 1)。

量化函数 Q 存在多种选择。该研究使用统一的量化函数，张量中的浮点值范围平分 [12, 42]，可以用 0, . . . , 2^k − 1 中的无符号整数来表示。但是，研究者选择使用统一的量化函数，以便获得更高效、简便的硬件实现。为了将梯度通过不可微函数 Q 进行反向传播，研究者使用了直通估计器（Straight-through Estimator，STE）。

混合精度量化

研究者探索了混合精度量化，即对敏感度较高的层分配更多 bit，以保持性能。为此研究者开发了 Hessian AWare Quantization (HAWQ)。由于每个层 Hessian 矩阵的大小为 7M × 7M，因此存在一个常见的误解：计算二阶信息并不可行。但是，Hessian 谱可以通过矩阵无关的幂迭代方法来计算，该方法无需显式信息。

为方便读者理解，这里以第一个编码器层为例。将该层的梯度标注为 g_1，对于和 g_1 具备相同维度的随机向量 v，存在

其中 H_1 是第一个编码器层的 Hessian 矩阵。第二个方程来自于「v 独立于 W_1」这一事实。然后利用幂迭代计算 top 特征值，如附录中算法 1 所示。λ_i 表示第 i 个编码器层的 top 特征值。

幂迭代算法。

图 2 展示了 BERT_BASE 中不同层的 top Hessian 特征值的分布情况。不同层的特征值数量不同，尽管所有层的结构和大小均相同。

图 2：从 (a) 到 (d)：不同编码器层在 SST-2、MNLI、CoNNL-03、SQuAD 任务上的 top 特征值分布情况。

对于 top 特征值较小的层（图 1 中比较平坦的损失分布），会执行更具攻击性的量化。但是，研究者发现，仅基于平均 top 特征值来分配 bit 的做法不适用于很多 NLP 任务。

图 1：不同层在 MNLI 和 CoNNL-03 任务上的损失分布，该分布图是通过沿着 Hessian 矩阵的前两个主要特征值扰动参数绘制而成的。铜球表示参数空间中 BERT 模型的收敛点。这些层表明较平坦的曲率会被量化为较低精度。

为了解决此问题，研究者使用以下度量指标来替代仅使用均值的方法：

其中 λ_i 是 H_i top 特征值的分布，基于 10% 的训练数据集计算得到。接下来，研究者基于选择的精度设置执行量化感知的微调（quantization-aware fine-tuning）。

研究者强调了一个重要的技术点：该方法预计，在执行量化前，训练模型已收敛至局部极小值。必要的最优性条件是零梯度和正曲率（即正 Hessian 特征值）。

根据分析，研究者发现，在 MNLI、CoNLL-03 和 SST-2 这三项任务中，top Hessian 特征值确实为正值。但针对 SQuAD 微调后的 BERT 模型无法收敛至局部极小值，参见图 2d 中的 Hessian 特征值，那里存在非常大的负特征值。直接可视化损失分布也可以证明这一点，详见下图 3：