MetaQuant: Learning to Quantize by Learning to Penetrate Non-differentiable Quantization阅读

绒绒毛毛雨

已于 2024-10-31 14:20:26 修改

阅读量1.0k

点赞数 11

CC 4.0 BY-SA版权

分类专栏：论文阅读记录文章标签：人工智能架构

于 2024-10-25 18:01:27 首次发布

本文链接：https://blog.csdn.net/m0_56243424/article/details/143240283

论文阅读记录专栏收录该内容

21 篇文章

订阅专栏

Abstract

深层神经网络的巨大参数量使其在计算能力和存储空间有限的边缘设备上部署变得不切实际。现有研究通过学习量化深度模型，减少模型尺寸和能耗，即在监督训练过程中将全精度权重（r）转换为离散值（q），取得了一定进展。然而，量化的训练过程是不可微的，这导致相对于r的梯度（gr）要么无限大，要么为零。为了解决这个问题，大多数基于训练的量化方法使用相对于q的梯度（gq）并通过截断来近似gr，采用直接通过估计器（STE）或手动设计他们的计算。然而，这些方法仅凭经验使基于训练的量化适用，而没有进一步分析近似梯度如何辅助量化网络的训练。在本文中，我们提出通过神经网络学习gr。具体而言，一个元网络使用gq和r作为输入，并输出gr以进行后续权重更新。元网络与原始量化网络一起更新。我们提出的方法缓解了不可微性问题，并且可以端到端地训练。在CIFAR10/100和ImageNet上对各种深层网络进行了广泛的实验，以证明我们提出的方法在收敛速度和性能方面的优势。代码发布于：https://github.com/csyhhu/MetaQuant

1 Introduction

深层神经网络在各种计算机视觉任务中展示了令人鼓舞的结果。然而，现代深度学习模型通常包含许多层和大量参数，这限制了它们在边缘设备上的应用。为了减少参数冗余，已经在架构改进上做出了持续性的努力，例如使用小卷积核和重用特征。考虑一个完全训练好的非常深的模型，要用它进行预测，大部分计算涉及在前向传播过程中将一个实值权重与一个实值激活相乘。这些乘法操作非常昂贵，因为它们都是浮点数到浮点数的乘法运算。为了解决这个问题，已经提出了许多方法来量化深度模型。Courbariaux等人和Hubara等人提出将深度模型的权重二值化为+1或-1。为了在每层中提供更多的量化值灵活性，Rastegari等人引入了一个称为层缩放因子的浮点值λ，将二值化权重转换为λ结合的激活和梯度量化。

基于训练的量化方法在训练机制下生成量化神经网络。现有的基于训练的量化方法大致可以分为“STE”和“非STE”方法。“STE”方法包含一个不可微的离散量化函数，将全精度权重和量化权重连接起来。在反向传播过程中，使用STE穿透这个不可微函数。例如，Hubara等人的方法和Rastegari等人的方法。非STE方法指的是直接在全精度权重上使用正则化器进行量化（例如），或使用近端梯度方法进行权重投影。这些非STE量化方法在训练过程中需要大量的超参数调优，如每一步中的权重分区比例和惩罚设置。

具体来说，STE量化方法遵循一个相当简单和标准的训练协议。给定一个具有全精度权重的神经网络，一个量化函数Q()（不失一般性，设Q®为从r映射到+1如果r≥0，否则为-1），以及标注数据(x, y)，其目标是最小化训练损失。但是，由于Q的不可微性，相对于W的梯度不能使用链式法则进行计算。为使量化训练稳定，Hubara等人提出了Straight-Through-Estimator (STE) 来重新定义Q® / r:

$\frac{\partial Q(r)}{\partial r} = \begin{cases} 1 & \text{if } r \geq 1 \\ 0 & \text{otherwise} \end{cases}$

STE在基于训练的量化方法中被广泛使用，因为它提供了一个穿透Q的近似梯度，并且实现简单。然而，它不可避免地带来了梯度不匹配的问题：权重的梯度不是使用权重的值生成的，而是使用其量化值生成的。虽然STE在离散约束下提供了端到端训练的好处，但很少有工作致力于研究如何获得更好的量化训练梯度。在HWGQ和Bi-real的方法中，Q® / r是手动定义的，但它们主要关注激活量化。

