Interpreting and Steering State-Space Models via Activation Subspace Bottlenecks

Vamshi Sunku Mohan; Kaustubh Gupta; Aneesha Das; Chandan Singh

cs.LGFeb 26, 2026

Interpreting and Steering State-Space Models via Activation Subspace Bottlenecks

Authors: Vamshi Sunku Mohan, Kaustubh Gupta, Aneesha Das, Chandan Singh

Organizations: AppViewX, Bangalore · Independent Researcher · Microsoft Research, Redmond

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.

Figures & tables

Panel of Figure 1 — Figure 1: Entropy across layers measured using Stochastic Parameter Decomposition ( Bushnaq et al., 2025 ) . (a) Vanilla Mamba exhibits a sharp entropy spike at Layer 20, indicating a parameter-level routing bottleneck where diverse information is forced through a narrow subset of parameters. After (b) steering and (c) architectural modifications, the spike is removed and entropy becomes smoother, indicating restored information flow. These modifications yield improved performance across three standard benchmarks; NIAH: Needle in a haystack, QA: question-answering and Pathfinder: long-context benchmark.

Figure 2: Workflow for identifying Activation Subspace Bottlenecks and using them to conduct post-hoc steering or to make architectural modifications to Mamba.

Relation	PPL (Transformer)		PPL (Mamba)
Relation	Before	After	Before	After
P264 (record label)	189.60	188.07	77.30	176.15
P449 (original network)	77.06	74.41	44.60	220.37
P413 (position played on team)	142.04	170.06	29.68	63.08
P463 (member of)	15.54	9.76	9.46	24.95
P530 (diplomatic relation)	31.29	24.81	15.39	52.41

Table 1: Perplexity in the Transformer and Mamba before and after ablating different knowledge-related subspaces . Cloze-style prompts on masked answer tokens are used ( Dai et al., 2022 ) , yielding higher values than standard next-token perplexity (often

>

30). Mamba shows substantially larger values, suggesting knowledge is encoded in relatively few overlapping subspaces.

Information Type	Result	Analysis
Activation Strength	0.87	High correlation - Model explicitly tracks activation strength
Token Index	0.65	Moderate correlation - Positional information distributed across subspaces 128, 256, 384, 512 and 640 at regular intervals of 128
Sparsity	0.45	Lower sparsity - Information is encoded nonlinearly and distributed across dimensions

Table 2: Correlation analysis of information types in SAE-derived features with encoded information strength.

Metric	Vanilla Mamba	Steered Mamba	Stable-Mamba
Reconstruction error	9404.5	238.2	267.6
Sparsity (%)	98.12	62.05	58.37
Active features (%) (Activation $>$ 0.1)	18.75	60.00	68.95
Perplexity (%)	36.77	28.83	27.50

Table 3: SAE dictionary statistics for Vanilla, Steered and Stable Mamba. Metrics report dictionary size, reconstruction error, sparsity, feature usage, and active features Steered and Stable-Mamba substantially reduce sparsity while lowering reconstruction error, indicating denser and more distributed representations.

Phase	Layers	Entropy	Variance ( $\sigma$ )	Rank
1	0-18	1.04	2.0	5.85
2	19	1.02	2.30	6.25
3	20	1.19	2.63	7.59
4	21	0.92	3.59	6.49
5	22-23	0.70	11.66	6.43

Table 4: Layer-level functional roles based on SPD statistics.

Metric	Mamba			Steered Mamba			Stable-Mamba
	L19	L20	L21	L19	L20	L21	L19	L20	L21
SPD Coefficient of Variation	0.346	0.067	0.526	0.031	0.165	0.019	0.0017	0.0098	0.0018
Gradient Sensitivity Norm	0.811	0.072	0.790	0.076	0.076	0.075	0.153	0.101	0.096
Post-Ablation KL Divergence	1.0	813.0	1.0	1.0	65.04	1.0	1.8	51.87	1.0

