Mamba-2 compression preserves jointly useful state directions
Balanced State Subspaces for Mamba-2 State Compression · arXiv:2610.36617
In plain words
Mamba-2 language models keep information in a changing internal state as they read text. Earlier compression methods choose individual parts of that state, but this paper argues that useful information can depend on how several parts work together.
Its BASS method finds and keeps a smaller set of directions that are both written to and read from. Across five Mamba-2 checkpoints, at half state size, it caused a smaller average increase in WikiText-2 perplexity than the GHOST compressor; results varied by checkpoint.
The research
A recurrent state can be smaller only if it keeps information that still shapes the model’s outputs.
The paper asks whether choosing individual state coordinates overlooks useful combinations, and whether a compressed model can preserve language prediction and retrieval quality by retaining directions that work together.
BASS compresses trained Mamba-2 states by retaining directions that are jointly written and read, while HARP allocates those directions across layers.
The paper asks whether trained Mamba-2 states can be compressed without discarding useful information that depends on several state coordinates together. BASS selects a subspace using how strongly directions are written and read, while HARP distributes the retained modes across layers.
Across five checkpoints, at half state retention, BASS with HARP had a smaller average WikiText-2 perplexity increase over GHOST.
In a matched average-rank test on Mamba2-130M, HARP’s retrieval floor also improved perplexity and LAMBADA last-word accuracy over uniform allocation, suggesting that layer allocation affects both measures.
The results cover the tested Mamba-2 checkpoints and language-modeling evaluations. The paper does not test length generalization, out-of-distribution or zero-shot tasks, vision models, or hybrid models.
What the evidence shows
- The compression question
BASS is a post-training method that compresses the recurrent state in trained Mamba-2 models. The paper asks whether compression can preserve how state directions work together, since coordinate-based methods score individual channels even when those scores use output information.
The authors argue that such scores miss joint write-and-read structure, which shapes how the state affects outputs.
- Keep joint state directions
BASS estimates how strongly inputs write into state directions and how strongly those directions are read into outputs. It uses both measurements to select a smaller subspace that can combine native coordinates.
The full-rank transform preserves the recurrence exactly; keeping fewer modes approximates it. HARP then allocates the retained modes across layers using their spectra and retrieval sensitivity.
- A structural comparison
A diagnostic compared rank-32 selections in Mamba2-130M. Across layers, the balanced subspace retained a median 96.6% of empirical Hankel energy, versus 70.0% for the strongest marginal-diagonal coordinate ordering.
This measures preservation of the joint write-and-read structure, not language-model quality. It provides evidence for why a mixed subspace may be more suitable than selecting coordinates individually.
- Quality at half retention
On WikiText-2, a text-prediction test where perplexity measures prediction difficulty, BASS with HARP increased perplexity by an average of 0.3 points across five checkpoints at half state retention, compared with 1.6 points for GHOST.
In a separate Mamba2-130M comparison at the same average rank, HARP’s retrieval floor improved both perplexity and last-word accuracy over uniform BASS allocation on the reported evaluations.
- Basis choice and boundaries
A control at 25% retention on Mamba2-130M compared BASS with a random orthogonal basis at the same rank. BASS raised WikiText-2 perplexity by 25.4%, while the random-basis run diverged.
This supports the interpretation that basis choice matters beyond the number of retained modes, within that tested model and retention level. Other task families and model types remain untested.
Deep
- Original paper authors
- Shang Tingyu, Cao Wenfei
- Submitted to arXiv
- License
- arXiv non-exclusive