IMSSA: Deploying modern state-space models on memristive in-memory compute hardware

TL;DR

This paper presents IMSSA, a method for deploying state-space models on memristive in-memory hardware, reducing computational demands and enabling efficient long sequence processing on edge devices.

Contribution

It introduces a novel approach to implement S4 models on memristive hardware using quantization-aware training and in-memory compute techniques.

Findings

Successfully deployed S4 kernels on memristive hardware.

Achieved quantization, including ternary weights, for efficient model deployment.

Demonstrated real-world application on edge hardware.

Abstract

Processing long temporal sequences is a key challenge in deep learning. In recent years, Transformers have become state-of-the-art for this task, but suffer from excessive memory requirements due to the need to explicitly store the sequences. To address this issue, structured state-space sequential (S4) models recently emerged, offering a fixed memory state while still enabling the processing of very long sequence contexts. The recurrent linear update of the state in these models makes them highly efficient on modern graphics processing units (GPU) by unrolling the recurrence into a convolution. However, this approach demands significant memory and massively parallel computation, which is only available on the latest GPUs. In this work, we aim to bring the power of S4 models to edge hardware by significantly reducing the size and computational demand of an S4D model through quantization-aware training, even achieving ternary weights for a simple real-world task. To this end, we extend conventional quantization-aware training to tailor it for analog in-memory compute hardware. We then demonstrate the deployment of recurrent S4D kernels on memrisitve crossbar arrays, enabling their computation in an in-memory compute fashion. To our knowledge, this is the first implementation of S4 kernels on in-memory compute hardware.