TLDR: A new method called Hybrid Projection Decomposition (HPD) significantly improves the robustness and accuracy of State Space Models (SSMs) when run on noisy analog Compute-in-Memory (CIM) hardware. HPD achieves this by strategically splitting the output projection layer’s computation, performing the noise-sensitive part on precise digital hardware while keeping the rest on efficient analog CIM, leading to up to 99.57% perplexity reduction and 96.67% accuracy gains.
In the rapidly evolving world of artificial intelligence, State Space Models (SSMs) are emerging as powerful alternatives to traditional sequence models like Transformers. These models are particularly good at handling long sequences of data with much less computational effort. A key reason for their efficiency is their heavy reliance on matrix multiplications, which makes them highly compatible with a new type of computing architecture called Compute-in-Memory (CIM).
CIM architectures are designed to perform computations directly within memory units, drastically reducing the energy consumed by moving data between processors and memory. This can lead to significant efficiency improvements. However, there’s a catch: analog computing elements, which are central to CIM, are prone to “non-idealities.” These are imperfections that introduce errors or “weight perturbations” into the calculations, which can severely degrade the accuracy of neural networks and undermine the benefits of CIM hardware.
The challenge is particularly acute for SSMs because of their recurrent nature. Small errors can accumulate over time, leading to significant performance drops. Despite the growing interest in deploying SSMs on efficient hardware, how robust these models are to such errors has been largely unexplored until now.
A recent research paper, “HPD: Hybrid Projection Decomposition for Robust State Space Models on Analog CIM Hardware,” by Yuannuo Feng, Wenyong Zhou, Yuexi Lyu, Hanjie Liu, Zhengwu Liu, Ngai Wong, and Wang Kang, addresses this critical gap. The researchers systematically analyzed the robustness of SSMs under noisy conditions, specifically identifying that the final block and output projection layers are the most vulnerable parts of the model to these hardware-induced perturbations.
Building on this crucial insight, the team proposed a novel solution called Hybrid Projection Decomposition (HPD). This strategy focuses on the last output projection layer, which was found to be highly susceptible to noise. Instead of using the original weight matrix in this layer, HPD replaces it with a decomposed version based on Singular Value Decomposition (SVD). Specifically, the weight matrix is broken down into three components: U, Σ, and Vᵀ. The genius of HPD lies in how these components are handled.
The first part of the computation, involving UΣ, is designed to be performed on the analog CIM hardware. This is possible because the product UΣ maintains the same structure and dimensions as the original weight matrix, meaning it doesn’t require any changes to existing CIM hardware designs. The second part, involving Váµ€, is offloaded to digital hardware. Digital hardware, while not as energy-efficient for certain tasks, offers much higher precision and is immune to the analog noise that plagues CIM. By splitting the computation this way, HPD ensures that the most sensitive part of the calculation benefits from the precision of digital processing, while the bulk of the computation still leverages the efficiency of CIM.
The researchers conducted extensive tests using Mamba models, a popular variant of SSMs, under various noise conditions, including Gaussian and lognormal noise. The results were highly impressive. HPD significantly reduced perplexity, a measure of how well a language model predicts a sample, by up to 99.57% compared to baseline models under noisy conditions. Furthermore, on the PIQA benchmark, which tests commonsense reasoning, HPD achieved accuracy gains of up to 96.67%.
Also Read:
- Enhancing Neural Network Robustness on Analog Compute-In-Memory Hardware with Extended Straight-Through Estimation
- Advancing Spiking Neural Networks with Single-Timestep Processing and Adaptive Optimization
These findings demonstrate that HPD offers a practical and effective way to make SSMs robust against the inherent non-idealities of analog CIM hardware. By strategically decomposing and distributing computations, this method allows for the continued pursuit of energy-efficient AI hardware without sacrificing model accuracy. This work paves the way for more reliable and high-performing AI systems on next-generation computing platforms. You can read the full research paper for more details: HPD: Hybrid Projection Decomposition for Robust State Space Models on Analog CIM Hardware.
