AdaRound：训练后量化的自适应舍入

本文是高通AI研究院发表在ICML 2020上的一篇 PTQ 文章。根据泰勒展开的公式推导证明了业界常用的round策略并不是最优的，并提出了可学习的自适应任务损失的量化策略AdaRound。本文实现了在少量数据无需微调的情况下，将 Resnet18 和 Resnet50 的权重量化为 4 位，同时保持 1％ 的精度损失。

• 论文题目：Up or Down? Adaptive Rounding for Post-Training Quantization
• 论文链接：https://arxiv.org/abs/2006.10518
• 代码链接：https://github.com/itayhubara/CalibTIP

摘要

在对神经网络进行量化时，主要方法是将每个浮点权重分配给其最接近的定点值。本文发现，这不是最佳的量化策略。本文提出了 AdaRound，一种用于训练后量化的更好的权重舍入机制，它可以适应数据和任务损失。 AdaRound 速度很快，不需要对网络进行微调，仅需要少量未标记的数据。本文首先从理论上分析预训练神经网络的舍入问题。通过用泰勒级数展开来逼近任务损失，舍入任务被视为二次无约束二值优化问简化为逐层局部损失，并建议通过软松弛来优化此损失。 AdaRound 不仅比舍入取整有显著的提升，而且还为几种网络和任务上的训练后量化建立了新的最新技术。无需进行微调，本文就可以将 Resnet18 和 Resnet50 的权重量化为 4 位，同时保持 1％ 的精度损失。

贡献

• 建立了一个理论框架，以一种既考虑数据本身又考虑任务损失的特征的方式来分析舍入的影响。 使用此框架，可以将舍入公式转化为二次无约束二进制优化（QUBO）问题。

• 提出了 AdaRound，这是一种新颖的方法，它可以通过连续松弛来为这种每层配置找到一个好的解决方案。 AdaRound 仅需要少量未标记的数据，计算效率高，并且适用于具有卷积层或完全连接层的任何神经网络体系结构。

• 在全面的研究中，表明 AdaRound 为几个网络和任务（包括 ResNet18，ResNet50，MobilenetV2，InceptionV3 和 DeeplabV3）提供了最新的训练后量化新技术。

动机

为了直观地了解为什么四舍五入可能不是最佳选择，让我们看一下在扰动预训练模型的权重时会发生什么。 $w$ 表示神经网络展开的权重参数。 ∆w 表示很小的扰动，而 $L(x，y，w)$ 表示要最小化的任务损失。 然后

$$\begin{array}{rl}& \mathbb{E}[\mathcal{L}(\mathbf{x},\mathbf{y},\mathbf{w}+\Delta \mathbf{w})-\mathcal{L}(\mathbf{x},\mathbf{y},\mathbf{w})]\\ & \stackrel{(a)}{\approx }\mathbb{E}[\Delta {\mathbf{w}}^T\cdot {\nabla }_{\mathbf{w}}\mathcal{L}(\mathbf{x},\mathbf{y},\mathbf{w})+\frac{1}{2}\Delta {\mathbf{w}}^T\cdot {\nabla }_{\mathbf{w}}^2\mathcal{L}(\mathbf{x},\mathbf{y},\mathbf{w})\cdot \Delta \mathbf{w}]\\ & =\Delta {\mathbf{w}}^T\cdot {\mathbf{g}}^{(\mathbf{w})}+\frac{1}{2}\Delta {\mathbf{w}}^T\cdot {\mathbf{H}}^{(\mathbf{w})}\cdot \Delta \mathbf{w}\end{array}$$

这里，采用二阶泰勒展开。${\mathbf{g}}^{(\mathbf{w})}$和${\mathbf{H}}^{(\mathbf{w})}$分别代表任务损失的梯度与 Hessian 矩阵：

$$\begin{array}{rl}{\mathbf{g}}^{(\mathbf{w})}& =\mathbb{E}\left[{\nabla }_{\mathbf{w}}\mathcal{L}(\mathbf{x},\mathbf{y},\mathbf{w})\right]\\ {\mathbf{H}}^{(\mathbf{w})}& =\mathbb{E}\left[{\nabla }_{\mathbf{w}}^2\mathcal{L}(\mathbf{x},\mathbf{y},\mathbf{w})\right].\end{array}$$

相对于指定变量，本文中的所有梯度项和Hessian项的任务损失均为L。 只要 $\Delta w$ 不太大，就可以忽略泰勒级数展开式中的高阶项。 假设网络经过训练可以收敛，我们也可以忽略梯度项，因为它接近于0。因此，H(w) 定义了不同扰动权重之间的相互作用，它共同影响着权重与任务损失$L(x，y，w+∆w)$。 下面的简单示例说明了舍入到最近的取舍可能不是最佳的方法：

