Interpreting and Steering State-Space Models via Activation Subspace Bottlenecks
Abstract
State-space models (SSMs) have emerged as an efficient strategy for building powerful language models, avoiding the quadratic complexity of computing attention in transformers. Despite their promise, the interpretability and steerability of modern SSMs remain relatively underexplored. We take a major step in this direction by identifying activation subspace bottlenecks in the Mamba family of SSM models using tools from mechanistic interpretability. We then introduce a test-time steering intervention that simply multiplies the activations of the identified bottlenecks by a scalar. Across 7 SSMs and 6 diverse benchmarks, this intervention improves performance by an average of 8.27%, without requiring any task-specific tuning. Finally, we validate that the identified bottlenecks are indeed hindering performance by modifying them to yield an architecture we call Stable-Mamba, which achieves long-context performance gains when retrained from scratch.
Lay Summary
Modern AI systems often rely on transformer models, which can be expensive to run on long text because their memory and computation grow quickly with input length. State-space models (SSMs), such as Mamba, are a newer class of AI models designed to process long sequences more efficiently. However, research has not progressed much in understanding how they store and route information, which makes them difficult to improve and control. In this work, we study how information flows inside SSMs using mechanistic interpretability. We discover that SSMs rely heavily on a small number of internal activation subspaces that act as bottlenecks, restricting the model’s ability to effectively use information over long contexts. Using this insight, we introduce a steering method that improves model behavior at test time without retraining by selectively amplifying important internal activations. This method improves performance across several language understanding and long-context benchmarks. We also design a modified architecture called Stable-Mamba that reduces these bottlenecks directly within the model design. Stable-Mamba achieves stronger long-context performance while adding very little computational overhead. Our work provides a method to understand, steer and improve SSMs, helping make them both more interpretable and reliable.