[论文解读] PIM-QAT: Neural Network Quantization for Processing-In-Memory (PIM) Systems

Qing Jin, Zhiyu Chen+4 more|arXiv (Cornell University)|Sep 18, 2022

Advanced Memory and Neural Computing被引用 4

一句话总结

本文提出PIM-QAT，一种面向存内计算（PIM）系统的量化感知训练方法，考虑了PIM特有的硬件约束，包括低分辨率模数转换以及噪声和非线性等非理想效应。通过在前向和反向传播中引入重缩放，并结合批量归一化校准与调整精度训练，PIM-QAT在CIFAR10和CIFAR100数据集上，针对不同网络深度，实现了与传统数字硬件上量化模型相当的推理精度。

ABSTRACT

Processing-in-memory (PIM), an increasingly studied neuromorphic hardware, promises orders of energy and throughput improvements for deep learning inference. Leveraging the massively parallel and efficient analog computing inside memories, PIM circumvents the bottlenecks of data movements in conventional digital hardware. However, an extra quantization step (i.e. PIM quantization), typically with limited resolution due to hardware constraints, is required to convert the analog computing results into digital domain. Meanwhile, non-ideal effects extensively exist in PIM quantization because of the imperfect analog-to-digital interface, which further compromises the inference accuracy. In this paper, we propose a method for training quantized networks to incorporate PIM quantization, which is ubiquitous to all PIM systems. Specifically, we propose a PIM quantization aware training (PIM-QAT) algorithm, and introduce rescaling techniques during backward and forward propagation by analyzing the training dynamics to facilitate training convergence. We also propose two techniques, namely batch normalization (BN) calibration and adjusted precision training, to suppress the adverse effects of non-ideal linearity and stochastic thermal noise involved in real PIM chips. Our method is validated on three mainstream PIM decomposition schemes, and physically on a prototype chip. Comparing with directly deploying conventionally trained quantized model on PIM systems, which does not take into account this extra quantization step and thus fails, our method provides significant improvement. It also achieves comparable inference accuracy on PIM systems as that of conventionally quantized models on digital hardware, across CIFAR10 and CIFAR100 datasets using various network depths for the most popular network topology.

研究动机与目标

• 解决由低分辨率模数转换和硬件非理想性引起的PIM特异性量化导致的显著精度下降问题。
• 开发一种考虑PIM系统独特计算流程的训练方法，其中矩阵乘法被分解为更小的子集。
• 通过在前向与反向传播中建模PIM量化动态，提升训练的收敛性与稳定性。
• 通过软件级校准与精度调整，抑制真实PIM芯片中非理想线性与随机热噪声的不利影响。
• 实现PIM系统上的推理精度，达到或超过传统量化模型在数字硬件上的表现。

提出的方法

• 提出PIM量化感知训练（PIM-QAT）算法，在反向传播与前向传播中显式建模PIM量化步骤。
• 在前向与反向传播中引入重缩放技术，以稳定训练动态并提升收敛性，尤其在低PIM位宽（如3–6位）下表现显著。
• 采用批量归一化（BN）校准，缓解真实PIM芯片中非线性与噪声导致的性能损失，提升鲁棒性，且无需额外训练。
• 应用调整精度训练以补偿硬件引起的失真，尤其在PIM传输函数存在增益与偏移变化时效果显著。
• 在三种PIM分解方案及物理原型芯片上验证该方法，确保其在真实场景中的适用性。
• 使用理想化与真实PIM量化曲线，模拟并评估在非理想条件（如随机噪声与非线性）下的性能表现。

实验结果

研究问题

• RQ1 如何在低分辨率量化与非理想模拟行为严重降低模型精度的PIM系统中稳定训练？
• RQ2 在PIM量化步骤不可逆且受硬件限制的条件下，哪些训练技术能有效维持收敛性与精度？
• RQ3 软件级校准技术（如BN校准）在多大程度上可补偿PIM系统中噪声与非线性等硬件非理想性？
• RQ4 统一的训练方法能否在PIM系统上实现与传统数字硬件上量化模型相当或更优的精度？
• RQ5 前向与反向传播中的重缩放对低PIM位宽下的训练稳定性与最终精度有何影响？

主要发现

• 在使用8位PIM量化时，PIM-QAT在CIFAR10上ResNet20模型达到90.8%的top-1精度，与数字硬件上传统量化模型性能相当。
• 在7位PIM量化下，PIM-QAT实现90.8%精度，而基线方法完全失效（仅10.0%精度），证明PIM感知训练的必要性。
• 前向与反向传播中的重缩放至关重要：若无此机制，低比特位宽（如3–5位）下训练极不稳定，精度降至20%以下。
• BN校准显著提升真实PIM系统上的性能，缩小理想化与真实情况下的推理差距，有效恢复因非线性与噪声导致的性能损失。
• 在72和144条存在增益与偏移变化的PIM传输曲线中，BN校准将精度从接近随机（10.0%）恢复至90%以上，充分证明其在真实校准中的有效性。
• 该方法在不同PIM分解方案与网络深度间具有良好泛化性，在CIFAR10与CIFAR100上的ResNet20与ResNet56上均实现一致的高性能表现。