PaperlayerView exact revision on arXiv

Source-linked AI summary

Balanced State Subspaces for Mamba-2 State Compression

Shang Tingyu, Cao Wenfei

arXiv:2610.36617v1cs.LGcs.AIcs.CLstat.APstat.ME

TL;DR

Mamba-2’s fixed-size recurrent state dominates decoding memory, while existing compression approaches score state values or coordinates rather than the joint input–output subspace. The paper proposes BASS, a post-training balanced-subspace transform, and HARP, which allocates ranks across layers. With HARP at 50% retention, BASS adds 0.3 average WikiText-2 perplexity points across five checkpoints, versus 1.6 for GHOST.

  • Problem

    Mamba-2’s fixed-size recurrent state dominates decoding memory, and prior state compressors score native coordinates rather than jointly measuring the write–read structure of trained states.

    Source: §1 · pp. 1–2

  • Method

    BASS applies balanced truncation using full empirical controllability and observability Gramians, while HARP allocates layer ranks from Hankel spectra and protects retrieval-sensitive layers.

    Source: §1 · p. 2, §3 · p. 4, §5.2 · p. 5

  • Results

    At 50% retention, HARP-assisted BASS adds 0.3 average WikiText-2 perplexity points across five checkpoints, versus 1.6 for GHOST.

    Source: §1 · p. 2, §8 · p. 8

  • Takeaways & Limitations

    The results support state-level, system-theoretic work on efficient recurrent models, with the performance gap widening at 25% retention.

    Source: §8 · p. 8

    In the unoptimized loop, decoding at batch size 1 is 6–10% slower after compression.

    Source: Limitations · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Mamba-2’s fixed recurrent state dominates decoding memory, but existing compressors score values or native coordinates rather than the joint write–read structure. BASS instead selects a balanced state subspace, while HARP adapts retained ranks across layers.

1 Introduction

  • At batch size 128, Mamba2-2.7B recurrent states account for roughly 80% of decoding memory, making state size a major inference cost.Each layer keeps a matrix state per head, and each generated token must stream the states from memory.Source: §1 · p. 1
  • Existing quantization and coordinate-pruning methods score state values or channels in the native basis, missing joint write–read structure.Even output-aware criteria score coordinates independently from calibration statistics.Source: §1 · p. 1
  • Trained Mamba-2 states are almost rank-one in energy, but their functional subspace is not coordinate-aligned and varies across layers rather than heads.The leading energy direction is bias-like, while the output reads low-energy directions.Source: §1 · p. 2
  • At 50% retention, HARP lowers BASS’s average WikiText-2 perplexity cost to 0.3 points, versus 1.6 for GHOST.The gap widens at 25% retention, where coordinate-space selection collapses.Source: §1 · p. 2
  • BASS uses full empirical controllability and observability Gramians for balanced truncation, applying the rotation exactly through Mamba-2’s scalar-identity transition.It is post-training and does not require retraining.Source: §1 · p. 2
  • HARP allocates ranks per layer from the Hankel spectrum and protects retrieval-sensitive layers, improving WikiText-2 perplexity and LAMBADA accuracy over uniform budgets.The paper examines Mamba-2 checkpoints from 130M to 2.7B.Source: §1 · pp. 2–3

2 Related Work

Prior state-compression methods quantize values or prune native coordinates, while balanced truncation offers a subspace-based alternative. For trained Mamba-2, the paper targets the joint write–read geometry using full empirical Gramians.

2 Related Work

  • Mamba-2’s scalar-times-identity transition is equivalent to masked linear attention, while Mamba-3’s architectural state reduction is distinct from post-training compression of trained checkpoints.The related work also places SSMs in the lineage from S4 through selective Mamba.Source: §2 · p. 3
  • Quantization methods compress state values, while GHOST prunes native channels using diagonal controllability–observability scores; these approaches do not capture joint write–read geometry.Other pruning approaches target low gate activity, blocks, sub-modules, or key/query channels; prior analyses did not use joint write–read geometry for compression.Source: §2 · p. 3
  • Classical balanced truncation retains simultaneously reachable and observable directions ranked by Hankel singular values, but its guarantees do not directly apply to time-varying Mamba-2.BASS instead estimates full empirical Gramians and exploits the scalar-identity transition to apply the transform without altering the recurrence.Source: §2 · p. 3

3 Preliminaries

Mamba-2 uses a recurrent state with input-dependent writes and reads. The section introduces balanced truncation and explains how Mamba-2’s scalar-identity transition allows state-coordinate changes.

