X-VMamba: Explainable Vision Mamba

AI summary

Overview

Research area: Interpretability and explainability for State Space Models (SSMs), specifically the Mamba/Vision Mamba family, applied to medical image classification.

Technical level: Advanced. The paper builds on classical control theory (controllability, controllability Gramian, Lyapunov equations) and requires comfort with linear dynamical systems, Jacobians, and matrix norms. The conceptual goal is explained accessibly, but the derivations are mathematically dense.

Scope: The paper proposes and validates a controllability-based framework, X-VMamba, that quantifies how individual input patches influence the internal state dynamics of vision SSMs, demonstrated on three medical imaging datasets using the MedMamba architecture.

What This Paper Is About

State Space Models such as Mamba offer linear computational complexity as an alternative to Transformers, but they lack the attention matrices that make Transformers relatively easy to inspect. Because SSMs propagate information implicitly through hidden state transitions, there is no direct way to answer "which input patches influenced this prediction and how strongly?" This paper addresses that gap by repurposing controllability from classical control theory as an intrinsic, single-forward-pass measure of input influence, and applying it to reveal how vision SSMs refine spatial information across layers.

Key Contributions

1. A controllability-based interpretability framework for SSMs. The authors reframe interpretability as a control problem: rather than approximating influence through gradients or perturbations, they measure how strongly an input token or patch steers the hidden state trajectory, weighted by how observable that state is in the output.

2. Two complementary formulations of an influence score. A Jacobian-based method that applies to any SSM architecture and traces influence through the full chain of state propagation, and a Gramian-based method for diagonal SSMs that uses a closed-form analytical solution (element-wise, based on the relationship involving the input matrix and the discretized state transition) for superior speed.

3. Efficiency and architectural neutrality. Both methods operate in a single forward pass with linear complexity, require no architectural modifications, no task-specific tuning, and no hyperparameters. Influence scores are computed in O(N) complexity per sequence element, where N is the state dimension, typically 8-64, versus perturbation methods that require O(L) forward passes for sequence length L.

4. Empirical validation across three medical imaging modalities, together with an articulation of applications spanning computer vision, natural language processing, time-series forecasting, audio processing, genomics, and reinforcement learning.

Main Findings

• Hierarchical feature refinement is a natural property of vision SSMs. Across all three datasets, controllability maps become progressively more focused with depth. Early layers produce diffuse scores capturing low-level texture such as edges, intensity gradients, and local contrast; deeper layers concentrate on semantically and clinically meaningful regions.

• Modality-specific controllability signatures emerge. Dermatology heatmaps weight color heterogeneity and border irregularity, with layer-3 maps showing sharp demarcation around the lesion perimeter. Mammography heatmaps shift from diffuse fibroglandular tissue to fat-tissue interfaces and then to fibroglandular regions surrounding mass abnormalities. Blood cell images start with uniform attention across many cells and later selectively enhance individual cells or cell clusters.

• Influence becomes sparser but stronger with depth. As layers deepen, the number of high-influence patches decreases while their maximum influence scores increase substantially, indicating that decision-making authority concentrates on a small subset of patches.

• Scanning strategy substantially changes the explanations. Influence score distributions are strongly modulated by which of the four 2D scanning patterns (forward, backward, transposed-forward, transposed-backward) is used, indicating that directional sequencing shapes the model's spatial attention.

• A naive controllability formulation fails. The initial approach—measuring influence on the final hidden state rather than the aggregated output—produces "vanishing influence" and a severe recency bias, because the long product of stable state matrices decays exponentially with chain length. Empirically this yields heatmaps concentrated in image corners, corresponding to the start and end points of the 2D scanning paths, an artifact of the method rather than content-based reasoning.

• Global Average Pooling motivates the correct target. Because the classifier consumes the mean of all output vectors rather than the final state, the Jacobian formulation sums direct influence plus propagated influence, giving early patches a baseline score and removing the corner artifact.

• The paper reports no quantitative classification accuracy numbers in the content provided; evaluation is presented qualitatively through controllability heatmaps, influence score distributions across layers, and comparisons across scanning schemes.

Methodology in Plain English

The authors treat a vision SSM as a dynamical system with an internal state that is nudged by each input patch and passed forward over time. Instead of asking how the output changes when you perturb the input (the gradient/saliency approach), they ask how much each patch can steer the state, and how much that steering is visible in the output.

Two ways of computing this:

• Jacobian method. Unroll the recurrence and compute how a change in one input patch propagates to every later output term, then sum the magnitudes of those contributions. This separates an immediate, direct influence at the patch's own step from a propagated influence on all future steps. It applies to any SSM, including non-diagonal ones, and is exact rather than approximate.

• Gramian method. For SSMs with a diagonal state matrix—the common case—the controllability Gramian has a closed-form, element-wise solution. Multiply this by the square of the output matrix entries to account for observability, then sum over state dimensions and average over inner channels. This is faster but requires the diagonal assumption; a non-diagonal state matrix would require solving a full Lyapunov equation.

