[论文解读] Mirror Descent View for Neural Network Quantization

Thalaiyasingam Ajanthan, Kartik Gupta+3 more|arXiv (Cornell University)|Oct 18, 2019

Advanced Neural Network Applications参考文献 51被引用 4

一句话总结

本文通过将连续参数解释为量化参数的对偶变量，提出了一种用于神经网络量化的一致优化方法——镜像下降（MD）框架。该方法从投影函数中推导出有效的镜像映射，通过辅助变量确保数值稳定性，并揭示了MD与梯度直通估计器（Straight-Through Estimator, STE）之间的理论等价性，从而在CIFAR-10/100、TinyImageNet和ImageNet上使用VGG-16、ResNet-18和MobileNetV2模型实现了最先进（SOTA）的准确率。

ABSTRACT

Quantizing large Neural Networks (NN) while maintaining the performance is highly desirable for resource-limited devices due to reduced memory and time complexity. It is usually formulated as a constrained optimization problem and optimized via a modified version of gradient descent. In this work, by interpreting the continuous parameters (unconstrained) as the dual of the quantized ones, we introduce a Mirror Descent (MD) framework for NN quantization. Specifically, we provide conditions on the projections (i.e., mapping from continuous to quantized ones) which would enable us to derive valid mirror maps and in turn the respective MD updates. Furthermore, we present a numerically stable implementation of MD that requires storing an additional set of auxiliary variables (unconstrained), and show that it is strikingly analogous to the Straight Through Estimator (STE) based method which is typically viewed as a "trick" to avoid vanishing gradients issue. Our experiments on CIFAR-10/100, TinyImageNet, and ImageNet classification datasets with VGG-16, ResNet-18, and MobileNetV2 architectures show that our MD variants obtain quantized networks with state-of-the-art performance. Code is available at https://github.com/kartikgupta-at-anu/md-bnn.

研究动机与目标

• 解决在资源受限设备上部署深度神经网络时的量化挑战，同时保持模型准确率。
• 通过将量化问题置于镜像下降优化框架中，为量化提供理论基础。
• 建立广泛使用的梯度直通估计器（STE）与镜像下降之间的一致性联系，解释其经验成功的原因。
• 设计一种数值稳定的MD实现，适用于训练量化网络。
• 通过所提出的基于MD的量化方法，在多种数据集和网络架构上展示最先进性能。

提出的方法

• 将网络量化建模为在离散集合上的约束优化问题，随后通过在无约束参数的对偶空间中使用镜像下降（MD）进行重述。
• 在投影函数P满足特定条件的前提下，从P函数中解析推导出镜像映射，该P函数用于将无约束参数映射到量化值。
• 通过使用单调递增的退火超参数，引入时变镜像映射，以确保即使在非凸量化集合下也能收敛到离散解。
• 提出一种数值稳定的MD变体，通过保持辅助无约束变量，其更新规则与梯度直通估计器（STE）高度相似。
• 使用标准优化器（Adam、SGD）配合学习率调度和权重衰减实现该方法，最终通过符号函数或四舍五入操作对权重进行离散化。
• 针对具有时变镜像映射的凸情况，分析了理论收敛性，表明其收敛速率与标准MD相同。

实验结果

研究问题

• RQ1 能否通过从投影函数中推导镜像映射，系统性地将镜像下降应用于神经网络量化？
• RQ2 所提出的具有时变镜像映射的MD框架如何在非凸量化设置下确保收敛到离散解？
• RQ3 在训练量化网络的背景下，梯度直通估计器（STE）与镜像下降之间存在何种理论关系？
• RQ4 所提出的基于MD的量化方法是否在多种数据集和架构上均实现了最先进性能？
• RQ5 能否设计一种数值稳定的MD实现，使其在实际效果上与STE相当？

主要发现

• 所提出的基于MD的量化方法在使用VGG-16的CIFAR-10和CIFAR-100上实现了最先进准确率，无论在二值化还是三值化设置下均优于先前方法。
• 在使用ResNet-18的TinyImageNet上，MD变体的Top-1准确率高于基线方法，表明其在更大规模数据集上的鲁棒性。
• 在使用ResNet-18的ImageNet上，从零开始训练的MD-tanh-s*变体达到了76.8%的Top-1准确率，与或超过先前最先进结果。
• 带有辅助变量的数值稳定MD变体在经验上与梯度直通估计器（STE）等价，从而将STE的实际成功解释为MD的一个特例。
• 理论分析证实，在凸设置下，使用时变镜像映射的MD收敛速率与标准MD相同，验证了该框架的收敛性特性。
• 超参数调优表明，学习率调度和beta缩放参数的退火显著提升了性能，尤其在二值化设置中效果明显。