Table 5: Parameter sensitivity and information compression for Layers 19-21 using SPD. A higher coefficient of variation indicates gradient stability, a high gradient norm indicates high optimization sensitivity and KL divergence quantifies information loss. Vanilla Mamba shows a Layer 20 bottleneck (low gradient, high KL). Steering partially redistributes sensitivity but retains elevated KL, while Stable-Mamba exhibits more uniform sensitivity, higher gradients, and reduced KL, indicating bottleneck mitigation.

Steering Target	Accuracy (%)
Layer 18	65.0
Layer 19	69.0
Layer 20	72.5
Layer 21	59.0
Layer 22	61.0
Delta-sensitive Subspaces	72.5

Table 6: Ablations of steering target selection by layer and by subspace . Values show next-token prediction accuracy for Mamba on IFEval. We find that both the selection of layer (top section) and filtering by delta-sensitive subspaces (bottom section) is important for achieving high accuracy when steering. Baseline accuracy for Vanilla Mamba is 72.50%.

Figure 3: Performance comparison of SSM models with and without transferred steering parameters. Models are trained on The Pile and evaluated on test sets of task-specific benchmarks: SQuAD (recall), IFEval (instruction following), RULER (long context), MuSiQue (multi-hop reasoning), DROP (basic reasoning) and TriviaQA (QA), representing query types where SSMs show strong performance ( Guan et al., 2025 ; Wang et al., 2025 ) . Solid and striped bars denote performance without and with steering, respectively.

	RULER				Long Range Arena					LongBench v2
Model	NIAH	Aggrega-	QA	AVG	ListOps	Retrie-	Image	Pathfi-	AVG	Accur-	Confide-	Calibra-	Faithf-	Robust-	AVG
	(%)	tion (%)	(%)	(%)	(%)	val (%)	(%)	nder (%)	(%)	acy (%)	nce (%)	tion (%)	ulness (%)	ness (%)	(%)
Mamba-130M	84.00	67.00	66.70	72.57	62.00	61.00	25.00	60.00	52.00	100.00	21.99	96.86	100.00	100.00	83.77
Mamba-1.4B	87.71	71.00	94.00	84.24	65.00	68.00	26.50	65.00	56.13	100.00	24.22	88.95	100.00	100.00	82.63
Mamba-2.8B	90.43	70.00	95.00	85.14	74.00	69.00	49.50	70.00	65.63	100.00	26.55	95.63	100.00	100.00	84.44
Steered Mamba-130M	92.00	67.00	76.00	78.33	73.00	61.00	35.50	74.00	60.88	100.00	36.61	99.73	100.00	100.00	87.27

Table 7: Performance comparison across long-context benchmarks . Column names follow the original benchmark definitions. RULER uses a synthetic context length of 1000 tokens, LongBench v2 truncates inputs to 256 tokens and LRA uses a fixed sequence length of 100 tokens. RULER and LRA report task-level performance, while LongBench v2 reports benchmark-defined metrics aggregated across tasks. Stable-Mamba variants achieve the strongest performance, with Steered Mamba consistently outperforming their vanilla counterparts.

Appendix figures & tables38 assets

Supplementary material from the paper’s appendix.

Appendix

Figure A1 : Universality score distribution of subspaces. Mamba vs. Transformer

Figure A2 : 2D PCA Projection of Universality subspaces. Mamba representations are less tightly clustered than Transformer representations. PC1 corresponds to task type; PC2 corresponds to model type (Mamba vs. Transformer)

Panel of Figure 13 — Figure A3 : subspace-level causal influence in Mamba and Transformer . (a) Cumulative causal impact distribution across subspaces. (b) Layer-wise activation contributions across the network.

Figure A4 : Cross-Layer Causal Influence in Mamba.

Figure A5 : Layer-wise subspace sparsity using Gini Index . Mamba vs. Transformer

Figure A6 : Cumulative activation variance determining sparsity thresholds in Mamba and Transformer

Panel of Figure 17 — Figure A7 : Bias heatmaps for delta-sensitive subspaces . Mamba vs. Transformer

Table A1 : Knowledge subspace overlap statistics for the topic: “Capital of France”

Panel of Figure 19 — Figure A8 : Dead subspace analysis across Mamba model variants and Transformer