为了克服梯度不匹配的问题并探讨在基于训练的方法中获得更好梯度的可能性，受启发于，我们提出通过神经网络(M)来学习Q(W) / W在量化训练中的表现。这个额外的神经网络称为元量化器，并与基础量化模型一起训练。整个过程称为元量化（MetaQuant）。特别地，在每次反向传播中，M以坐标方式将Q(W)和W作为输入，然后其输出用于计算W，以使用常见的优化方法如SGD或Adam更新权重W。在一次前向传播中，推理是使用更新后的量化权重进行的，其产生的最终输出与真实标签进行比较以进行反向计算。在这个过程中，量化权重到全精度权重的梯度传播由M处理，从而避免了不可微性和梯度不匹配的问题。此外，由元量化器生成的梯度是感知损失的，有助于量化训练的更好表现。

与常用的STE和手动设计的量化训练中的梯度传播相比，元量化器无需任何手动设计的知识就可以学习生成适当的梯度。整个过程是端到端的。元量化器可以视为任何基础模型的插件，使其在现代架构中易于实现且具有通用性。在量化训练完成后，元量化器可以被移除，并且不会占用额外的推理空间。我们在不同的量化函数（如dorefa和BWN）和优化技术（如SGD和Adam）下，与CIFAR10/100和ImageNet上的各种基础模型进行比较，以验证元量化的泛化能力。广泛的实验表明，元量化在SGD下实现了更快的收敛速度，在SGD/Adam下实现了更好的性能。

2 Related Work

Courbariaux等人提出通过在反向传播后对参数更新进行确定性和随机舍入来训练二值化网络。这一想法在后续研究中得到了进一步扩展，通过引入二值化激活。然而，这些开创性的尝试在反向传播过程中面临不可微舍入运算符的问题，这通过直接穿透舍入操作并保持梯度不变来解决。为了绕过不可微性，Leng等人通过使用交替方向乘子法（ADMM）修改了量化训练的目标函数，将训练实值参数和量化更新参数的过程分开。Zhou等人提出基于权重分区逐步量化部分参数，并通过正常训练更新未量化参数。然而，这类方法引入了更多的超参数调优，例如确定部分量化的过程，从而使量化复杂化。Bai等人在量化训练中添加了一个正则化器，将全精度权重转换为量化值。尽管这种方法简化了量化训练过程，但其优化过程涉及近端方法，使得训练成本昂贵。

3 ProblemStatement

给定一个包含n个带标签实例(x, y)的训练集，以及一个预训练的全精度基础模型f，该模型由L层构成，并由参数W = [W1, …, WL] 表示。我们定义一个预处理函数A()和一个量化函数Q()。A()将W转换为(\widetilde{W})，使其重新缩放并中心化，以便于量化。Q()使用k位离散化(\widetilde{W})到(\widehat{W})。特别地，本文研究了两种预处理函数及对应的量化方法（dorefa2，BWN）：

dorefa:
$\widetilde{W} = A(W) = \frac{\tanh(W)}{2 \max(\tanh(W))} + \frac{1}{2}$

$W^=Q(W~)=2⋅round((2k−1)W~)2k−1 \widehat{W} = Q(\widetilde{W}) = \frac{2 \cdot \text{round}((2^k - 1) \widetilde{W})}{2^k - 1}$

BWN:
$\widetilde{W} = A(W) = W$
$W^=Q(W)=1n∥W∥1sign(W) \widehat{W} = Q(W) = \frac{1}{n \|W\|_1} \text{sign}(W)$

基于训练的量化旨在训练W的量化版本，即(\widehat{W})，使得量化后的模型f的损失最小化：

$min⁡W^L(f(W^;x),y) \min_{\widehat{W}} \mathcal{L}(f(\widehat{W}; x), y)$

4 MetaQuantization

4.1 元梯度的生成

我们提出的MetaQuant将一个跨层共享的元量化器M引入到量化训练中。将W量化为(\widehat{W})（为了简化记号省略下标l）后，通过比较(f(\widehat{W}; x))与真实标签生成损失。

在反向传播中，(\widehat{W})相对于W的梯度通过链式法则计算，记为

$gW=∂L∂W^∂W^∂W g_W = \frac{\partial \mathcal{L}}{\partial \widehat{W}} \frac{\partial \widehat{W}}{\partial W}$

元量化器M接收(g_W)和W作为输入，并输出(\widehat{W})相对于W的梯度，记为(\widetilde{g}_W)，表示为：

$g~W=∂L∂W=M(gW⋅∂W^∂W) \widetilde{g}_W = \frac{\partial \mathcal{L}}{\partial W} = M(g_W \cdot \frac{\partial \widehat{W}}{\partial W})$

梯度(\widetilde{g}_W)进一步用于计算(\mathcal{L})相对W的梯度，记为(g_W)，其中(g_W)通过以下方式计算：

$g_W = \frac{\partial \mathcal{L}}{\partial W}$