示例：假设 $∆w^T = [∆w_1, ∆w_2] $ 并且 ${\mathbf{H}}^{(\mathbf{w})}=\left[\begin{array}{cc}1& 0.5\\ 0.5& 1\end{array}\right]$。那么由于扰动导致的任务损失增加大约成正比：
$$\Delta {\mathbf{w}}^T\cdot {\mathbf{H}}^{(\mathbf{w})}\cdot \Delta \mathbf{w}=\Delta {\mathbf{w}}_1^2+\Delta {\mathbf{w}}_2^2+\Delta {\mathbf{w}}_1\Delta {\mathbf{w}}_2$$
对于对角线项$∆w_1^2$和$∆w_2^2$，仅扰动的大小重要。因此，在本示例中，如果仅考虑这些对角线项时，rounding-to-nearest 是最佳的。 但是，对于非对角线项$∆w_1∆w_2$，扰动的符号很重要，其中两个扰动的相反符号会改善损失。 为了使量化对任务损失的总体影响最小，需要在对角项和非对角项的贡献之间进行权衡。 Rounding-to-nearest 忽略了对角线以外的贡献，使其通常不是最佳的。

先前的分析对于任何参数系统的量化都是有效的。本文证明了这种效果也适用于神经网络。为了说明这一点，本文为Resnet18的第一层生成了100个随机舍入，并仅对第一层进行了量化就评估了网络的性能。结果列于下表。在100次运行中，发现48个随机抽样的舍入选择比 Rounding-to-nearest 有更好的性能。这意味着存在许多舍入解决方案要比四舍五入法更好。此外，在这100个随机样本中，最好的样本可将网络的准确性提高10％以上。并且，还看到刻意地将所有值向上或向下取整会产生灾难性的影响。这意味着在进行训练后量化时，通过仔细舍入权重，可以获得很多收益。本文的其余部分旨在设计一种有充分根据和计算效率的舍入机制。

方法

在本节中，提出AdaRound，这是一种用于训练后量化的新舍入程序，在理论上是有充分根据的，并且在实践中显示出显着的性能改进。 本文从理论上分析由于量化引起的损失。 然后，制定了有效的每层算法来对其进行优化。

基于任务损失的 Rounding

在量化预训练的NN时，我们的目标是最大程度地减少由于量化引起的性能损失。 假设遵循 per-layer 权重量化，则量化后的权重表示为：

$${\hat{\mathbf{w}}}_i^{(\ell )}\in \left\{{\mathbf{w}}_i^{(\ell ),floor},{\mathbf{w}}_i^{(\ell ),ceil}\right\}(1)$$
其中，
$${\mathbf{w}}_i^{(\ell ),floor}={\mathrm{s}}^{(\ell )}\cdot \mathrm{clip}\left(\lfloor \frac{{\mathbf{w}}_i^{(\ell )}}{{\mathrm{s}}^{(\ell )}}],\mathrm{n},\mathrm{p}\right)(2)$$
另外${\mathbf{w}}_i^{(\ell ),ceil}$跟上述定义类似，只需要将符号$\lfloor \cdot \rfloor$替换为$\lceil \cdot \rceil$。$\Delta {\mathbf{w}}_i^{(\ell )}={\mathbf{w}}^{(\ell )}-{\hat{\mathbf{w}}}_i^{(\ell )}$表示由于量化产生的扰动。在这项工作中，假设在优化舍入过程之前先将${\mathrm{s}}^{(\ell )}$固定好。最后，每当优化损失函数超过$\Delta {\mathbf{w}}_i^{(\ell )}$时，$\mathrm{s}^{(\ell)} $ 只能采用公式(1)中指定的两个值。