Panel of Figure 20 — Figure A9 : Layer-wise dynamics of token-detecting subspaces across Mamba model variants and Transformer . Highlights size-dependent differences in subspace allocation, token coverage and the depth at which new token representations emerge

Panel of Figure 21 — Figure A10 : N-Gram subspaces across layers and model scales . Each plot shows the number of subspaces in each layer that respond to unique unigrams, grouped by specificity. Dark green bins indicate highly specific subspaces, lighter green/yellow indicate moderately specific subspaces and orange/red indicate generalizt subspaces. GPT-2 exhibits relatively consistent N-Gram detection patterns across scales, reflecting attention-enabled parallel specialization. In contrast, Mamba models show a size-dependent shift in N-Gram selectivity, consistent with hierarchical recursive buildup through sequential state-space processing

Figure A11 : Receptive field sensitivity across layers for 1-gram, 2-gram and 3-gram in Mamba model variants and Transformer . Larger Mamba models exhibit a delayed emergence of higher order N-gram sensitivity while GPT-2 shows strong, depth-invariant specialisation

Panel of Figure 23 — Figure A12 : Cross-layer recursive dynamics in Mamba . (a) Temporal variance across layers shows a three-stage pattern: high variance in early layers (Layer 0 $\approx$ 0.102), a sharp reduction in mid-layers (Layer 2 $\approx$ 0.012) consistent with selective SSM gating and a gradual increase in deeper layers (Layer 10 $\approx$ 0.05) indicating accumulation of structured temporal representations through recursion. (b) Cross-layer activation changes measure the magnitude of feature transformation between layers, with larger changes in early transitions reflecting rapid state updates and convergence in long-range transitions indicating cumulative state integration across depth

Figure A13 : Cross-layer activation correlations between consecutive network layers . Mean correlations are near zero while maximum correlations approach unity, indicating sparse but highly selective inter-layer information transfer

Panel of Figure 25 — Figure A14 : Induction behavior in Mamba

Method	Description	Key Findings
Sparse Autoencoders (SAEs)	Reconstruct hidden activations using sparse latent representations to identify activation-level features and compact summaries of model activations	Identified sparse task-relevant features carrying most of the useful signal. SAE analysis revealed extreme sparsity and emergent specialized features responsible for key behaviors
Dictionary Learning	Applied on SAE latent representations to decompose activations into interpretable components and quantify sparsity, reconstruction error and feature utilization	Showed that sparsity, reconstruction error and feature usage occupy distinct latent dimensions, supporting functional specialization within representations
SPD (Stochastic Parameter Decomposition)	Decomposes parameters into sparsely used vectors and computes entropy, variance, effective rank, gradient sensitivity and post-ablation KL divergence	Identified a major Layer 20 information bottleneck characterized by high entropy, high variance and extreme KL divergence after ablation
Hypothesis Probing	Probe-based analysis examining whether task-relevant information is concentrated in sparse SAE-discovered features and circuits	SAE-based circuits showed strongest task alignment and predictive structure with structured positional selectivity
Polysemanticity Analysis	Measures feature interference, effective dimensionality and activation sparsity to analyze feature superposition	Found moderate superposition and low effective dimensionality, indicating substantial feature reuse across dimensions
Grokking Analysis	Analyzes emergent feature organization, induction behavior and transition dynamics during training	Observed sparse feature specialization, gradual learning dynamics and increasing induction strength during recursive integration

Appendix

Table A1 : Summary of mechanistic interpretability methods, analyses and key findings

Panel of Figure 27 — Figure B1 : Hypothesis probing of SAE-discovered features . (a) Activation strength distribution for a representative SAE feature ( $\rho=0.479$ ) follows a sparse, power-law pattern, where most activations are near zero and a small subset carries task-relevant signal, (b) Correlation between SAE features and token position varies non-uniformly across the sequence, peaking at mid-sequence positions (Layers 15–18), indicating recurrence and structured positional selectivity

Panel of Figure 28 — Figure B2 : Polysemanticity in Mamba

Panel of Figure 29 — Figure B3 : Grokking dynamics in Mamba . The model exhibits structured feature emergence and a discrete generalization transition following stable memorization