3 Preliminaries

  • Each Mamba-2 head maintains a state that receives input-dependent writes and is read through output couplings.The scalar decay is shared across the state elements, and the state is the layer’s context memory.Source: §3 · p. 3
  • Balanced truncation ranks state directions by how strongly inputs reach them and how strongly they affect outputs.For an LTI system, the eigenvalues of the controllability-observability Gramian product are squared Hankel singular values.Source: §3 · pp. 3–4
  • Selective SSM transitions vary over time, so LTI guarantees do not directly apply; empirical Gramians are evaluated end to end.The empirical Gramians are estimated from calibration trajectories as analogues of the LTI Gramians.Source: §3 · p. 4
  • Because Mamba-2’s transition is a scalar multiple of the identity, changing state coordinates commutes with the recurrence when input and output couplings are transformed.This property enables the truncated transform to be applied exactly.Source: §3 · p. 4

4 The Structure of Trained Mamba-2 States

Trained Mamba-2 states have highly concentrated energy, but their useful input-output subspace is not coordinate-aligned. The useful rank varies chiefly across layers, motivating per-layer budgets.

Measurement setup

  • The state analysis captures per-layer, per-head states from 130M, 370M, and 1.3B models on four WikiText-2 training sequences.Energy denotes mean squared state norm; anisotropy denotes mean pairwise cosine between cross-sequence states.Source: §4 · p. 4

Energy concentration

  • 89.6–91.9% of heads have effective rank at most two, and the leading direction carries 89–93% of energy.Median effective rank is 1.16–1.28 across the studied model scales.Source: §4 · p. 4, Figure 2 · p. 4

Functional content and coordinates

  • The leading energy direction aligns with the mean write direction, but individual writes vary across tokens.Median absolute cosine with the layer’s mean write direction is 0.88–0.95; consecutive write cosine is 0.51–0.54.Source: §4 · p. 4
  • Four randomly retained coordinates out of 128 preserve at most 8.5% of output energy; at r=64, even the best native-coordinate subset captures only 0.69–0.72 of balanced-subspace row mass.At r=32 on 130M, the balanced subspace retains 96.6% of empirical Hankel energy versus 70.0% for the strongest diagonal ordering.Source: §4 · p. 4

Layer-wise rank variation

  • Across layers, effective rank ranges up to 12.8 in some early-to-middle layers while many others are effectively rank-one; heads within a layer are similar.This heterogeneity means uniform budgets over-provision many near-rank-one layers, while HARP reallocates the same average budget and improves perplexity and retrieval.Source: §4 · pp. 4–5

5 Method: BASS

BASS estimates empirical write-read Gramians, retains a balanced state subspace, and applies the transform through Mamba-2’s scalar-identity transition. HARP allocates ranks by layer and protects layers sensitive to retrieval.

5.2 Balanced subspace and exact transform

  • BASS projects the state and write/read couplings into the retained basis; the scalar-identity transition makes the rotation commute with the recurrence.At full rank the transform is exact, while truncation approximates; a runtime linear map adds 2Nr multiply–accumulates per token per layer.Source: §5 · p. 5, §5.2 · p. 5

5.1 Calibrating the state Gramians

  • BASS estimates each layer’s empirical Gramians from write and read couplings collected over four 1,024-token calibration sequences.For the selective case, the Gramians are empirical analogues rather than exact infinite-horizon Gramians; estimation uses a single forward pass over 4,096 tokens.Source: §5.1 · p. 5

5.2 Balanced subspace and exact transform

  • Balanced truncation retains directions that are both strongly written and strongly read, ranked by Hankel singular values.At fixed rank, it minimizes discarded Hankel energy; coordinate selection is the special case that keeps only native axes.Source: §5.2 · p. 5

5.3 HARP: Hankel allocation with retrieval protection

  • HARP first allocates ranks by greedily assigning modes to layers with the largest marginal reduction in Hankel-tail energy.The diminishing marginal reductions make this greedy procedure optimal for the discrete allocation problem.Source: §5.3 · p. 6
  • HARP probes retrieval sensitivity by compressing one layer at a time, then keeps the most sensitive layers at the uniform rank before reallocating the remaining budget.The retrieval floor adds no parameters or FLOPs beyond BASS at a matched average rank; the probe uses no gradients.Source: §5.3 · p. 6

6 Experiments

Experiments compare state-compression methods on WikiText-2 and LAMBADA across Mamba-2 checkpoints, testing uniform and layer-adaptive ranks as well as exactness controls. BASS+HARP gives the strongest reported perplexity results, while retrieval allocation and exactness are evaluated separately.