这里，(\frac{\partial \widehat{W}}{\partial W})取决于W和(\widehat{W})之间的预处理函数：

对于dorefa：

$∂W^∂W=1max⁡(tanh⁡(W))⋅(1−tanh⁡2(W)) \frac{\partial \widehat{W}}{\partial W} = \frac{1}{\max(\tanh(W)) \cdot (1 - \tanh^2(W))}$

对于BWN：

$∂W^∂W=1 \frac{\partial \widehat{W}}{\partial W} = 1$

这个过程被称为校准。使用(g_W)更新W之前，(g_W)首先根据不同的优化方法进行处理，以产生每个权重的最终更新值。这个过程被称为梯度精炼，记为(\Delta(g_W))。特别地，对于SGD，(\Delta(g_W) = g_W)。对于其他优化方法如Adam，可以实现为(\Delta(g_W) = g_W + \text{residual})，其中“residual”根据不同的梯度精炼方法计算。最后，全精度权重W更新为：

$W_{t+1} = W_t - \eta (\Delta(g_W))$

其中t表示第t次训练迭代，(\eta)是学习率。图1展示了MetaQuant的整体过程。

与[1]相比，直接学习(g_W)，MetaQuant构建了一个神经网络来学习(\widetilde{g}_W)，由于量化函数的不可微性，在量化训练中无法直接计算。我们的工作解决了不可微性问题，并且对不同的优化方法通用。MetaQuant的工作原理和原因在附录7.2中进行了研究。

4.2 元量化器的训练

类似于[1]，我们提出的元量化器是一个坐标级的神经网络，这意味着每个权重参数都是独立处理的。对于梯度(\widetilde{g}W)中的单个权重索引i，(W_i)接收其对应的梯度(g{W_i}) 通过以下方式:

$\widetilde{g}_{W_i} = M(g_{W_i}, W_i)$

为了高效处理，在推理过程中，将(3)中的输入按大小为1的批次排列。特别地，假设(\widehat{W})来自一个形状为(R^{o \times i \times k \times k})的卷积层，其中o，i和k分别表示输出通道数，输入通道数和内核大小。然后(\widehat{W})，W和对应的梯度共享相同的形状，这是对(3)中输入的重塑为(R^{o \times i \times k^2 \times 1})。回顾(5)和(4)，M的输出被结合到更新的(W_t)的值中，然后在下一次迭代的推理中量化。因此，M与最终的量化训练损失相关，通过从最终损失反向传播的梯度更新。通过引入元量化器来生成(\widetilde{g}_W)，MetaQuant不仅解决了基础模型中参数的不可微性问题，还提供了整个网络的端到端训练的好处。此外，元量化器是损失感知的，因此它被训练生成更准确的权重更新以减少最终损失，这探索了如何修改梯度以辅助量化训练。图2展示了将元量化器结合到基础模型的量化训练中的详细过程，这形成了从最终损失到(\mathcal{L})的可微路径。在W的量化训练同时，(\mathcal{L})也在每次训练迭代t中被学习：

前向传播：

$W_t = A(W_t) = A(W_{t-1})$

损失：

$L=f(Q(Wt^);x,y) \mathcal{L} = f(Q(\widehat{W_t}); x, y)$

反向传播：

$\frac{\partial \mathcal{L}}{\partial W_t} = \frac{\partial \mathcal{L}}{\partial W_t} = M(g_{t-1}, W_{t-1}) \frac{\partial W_{t-1}}{\partial W_t}$

在前向传播中，我们使用(W_{t-1})和元梯度的组合来表示(W_t)，以便结合M。特别是在(6)中，元梯度来自M的输出，首先通过乘法计算W的梯度，然后通过优化进行精炼。最后，通过学习率调整成为元梯度。(8)计算元梯度的梯度，在这里W是可微的，因为A是可微的。此外，选择了一个可微的元神经网络。W_t将在反向传播后实际更新，这可以视为权重的延迟更新。

5 Experiment

4.3 元量化器的设计

元量化器 (M) 是一个参数化的、可微分的神经网络，用于生成元梯度。它可以被视为STE（Straight-Through Estimator）的推广。例如，如果 (M) 根据 (W) 的绝对值幅度剪辑 (g_W)，则 (M) 退化为STE：(g_W = M(g_W, W) = g_W \cdot \mathbf{1}_{|W|})。

我们设计了三种不同架构的元量化器。第一种架构仅使用一个由两层或多层全连接层组成的神经网络。它只需要 (g_W) 作为输入：

$\text{FCGrad} : M(g_W) = \text{FCs}(g_W) \quad (9)$