寻找最佳舍入程序可以公式化为以下二进制优化问题：
$$\underset{\Delta \mathbf{w}}{\mathrm{argmin}}\mathbb{E}[\mathcal{L}(\mathbf{x},\mathbf{y},\mathbf{w}+\Delta \mathbf{w})-\mathcal{L}(\mathbf{x},\mathbf{y},\mathbf{w})](3)$$
在公式（3）中评估成本需要在优化过程中对每个新$\Delta w$的输入数据样本进行前向传递。为了避免重复的正向遍历数据的计算开销，本文利用二阶泰勒级数逼近。 此外，本文忽略了属于不同层的权重之间的相互作用。 反过来，这意味着，假设一个块对角线 ${\mathbf{H}}^{\left(\mathbf{w}\right)}$，其中每个非零块仅对应一层。 因此，最终简化为以下每层优化问题：
$$\underset{\Delta \mathbf{w}(\ell )}{\mathrm{argmin}}\mathbb{E}\left[{\mathbf{g}}^{{\left({\mathbf{w}}^{(\ell )}\right)}^T}\Delta {\mathbf{w}}^{(\ell )}+\frac{1}{2}\Delta {\mathbf{w}}^{(\ell {)}^T}{\mathbf{H}}^{\left({\mathbf{w}}^{(\ell )}\right)}\Delta {\mathbf{w}}^{(\ell )}\right](4)$$
如之前的示例所示，我们需要二阶项来利用权重扰动之间的联合相互作用。 公式（4）是一个二值优化问题，因为 $\Delta {\mathbf{w}}^{(\ell )}$ 是二进制变量。 对于收敛的预训练模型，可以安全地忽略梯度项对（4）中的优化的影响。进一步简化为：
$$\underset{\Delta \mathbf{w}(\ell )}{\mathrm{argmin}}\mathbb{E}\left[\Delta {\mathbf{w}}^{(\ell {)}^T}{\mathbf{H}}^{\left({\mathbf{w}}^{(\ell )}\right)}\Delta {\mathbf{w}}^{(\ell )}\right](5)$$

为验证公式（5）可以很好地优化由于量化而导致的任务损失，本文在仅对 Resnet18 的第一层进行量化时，将（5）中的损失与100个随机舍入向量的验证精度进行了比较。 上图显示了两个量之间的显著相关性。 这证明了本文对优化进行了合理的近似值，即使是4位量化也是如此。 优化公式（5）能显着提升量化性能，但是其应用受到两个问题的限制：

• 即使对于中等大小的层，${\mathbf{H}}^{\left({\mathbf{w}}^{(\ell )}\right)}$ 也会遇到计算和空间复杂性大的问题。
• 公式（5）是一个NP难的优化问题。 解决它的复杂性随 $\Delta {\mathbf{w}}^{(\ell )}$ 的大小迅速扩展，再次阻止了公式（5）应用于中等大小的层。

从泰勒展开到局部损失

为了了解造成 ${\mathbf{H}}^{\left({\mathbf{w}}^{(\ell )}\right)}$ 复杂的原因，让我们看一下它的元素。 对于同一全连接层中的两个权重：
$$\begin{array}{rl}\frac{{\partial }^2\mathcal{L}}{\partial {\mathbf{W}}_{i,j}^{(\ell )}\partial {\mathbf{W}}_{m,o}^{(\ell )}}& =\frac{\partial }{\partial {\mathbf{W}}_{m,o}^{(\ell )}}\left[\frac{\partial \mathcal{L}}{\partial {\mathbf{z}}_i^{(\ell )}}\cdot {\mathbf{x}}_j^{(\ell -1)}\right]\\ & =\frac{{\partial }^2\mathcal{L}}{\partial {\mathbf{z}}_i^{(\ell )}\partial {\mathbf{z}}_m^{(\ell )}}\cdot {\mathbf{x}}_j^{(\ell -1)}{\mathbf{x}}_o^{(\ell -1)}\end{array}(6)$$
其中${\mathbf{z}}^{(\ell )}={\mathbf{W}}^{(\ell )}{\mathbf{x}}^{(\ell -1)}$是层 $\ell$ 的预激活(preactivations)，而 ${\mathbf{x}}^{(\ell -1)}$ 表示层 $\ell$ 的输入。 将其写成矩阵公式（ 对于全展开的${\mathbf{w}}^{(\ell )}$ ），我们有:
$${\mathbf{H}}^{\left({\mathbf{w}}^{(\ell )}\right)}=\mathbb{E}\left[{\mathbf{x}}^{(\ell -1)}{\mathbf{x}}^{(\ell -1{)}^T}\otimes {\nabla }_{{\mathbf{z}}^{(\ell )}}^2\mathcal{L}\right](7)$$
其中 $\otimes$ 表示两个矩阵的 Kronecker 乘积，${\nabla }_{{\mathbf{z}}^{(\ell )}}^2\mathcal{L}$是任务损失w.r.t. ${\mathbf{z}}^{(\ell )}$ 的Hessian。从公式（7）可以清楚地看出，复杂性问题主要是由 ${\nabla }_{{\mathbf{z}}^{(\ell )}}^2\mathcal{L}$ 引起的，它要求二阶导数通过网络的后续层进行反向传播。 为了解决这个问题，我们假设任务损失为 w.r.t. 预激活，比如， ${\nabla }_{{\mathbf{z}}^{(\ell )}}^2\mathcal{L}$ 是对角矩阵，由 $\mathrm{diag}\left({\nabla }_{{\mathbf{z}}^{(\ell )}}^2{\mathcal{L}}_{i,i}\right)$ 表示。 这将导致：
$${\mathbf{H}}^{\left({\mathbf{w}}^{(\ell )}\right)}=\mathbb{E}\left[{\mathbf{x}}^{(\ell -1)}{\mathbf{x}}^{(\ell -1{)}^T}\otimes \mathrm{diag}\left({\nabla }_{{\mathbf{z}}^{(\ell )}}^2{\mathcal{L}}_{i,i}\right)\right](8)$$
注意，在公式（8）中表示的 ${\mathbf{H}}^{\left({\mathbf{w}}^{(\ell )}\right)}$ 的近似值不是对角线。 将公式（8）插入方程中以找到优化损失（5）的舍入向量，得到:
$$\begin{array}{l}\arg {\min }_{\Delta {\mathbf{W}}_{k,:}^{(\ell )}}\mathbb{E}\left[{\nabla }_{{\mathbf{z}}^{(\ell )}}^2{\mathcal{L}}_{k,k}\cdot \Delta {\mathbf{W}}_{k,:}^{(\ell )}{\mathbf{x}}^{(\ell -1)}{\mathbf{x}}^{(\ell -1{)}^T}\Delta {{\mathbf{W}}_{k,:}^{(\ell )}}^T\right]\\ \stackrel{(a)}{=}\underset{\Delta {\mathbf{W}}_{k,:}^{(\ell )}}{\mathrm{argmin}}\Delta {\mathbf{W}}_{k,:}^{(\ell )}\mathbb{E}\left[{\mathbf{x}}^{(\ell -1)}{\mathbf{x}}^{(\ell -1{)}^T}\right]\Delta {\mathbf{W}}_{k,:}^{(\ell {)}^T}\\ =\underset{\Delta {\mathbf{W}}_{k,:}^{(\ell )}}{\mathrm{argmin}}\mathbb{E}\left[{\left(\Delta {\mathbf{W}}_{k,:}^{(\ell )}{\mathbf{x}}^{(\ell -1)}\right)}^2\right]\end{array}(9)$$

