[论文解读] Sharpness-aware Quantization for Deep Neural Networks

Jing Liu, Jianfei Cai+1 more|arXiv (Cornell University)|Nov 24, 2021

Advanced Neural Network Applications参考文献 61被引用 8

一句话总结

本文提出了一种新颖的方法——清晰度感知量化（SAQ），通过将量化噪声与对抗性扰动统一建模为模型权重上的噪声源，联合优化模型量化与损失曲面平滑化。SAQ 在低比特量化模型中提升了泛化性能，实现了当前最优（SOTA）结果，包括在 ImageNet 上 4-bit ViT-B/16 实现 1.2% 的 Top-1 准确率提升，以及在 4-bit ResNet-50 上实现 0.9% 的提升，同时保持了与 AdamW 等标准优化器相当的训练效率。

ABSTRACT

Network quantization is a dominant paradigm of model compression. However, the abrupt changes in quantized weights during training often lead to severe loss fluctuations and result in a sharp loss landscape, making the gradients unstable and thus degrading the performance. Recently, Sharpness-Aware Minimization (SAM) has been proposed to smooth the loss landscape and improve the generalization performance of the models. Nevertheless, directly applying SAM to the quantized models can lead to perturbation mismatch or diminishment issues, resulting in suboptimal performance. In this paper, we propose a novel method, dubbed Sharpness-Aware Quantization (SAQ), to explore the effect of SAM in model compression, particularly quantization for the first time. Specifically, we first provide a unified view of quantization and SAM by treating them as introducing quantization noises and adversarial perturbations to the model weights, respectively. According to whether the noise and perturbation terms depend on each other, SAQ can be formulated into three cases, which are analyzed and compared comprehensively. Furthermore, by introducing an efficient training strategy, SAQ only incurs a little additional training overhead compared with the default optimizer (e.g., SGD or AdamW). Extensive experiments on both convolutional neural networks and Transformers across various datasets (i.e., ImageNet, CIFAR-10/100, Oxford Flowers-102, Oxford-IIIT Pets) show that SAQ improves the generalization performance of the quantized models, yielding the SOTA results in uniform quantization. For example, on ImageNet, SAQ outperforms AdamW by 1.2% on the Top-1 accuracy for 4-bit ViT-B/16. Our 4-bit ResNet-50 surpasses the previous SOTA method by 0.9% on the Top-1 accuracy.

研究动机与目标

• 为解决由于离散权重更新导致的损失曲面尖锐性所引发的低精度量化模型不稳定与性能下降问题。
• 探索将清晰度感知最小化（SAM）与网络量化相结合，这一问题在模型压缩领域此前尚未被研究。
• 构建一个统一框架，将量化与对抗性扰动视为模型权重上的噪声源，从而系统性地分析其相互作用。
• 设计一种高效的训练策略，在保持 SAM 在量化设置中优势的同时，将额外计算开销降至最低。
• 在多种架构（CNN 与 Transformer）和数据集上进行实证验证，证明 SAQ 在标准与迁移学习设置下均能持续提升性能。

提出的方法

• 将量化与 SAM 分别建模为向模型权重引入量化噪声与对抗性扰动，从而实现对二者交互作用的统一视角。
• 基于量化噪声与对抗性扰动之间的依赖关系，定义三种 SAQ 情况：独立、依赖与混合，并进行理论与实证比较。
• 提出一种高效的训练策略，通过重用梯度信息来减少对抗性扰动计算的开销，使训练速度接近 SGD 或 AdamW。
• 在量化模型的端到端训练过程中应用 SAQ，联合优化权重量化与基于扰动的清晰度最小化。
• 采用两步优化循环：首先在半径内计算使损失最大化的扰动，然后更新权重以最小化扰动后的损失，该方法适配于量化权重空间。
• 使用标准优化器（如 AdamW）实现该方法，并将其应用于多个基准测试中的卷积网络与视觉 Transformer。

实验结果

研究问题

• RQ1 清晰度感知最小化（SAM）能否有效应用于量化神经网络？量化与对抗性扰动之间的相互作用会带来哪些挑战？
• RQ2 量化噪声与对抗性扰动之间的依赖关系如何影响量化模型的泛化性能？
• RQ3 SAQ 的最优公式是什么？即在三种推导出的情况（独立、依赖、混合）中，哪一种在性能与训练效率之间实现了最佳平衡？
• RQ4 SAQ 是否能在均匀量化中实现 SOTA 性能，同时仅引入极少的额外训练成本？
• RQ5 与先应用 SAM 再进行量化（SAM → SGD）的串行方法相比，联合优化损失曲面平坦性与量化是否能带来更好的泛化与迁移学习性能？

主要发现

• 在 ImageNet 上，SAQ 在 4-bit ViT-B/16 上相比 AdamW 实现了 1.2% 的 Top-1 准确率提升，证明了其在均匀量化中的 SOTA 性能。
• 对于 4-bit ResNet-50，SAQ 在 Top-1 准确率上比之前 SOTA 方法高出 0.9%，证实了其在卷积架构上的有效性。
• 在迁移学习中，SAQ 相比 SGD 训练的量化模型，在 Oxford-IIIT Pets 上提升了 1.0% 的 Top-1 准确率，在 CIFAR-100 上提升了 0.7%。
• 高效的训练策略使 SAQ 的训练时间控制在 SGD 的 11% 以内，即使增加了扰动计算，仍具备扩展至大模型的能力。
• 在 2-bit ResNet-18 上，SAQ 相比先对全精度模型应用 SAM 再量化（SAM → SGD）的串行方法，Top-1 准确率高出 0.4%，证明了联合优化的优势。
• 全面的消融实验表明，混合情况（即扰动依赖于量化噪声）的 SAQ 表现最佳，凸显了建模二者交互作用的重要性。