To handle the four scanning directions used by VMamba's 2D Selective Scan, the algorithm computes a score map for each direction, re-orients each 1D sequence back into a 2D map, sums them, and divides by four. The core computation iterates backwards through the sequence while maintaining a running "future influence propagator" that is updated with the output matrix and state transition matrix at each step.

Validation used MedMamba trained on three datasets chosen to stress different visual characteristics: CMMD grayscale high-resolution mammography, DermaMNIST RGB skin lesion images, and BloodMNIST blood cell microscopy images. The authors report evaluating the Jacobian approach because it is more general and does not require the diagonal state matrix assumption of the Gramian method.

Why This Matters

SSMs are increasingly positioned as foundational-model alternatives to Transformers, yet interpretability tooling has lagged behind. Gradient-based methods (Grad-CAM, Integrated Gradients, SmoothGrad) inherit problems of gradient noise, saturation artifacts, and sensitivity to initialization—problems that are worse in SSMs because of exponential discretization and complex recurrent dynamics. Attention-based interpretation is simply unavailable, since SSMs have no attention weights. This work argues that controllability is a faithful, intrinsic property derived from first principles rather than a post-hoc approximation, and that layer-wise controllability analysis reveals information refinement in a way attention rollout cannot.

Real-world applications:

• Clinical decision support. If an SSM-based diagnostic model flags a scan as high-risk, clinicians need to see which image regions drove the prediction and how strongly, which matters for regulatory demands for mechanistic explanations in medical AI.
• Mammography and dermatology screening. The method localizes the tissue interfaces, lesion boundaries, and pigmentation patterns that align with how radiologists and dermatologists actually read images, supporting radiologist-in-the-loop workflows.
• Hematology and pathology. Dense, repetitive cell structures with subtle morphological differences benefit from a method that can show which individual cells or clusters the model prioritized over redundant normal cells.
• Autonomous driving and other high-stakes vision. The framework is domain-independent in its mathematics; only the interpretation of "influence" changes, so the same machinery applies to any spatial decision that must be justified.

Industry relevance: Organizations deploying SSM-based models at scale get an explanation tool that runs in a single forward pass at linear cost, requires no retraining or architecture changes, and emits spatially interpretable maps. That combination—cheap, faithful, and non-invasive—is what makes it plausible for regulated production pipelines rather than research demos.

Future Directions

• Extending beyond vision. The authors explicitly articulate applications to language modeling (token-level influence), time-series forecasting (temporal importance), audio processing (spectro-temporal analysis), genomics (base-pair relevance), and reinforcement learning. These remain proposed, not demonstrated.
• Comparing the two formulations head-to-head. The paper evaluates only the Jacobian method on real data, leaving open how the much faster Gramian method performs empirically on diagonal SSMs.
• Understanding the scanning-strategy effect. Since scanning direction substantially modulates influence distributions, an open question is whether scanning schemes can be designed or selected using controllability as an objective.
• Replacing the diagonal assumption. The Gramian method requires a diagonal state matrix; relaxing this would need a scalable approach to the full Lyapunov equation.
• The authors state that code and analysis tools will be made available upon publication, which would enable independent reproduction of the layer-wise and scan-direction analyses.

Target Audience

Researchers and engineers working on model interpretability, state space models, or efficient vision architectures will get the most from this paper, particularly those who need explanations they can defend in regulated settings. Practitioners applying Mamba-family models to medical imaging, scientific imaging, or any long-sequence spatial task will find the influence-score machinery directly usable. Readers without a control theory background will need to invest in the mathematical sections, but the experimental narrative—hierarchical refinement, modality-specific signatures, and scan-direction sensitivity—is accessible on its own.

Authors’ abstract

State Space Models (SSMs), particularly the Mamba architecture, have recently emerged as powerful alternatives to Transformers for sequence modeling, offering linear computational complexity while achieving competitive performance. Yet, despite their effectiveness, understanding how these Vision SSMs process spatial information remains challenging due to the lack of transparent, attention-like mechanisms. To address this gap, we introduce a controllability-based interpretability framework that quantifies how different parts of the input sequence (tokens or patches) influence the internal state dynamics of SSMs. We propose two complementary formulations: a Jacobian-based method applicable to any SSM architecture that measures influence through the full chain of state propagation, and a Gramian-based approach for diagonal SSMs that achieves superior speed through closed-form analytical solutions. Both methods operate in a single forward pass with linear complexity, requiring no architectural modifications or hyperparameter tuning. We validate our framework through experiments on three diverse medical imaging modalities, demonstrating that SSMs naturally implement hierarchical feature refinement from diffuse low-level textures in early layers to focused, clinically meaningful patterns in deeper layers. Our analysis reveals domain-specific controllability signatures aligned with diagnostic criteria, progressive spatial selectivity across the network hierarchy, and the substantial influence of scanning strategies on attention patterns. Beyond medical imaging, we articulate applications spanning computer vision, natural language processing, and cross-domain tasks. Our framework establishes controllability analysis as a unified, foundational interpretability paradigm for SSMs across all domains. Code and analysis tools will be made available upon publication