6.1 Setup

  • Four 1024-token WikiText-2 training sequences calibrate the methods, while evaluation uses chunked WikiText-2 perplexity and LAMBADA last-token accuracy.Calibration sequences are disjoint from evaluation data; LAMBADA uses the first 512 examples.Source: §6.1 · p. 6

6.2 Main results

  • At 50% retention, BASS+HARP adds 0.3 WikiText-2 perplexity points on average versus 1.6 for GHOST; its advantage widens at 25% retention.The averages cover five checkpoints; Table 2 reports perplexity at 50% and 25% retention.Source: §6.2 · p. 6, Table 2 · p. 7
  • Magnitude scoring is catastrophic and nonmonotone in budget, while Taylor-2 is numerically explosive at 50% retention.These are reruns of criteria compared in the GHOST paper under a shared protocol.Source: §6.2 · p. 7

6.3 Layer-adaptive budgets and retrieval

  • At 130M and mean rank 64, a retrieval floor restores most LAMBADA accuracy: +2.9% at 0.404 versus +7.2% at 0.371 for uniform allocation.Hankel allocation nearly closes the perplexity gap but loses LAMBADA accuracy; sensitivity-weighted allocation is unstable across probe draws.Source: §6.3 · p. 7

6.4 Exactness controls

  • At full rank, the identity and full-rank balanced transform match baseline perplexity exactly; offline state recomputation matches the reference within 2.2×10^-6 relative error.At reduced rank, only the retained subspace changes.Source: §6.4 · p. 7

7 Discussion

The discussion attributes the advantage of subspace compression to jointly modeling state writes and reads, which coordinate-wise criteria miss. It also considers how state quantization and subspace reduction interact across compression budgets.

Interaction with quantization

  • Table 3 reports LAMBADA last-word accuracy at the Table 2 budgets, with final columns giving average changes versus dense over five checkpoints.HARP selects layers using held-out examples disjoint from evaluation.Source: Table 3 · p. 8

Why balanced truncation wins

  • Coordinate-wise scores ignore cross-channel covariance; rotating into a basis that diagonalizes the input–output map lets selection target the jointly defined write–read subspace.In the LTI limit, the subspace is optimal; empirically, it is best or within 0.2 points of the best control at every budget.Source: §7 · p. 8

Interaction with quantization

  • At 2× compression, 16-bit quantization is close to lossless on smaller models; at 8× and 16×, only subspace reduction plus int8 remains functional.The best 4× mix varies by scale: HARP with int16 leads on 130M/370M, while dense int8 leads on 1.3B.Source: §7 · p. 8

8 Conclusion

The conclusion presents BASS as balanced truncation for Mamba-2 states whose functional content is not coordinate-aligned, with HARP adding layer-adaptive allocation and a retrieval floor. At 50% retention, BASS+HARP has a lower average WikiText-2 perplexity cost than GHOST, and HARP improves LAMBADA accuracy on the probed scales.

8 Conclusion

  • BASS applies balanced truncation through Mamba-2’s scalar-identity transition, while HARP uses per-layer Hankel allocation and a retrieval floor.The conclusion states that the characterization and released toolkit support state-level, system-theoretic work on efficient recurrent models.Source: §8 · p. 8
  • At 50% retention, BASS adds 0.3 WikiText-2 perplexity points on average across five checkpoints, versus 1.6 for GHOST.On the three probed scales, HARP recovers 0.02–0.03 LAMBADA accuracy, and the gap widens at 25% retention.Source: §8 · p. 8

Limitations

BASS’s compression results depend on calibration choices, evaluation coverage is limited to WikiText-2 and LAMBADA on 130M–2.7B checkpoints, and the reported decode trade-off depends on batch size and implementation. The state-capture protocol reconstructs recurrent states offline because fused kernels do not expose them.

Limitations

  • Calibration sources change the 50% cost by up to 1.5 points, while 128-token sequences raise it by up to 1.0/9.9 points at r=64/32.At 25% retention, source variation reaches 7.7 points on 130M.Source: Limitations · p. 8
  • Evaluation covers WikiText-2 and LAMBADA on 130M–2.7B checkpoints, leaving length generalization, out-of-distribution and zero-shot suites, and vision and hybrid models untested.Source: Limitations · p. 8
  • The runtime rotation adds 0.6–0.8% prefill latency; decode gains at large batches contrast with a 6–10% batch-1 loss in the unoptimized loop.State quantization uses simulated per-mode integers rather than production kernels.Source: Limitations · p. 8
  • Fusing the rotation into decode kernels or using recovery fine-tuning are proposed to remove batch-1 overhead; extending BASS to Mamba-1, Mamba-3, and long-context recall remains future work.Source: Limitations · pp. 8–9

