Mamba Knockout for Unraveling Factual Information Flow

TL;DR

This paper explores how factual information flows within Mamba language models by adapting attention interpretability techniques, revealing both model-specific and universal patterns of information transmission across tokens and layers.

Contribution

It introduces a novel application of attentional interpretability to Mamba models, uncovering internal information dynamics and their relation to Transformer architectures.

Findings

Identifies patterns of information emergence and layer-wise dynamics.

Reveals both model-specific and universal information flow phenomena.

Disentangles feature roles in token exchange and enrichment.

Abstract

This paper investigates the flow of factual information in Mamba State-Space Model (SSM)-based language models. 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.