[论文解读] BitPruning: Learning Bitlengths for Aggressive and Accurate Quantization

Miloš Nikolić, Ghouthi Boukli Hacene+6 more|arXiv (Cornell University)|Feb 8, 2020

Advanced Neural Network Applications参考文献 36被引用 15

一句话总结

BitPruning 引入了一种可微分的正则化方法，在训练过程中联合学习深度神经网络各层的最优比特长度，实现极强的量化并仅造成极小的精度损失。在 ImageNet 上，AlexNet、ResNet18 和 MobileNet V2 的平均比特长度分别为 4.13、3.76 和 4.36，精度损失分别控制在全精度基线的 0.5%–2.0% 以内。

ABSTRACT

Neural networks have demonstrably achieved state-of-the art accuracy using low-bitlength integer quantization, yielding both execution time and energy benefits on existing hardware designs that support short bitlengths. However, the question of finding the minimum bitlength for a desired accuracy remains open. We introduce a training method for minimizing inference bitlength at any granularity while maintaining accuracy. Namely, we propose a regularizer that penalizes large bitlength representations throughout the architecture and show how it can be modified to minimize other quantifiable criteria, such as number of operations or memory footprint. We demonstrate that our method learns thrifty representations while maintaining accuracy. With ImageNet, the method produces an average per layer bitlength of 4.13, 3.76 and 4.36 bits on AlexNet, ResNet18 and MobileNet V2 respectively, remaining within 2.0%, 0.5% and 0.5% of the base TOP-1 accuracy.

研究动机与目标

• 为解决在低精度推理中寻找每层最小比特长度以维持模型精度的挑战。
• 实现各层异构比特长度的端到端学习，而非依赖固定或人工调优的比特宽度。
• 通过灵活的正则化器，不仅最小化比特长度，还最小化可度量的指标，如内存占用和计算负载。
• 通过学习与硬件优化数据类型对齐的比特宽度，支持在现有及专用硬件上的高效部署。
• 在无需为每种比特宽度重新训练的情况下，降低能耗并提升推理效率。

提出的方法

• 引入一种可微分的正则化器，对网络中所有层的大比特长度表示施加惩罚，促使模型学习每层最小但准确的比特宽度。
• 该正则化器可加权调节，以优先最小化比特长度、内存占用或计算操作数。
• 该方法在标准反向传播过程中应用，支持权重与比特长度的联合端到端训练。
• 通过直通估计器实现比特长度的可微性，即使在离散选择下也能实现梯度流动。
• 该方法支持任意粒度的比特宽度分配，包括每层、每卷积核或每张量级别。
• 通过调整正则化器的系数，可将该方法适配至不同硬件约束下的优化。

实验结果

研究问题

• RQ1 深度学习模型能否学习到在保持高精度的同时最小化比特宽度的各层比特长度？
• RQ2 与推理后或手工设计的比特宽度选择相比，所提出的正则化器在精度和效率方面表现如何？
• RQ3 该方法能否扩展以最小化其他硬件感知指标，如内存占用或 FLOPs？
• RQ4 使用学习到的比特宽度时，训练成本与推理效率之间的权衡如何？
• RQ5 该方法在不同网络架构和数据集上是否具备泛化能力？

主要发现

• 在 ImageNet 上，BitPruning 为 AlexNet 实现平均 4.13 比特，ResNet18 为 3.76 比特，MobileNet V2 为 4.36 比特，与全精度基线相比，精度损失分别为 2.0%、0.5% 和 0.5%。
• 在 CIFAR-10 上，该方法为 AlexNet 和 ResNet18 分别实现平均 3.92 和 2.17 比特，精度损失在基线的 0.9% 以内。
• 对于 AlexNet 和 ResNet18，该方法在 CIFAR-10 上将内存占用减少 10%–24%，计算负载减少 4%–8%，展现出显著的效率提升。
• 训练成本较高，需 180 个周期（为基线的 2 倍），且耗时为基线的 4.6×–5.4×，但这一代价由推理阶段的显著收益所合理化。
• 该方法支持端到端量化，包括首层和末层，实现无需架构修改的完整模型量化。
• 正则化器可通过重新加权以优化内存或计算，使其可灵活适配多种硬件目标。