Metric	Value	Interpretation
Emerging features	77/768 (10%)	Sparse specialization
Grokking score	0.023	Gradual learning
Transition points	58	Frequent micro-shifts
Average stability	0.163	Moderate consistency
Gradient range	5.4 - 14.0	Stable convergence
Feature coherence	0.15	Low co-activation

Appendix

Table B2 : Summary of grokking-related training metrics.

Model	Layers	Bottleneck Layer	Depth Ratio (%)
Mamba-130M	24	20	87.5
Mamba-370M	48	40	83.3
Mamba-790M	48	41	85.4

Appendix

Table B3 : Scale-invariant bottleneck location across Mamba model variants

Model Size	Modularity	Circuit Complexity
130M	1.0061	1.0
370M	1.0028	2.0
Correlation	$-1$	$+1$

Appendix

Table B4 : Modularity–complexity trade-off across Mamba model variants

Layer	Entropy	Variance	Effective Rank
19	1.02	2.30	6.25
20	1.19	2.63	7.59
21	0.92	3.59	6.49

Appendix

Table B5 : Entropy–variance dynamics around the bottleneck layer

Metric	Value	Relative	Interpretation
APD Attribution	0.00758	5.0x L19	Causal impact
SPD Stability	CoV = 0.001	1000x avg	Parameter rigidity
Gradient Magnitude	0.072	0.09x avg	Frozen dynamics
Ablation KL	813.0	160x avg	Catastrophic loss
Prediction $\Delta$	45%	–	Output control
Susceptibility	-41.54	0.13 of max CoV	Perturbation resistance

Appendix

Table B6 : Causal and stability metrics for dt_proj.bias at Layer 20

Circuit Type	Intervention Type	Metric Affected	Effect Size	Interpretation
SAE-based Circuits	Feature Ablation	Perplexity	$\Delta$ PPL $\approx$ +394.1%	Circuits are causally necessary for factual recall
Probe-based Circuits	Activation Amplification	Token Probability	$\Delta$ Prob $\approx$ +218.0%	Circuits contribute directly to retrieval behaviour
Temporal Circuits	Information Gating ( $\Delta$ )	Accuracy	$\Delta$ Acc $\approx$ -42%	Circuits control long-range memory and sequence integration

Appendix

Table B7 : Causal Validation of Identified Mamba Circuits

Figure D1 : Steerable Delta-Sensitive subspaces causing recurrence in Mamba.

Layer	Role	Accuracy After Ablation (%)
18	Pre-bottleneck	56.5
19	Pre-compression	33.0
20	Bottleneck	77.0
21	Post-bottleneck	32.0
22	Output projection	8.0

Appendix

Table D1 : Layer-wise ablation analysis . Shows the impact of removing activations at different depths. Measured wrt baseline accuracy of 63.5%

Subspace Group	Accuracy After Ablation (%)
Delta-sensitive subspaces	78.0
Random Selection	50.0
High-Variance subspaces	42.0

Appendix

Table D2 : Subspace ablation analysis within the bottleneck layer (Layer 20) . Targeted ablation of Delta-sensitive subspaces improves performance, indicating that

\Delta

-controlled dynamics are the primary source of over-compression in this layer. In contrast, ablating high-variance or randomly selected subspaces degrades accuracy, showing that much of the remaining representational subspace carries task-relevant information. Results are measured on The Pile dataset relative to a baseline accuracy of 63.5%.

Figure D2 : Structural overview of the Mamba model hierarchy highlighting the steering location . Steering is applied within the SSM mixer on the bottleneck layer, targeting parameters in the dt_proj projection

Group	Change in	Total Subspaces	Subspace Indices
	Performance (%)