其中 (\text{FCs}) 表示非线性激活函数。由于先前成功的实验结果表明，通过考虑 (W) 的值生成一个好的 (g_W) 是可行的。基于这一观察，我们构建了另外两种以 (W) 作为输入的元量化器架构，并将这些神经网络的输出与 (g_W) 相乘，以结合其后续步骤中的梯度信息。具体来说，其中一种基于全连接（FC）层：

$\text{MultiFC} : M(g_W, W) = g_W \cdot \text{FCs}(W) \quad (10)$

另一种网络结合了LSTM和FC层来构建 (M)，这是受到[1]中使用基于记忆的神经网络作为元学习器的启发：

$\text{LSTMFC} : M(g_W, W) = g_W \cdot \text{FCs}(\text{FCs}(\text{LSTM}(\text{LSTM}(W)))) \quad (11)$

当使用LSTM作为元量化器时，每个权重的坐标都跟踪由LSTM生成的隐藏状态，这些隐藏状态包含了 (g_W) 和 (W) 的历史信息。元量化器的内存消耗和详细的超参数研究在附录7.1和7.3中讨论。

4.4 算法和实现细节

MetaQuant的详细过程如算法1所示。首先构建并随机初始化一个共享的元量化器 (M)。在每次训练迭代中，第2-6行描述了前向过程：对于每一层，从上一迭代获得的 (g_W) 和 (W) 被输入到 (M) 中以生成元梯度 (g_W) 进行推理，如第3-5行所示。由于在第一次迭代中没有计算 (g_W)，正常的量化训练在第一次迭代中进行：(W = Q(\widetilde{W}) = Q[A(W)]) 以替代第4行。第7-9行展示了反向过程：通过误差反向传播可以获得 (W) 的梯度，如第7行所示。在反向过程中，获得当前迭代的 (g_W) 和 (W) 并将其输出保存以供下一次迭代计算，如第7-8行所述，用 (g_W^{t+1}) 表示。通过将 (M) 合并到推理图中，其梯度在第9行中获得。最后，(g_W) 用于计算 (g_W)，然后通过不同的优化方法进行处理，导致第10-12行中 (W) 的更新。在第一次迭代中，由于缺少 (g_W)，(W) 的权重更新未进行。注意元量化器在第13行中进行更新。

算法 1: MetaQuant 量化训练
输入: 训练集 (x, y), 初始权重 W, 预处理函数 A(), 量化函数 Q(), 元量化器 M, 学习率 η
输出: 量化后的权重 W
初始化: 随机初始化元量化器 M

1:  for t = 1 to 最大迭代次数 do
2:      for l = 1 to L do
3:          if t == 1 then
4:              W_l = Q[A(W_l)]
5:          else
6:              g_W, W = M(g_W, W)
7:          end if
8:      end for
9:      计算损失 \(\mathcal{L}\) 并反向传播
10:     for l = 1 to L do
11:         计算 W_l 的梯度 g_W
12:         根据优化方法更新 W_l
13:     end for
14:     更新元量化器 M
15: end for

5 Experiment

5.1 实验设置

MetaQuant 关注于在训练过程中穿透不可微的量化函数的方法。我们使用以下两种前向量化方法和两种优化方法与STE进行对比实验：

前向量化方法：

dorefa [17]
BWN [13]

优化方法：

SGD
Adam [8]

当使用dorefa或BWN作为前向量化函数并使用STE作为后向方法进行量化训练时，分别成为[17]或[13]中提出方法的权重量化版本。我们使用了三个基准数据集，包括ImageNet ILSVRC-2012和CIFAR10/100。对于深度架构，我们在CIFAR10上实验了ResNet20/32/44。由于CIFAR10/100共享相同的输入维度，我们在CIFAR100中将ResNet56/110的最后一个全连接层的输出维度从10修改为100。对于ImageNet，使用ResNet18进行比较。

在所有实验中，我们对网络中的所有层进行1位量化：每层仅包含两个值。在CIFAR10/100的实验中，我们将基础模型的初始学习率设为(\eta = 1 \times 10^{-3})，元量化器的初始学习率也设为(\eta = 1 \times 10^{-3})。为了公平比较，我们将所有实验的总训练轮数设为100，每30轮后学习率减小10倍。对于ImageNet，基础模型使用dorefa和BWN的初始学习率设为(\eta = 1 \times 10^{-4})，元量化器的初始学习率也设为(\eta = 1 \times 10^{-3})。当训练进行到10/20轮时，学习率分别减小到(\eta = 1 \times 10^{-5})和(\eta = 1 \times 10^{-6})。总轮数设为30。在基础模型的学习率变化时，元量化器的学习率也相应减小到(\eta = 1 \times 10^{-4})和(\eta = 1 \times 10^{-5})。批次大小为128（适用于CIFAR和ImageNet）。所有实验进行了5次，并报告了最后10/5轮测试准确性的统计数据，作为CIFAR/ImageNet数据集中提出方法和基线方法的表现。我们还通过训练损失曲线展示了不同方法之间的经验收敛速度。

