Balanced State Subspaces for Mamba-2 State Compression

Shang Tingyu; Cao Wenfei

cs.LGOct 27, 2026

Balanced State Subspaces for Mamba-2 State Compression

Authors: Shang Tingyu, Cao Wenfei

Organizations: School of Mathematics and Statistics, Shaanxi Normal University, Xi’an 710119, China

Abstract

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-

\Delta

write

\tilde{B}_{t}

, read

C_{t}

, and step

\Delta_{t}

, where

B_{t}=\Delta_{t}\odot\tilde{B}_{t}

; state and output follow Eqs. ( 1 )–( 2 ). Bottom: BASS estimates

W_{c},W_{o}

, keeps the balanced basis

Q

, and applies

H^{\prime}=HQ

B^{\prime}=Q^{\top}B_{t}

C^{\prime}=Q^{\top}C_{t}

; the compressed layer reruns update and read with

H^{\prime},B^{\prime},C^{\prime}

. 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)	$\leq 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; “

\leq 2

”: fraction of heads with effective rank

\leq 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 $\downarrow$										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 $\uparrow$										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	$\mathbf{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{\cdot}10^{16}$	$4.3{\cdot}10^{18}$	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

\Delta

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 $\downarrow$
	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{\cdot}10^{18}$	$4.5{\cdot}10^{34}$	13851	11654	$1.7{\cdot}10^{5}$	$1.2{\cdot}10^{6}$	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

\mathbb{E}[H^{2}]

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 $\times$ 128	4 $\times$ 1024	128 $\times$ 128	4 $\times$ 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 $\bar{r}{=}64$ , floor 2	$+2.78\%$	0.320
Adaptive $\bar{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 $\bar{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 (

r_{\mathrm{p}}{=}8

for the probed layer,

\bar{r}{=}64

otherwise; held-out LAMBADA examples 512–767, disjoint from the evaluation set).

	BASS	BASS	GHOST	GHOST
	4 $\times$ 1024	128 $\times$ 128	128 $\times$ 128	4 $\times$ 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

\times

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	$\mathbf{+7.2}$	$+7.4$	$+10.0$	$+78.7$
130M	32	$\mathbf{+25.4}$	$+28.6$	$+34.7$	$\infty^{*}$
370M	64	$\mathbf{+5.7}$	$+6.2$	$+6.1$	$+135.4$
370M	32	$\mathbf{+20.3}$	$+22.5$	$+60.1$	$+239.9$
1.3B	64	$+5.0$	$+6.1$	$+4.9$	$+31.3$
1.3B	32	$\mathbf{+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\times 10^{21}

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

\operatorname{tr}(W_{c}W_{o})

retained by an

r

-dimensional subspace (median over layers). Coordinate columns select native coordinates by the layer-mean score

\mathbb{E}[H^{2}]

(energy),

\mathbb{E}[H^{2}C^{2}]

(saliency), or the strongest marginal-diagonal ordering

\operatorname{diag}(W_{c})\odot\operatorname{diag}(W_{o})

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

\times

1024 tokens): perplexity of the uncompressed model and BASS relative change; CIs are for the

r{=}64

change (the

r{=}32

CIs are within

\pm 1.3

points, [23.5, 26.0]/[19.4, 21.1]/[16.0, 17.4]).

4096 bits/head (1 $\times$ )
Configuration	130M	370M	1.3B
Dense, $r{=}128$ , fp32	25.21	17.97	12.39
2048 bits/head (2 $\times$ )
Dense, $r{=}128$ , int16	25.21	17.97	13.36
BASS, $r{=}64$ , fp32	27.02	19.00	13.01
BASS+HARP, $\bar{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

\times

bits per value).

r

: modes per layer;

\bar{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.

LLM Compression

Large language model compression reduces the memory, compute, and context costs of LLMs while aiming to preserve their capabilities.

See all 283 papers

Lost in State Space: Probing Frozen Mamba Representations
Apr 30, 2026Bhagyashree Wagh, Akash SinghRepresentation Geometry Representation Collapse
University of Washington · University of Florida
Beyond State Capacity: Why Retrieval Breaks in Mamba
Oct 14, 2026Nir Yarden, Amit Ben-Artzy, Roy SchwartzIn-Context Retrieval Selective SSMs
The Hebrew University of Jerusalem
2Mamba2Furious: Linear in Complexity, Competitive in Accuracy
Feb 19, 2026Gabriel Mongaras, Eric C. LarsonLinear Attention Memory-Efficient Attention
Lyle School of Engineering Southern Methodist University Dallas, TX 75205

Apr 30, 2026J/K move · Enter open · S save

Lost in State Space: Probing Frozen Mamba Representations

Bhagyashree Wagh, Akash Singh

University of Washington · University of Florida

Representation Geometry Representation Collapse Sentence Embeddings State Representation Learning State Space Models Semantic Representation Learning

Open paper AskPDF arXiv

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

Abstract

LLM Compression

Explore similar work

Lost in State Space: Probing Frozen Mamba Representations

Beyond State Capacity: Why Retrieval Breaks in Mamba

2Mamba2Furious: Linear in Complexity, Competitive in Accuracy