QS4D: Quantization-aware training for efficient hardware deployment of structured state-space sequential models

TL;DR

This paper introduces quantization-aware training for structured state-space models, significantly reducing their complexity and enabling efficient deployment on analog in-memory computing hardware, with improved robustness and pruning capabilities.

Contribution

It demonstrates how QAT can drastically simplify SSMs, improve robustness to noise, and facilitate deployment on specialized edge hardware, which was not previously addressed.

Findings

QAT reduces SSM complexity by up to 100x.

QAT improves robustness to analog noise.

Structural pruning is enabled for SSMs.

Abstract

Structured State Space models (SSM) have recently emerged as a new class of deep learning models, particularly well-suited for processing long sequences. Their constant memory footprint, in contrast to the linearly scaling memory demands of Transformers, makes them attractive candidates for deployment on resource-constrained edge-computing devices. While recent works have explored the effect of quantization-aware training (QAT) on SSMs, they typically do not address its implications for specialized edge hardware, for example, analog in-memory computing (AIMC) chips. In this work, we demonstrate that QAT can significantly reduce the complexity of SSMs by up to two orders of magnitude across various performance metrics. We analyze the relation between model size and numerical precision, and show that QAT enhances robustness to analog noise and enables structural pruning. Finally, we integrate these techniques to deploy SSMs on a memristive analog in-memory computing substrate and highlight the resulting benefits in terms of computational efficiency.