[论文解读] Learning low-precision neural networks without Straight-Through Estimator(STE)

Zhigang Liu, Matthew Mattina|arXiv (Cornell University)|Mar 4, 2019

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

一句话总结

本文提出了一种名为 alpha-blending（AB）的新方法，用于在不依赖直通估计器（Straight-Through Estimator, STE）的情况下训练低精度神经网络。AB 在损失函数中用全精度权重与量化权重的仿射组合替代了量化权重，逐步将混合系数 α 从 0 增加到 1，从而实现通过全精度路径的稳定梯度传播。该方法实现了当前最优的精度表现，在 4 位权重和 8 位激活的 ImageNet 模型上，相比以往基于 STE 的量化方法，top-1 准确率提升了 2.93%。

ABSTRACT

The Straight-Through Estimator (STE) is widely used for back-propagating gradients through the quantization function, but the STE technique lacks a complete theoretical understanding. We propose an alternative methodology called alpha-blending (AB), which quantizes neural networks to low-precision using stochastic gradient descent (SGD). Our method (AB) avoids STE approximation by replacing the quantized weight in the loss function by an affine combination of the quantized weight w_q and the corresponding full-precision weight w with non-trainable scalar coefficient $α$ and $1-α$. During training, $α$ is gradually increased from 0 to 1; the gradient updates to the weights are through the full-precision term, $(1-α)w$, of the affine combination; the model is converted from full-precision to low-precision progressively. To evaluate the method, a 1-bit BinaryNet on CIFAR10 dataset and 8-bits, 4-bits MobileNet v1, ResNet_50 v1/2 on ImageNet dataset are trained using the alpha-blending approach, and the evaluation indicates that AB improves top-1 accuracy by 0.9%, 0.82% and 2.93% respectively compared to the results of STE based quantization.

研究动机与目标

• 解决直通估计器（Straight-Through Estimator, STE）在量化神经网络反向传播中缺乏理论基础的问题。
• 开发一种训练方法，避免使用 STE 的经验近似，通过在量化过程中保持全精度权重的稳定梯度更新。
• 评估所提方法在低精度模型上的有效性，包括 1 位 BinaryNet 以及 4 位/8 位 MobileNet 和 ResNet 架构。
• 证明 alpha-blending 在推理精度上优于当前最先进的基于 STE 的量化技术。

提出的方法

• 在损失函数中，用仿射组合 w_ab = (1 - α)w + αwq 替换量化权重 wq，其中 w 为全精度权重，α 为不可训练的标量。
• 在训练过程中，逐步将 α 从 0.0 增加到 1.0，使模型逐步从全精度行为过渡到低精度行为。
• 对全精度部分 (1 - α)w 执行梯度更新，确保反向传播过程中梯度通过可微路径流动。
• 采用逐层或逐通道的量化策略，将 AB 应用于不同网络组件，包括权重和激活。
• 使用随机梯度下降（SGD）训练模型，结合修改后的损失函数，全程保持全精度权重作为可训练参数。
• 训练完成后，将最终模型转换为低精度推理格式，权重和激活根据最终 α = 1.0 进行量化。

实验结果

研究问题

• RQ1 能否设计一种不使用 STE 的训练方法，同时保持或提升量化精度？
• RQ2 全精度权重与量化权重的渐进式混合是否能实现低精度网络中稳定且有效的反向传播？
• RQ3 在 CIFAR10 和 ImageNet 等标准基准上，alpha-blending 与基于 STE 的量化在 top-1 准确率上的表现如何比较？
• RQ4 AB 是否能在不依赖 STE 的前提下，实现与混合精度量化（如 4 位权重和 8 位激活）相当的性能？

主要发现

• 与基于 STE 的量化相比，alpha-blending 在 CIFAR10 的 1 位 BinaryNet 上将 top-1 准确率提升了 0.9%。
• 在 8 位权重和 8 位激活量化设置下，AB 在 ImageNet 上比当前最先进的基于 STE 的方法高出 0.82% 的 top-1 准确率。
• 在 4 位权重和 8 位激活模型上，AB 在 MobileNet 和 ResNet 架构上的平均 top-1 准确率相比先前工作提升了 2.93%。
• AB 在 4 位权重和 4 位激活量化设置下的精度与当前最佳的 4 位权重和 8 位激活方法相当，证明了其在混合精度设置下的有效性。
• AB 的逐通道量化变体在 4 位权重和 8 位激活设置下平均仅出现 0.9% 的准确率下降，相比先前结果提升了 2.93%。
• 该方法成功实现了低精度模型的训练，即使在存在不可微量化函数的情况下，也未出现梯度消失问题。