A flexible and fast PyTorch toolkit for simulating training and inference on analog crossbar arrays

TL;DR

The paper presents the IBM Analog Hardware Acceleration Kit, an open-source PyTorch toolkit for simulating analog crossbar arrays, enabling flexible, fast, and hardware-aware neural network modeling with non-idealities and noise considerations.

Contribution

It introduces the first comprehensive, GPU-accelerated PyTorch toolkit for simulating analog hardware effects in neural network training and inference, including customizable analog tiles and noise models.

Findings

Enables simulation of various analog hardware characteristics and non-idealities.

Supports hardware-aware training with ideal backward and update behaviors.

Provides calibrated noise and drift models for phase-change memory hardware.

Abstract

We introduce the IBM Analog Hardware Acceleration Kit, a new and first of a kind open source toolkit to simulate analog crossbar arrays in a convenient fashion from within PyTorch (freely available at https://github.com/IBM/aihwkit). The toolkit is under active development and is centered around the concept of an "analog tile" which captures the computations performed on a crossbar array. Analog tiles are building blocks that can be used to extend existing network modules with analog components and compose arbitrary artificial neural networks (ANNs) using the flexibility of the PyTorch framework. Analog tiles can be conveniently configured to emulate a plethora of different analog hardware characteristics and their non-idealities, such as device-to-device and cycle-to-cycle variations, resistive device response curves, and weight and output noise. Additionally, the toolkit makes it possible to design custom unit cell configurations and to use advanced analog optimization algorithms such as Tiki-Taka. Moreover, the backward and update behavior can be set to "ideal" to enable hardware-aware training features for chips that target inference acceleration only. To evaluate the inference accuracy of such chips over time, we provide statistical programming noise and drift models calibrated on phase-change memory hardware. Our new toolkit is fully GPU accelerated and can be used to conveniently estimate the impact of material properties and non-idealities of future analog technology on the accuracy for arbitrary ANNs.