Mamba-2 compresses all context into a fixed-size recurrent state that dominates decoding memory and bandwidth. Existing compressors quantize state values or prune state coordinates, scoring each channel independently in the native basis. We examine trained Mamba-2 states and find that energy is almost rank-one, but the output does not read the energetic directions: energy concentrates in a bias-like direction accumulated by the recurrence, while the content the output reads lives in a low-dimensional subspace that is not aligned with the state's coordinates. Even output-aware scores are computed coordinate by coordinate, so they miss the directions that determine the input--output map. Compression should therefore target the write--read subspace rather than individual coordinates. We propose \method{}, a post-training transform that instantiates balanced truncation with full controllability and observability Gramians and applies the rotation exactly via Mamba-2's scalar-identity transition. Because the useful rank varies widely across layers, we further introduce HARP, which allocates ranks per layer from the Hankel spectrum and protects retrieval-sensitive layers, improving both perplexity and retrieval over uniform budgets. Calibration takes a single forward pass, an order of magnitude faster than GHOST's released protocol, with lower peak memory. With HARP, \method{} halves the state dimension at an average WikiText-2 perplexity increase of 0.3 points (1.8%) across 130M--2.7B, versus 1.6 (9.3%) for GHOST, and the gap widens at 25% retention.
Figures & tables
Figure 1: BASS overview. Top: one Mamba-2 layer with pre- Δ write B~t , read Ct , and step Δt , where Bt=Δt⊙B~t ; state and output follow Eqs. ( 1 )–( 2 ). Bottom: BASS estimates Wc,Wo , keeps the balanced basis Q , and applies H′=HQ , B′=Q⊤Bt , C′=Q⊤Ct ; the compressed layer reruns update and read with H′,B′,C′ . The rotation is exact for any r ; lossless at r=N ; only truncation r<N approximates. Post-training; all other parameters untouched.
Model
L
H
erank med. (p90)
≤2
top-1
aniso
130M
24
24
1.16 (2.10)
89.6%
0.93
0.43
370M
48
32
1.17 (1.76)
91.9%
0.92
0.39
1.3B
48
64
1.28 (1.93)
90.7%
0.89
0.32
Table 1: State statistics per model. L : layers; H : heads. “erank”: median effective rank of a head; “ ≤2 ”: fraction of heads with effective rank ≤2 ; “top-1”: median energy share of the leading direction; “aniso”: median per-layer mean pairwise cosine of cross-sequence states.
Figure 2: Empirical CDF of per-head effective rank. About 90% of state heads live in at most two directions, across all model scales.
WikiText-2 perplexity ↓
Avg. increase vs. dense
130M
370M
780M
1.3B
2.7B
Method
r=64
r=32
r=64
r=32
r=64
r=32
r=64
r=32
r=64
r=32
50%
25%
Dense
25.21
17.97
14.78
12.39
10.93
ref.
ref.
GHOST
28.35
38.67
19.69
459.6
16.17
101.2
13.47
22.75
11.65
14.56
+1.6 (9.3%)
+111 (643%)
Magnitude
1972
105.8
56.81
2771
26.77
91.71
21.72
33.49
16.03
24.17
+402 (1628%)
+589 (3290%)
Saliency
28.30
37.11
19.54
62.65
15.96
43.91
13.49
19.78
11.48
13.71
+1.5 (8.6%)
+19.2 (116%)
Table 2: WikiText-2 perplexity at r=64/32 (50%/25% retention). Last two columns: average increase vs. dense in ppl points, with relative change in parentheses (mean over five checkpoints). Bold marks the best compressed method in each column and underline indicates the second best; Dense is a reference and is not compared. Official rows use released implementations; ours use BASS .
Figure 3: (a) Perplexity change vs. retained rank for Mamba2-130M (log scale). (b) At 50% state reduction for the three checkpoints with complete method sweeps (log scale); 780M/2.7B are in Table 2 .
LAMBADA last-word accuracy ↑
Avg. change vs. dense
130M
370M
780M
1.3B
2.7B
Method
r=64
r=32
r=64
r=32
r=64
r=32
r=64
r=32
r=64
r=32
50%
25%
Dense
0.449
0.557
0.607
0.648
0.703
ref.
ref.
GHOST
0.355
0.193
0.463
0.246
0.471
0.307
0.520
0.332
0.609
0.459
−10.9 ( −18.7% )
−28.6 ( −49.2% )
Magnitude
0.092
0.037
0.102
0.006
0.164
0.027
0.154
0.039
0.295
0.088
−43.2 ( −73.7% )
−55.4 ( −93.5% )
Saliency
0.330
0.139
0.449
0.279
0.461
0.299
0.520
0.318
0.596
0.443
−12.2 ( −21.0% )
−29.7 ( −51.5% )
Table 3: LAMBADA last-word accuracy at the budgets of Table 2 (512 examples). Last two columns: average change vs. dense in accuracy points, with relative change in parentheses (unweighted mean over the five checkpoints). Official rows use released implementations; ours use BASS . HARP selects layers on held-out examples (512–767), disjoint from evaluation. Bold marks the best compressed method in each column and underline indicates the second best; Dense is a reference and is not compared. Magnitude is a failure floor.
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
Method
Model / budget
Released
Ours
Metric
Reading
GHOST-2
130M / 50%
+3.21
+3.14
word / token ppl
match
Random
130M / 50%
45.6
47.5
token ppl
match
Taylor-2
370M / 50%
12647
13851
word / token ppl
same order
Taylor-2
1.3B / 50%
13.94
539
word / token ppl
released likely no-op
Taylor-2
130M / 50%
1.0⋅1016
4.3⋅1018
word / token ppl
both catastrophic
Magnitude
130M / 50%
910
1972
word / token ppl
metric gap
Appendix
Table 4: Reproduction audit. Released values are word-level perplexities from the GHOST paper or repository; ours are reruns with the released compressor under our evaluation. No Δ is reported because the metrics differ. “Reading” states the interpretation defended in the text.
Setting
130M
370M
1.3B
Replica (Table 2 )
+53.4
+2458
+83.7
Replica, 32k window
+47.3
+1719.7
+106.3
Official calib., 32k window
+47.3
+1719.3
+61.7
Official calib., draw 2, 32k
–
+1748.9
+62.5
Replica calib., draw 2, 32k
–
+2063.6
–
Hard truncation
+53.4
+2458
+83.7
Appendix
Table 5: GHOST-2 perplexity change at 25% under different calibration protocols, calibration draws and evaluation windows (all relative to the same uncompressed model). The collapse persists in every setting. Hard truncation removes channels physically and retains 31–32 of 128 states per layer. “–”: not run.
Sparsity
0.1
0.3
0.5
0.7
0.9
Dense
13.17
GHOST-2
13.24
13.51
14.23
16.16
25.07
Magnitude
427.7
22.39
29.34
39.96
69.49
Taylor-2
13.18
13.26
13.94
4255
4690
Random
13.52
15.04
17.77
30.14
64.21
SparseGPT
13.17
13.19
13.25
13.52
15.51
Appendix
Table 6: Word-level WikiText-2 perplexity from the authors’ released sweep on Mamba2-1.3B. SparseGPT is unstructured and does not reduce the state dimension; it is shown for reference.
WikiText-2 perplexity ↓
130M
370M
780M
1.3B
2.7B
Method
r=64
r=32
r=64
r=32
r=64
r=32
r=64
r=32
r=64
r=32
Taylor-2 (released code, our run)
4.3⋅1018
4.5⋅1034
13851
11654
1.7⋅105
1.2⋅106
539.0
813.4
511.7
542.9
Energy (ours)
28.16
38.64
20.07
24.47
16.11
21.04
13.63
17.32
11.67
13.74
Random (ours)
39.23
84.54
28.21
–
26.20
–
18.38
–
14.02
–
Appendix
Table 7: Additional baselines moved from Table 2 to keep the main comparison focused. Taylor-2 entries for 780M/2.7B are evaluated in bf16 (fp32 exceeds 24 GB); energy selects the r largest E[H2] coordinates, random draws coordinates uniformly, and both use the BASS protocol and the same evaluation as Table 2 . “–”: not run.
130M
1.3B
Calibration source
r=64
r=32
r=64
r=32
Train, 4 seq. (reported)
+7.18
+25.38
+5.06
+22.75
Validation, 4 seq.
+8.18
+33.06
+5.00
+22.34
Train, 8 seq.
+7.38
+27.09
+4.89
+22.02
Eval. text (leaky, excluded)
+3.50
+13.96
+1.02
+8.95
Appendix
Table 8: WikiText-2 perplexity change of BASS under different calibration sources. The last row calibrates on the evaluation text as a leakage control; all reported results use train calibration.
r=64
r=32
Model
128 × 128
4 × 1024
128 × 128
4 × 1024
130M
+8.14
+7.18
+35.25
+25.38
370M
+6.24
+5.73
+23.97
+20.34
1.3B
+5.98
+5.06
+25.59
+22.75
Appendix
Table 9: Calibration-length control: BASS perplexity change when the Gramians are estimated from 128-token sequences (GHOST’s protocol, 16k tokens) instead of the reported 1024-token sequences.
Configuration (130M)
ppl change
LAMBADA
Uncompressed baseline
0.00%
0.449
Uniform r=64
+7.18%
0.371
Adaptive rˉ=64 , floor 2
+2.78%
0.320
Adaptive rˉ=64 , floor 32
+2.81%
0.313
Retrieval floor ( k=6 )
+2.92%
0.404
Retrieval floor, weighted
+3.24%
0.422
Appendix
Table 10: Allocation ablations at a matched average budget on Mamba2-130M. Hankel allocation improves perplexity but trades off retrieval; per-layer rank floors do not recover it, while the retrieval floor does.
Uniform rˉ=64
HARP (floor)
HARP (weighted)
Model
ppl
LAMB.
ppl
LAMB.
ppl
LAMB.
130M
+7.18
0.371
+2.92
0.404
+3.24
0.422
370M
+5.73
0.412
+2.75
0.438
+2.05
0.563
1.3B
+5.0
0.572
+1.92
0.592
+59.0
0.440
Appendix
Table 11: Retrieval-aware allocation across scales. The probe selects layers on held-out LAMBADA examples (512–767), disjoint from the evaluation set (0–511). The retrieval floor ( k=6 ) protects the six most probe-sensitive layers; the weighted variant is unstable across probe draws.
Figure 4: Uniform, Hankel-adaptive, retrieval-floor ( k=6 ), sensitivity-weighted and anti-allocated budgets at 50% reduction (130M): the retrieval-aware budgets dominate uniform allocation on both perplexity and LAMBADA.
Model
probe-sensitive layers (drop, pts)
probe time
130M
20 (16.4), 19 (7.4), 10 (2.3), 17 (2.3)
218 s
370M
38 (15.6), 32 (3.5), 35 (2.7), 28 (2.3)
779 s
1.3B
33 (26.2), 6 (1.6), 25 (1.2), 37 (0.8)
1362 s
Appendix
Table 12: Leave-one-layer-out retrieval probe ( rp=8 for the probed layer, rˉ=64 otherwise; held-out LAMBADA examples 512–767, disjoint from the evaluation set).
BASS
BASS
GHOST
GHOST
4 × 1024
128 × 128
128 × 128
4 × 1024
130M
1.7 / 2.3
4.6 / 2.3
14.9 / 3.3
13.1 / 10.2
370M
3.7 / 4.0
11.3 / 4.0
37.2 / 5.7
33.3 / 15.1
1.3B
6.8 / 10.2
22.0 / 10.3
70.4 / 16.3
OOM
Appendix
Table 13: One-time calibration cost (median seconds / peak GB, RTX 3090, float32) under each method’s released protocol and under a token-matched swap. GHOST’s 4 × 1024 run exceeds 24 GB at 1.3B.
Model
N=128
r=64
r=32
130M
18.0
9.0
4.5
370M
48.0
24.0
12.0
1.3B
96.0
48.0
24.0
Appendix
Table 14: Recurrent-state size summed over layers (MB, float32).
Model
baseline
+map
r=64
r=32
130M
244.6 / 5.12
+0.8%
173.0 / 4.01
137.9 / 3.61
370M
657.2 / 12.63
+0.6%
460.4 / 10.20
367.2 / 9.18
1.3B
OOM
OOM
968.2 / 22.50
782.7 / 20.47
Appendix
Table 15: Prefill at L=1024 , batch 1, float32, one RTX 3090 (median ms / peak GB over five runs). All rows use the same chunked scan; “+map” applies an identity map to the uncompressed state. At 1.3B the uncompressed scan exceeds 24 GB.
decode tokens/s
peak GB
Model
B
dense
r=64
r=32
dense
r=64 / r=32
130M
1
66.9
69.2
71.9
2.53
2.52 / 2.51
130M
128
4050
6555
8717
6.32
4.74 / 3.95
370M
1
37.0
34.2
34.7
7.19
7.17 / 7.15
370M
32
1131
1052
1074
10.31
9.47 / 9.05
1.3B
1
37.7
33.8
35.2
5.24
5.19 / 5.17
Appendix
Table 16: Decode throughput and peak GPU memory (float32, one RTX 3090), measured over 128 steps after warmup with a cached-state reference loop; rotation included for BASS rows.
Model
r
BASS
write-only
read-only
rand-rot
130M
64
+7.2
+7.4
+10.0
+78.7
130M
32
+25.4
+28.6
+34.7
∞∗
370M
64
+5.7
+6.2
+6.1
+135.4
370M
32
+20.3
+22.5
+60.1
+239.9
1.3B
64
+5.0
+6.1
+4.9
+31.3
1.3B
32
+22.7
+27.7
+23.3
+206.2
Appendix
Table 17: WikiText-2 perplexity change (%) under matched surgery with single-sided Gramian bases and a random orthogonal rotation. ∗ Diverged to 4×1021 ppl.
Model
r
Balanced
Hankel-diag
Saliency
Energy
130M
8
89.6
42.7
37.1
23.8
130M
16
93.3
47.4
54.3
41.2
130M
32
96.6
70.0
69.5
61.2
370M
8
86.9
32.3
33.4
26.1
370M
16
91.5
44.3
42.5
39.8
370M
32
95.5
65.7
63.9
55.7
Appendix
Table 18: Fraction (%) of the empirical Hankel spectral energy tr(WcWo) retained by an r -dimensional subspace (median over layers). Coordinate columns select native coordinates by the layer-mean score E[H2] (energy), E[H2C2] (saliency), or the strongest marginal-diagonal ordering diag(Wc)⊙diag(Wo) .
Model
dense ppl
r=64
95% CI
r=32
130M
22.69
+6.8%
[6.5,7.2]
+24.7%
370M
15.87
+5.9%
[5.6,6.1]
+20.2%
1.3B
11.71
+4.9%
[4.7,5.1]
+16.7%
Appendix
Table 19: Full-test-file evaluation (280 × 1024 tokens): perplexity of the uncompressed model and BASS relative change; CIs are for the r=64 change (the r=32 CIs are within ±1.3 points, [23.5, 26.0]/[19.4, 21.1]/[16.0, 17.4]).
Configuration
130M
370M
1.3B
4096 bits/head (1 × )
Dense, r=128 , fp32
25.21
17.97
12.39
2048 bits/head (2 × )
Dense, r=128 , int16
25.21
17.97
13.36
BASS, r=64 , fp32
27.02
19.00
13.01
BASS+HARP, rˉ=64 , fp32
25.95
18.46
12.63
Appendix
Table 20: WikiText-2 perplexity (lower is better) at matched per-head state budgets (bits = retained modes × bits per value). r : modes per layer; rˉ : average for HARP, whose per-layer ranks vary; compression is relative to the 4,096-bit dense state. Bold marks the best entry for each model within a budget. All rows use the same token-by-token loop; integer rows use per-mode absmax scales without fused kernels or clipping calibration.
Mamba's recurrent state h_t is, by construction, a compressed summary of every token seen so far. This raises a tempting hypothesis: if we extract token-level outputs y_t at fixed patch boundaries, we obtain semantic sentence summaries for free, with no pooling head, no fine-tuning, and no [CLS] token. We test this hypothesis carefully. Across five benchmarks (SST-2, CoLA, MRPC, STS-B, IMDb), we compare four strategies for extracting frozen sentence representations from a pretrained Mamba-130M backbone under a strict frozen-feature probing protocol, using three random seeds where computationally feasible. The results do not support the hypothesis: patch boundary readouts do not consistently outperform simple mean pooling. We identify and quantify two structural pathologies: severe anisotropy (mean pairwise cosine similarity 0.9999, std 0.000044) and representational collapse in the raw final SSM state (MCC = 0.000 on CoLA across all three seeds, confirmed via confusion matrix). We further propose orthogonal injection, a modified recurrence that constrains new information per
Mamba models underperform Transformers on in-context retrieval. This gap is usually attributed to Mamba's fixed-size recurrent state, which loses information as the context grows. We challenge this assumption by focusing on query-first retrieval, a setup in which the query precedes the context and state capacity should not be the binding constraint. We show that even in this setting Mamba lags substantially behind Transformers, despite its selection mechanism, which was designed to distinguish relevant content from distractors. We then demonstrate that the Mamba selection gates have sufficient control over the state updates to influence the model output: intervening on specific gates alone can swing accuracy between extremes. However, as gates are input-dependent, they require the incoming token's representation to encode its relevance to the query. We therefore probe token representations across layers, and find that the relevance signal emerges only gradually with model depth, and that the gates begin to reflect it only several layers deeper. Crucially, we find that gates in retrieval-critical layers are not yet aligned with relevance. The selection mechanism is thus expressive enough to solve the task, but it commonly operates on a relevance signal that has not yet arrived. Mamba's retrieval bottleneck therefore extends beyond state capacity: even when capacity is sufficient, the delayed emergence of relevance information across depth often prevents successful retrieval.
Linear attention transformers have become a strong alternative to softmax attention due to their efficiency. However, linear attention tends to be less expressive and results in reduced accuracy compared to softmax attention. To bridge the accuracy gap between softmax attention and linear attention, we manipulate Mamba-2, a very strong linear attention variant. We first simplify Mamba-2 down to its most fundamental and important components, evaluating which specific choices make it most accurate. From this simplified Mamba variant (Mamba-2S), we improve the A-mask and increase the order of the hidden state, resulting in a method, which we call 2Mamba, that is nearly as accurate as softmax attention, yet much more memory efficient for long context lengths. We also investigate elements to Mamba-2 that help surpass softmax attention accuracy. Code is provided for all our experiments.
Gabriel Mongaras, Eric C. Larson
Lyle School of Engineering Southern Methodist University Dallas, TX 75205