On-Device Continual Learning with STT-Assisted-SOT MRAM Based In-Memory Computing

Due to the separate memory and computation units in traditional von Neumann architecture, massive data transfer dominates the overall computing system's power and latency, known as the "Memory-Wall" issue. Especially with ever-increasing deep learning-based AI model size and computing complexity, it becomes the bottleneck for state-of-the-art AI computing systems. To address this challenge, in-memory computing (IMC)-based Neural Network accelerators have been widely investigated to support AI computing within memory. However, most of those works focus only on inference. The on-device training and continual learning have not been well explored yet. In this work, for the first time, we introduce on-device continual learning with STT-assisted-SOT (SAS) magnetoresistive random-access memory (MRAM)-based IMC system. On the hardware side, we have fabricated a STT-assisted-SOT MRAM (SAS-MRAM) device prototype with 4 magnetic tunnel junctions (MTJs, each at 100 nm x50 nm) sharing a common heavy metal layer, achieving significantly improved memory writing and area efficiency compared to traditional SOT-MRAM. Next, we designed fully digital IMC circuits with our SAS-MRAM to support both neural network inference and on-device learning. To enable efficient on-device continual learning for new task data, we present an 8-bit integer (INT8)-based continual learning algorithm that utilizes our SAS-MRAM IMC-supported bit-serial digital in-memory convolution operations to train a small parallel reprogramming network (Rep-Net) while freezing the major backbone model. Extensive studies have been presented based on our fabricated SAS-MRAM device prototype, cross-layer device-circuit benchmarking and simulation, as well as the on-device continual learning system evaluation.