其中（8）中的优化问题,现在分解为（9）中的独立子问题。每个子问题都处理一行$\Delta {\mathbf{W}}_{k,:}^{(\ell )}$ 和（a）是进一步假设 ${\nabla }_{{\mathbf{z}}^{(\ell )}}^2{\mathcal{L}}_{i,i}={\mathbf{c}}_i$ 是与输入数据无关的常数的结果。值得注意的是，优化（9）不需要了解后续层和任务损失。在（9）中，我们只是最小化了由量化而在预激活 ${\mathbf{z}}^{(\ell )}$ 中引入的均方误差（MSE）。这与在多个神经网络压缩论文中优化的逐层目标相同。但是，与这些工作不同的是，本文以有原则的方式实现了这一目标，并得出结论，如假设（9）中所述，优化MSE 是我们在不了解超出层的网络其余部分的知识的情况下可以做的最好的事情。

公式（9）中的优化问题可以通过预先计算 $\mathbb{E}\left[{\mathbf{x}}^{(\ell -1)}{\mathbf{x}}^{(\ell -1{)}^T}\right]$ 来解决，如在（8）中所做的，然后在 $\Delta {\mathbf{W}}_{k,:}^{(\ell )}$ 上执行优化，或在优化过程中，执行每层前向传播的 $\Delta {\mathbf{W}}_{k,:}^{(\ell )}$ 。

AdaRound

解公式（9）不会产生与 ${\mathbf{H}}^{\left({\mathbf{w}}^{(\ell )}\right)}$ 相关的复杂性问题。 然而，它仍然是一个 NP难 的离散优化问题。 对于大量优化变量而言，找到具有合理计算复杂度的良好（次优）解决方案可能是一项挑战。 为了解决这个问题，本文将公式（9）放宽到以下基于软量化变量的连续优化问题。

$$\underset{\mathbf{V}}{\mathrm{argmin}}\Vert \mathbf{W}\mathbf{x}-\tilde{\mathbf{W}}\mathbf{x}{\Vert }_F^2+\lambda f_{reg}(\mathbf{V})(10)$$

其中 $\Vert \cdot {\Vert }_F^2$ 表示Frobenius范数，$\tilde{\mathbf{W}}$ 是要优化的软量化权重：