Critical Beneficial	$>10.0$	27	43, 53, 54, 59, 66, 75, 85, 127, 130, 134, 138, 142, 166, 173, 176, 191, 200, 202, 206, 207, 211, 214, 254, 259, 279, 284, 286, 303, 305, 309, 313, 314, 324, 341, 354, 355, 356, 369, 370, 374, 377, 379, 385, 402, 406, 408, 411, 416, 419, 430, 431, 435, 437, 442, 444, 455, 461, 465, 467, 476, 487, 502, 514, 520, 547, 553, 555, 566, 574, 575, 577, 578, 581, 583, 593, 597, 613, 625, 626, 646, 650, 657, 662, 663, 667, 672, 684, 693, 700, 703, 704, 710, 718, 729, 740, 757, 20, 33, 36, 76, 83, 93, 101, 149, 153, 159, 208, 223, 232, 239, 243, 340, 347, 361, 375, 400, 432, 458, 562, 572, 604, 605, 679, 694, 698, 706, 717
Very Beneficial	$5.0$ to $10.0$	82	0, 18, 40, 58, 79, 88, 117, 120, 124, 133, 139, 144, 145, 148, 152, 161, 168, 169, 201, 213, 225, 237, 245, 251, 273, 275, 287, 304, 367, 449, 475, 480, 486, 490, 493, 512, 521, 536, 538, 552, 556, 563, 567, 585, 590, 622, 624, 636, 649, 719, 730, 744, 763, 2, 21, 29, 103, 118, 128, 137, 170, 216, 271, 285, 288, 308, 393, 423, 434, 436, 443, 445, 529, 560, 616, 632, 656, 676, 680, 682, 702, 739
Beneficial	$2.0$ to $5.0$	326	1, 3, 5, 6, 7, 8, 9, 10, 12, 15, 17, 24, 25, 26, 27, 31, 32, 34, 35, 37, 38, 45, 57, 62, 64, 67, 71, 73, 77, 78, 82, 90, 91, 92, 96, 97, 98, 105, 106, 109, 110, 113, 116, 126, 129, 132, 135, 136, 141, 147, 150, 151, 154, 162, 163, 165, 171, 174, 177, 183, 184, 186, 192, 193, 194, 196, 197, 198, 205, 209, 212, 217, 218, 234, 235, 236, 242, 246, 247, 248, 250, 253, 255, 256, 257, 258, 260, 262, 263, 264, 265, 268, 269, 270, 272, 276, 283, 291, 292, 293, 295, 297, 298, 301, 302, 306, 307, 315, 316, 318, 319, 321, 322, 327, 331, 338, 342, 344, 345, 346, 349, 358, 360, 362, 363, 365, 371, 372, 373, 378, 382, 387, 388, 389, 390, 392, 399, 401, 403, 407, 412, 413, 417, 420, 421, 424, 425, 427, 428, 440, 451, 454, 462, 463, 466, 468, 472, 474, 477, 478, 479, 481, 484, 489, 491, 494, 498, 499, 500, 501, 505, 508, 510, 516, 519, 524, 526, 527, 530, 535, 551, 554, 557, 558, 568, 571, 576, 580, 582, 584, 586, 588, 592, 594, 596, 598, 600, 603, 608, 612, 614, 617, 620, 635, 638, 647, 651, 654, 658, 659, 660, 664, 668, 671, 677, 678, 685, 686, 687, 692, 695, 696, 699, 705, 708, 711, 712, 722, 724, 726, 727, 732, 733, 734, 735, 737, 741, 742, 745, 749, 750, 751, 752, 753, 758, 764, 765, 766, 767, 11, 16, 28, 41, 46, 47, 65, 81, 89, 95, 111, 131, 140, 146, 157, 185, 188, 190, 219, 281, 282, 299, 310, 326, 333, 334, 335, 336, 343, 348, 357, 364, 376, 383, 395, 398, 404, 415, 429, 438, 439, 460, 469, 473, 483, 488, 495, 496, 503, 541, 542, 545, 546, 548, 561, 569, 599, 601, 602, 611, 623, 627, 628, 631, 634, 639, 641, 644, 653, 665, 666, 670, 683, 701, 728, 746, 754
Neutral	$-2.0$ to $2.0$	155	4, 14, 19, 22, 23, 39, 42, 50, 51, 52, 56, 60, 61, 63, 70, 72, 84, 86, 87, 94, 100, 102, 104, 107, 112, 115, 119, 122, 123, 125, 143, 155, 156, 158, 160, 164, 167, 172, 178, 180, 182, 187, 210, 215, 222, 226, 230, 233, 240, 249, 252, 261, 266, 267, 274, 277, 278, 289, 290, 311, 312, 317, 323, 325, 329, 330, 332, 337, 339, 350, 351, 352, 366, 368, 380, 391, 396, 397, 405, 409, 410, 418, 422, 433, 441, 447, 448, 450, 452, 453, 456, 457, 459, 470, 492, 497, 506, 507, 509, 511, 517, 522, 523, 534, 537, 539, 543, 549, 550, 559, 565, 570, 587, 591, 595, 606, 607, 609, 610, 615, 618, 619, 629, 630, 633, 637, 640, 642, 643, 648, 652, 655, 661, 673, 674, 675, 688, 689, 690, 691, 697, 707, 713, 715, 716, 720, 721, 723, 736, 738, 748, 755, 759, 760, 761
Slightly Detrimental	$-5.0$ to $-2.0$	55	55, 30, 44, 69, 74, 80, 114, 195, 199, 203, 204, 220, 221, 224, 228, 229, 231, 238, 244, 294, 296, 328, 353, 386, 394, 414, 471, 482, 485, 504, 513, 515, 518, 525, 528, 531, 532, 533, 540, 544, 564, 573, 579, 621, 645, 669, 681, 709, 714, 731, 747, 756, 762, 320, 359

