[论文解读] Position-based Scaled Gradient for Model Quantization and Pruning

Jangho Kim, KiYoon Yoo+1 more|arXiv (Cornell University)|May 22, 2020

Advanced Neural Network Applications被引用 14

一句话总结

本文提出了一种基于位置的缩放梯度（PSG）方法，这是一种新颖的正则化技术，通过根据权重位置缩放梯度，引导训练过程朝向更适合压缩的极小值点。PSGD通过使全精度模型更易于在不微调的情况下进行推理后量化和剪枝，减少了全精度与低精度权重分布之间的差距，在CIFAR-100和ImageNet上实现了4比特量化下的最先进性能。

ABSTRACT

We propose the position-based scaled gradient (PSG) that scales the gradient depending on the position of a weight vector to make it more compression-friendly. First, we theoretically show that applying PSG to the standard gradient descent (GD), which is called PSGD, is equivalent to the GD in the warped weight space, a space made by warping the original weight space via an appropriately designed invertible function. Second, we empirically show that PSG acting as a regularizer to a weight vector is favorable for model compression domains such as quantization and pruning. PSG reduces the gap between the weight distributions of a full-precision model and its compressed counterpart. This enables the versatile deployment of a model either as an uncompressed mode or as a compressed mode depending on the availability of resources. The experimental results on CIFAR-10/100 and ImageNet datasets show the effectiveness of the proposed PSG in both domains of pruning and quantization even for extremely low bits. The code is released in Github.

研究动机与目标

• 解决全精度与量化权重分布之间的根本性分布差距，该差距阻碍了有效的推理后量化。
• 实现在无需微调或数据访问的情况下，单一模型在全精度和压缩模式下灵活部署。
• 开发一种训练方法，通过正则化使收敛过程趋向特定目标点，从而内在地生成适合低比特量化和剪枝的权重。
• 提供一种理论基础坚实的替代方案，以应对现有量化感知训练（QAT）和推理后量化（PTQ）方法所依赖的额外微调或数据。
• 证明PSGD能够找到更兼容量化与剪枝所用离散权重网格的更尖锐极小值点。

提出的方法

• 提出基于位置的缩放梯度（PSG），其中梯度根据当前权重与目标量化值（例如最近的4比特网格点）之间的L2距离进行缩放。
• 将PSGD定义为使用PSG而非标准梯度的随机梯度下降，从而有效将权重正则化至适合压缩的位置。
• 理论上证明PSGD等价于在经扭曲的权重空间中的标准SGD，其中扭曲函数为可逆函数，且设计用于将权重拉向目标网格点。
• 使用可逆且可微的函数对原始权重空间进行扭曲，使得在原始空间中的梯度更新等价于在扭曲空间中的标准梯度下降。
• 在剪枝中，目标点设为零，通过基于位置的梯度缩放实现稀疏训练，使权重被正则化趋向于零。
• 在初始训练中应用PSGD，以生成一个保持高精度的全精度模型，同时可直接量化为低比特格式。

实验结果

研究问题

• RQ1 基于权重位置的梯度缩放是否能减少全精度与低精度模型之间的分布差距？
• RQ2 PSGD是否能找到比标准SGD更兼容量化与剪枝的极小值点？
• RQ3 使用PSGD训练的模型是否能在无需额外微调的情况下，在全精度和低精度模式下均保持高精度？
• RQ4 在精度与部署灵活性方面，PSGD相较于现有QAT与PTQ方法表现如何？
• RQ5 PSGD在极低比特量化（如4比特）和结构化剪枝中是否有效？

主要发现

• PSGD显著减少了全精度与4比特量化权重之间的分布差距，如图1(b)所示，PSGD训练模型的权重分布与低精度模型高度相似。
• 在CIFAR-100上，PSGD在4比特权重与激活量化（W4A4）下实现了66.5%的top-1准确率，优于标准SGD和同期的PTQ方法，尤其在作为预训练模型使用时表现更优。
• 当与LAPQ PTQ方法结合时，PSGD预训练模型在ImageNet W4A4上达到66.5%的准确率，分别超出LAPQ-only（60.3%）和PSGD-only（63.3%）3.1和6.2个百分点。
• PSGD训练的模型表现出更尖锐的极小值点，其Hessian矩阵中具有9倍于标准SGD的大于10⁻³的正特征值（λ > 10⁻³），表明其收敛于更陡峭的谷底。
• 该方法实现了单一模型在全精度与压缩模式下的直接部署，无需在资源受限环境中进行微调或访问数据。
• PSGD在权重量化与基于幅值的剪枝中均表现有效，证明其作为统一正则化方法在模型压缩中的多功能性。