[Paper Review] 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 Computing4 citations

TL;DR

This paper proposes PIM-QAT, a quantization-aware training method for processing-in-memory (PIM) systems that accounts for the unique hardware constraints of PIM, including low-resolution analog-to-digital conversion and non-ideal effects like noise and nonlinearity. By integrating rescaling in forward and backward passes and using BN calibration and adjusted precision training, PIM-QAT achieves inference accuracy on PIM systems comparable to conventionally quantized models on digital hardware across CIFAR10 and CIFAR100 with various network depths.

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.

Motivation & Objective

• To address the significant accuracy degradation caused by PIM-specific quantization, which results from low-resolution analog-to-digital conversion and hardware non-idealities.
• To develop a training method that accounts for the unique computing flow of PIM systems, where matrix multiplication is decomposed into smaller subsets.
• To improve convergence and stability during training by modeling the dynamics of PIM quantization in both forward and backward propagation.
• To suppress the adverse effects of non-ideal linearity and stochastic thermal noise in real PIM chips through software-level calibration and precision adjustment.
• To achieve inference accuracy on PIM systems that matches or exceeds that of conventionally quantized models on digital hardware.

Proposed method

• Proposes a PIM quantization-aware training (PIM-QAT) algorithm that explicitly models the PIM quantization step during backpropagation and forward pass.
• Introduces rescaling techniques in both forward and backward propagation to stabilize training dynamics and improve convergence, especially at low PIM bit-widths (e.g., 3–6 bits).
• Employs batch normalization (BN) calibration to mitigate performance loss caused by nonlinearity and noise in real PIM chips, improving robustness without additional training.
• Applies adjusted precision training to compensate for hardware-induced distortions, particularly in the presence of gain and offset variations in PIM transfer functions.
• Validates the method across three PIM decomposition schemes and on a physical prototype chip, ensuring real-world relevance.
• Uses idealized and real PIM quantization curves to simulate and evaluate performance under non-ideal conditions, including stochastic noise and nonlinearity.

Experimental results

Research questions

• RQ1 How can we stabilize training in PIM systems where low-resolution quantization and non-ideal analog behavior severely degrade model accuracy?
• RQ2 What training techniques are effective for maintaining convergence and accuracy when the PIM quantization step is irreversible and hardware-constrained?
• RQ3 To what extent can software-level calibration techniques like BN calibration compensate for hardware non-idealities such as noise and nonlinearity in PIM systems?
• RQ4 Can a unified training method achieve accuracy on PIM systems that matches or exceeds that of conventionally quantized models on digital hardware?
• RQ5 How do rescaling in forward and backward passes affect training stability and final accuracy at low PIM bit-widths?

Key findings

• PIM-QAT achieves 90.8% top-1 accuracy on CIFAR10 with ResNet20 using 8-bit PIM quantization, matching the performance of conventionally quantized models on digital hardware.
• For 7-bit PIM quantization, PIM-QAT achieves 90.8% accuracy, while the baseline method fails completely (10.0% accuracy), demonstrating the necessity of PIM-aware training.
• Rescaling in both forward and backward passes is essential: without it, training becomes unstable at low bit-widths (e.g., 3–5 bits), with accuracy dropping below 20%.
• BN calibration improves performance significantly on real PIM systems, reducing the gap between idealized and real-case inference and enabling recovery from nonlinearity and noise effects.
• With 72 and 144 PIM transfer curves exhibiting gain and offset variations, BN calibration recovers accuracy from near-random (10.0%) to over 90%, proving its effectiveness in real-world calibration.
• The method generalizes across different PIM decomposition schemes and network depths, achieving consistent performance on ResNet20 and ResNet56 across CIFAR10 and CIFAR100.