Appendix

Table D3 : Subspace-level ablation categorization based on change in performance accuracy . Subspaces are grouped by ablated causal impact.

Steering Factor	Accuracy (%)	Steering Factor	Accuracy (%)
0.1	37.0	20.0	41.6
1.5	43.0	30.0	37.0
2.0	47.4	40.0	35.0
3.0	46.0	50.0	31.6
4.0	47.2	60.0	31.0
5.0	48.5	70.0	29.5

Appendix

Table D4 : Effect of steering factor on next-token prediction accuracy. Baseline accuracy is 47.0%.

Panel of Figure 42 — Figure E1 : Stable-Mamba architecture illustrating ensemble gating, multi-timescale SSMs and sparse attention at the bottleneck

Component	Mamba	Stable Mamba	Performance Improvement in Tweaked Mamba
State transition	$h_{t}=\bar{A}h_{t-1}+\bar{B}x_{t}$ (single time-scale)	$h_{t}^{(k)}=\bar{A}_{k}h_{t-1}^{(k)}+\bar{B}_{k}(g\odot x_{t})$ (3 time-scales: slow, medium, fast)	1) +23% perplexity improvement on long sequences ( $>$ 1024 tokens) 2) +15% accuracy on multi-scale reasoning tasks 3) Entropy variance reduced by 31%
Tweaked Mamba	$y_{t}=Ch_{t}+Dx_{t}$ (linear)	$y_{t}=\sum_{k}w_{k}(C_{k}h_{t}^{(k)}+D_{k}x_{t})$ (weighted ensemble)	1) Adaptive temporal resolution via learned weighting 2) Top features show $0.16+$ importance vs. $0.02$ in Mamba
Tweaked Mamba	Not present	$h_{\text{global}}=\text{SparseAttn}(h_{\text{local}})$ $\text{output}=\alpha h_{\text{global}}+(1-\alpha)h_{\text{local}}$	1) 22% reduction in bottleneck entropy (1.21 $\rightarrow$ 0.94) 2) Smoother information flow in layers 19–21 3) $\sim$ 0.4 $\times$ cost of full Transformer attention
Gating	Uniform processing	$g=\sum_{i}\alpha_{i}\sigma(W_{i}\cdot\text{LN}(x))$ (ensemble gates, $n=3$ )	1) +683% feature usage efficiency (1.88% $\rightarrow$ 14.7%) 2) Sparsity reduced from 98.1% to 85.3%
Compression	Not present	$c=0.5+0.5\cdot\sigma(\text{MLP}(\bar{x}))$	1) Adaptive signal preservation with reduced noise 2) Effective compression strength in range 0.2-0.3
Gradient control	Standard backpropagation	$\frac{\partial L}{\partial\theta}=\frac{\partial L}{\partial\text{out}}\cdot\lambda_{\text{comp}}\cdot\frac{\partial\text{out}}{\partial\theta}$ (5 learnable scales)	1) 46% reduction in gradient instability (CoV: 17.3% $\rightarrow$ 9.4%) 2) Stable convergence without manual tuning

