TLDR: This paper introduces QS4D, a method using Quantization-Aware Training (QAT) to significantly reduce the computational complexity and memory footprint of Structured State Space Models (SSMs) for efficient deployment on resource-constrained edge hardware. QAT enables aggressive quantization, enhances noise robustness, and allows for structural pruning, culminating in a successful demonstration of SSMs on memristive analog in-memory computing substrates, showcasing substantial energy efficiency gains.
Structured State Space models (SSMs) are emerging as a powerful new class of deep learning models, particularly adept at handling long sequences of data. Unlike traditional Transformers, SSMs maintain a constant memory footprint, making them highly attractive for deployment on resource-constrained edge-computing devices. This characteristic positions them as promising candidates for future Large Language Models (LLMs) and crucial for real-world applications such as processing physiological signals in biomedical devices or enabling long-term planning in autonomous vehicles, where power and latency are critical considerations.
A common strategy for deploying large machine learning models on edge hardware involves quantization, which reduces memory usage and computational load by lowering the precision of model weights and activation signals. While quantization-aware training (QAT) has proven effective for Transformers, recurrent neural networks like SSMs are known to be highly sensitive to it. Previous research on SSM quantization often focused on general-purpose GPUs, overlooking the specific implications for specialized edge hardware, such as analog in-memory computing (AIMC) chips.
This research introduces QS4D, a method demonstrating that QAT can dramatically reduce the complexity of SSMs, specifically the S4D model, by up to two orders of magnitude across various performance metrics. The study delves into the relationship between model size and numerical precision, revealing that QAT not only enhances robustness to analog noise but also facilitates structural pruning, allowing for the removal of entire sections of the model without significant performance loss. A key highlight of this work is the successful integration of these techniques to deploy SSMs on a memristive analog in-memory computing substrate, showcasing substantial benefits in computational efficiency.
The findings indicate that QAT enables significantly more aggressive quantization compared to post-training quantization (PTQ). For instance, in a sequential CIFAR10 benchmark, QAT allowed for a homogeneous quantization of 6 bits before the error increased by 1% over the baseline, whereas PTQ reached this threshold at 10 bits. This advantage of QAT is particularly evident for parameters involved in the recurrent state update, such as the transition matrix A, the state, and the time-step parameter (dt).
Beyond just reducing bit precision, aggressive quantization through QAT translates into substantial efficiency gains. It can decrease computational effort by factors ranging from 2 to 11.5, reduce the model’s memory footprint by factors up to 2.9, and lower the peripheral analog-to-digital conversion (ADC) complexity by factors up to 4. These improvements are especially pronounced for audio classification tasks, a likely application area for SSMs on edge hardware.
The study also explores the trade-off between model size and quantization. It shows that increasing model dimensions, such as state dimension or width, can compensate for some accuracy loss due to quantization. Furthermore, more aggressive quantization allows for more extensive structural pruning, which is highly beneficial for hardware implementations where kernels and states need to be physically materialized.
A crucial aspect addressed is noise robustness. Analog in-memory computing systems are inherently prone to noise. The research demonstrates that QAT makes models more resilient to transient read noise. This resilience is attributed to quantization itself acting as a form of noise during training. The study further shows that explicitly introducing Gaussian noise during training can make models even more robust and, in some cases, even improve the performance of highly quantized models.
Also Read:
- Racetrack Memory: A New Path for Energy-Efficient AI in Embedded Systems
- Revolutionizing mmWave MIMO: How AI is Making Intelligent Surfaces Smarter and Faster
Finally, the paper details the successful deployment of S4D kernels on memristive crossbar arrays. These arrays are well-suited for efficient vector-matrix multiplication through analog computation. The proposed In-Memory State Space Accelerator (IMSSA) architecture maps the S4D kernel computation to a single memristive crossbar array, enabling the entire state update and output to be computed in one step. A comparison with a commercial edge GPU, the NVIDIA Jetson Nano, highlights the significant energy efficiency gains (TOPS/W) achieved by this specialized hardware approach. This comprehensive work paves the way for efficient sequence processing with ultra-long contexts at the edge, addressing critical hardware limitations. For more details, you can refer to the full research paper here.