A State Capture Protocol

  • Because fused scans do not expose per-step states, the protocol reconstructs them offline from reference equations and verifies replay with a maximum relative error of 2.2×10−6.Calibration uses four 1,024-token WikiText-2 training sequences; diagnostic analysis uses held-out windows and evaluation uses the test split.Source: A State Capture Protocol · p. 10

B Baseline Reproduction Details

The reproduction audit distinguishes released word-level perplexities from the authors’ reruns and treats reproduction as preserving order of magnitude and dense-model ordering. The 25% GHOST-2 collapse persists across tested protocols, draws, and windows.

B Baseline Reproduction Details

  • Released GHOST values are word-level perplexities, whereas reruns use token-level metrics, so the audit reports each metric per row without a Δ.A method counts as reproduced when it remains in the same order of magnitude and preserves its ordering against the dense model.Source: B Baseline Reproduction Details · p. 10, Table 4 · p. 11
  • At 370M/25%, the released pipeline is far worse than the joint re-implementation of the same diagonal criterion, and both remain behind BASS.The residual gap is consistent with sequential versus joint selection.Source: B Baseline Reproduction Details · p. 10
  • The 25% GHOST-2 collapse persists across calibration protocols, calibration draws, and evaluation windows; hard truncation reproduces the soft results with 31–32 of 128 states retained.Source: B Baseline Reproduction Details · p. 10, Table 5 · p. 11
  • The released 1.3B sweep shows GHOST degrading smoothly but severely at high sparsity, while Taylor-2 jumps between mid and high sparsity and Magnitude is unstable from the lowest sparsity.Source: B Baseline Reproduction Details · pp. 10–11, Table 6 · p. 11

C Calibration Sensitivity

Calibration-source and sequence-length changes affect BASS costs, but reported comparisons with GHOST remain favorable. Across the additional analyses, balanced subspaces outperform coordinate orderings, retrieval-aware allocation protects retrieval performance, and state-size reduction benefits larger-batch decoding.

C Calibration Sensitivity

  • Across train and validation calibration sources, 50% perplexity cost varies by at most 1.5 points, while 25% variation reaches 7.7 points on 130M; the weakest 130M/50% result (+8.18%) remains better than GHOST (+12.5%).At 1.3B, sources lie within 0.2 points; ridge shrinkage and doubled calibration data did not further improve results.Source: C Calibration Sensitivity · p. 11, Table 8 · p. 11

D Off-Diagonal Coupling and Retained-Subspace Stability

  • Balanced subspaces retain stable behavior despite tail-basis variability: different calibration draws can yield principal angles near 85° and O(1) tail-operator change, while rank-64 subspaces overlap by 0.76–0.91 mean squared principal cosines.The retained Hankel tail mass agrees to ≤0.004 at rank 64; native coordinates capture only 0.69–0.72 of the balanced subspace’s squared row mass.Source: C Calibration Sensitivity · p. 11, D Off-Diagonal Coupling and… · p. 12

C Calibration Sensitivity

E Layer-Adaptive Budgets and Retrieval

F Computational Cost

  • BASS calibration is 8.8–10.3× faster than GHOST under released protocols and peaks lower; compressed prefill is 1.4×–1.8× faster, with the 1.3B dense scan exceeding the 24 GB card.The runtime rotation adds only 0.6–0.8% to the uncompressed scan’s prefill cost.Source: F Computational Cost · pp. 13–14, Table 13 · p. 13, Table 15 · p. 14

G Decode Measurements

  • At batch 1, 370M/1.3B decoding is 6–10% slower from per-step rotation, but larger batches gain 1.5–2.2× throughput at 130M and 1.5–1.65× at 1.3B as peak memory falls with N/r.These measurements use an unoptimized reference loop.Source: G Decode Measurements · p. 14, Table 16 · p. 14

H Rotation and Single-Sided Controls

J Full-Test-Set Evaluation

On the full WikiText-2 test file, BASS costs agree with the main evaluation at r=64 and are slightly smaller at r=32. At matched budgets, quantization performs best at mild compression, while subspace reduction extends usable bit-widths at aggressive compression.

J Full-Test-Set Evaluation

K State Quantization at Matched Budgets

Related papers

Loading 2610.36617v1…