Appendix

Table E1: Component-wise architectural modifications introduced in Stable-Mamba and their impact on performance and stability

Component	Parameters	Memory	Frequency
SSM ( $\times$ 3)	$3d^{2}$ = 1.77M	$O(d\cdot s)$	Every layer
Gates ( $\times$ 3)	$3d^{2}$ = 1.77M	$O(d)$	Every layer
Sparse Attention	$0.3sd$ = 47K	$O(0.3s^{2})$	Every 5th layer
Total / layer (avg)	$\sim$ 4.95M	$O(d\cdot s+0.06s^{2})$	-

Appendix

Table E2: Average per-layer computational cost of Stable- Mamba

Mechanism	Metric	Mamba	DeepTrace	Performance Improvement
Ensemble gating	Sparsity	98.1%	85.3%	-12.8pp denser representations
	Feature usage	1.88%	14.7%	7.81x higher feature utilisation
Multi-timescale states	Long-context PPL	15.2	12.4	-18.4% better long-range modeling
	Multi-scale accuracy	73.5%	84.7%	+11.2pp improved multi-scale reasoning
Sparse global context	Bottleneck entropy	1.21	0.94	-22% smoother information flow
Learnable gradient scaling	Gradient CoV	0.173	0.094	-46% more stable optimisation

Appendix

Table E3 : Quantitative impact of the architectural modifications introduced in Stable-Mamba

Model	Latency (s)	Peak Memory (GB)	Throughput (tokens/s)	GPU Hours
Mamba-130M	0.709	1.668	1444	0.0068
Mamba-1.4B	1.475	2.307	694	0.0081
Mamba-2.8B	1.548	3.368	662	0.0104
Steered Mamba	0.713	1.724	1436	0.0070
Steered Mamba-1.4B	1.488	2.402	688	0.0081
Steered Mamba-2.8B	1.550	3.568	661	0.0104

Appendix

Table E4 : Latency, peak memory, throughput and GPU hours comparison on LongBench v2. Experiments are conducted on GeForce RTX 3090. Models are trained on The Pile and evaluated on WikiText-2-v1 .

Component	FLOPs	Explanation
SSM	$d^{2}=590$ K	Matrix multiply $(d\times d)$
in_proj + out_proj	$2d^{2}=1.18$ M	Two linear projections
Normalization	$2d=1.5$ K	Mean + std
Skip connections	$d=768$	Element-wise addition
Total	$\approx 1.77$ M

Appendix

Table E5 : FLOP calculation for Mamba and Steered Mamba

Component	FLOPs	Explanation
SSM $\times$ 3	$3d^{2}=1.77$ M	Three parallel SSMs
Gates $\times$ 3	$3d^{2}=1.77$ M	Three gating projections
Compression MLP	$(d^{2})/4+d/4=148$ K	Two-layer MLP
in_proj + out_proj	$2d^{2}=1.18$ M	Linear projections
Sparse attention (20%)	$\approx 47$ K	Partial attention
Normalization	$4d=3$ K	LayerNorm

Appendix

Table E6 : FLOP calculation for Stable-Mamba

State Space Language Models

See all 18 papers

Oct 8, 2025

Contextual Representation Flow Across Depth in State-Space and Transformer Architectures

