[论文解读] Training with Quantization Noise for Extreme Fixed-Point Compression

Angela Fan, Pierre Stock+5 more|arXiv (Cornell University)|Apr 15, 2020

Advanced Neural Network Applications被引用 9

一句话总结

本文提出通过在训练中引入量化噪声，实现超越标准 int8 的极端定点压缩，采用对随机子集的随机权重量化以实现无偏梯度。该方法在准确率-压缩率权衡上达到当前最优性能，使 14MB 的 RoBERTa 模型在 MNLI 上达到 82.5% 的准确率，使 3.3MB 的 EfficientNet-B3 在 ImageNet 上达到 80.0% 的 top-1 准确率。

ABSTRACT

We tackle the problem of producing compact models, maximizing their accuracy for a given model size. A standard solution is to train networks with Quantization Aware Training, where the weights are quantized during training and the gradients approximated with the Straight-Through Estimator. In this paper, we extend this approach to work beyond int8 fixed-point quantization with extreme compression methods where the approximations introduced by STE are severe, such as Product Quantization. Our proposal is to only quantize a different random subset of weights during each forward, allowing for unbiased gradients to flow through the other weights. Controlling the amount of noise and its form allows for extreme compression rates while maintaining the performance of the original model. As a result we establish new state-of-the-art compromises between accuracy and model size both in natural language processing and image classification. For example, applying our method to state-of-the-art Transformer and ConvNet architectures, we can achieve 82.5% accuracy on MNLI by compressing RoBERTa to 14MB and 80.0 top-1 accuracy on ImageNet by compressing an EfficientNet-B3 to 3.3MB.

研究动机与目标

• 为解决在极端压缩约束下（尤其是超越 int8 量化）实现高模型准确率的挑战。
• 克服在极端量化过程中，Straight-Through Estimator (STE) 由于梯度近似误差导致的性能下降问题。
• 通过引入受控的随机噪声，实现对低精度定点表示（包括 Product Quantization）的有效训练。
• 在自然语言处理和计算机视觉领域建立新的准确率-压缩率权衡的 SOTA。

提出的方法

• 在每次前向传播中，对权重的不同随机子集进行随机量化，以引入受控噪声并实现无偏梯度流动。
• 使用噪声控制机制，平衡梯度准确性与压缩效率之间的权衡。
• 在训练中应用该方法，结合考虑随机量化过程的修改损失函数。
• 将该方法扩展至极端压缩技术（如 Product Quantization），传统 STE 因高近似误差而失效的场景。
• 为未量化权重保持全精度梯度，同时允许量化权重通过注入噪声的梯度进行更新。
• 将该方法无缝集成到标准 Transformer 和 ConvNet 训练流水线中，无需架构修改。

实验结果

研究问题

• RQ1 在极端定点量化中，通过随机子集选择的随机权重量化是否能改善梯度估计？
• RQ2 引入受控量化噪声是否能实现如 Product Quantization 等压缩方法的有效训练？
• RQ3 所提出的方法是否能在自然语言处理和图像分类任务中均实现 SOTA 的准确率-压缩率权衡？
• RQ4 该方法在极端压缩下（如 14MB RoBERTa 或 3.3MB EfficientNet-B3）表现如何？

主要发现

• 在将 RoBERTa 压缩至 14MB 时，该方法在 MNLI 基准上实现 82.5% 的准确率，创下该模型尺寸下的新 SOTA。
• 通过将 EfficientNet-B3 压缩至仅 3.3MB，该方法在 ImageNet 上实现 80.0% 的 top-1 准确率，优于先前方法。
• 该方法即使在极端量化（如 Product Quantization）下也能实现稳定训练和高性能，而传统 STE 在此类场景下会失效。
• 通过采用随机子集量化，该方法提供无偏梯度，从而降低低精度训练中的近似误差。