[论文解读] Searching for Low-Bit Weights in Quantized Neural Networks

Zhaohui Yang, Wang, Yunhe+5 more|arXiv (Cornell University)|Sep 18, 2020

Advanced Neural Network Applications参考文献 52被引用 35

一句话总结

本文提出一种可微分的权重搜索方法，用于训练量化神经网络，通过将低位权重表示为离散值的概率分布，从而实现端到端优化，而无需对非微分的量化进行梯度估计。

ABSTRACT

Quantized neural networks with low-bit weights and activations are attractive for developing AI accelerators. However, the quantization functions used in most conventional quantization methods are non-differentiable, which increases the optimization difficulty of quantized networks. Compared with full-precision parameters (i.e., 32-bit floating numbers), low-bit values are selected from a much smaller set. For example, there are only 16 possibilities in 4-bit space. Thus, we present to regard the discrete weights in an arbitrary quantized neural network as searchable variables, and utilize a differential method to search them accurately. In particular, each weight is represented as a probability distribution over the discrete value set. The probabilities are optimized during training and the values with the highest probability are selected to establish the desired quantized network. Experimental results on benchmarks demonstrate that the proposed method is able to produce quantized neural networks with higher performance over the state-of-the-art methods on both image classification and super-resolution tasks.

研究动机与目标

• 通过低位权重提升内存与计算效率来激励量化网络。
• 解决来自不可微量化函数的优化挑战。
• 提出一个可微分的权重搜索框架，用于端到端学习量化权重。
• 降低训练（连续）与推理（离散）之间的量化差距。
• 在图像分类和超分辨率基准测试中展示了更高的准确性。

提出的方法

• 将每个权重表示为一组低位值的概率分布。
• 使用一个连续的辅助张量 A 来学习该分布，并对值执行 softmax（Eq. 5）。
• 在训练中，将期望的连续权重 W_c 计算为 W_c = sum_i P_i v_i (Eq. 6)，推理时选择最大概率值（Eq. 7）。
• 梯度相对于 A 计算，避免对量化函数进行梯度估计。
• 逐步降低温度以使分布变尖、减少量化差距（温度极限定理）。
• 引入状态批量归一化（SBN）以使连续训练输出和离散推理输出的统计量对齐。

实验结果

研究问题

• RQ1 可微分的离散低位权值放松是否能够在没有基于 STE 的梯度估计的情况下实现量化网络的端到端优化？
• RQ2 逐步降低温度并使用状态批量归一化是否能够降低量化差距并提升在分类和超分辨率任务上的准确性？
• RQ3 与最先进的量化方法相比，所提出的 SLB 方法在不同比特宽度和体系结构上的性能表现如何？
• RQ4 学习得到的权重分布在用于图像分类的 CNN 以及超分辨率模型中是否同样有效？

主要发现

• 在 CIFAR-10 上使用 ResNet 与 VGG-Small 主干，在多种比特宽度设置下实现比若干最先进的量化方法更高的准确率。
• 在 ImageNet（ILSVRC2012）上使用 ResNet18，SLB 采用 1-bit 权重和激活实现 Top-1 61.3% 和 Top-5 83.1%，优于若干基线。
• SLB 在 1/1、1/2、1/4、1/8、1/32 位宽配置中持续改进，并扩展到 2/2、2/4、2/8、2/32、4/4、4/8、4/32。
• 超分辨率实验表明，在放大因子 2 和 3 时，SLB 的 PSNR 高于 DoReFa，结果更接近全精度。
• 消融实验表明 State Batch Normalization 在使用离散权重时有助于保持性能。