State Space Models (SSMs) have emerged as efficient, linear-time alternatives to Transformer-Based Models (TBMs) for long-context processing. However, their different mechanisms raise an important question: how does contextual representation flow across depth differ between these model families and their hybrids? We examine this question across twenty-one pretrained models, combining token-wise spectral diversity, layer-wise variation on seven datasets, and answer-location probing on two retrieval tasks. This joint analysis reveals that TBMs retain most of their token diversity at the output, whereas SSMs reach richer intermediate diversity but retain less of it on average, with late-layer compression. Comparisons with reinitialized weights reveal that these profiles depend on pretrained weights, while hybrids combine strong answer-location accessibility with architecture-dependent geometry, showing that high accessibility need not require the steadiest representation flow. Finally, we conduct theoretical case studies connecting these contrasts to mechanism-level differences: query-time token mixing in TBMs, repeated recurrent updates in SSMs, and preserved attention updates in hybrids.

Nhat M. Hoang, Do Xuan Long, Cong-Duy Nguyen +3

Nanyang Technological University, Singapore · National University of Singapore

State Space Models State Space Language Models Transformer Representation Geometry in Language Models Long-Context Modeling

Ask PDF

Sep 21, 2026

Comparing Latent Concept Formation in State Space Models and Transformers via Sparse Autoencoders

The quadratic scaling of Transformer self-attention has driven the adoption of sub-quadratic Selective State Space Models (SSMs) like Mamba, which compress past context into a fixed-size recurrent hidden state. This strict informational bottleneck raises a foundational question for mechanistic interpretability: do SSMs and Transformers learn fundamentally distinct latent representations? In this work, we employ Sparse Autoencoders (SAEs) to conduct a large-scale, feature-level correspondence analysis between Mamba-130m and Pythia-70m over a 10-million token corpus. Contrary to hypotheses predicting widespread architectural divergence, we find no evidence of systematic representational divergence between architectures: across the observed Jaccard distribution, 99.98% of Mamba features cluster toward the upper alignment boundary, providing preliminary feature-level support for the Universality Hypothesis. We further identify and qualitatively characterize this microscopic fraction (0.02%) of diverging features, finding patterns consistent with the hypothesis that the recurrent bottleneck selectively limits the parsing of rigid syntax rather than broad semantic ontology. We demonstrate that while Pythia's unconstrained attention permits the monosemantic decomposition of distinct formatting edge-cases, Mamba is forced to compress unrelated syntactical anomalies into polysemantic "junk drawer" neurons to preserve state capacity. Collectively, these results suggest that architectural routing mechanisms may have negligible impact on core semantic understanding, with representational divergence confined to extreme structural margins.

Rithin Nagaraj, Rupa Laalasa Oruganti, Prerna Subhashchandra Kunder +1

Department of Computer Science and Engineering PES University Bangalore, India

Transformer Interpretability Language Modeling Representation Geometry in Language Models Representation Alignment State Space Models

Ask PDF

Sep 15, 2026

On the Importance of Gating: Memorization vs. In-Context Learning in State Space Models

State Space Models (SSMs) have emerged as a compelling alternative to Transformers, enabling sequence modeling with constant memory and linear compute. Although SSMs exhibit reasonable performance and favorable computational characteristics, they continue to lag behind Transformers on tasks that require in-context learning and precise retrieval, slowing their adoption for large-scale language modeling. In this work, we demonstrate that both the success and failure of SSMs in these domains can be explained by studying the role of the gating mechanism, a prevalent component in modern recurrent networks. Specifically, we show through theory and experiments that this gating mechanism causes SSMs to first learn an in-weights "memorization" solution, while delaying, or even preventing, convergence to a correct in-context learning solution. Importantly, this happens even in cases where there are no fundamental limitations due to the architecture or its memory capacity. On the other hand, we find that gating is often beneficial for improving generalization to long sequence lengths. Our results illuminate the crucial role of the gating mechanism in shaping both the training dynamics and generalization of SSMs, and provide a basis for understanding and improving linear-time models.

William L. Tong, Aryo Lotfi, Emmanuel Abbe +6

Harvard University · Apple

Memorization in Language Models Long-Context Modeling In-Context Learning State Space Models

Ask PDF

Abstract

State Space Language Models

Explore similar work

Contextual Representation Flow Across Depth in State-Space and Transformer Architectures

Comparing Latent Concept Formation in State Space Models and Transformers via Sparse Autoencoders

On the Importance of Gating: Memorization vs. In-Context Learning in State Space Models