A framework for analog-digital mixed-precision neural network training and inference

Recent advancements in AI hardware highlight the potential of mixed-signal accelerators, which integrate analog computation for matrix multiplications with reduced-precision digital operations to achieve superior performance and energy efficiency. In this paper, we present a framework designed to perform hardware-aware training of and to evaluate neural networks (NNs) on such accelerators. This framework extends an existing framework, the IBM Analog AI Hardware Kit (AIHWKit), using a quantization library, enabling flexible layer-wise deployment in either analog or digital units, the latter with configurable precision. Our combined framework supports simultaneous quantization- and analog-aware training. It can also evaluate the accuracy of NNs when deployed on mixed-signal accelerators. We demonstrate the effectiveness of this combined training approach through extensive ablation studies on a ResNet-based vision model and a BERT-based language model, highlighting its importance for maximizing accuracy. Our contribution will be open-sourced as part of the core code of AIHWKit.