$$\tilde{\mathbf{W}}=\mathbf{s}\cdot \mathrm{clip}\left(\lfloor \frac{\mathbf{W}}{\mathbf{s}}\rfloor +h(\mathbf{V}),\mathbf{n},\mathbf{p}\right)(11)$$

在卷积层的情况下，$\mathbf{W}\mathbf{x}$矩阵乘法被卷积代替。 ${\mathbf{V}}_{i,j}$ 是在其上优化的连续变量，并且 $h\left({\mathbf{V}}_{i,j}\right)$ 可以是任何可微函数，其值介于0和1之间，即 $h\left({\mathbf{V}}_{i,j}\right)\in [0,1]$。 附加项 $f_{\text{reg }}(\mathbf{V})$ 是可微分的正则化器，其引入是为了激励优化变量 ${\mathbf{V}}_{i,j}$ 向0或1方向收敛，即收敛为 h ( V i , j ) ∈ { 0 , 1 } h\left(\mathbf{V}_{i, j}\right) \in\{0,1\} h(Vi,j​)∈{0,1}。本文使用经整流的 Sigmoid 型曲线作为 $h\left({\mathbf{V}}_{i,j}\right)$，定义为:

$$h\left({\mathbf{V}}_{i,j}\right)=\mathrm{clip}\left(\sigma \left({\mathbf{V}}_{i,j}\right)(\zeta -\gamma )+\gamma ,0,1\right)(12)$$

其中 σ（·）是 Sigmoid 函数，ζ 和 γ 是拉伸参数，分别固定为 1.1 和 -0.1 。 当 $h\left({\mathbf{V}}_{i,j}\right)$ 接近 0 或 1 时，整流后的 Sigmoid 梯度不会消失，当${\mathbf{V}}_{i,j}$ 移至两端时，这有助于学习过程。 对于正则化，本文使用

$$f_{\text{reg }}(\mathbf{V})=\sum _{i,j}1-{\left|2h\left({\mathbf{V}}_{i,j}\right)-1\right|}^{\beta }(13)$$

在这里对参数 β 进行退火。这使大多数$h\left({\mathbf{V}}_{i,j}\right)$在初始阶段（较高的β）可以自由适应以改善 MSE，并鼓励其在优化的后期（较低的β）收敛到 0 或 1。图2说明了退火β的效果。

图3显示了经过整流的 Sigmoid 和 $f_{reg}$ 的这种组合如何导致许多权重学习舍入而不是舍入到最接近的舍入，以提高性能，同时最终收敛到接近0或1的水平。

这种优化方法（11）是用于二进制约束优化问题的 Hopfield 方法通用系列的特定实例。这些类型的方法通常用作大规模组合问题的有效逼近算法。

为了量化整个模型，我们逐层优化（11）。然而，这不能解决由于先前的层而引入的量化误差。为了避免更深层网络的量化误差累积以及考虑激活函数，我们使用以下不对称重建公式：

$$\underset{\mathbf{V}}{\mathrm{argmin}}{\Vert f_a(\mathbf{W}\mathbf{x})-f_a(\tilde{\mathbf{W}}\hat{\mathbf{x}})\Vert }_F^2+\lambda f_{reg}(\mathbf{V})(14)$$

其中 $\hat{\mathbf{x}}$ 是所有先前的层都已量化的该层的输入，而 $f_a$ 是激活函数。 尽管出于不同的目的，以前在中也使用了类似的损失表述方式。 （14）定义了我们的最终目标，即我们可以通过随机梯度下降进行优化。 并将此算法称为AdaRound，因为它适用于输入数据的统计以及任务损失（近似）。

实验

消融实验

From task loss to local loss

Design choices for AdaRound

Optimization using STE

Influence of quantization grid

Optimization robustness to data

对比试验

个人总结与思考

本文通过任务损失的泰勒展开形式推导说明了量化常用的round策略并不是最优的。进一步，为了实现自适应的量化，本文针对预训练的权重值施加了(0,1)的扰动学习以进一步优化量化损失。本文也首次将权重的PTQ做到了4bit几乎没有精度损失，给PTQ方法提供了新的研究方向与动力。但是针对PTQ的量化问题还有如下几个问题需要解决：

• PTQ的激活量化没法通过(0,1)的扰动学习进位取舍。
• (0,1)范围的扰动是否是最优的？是否可以增大范围比如(-1，1)？
• AdaRound最终优化公式与前期推导的结果之间存在太多的近似约束。
• 基于泰勒展开的二次项优化重点落在了海森矩阵的近似求解上。

---