上周四，加州大学伯克利分校 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：
图 3：不同层在 SQuAD 任务上的损失分布，该分布图是通过沿着 Hessian 矩阵的前两个主要特征值扰动参数绘制而成的。铜球表示参数空间中 BERT 模型的收敛点。
组量化
假设输入序列中有 n 个词，每个词有一个 d 维嵌入向量（BERT_BASE 中 d = 768）。在 Transformer 编码器中，每一个自注意力头具备 4 个密集矩阵，即
其中 N_h 是注意力头的数量。这里，W_k、W_q、W_v 和 W_o 分别表示 key、query、值和输出权重矩阵。每个自注意力头按以下公式计算加权和：
该研究针对基于注意力的模型提出组量化机制。将多头自注意力（MHSA）的密集矩阵中每个注意力头的矩阵 W 看作一个组，这样一共有 12 组。在每个组中，将多个顺序输出神经元看作一个子组，每个子组具备自己的量化范围。
W_v 的图示见下图 4，研究者将 N_h 值矩阵 W_v 连接为一个三维张量。
图 4：组量化方法概览。该图使用多头自注意力层的值矩阵绘制而成。
实验
研究者在四个下游任务中评估 Q-BERT，这些任务包括情感分类、自然语言推断、命名实体识别和机器阅读理解。Q-BERT 达到了 13 倍的权重压缩率，激活大小和嵌入大小仅为原版的 1/4，且准确率损失在 2.3% 以内。
据作者表示，这是首个将 BERT 以超低精度进行量化，且还能保持性能损失在可接受范围内的工作。
表 1：BERT_BASE 在自然语言理解任务上的量化结果，每个层有 128 个组。出于简洁性和效用考虑，除基线以外的所有模型都使用 8-bits 激活。此外，研究者对比了 Q-BERT 和未使用混合精度量化或组量化的直接量化方法（DirectQ）。
表 2：Q-BERT 组量化在三个任务上的效果。所有任务中量化 bit 的设置均为：权重 4、嵌入 8、激活 8。自上而下，组的数量逐渐增加。为了平衡准确率和硬件效率，其他实验中的组数量均为 128。
表 3：不同模块的量化效果。ew-bits、ep-bits、s-bits 和 f-bits 分别表示词嵌入、位置嵌入、多头注意力层和全连接层所使用的量化 bit。(a) (b) 中，权重和激活 bit 均设置为 8。
图 5：Q-BERT/DirectQ 和基线模型之间的注意力分布 KL 散度。Q-BERT 和基线之间的距离要比 DirectQ 和基线之间的距离短得多。