在CIFAR实验中，不同实现的MetaQuant的详细超参数如下：

在MultiFC中，使用2层全连接层，隐藏层大小为100，不使用非线性激活。
在LSTMFC中，使用1层LSTM和1层全连接层，隐藏维度设为100。
在FCGrad中，使用2层全连接元模型，隐藏层大小为100，不使用非线性激活。
在ImageNet实验中，我们使用MultiFC/FCGrad，分别使用2/1层全连接层，隐藏维度为100。

5.2 实验结果与分析

CIFAR100比CIFAR10更具挑战性，因为它包含更多的细粒度类别，总共有100个类别。表2显示了MetaQuant和STE在使用不同前向量化方法和优化方法时在CIFAR100上的总体实验结果。与CIFAR10类似，MetaQuant在所有情况下都比STE表现出显著的提升，这表明MetaQuant在比传统方法更具挑战性的任务中有显著的改进。

5.3 经验收敛分析

在这个实验中，我们比较了MetaQuant和STE在训练过程中的不同变体的性能，以展示它们的收敛速度。我们以使用dorefa的ResNet20为例。如图3所示，在相同任务和前向量化方法下，MetaQuant在使用SGD时显示出比STE更显著的收敛优势，包括损失下降速度显著更快和明显更低的损失值。在使用Adam时，尽管所有方法显示出相似的下降速度，MetaQuant方法最终达到了更低的损失值，这也反映在表1中的测试准确性报告中。总体而言，MetaQuant在使用不同的前向量化和优化方法时显示出比STE更好的收敛性。选择SGD时，改进更为明显。

我们推测，SGD和Adam之间的性能差异是由于以下原因：SGD仅使用从 (g_W) 校准的梯度更新全精度权重，这直接反映了元量化器的输出，而不是STE。Adam聚合了 (g_W) 的历史信息并归一化当前梯度，在某种程度上缩小了元量化器和STE之间的差异。更多关于训练准确性和在更多任务上的测试准确性比较见附录7.4。

5.4 与非STE训练量化方法的性能比较

MetaQuant旨在通过学习更好的梯度来提升基于训练的量化，以穿透不可微量化函数。一些先进的量化方法避免了离散量化。在本节中，我们将MetaQuant与非STE训练量化方法进行比较：ProxQuant ([2])，LAB ([5])，以证明传统的STE训练量化可以通过使用MetaQuant实现更好的性能。

由于所使用的初始全精度模型的差异，我们仅报告量化后测试准确性的性能下降（越小越好）。我们使用ResNet20/32/44与ProxQuant进行比较，使用其提出的架构3与LAB在CIFAR10上进行比较，所有层都量化为二进制值。如表4所示，MetaQuant显示出比两个基线更好的性能。

在表4的第3行中，我们使用ImageNet数据集比较了ELQ ([16]) 和TTQ ([18])。尽管ELQ表现优异，但它是组合了一系列先前的量化方法和增量量化技巧的结果。MetaQuant更关注于如何改进基于STE的训练量化，没有任何额外的损失和训练技巧。TTQ是一种非对称的三值化方法，使用0作为三值点。MetaQuant遵循dorefa，使用对称量化，这有助于提高推理效率。

5.5 MetaQuant训练分析

MetaQuant的训练涉及元量化器的计算。为了分析额外的训练时间，我们在CIFAR10中使用ResNet20比较了使用MultiFC和DoReFa的MetaQuant与使用STE的训练时间（在Intel Xeon CPU E5-1650和GeForce GTX 750 Ti上）。MetaQuant完成一次训练迭代需耗时51.15秒，而基线方法耗时38.17秒。然而，在实际部署中，元量化器被移除，MetaQuant能够在不增加任何额外推理时间的情况下提供更好的测试性能。

6 Conclusion

在本文中，我们提出了一种新方法（MetaQuant），通过元量化器在基于训练的量化中学习梯度，以穿透不可微的量化函数。这个元网络足够通用，可以融入各种基础模型中，并可以使用基础模型的损失进行更新。我们提出了三种类型的元量化器，并展示了通过这些模块生成的元梯度能够在不同的前向量化函数和优化方法下，提供更快的收敛速度和更好的最终量化性能。