Mamba Knockout for Unraveling Factual Information Flow

Abstract: This paper investigates the flow of factual information in Mamba State-Space Model (SSM)-based LLMs. We rely on theoretical and empirical connections to Transformer-based architectures and their attention mechanisms. Exploiting this relationship, we adapt attentional interpretability techniques originally developed for Transformers--specifically, the Attention Knockout methodology--to both Mamba-1 and Mamba-2. Using them we trace how information is transmitted and localized across tokens and layers, revealing patterns of subject-token information emergence and layer-wise dynamics. Notably, some phenomena vary between mamba models and Transformer based models, while others appear universally across all models inspected--hinting that these may be inherent to LLMs in general. By further leveraging Mamba's structured factorization, we disentangle how distinct "features" either enable token-to-token information exchange or enrich individual tokens, thus offering a unified lens to understand Mamba internal operations.

Mamba Knockout for Unraveling Factual Information Flow

This paper presents an examination of factual information dynamics within Mamba-based language models. The authors employ interpretability techniques adapted from Transformer architectures, leveraging similarities between State-Space Models (SSMs) and attention mechanisms to analyze information flow and localization through Mamba-1 and Mamba-2 models. The research reveals the intricate pathways of subject-token information transmission and layer-specific dynamics, demonstrating how features within Mamba models either mediate token-to-token information exchange or enhance individual tokens.

The study expands upon the existing Attention Knockout methodology, originally developed for Transformers. By applying this in Mamba models, the authors successfully dissect the flow of information at various levels of the architecture, identifying certain shared characteristics across all inspected models. These include the crucial role of subject tokens in directing information flow, aligning with analogous findings in Transformer models. Such patterns underscore the universal aspects of factual information processing in large language models, irrespective of their architectural distinctions.

The paper is noteworthy for two central contributions: Firstly, it extends Attention Knockout to SSMs, unveiling parallels and disparities in factual information dynamics between Mamba-based and Transformer-based models. Secondly, it introduces a novel 'feature knockout' mechanism, capitalizing on the unique structure of SSMs to facilitate nuanced insights into the contribution of distinct feature types to model behavior.

Several implications for this research are apparent. Practically, these findings could inform strategies for optimized training and deployment of Mamba models. The identification of redundant feature types and critical information pathways may prompt efficient pruning or targeted fine-tuning, enhancing computational performance without degrading accuracy. Theoretically, the study advances our understanding of the underlying principles governing factual information flow in LLMs, highlighting the role of Mamba’s factorized structure in token enrichment and information exchange.

Looking forward, this research paves the way for promising explorations in AI interpretability. The authors provide a methodological foundation that could guide future studies aiming to demystify the internal workings of SSM-based architectures, forming a basis for a more unified framework for understanding language model operations across different paradigms. The in-depth analysis served to not only reinforce known understandings of attention's role in LLMs but also to extend these insights into the less charted territory of structured state-space models.