AnalogAI: A Generalized Accuracy Recovery Framework for DNNs under Analog Computing

Analog In-Memory Computing (AIMC) offers a promising approach to accelerate Deep Neural Networks (DNNs) while reducing energy consumption. However, instabilities caused by manufacturing and device variations limit its accuracy and reliability. Existing methods typically focus on recovery on specific devices under ideal conditions, lacking the ability to handle complex realistic noise patterns and a unified framework to leverage various algorithms. To address these issues, we propose AnalogAI, a generalized framework for the accuracy recovery of analog DNNs. AnalogAI introduces a multi-noise fusion injection (FI) strategy working with multiple feedback loops (FL), which is designed to achieve multi-scale noise awareness by aggregating distributed weights while improving training stability. Furthermore, the framework supports custom quantization-aware training and integrates various noise-robust training algorithms for rapid accuracy recovery exploration. Extensive experiments on IBM Analog Hardware Acceleration Kit (IBM-aihwkit), a widely adopted open-source simulation platform, demonstrate that AnalogAI significantly improves inference accuracy across different datasets and network architectures, highlighting its applicability and effectiveness. Our code is released at https://github.com/PJLAB-CHIP/AnalogAI.