ssm-state-quantization

A small learning project on how rounding the recurrent state of Mamba-2 models affects accuracy, and why the damage depends on how fast each head forgets. It is a set of experiments with logs. It is not a paper or a library, and nothing here is a tested quantization method.

Summary. Rounding Mamba-2's recurrent state to 8 bits is nearly free with an asymmetric grid (+0.2 to +0.8 perplexity on 130M) but costly with a symmetric one, and at 4 bits both fail. The damage concentrates in heads with mid-range decay (0.9 < lambda <= 0.99), not the slowest-forgetting ones, and it depends on the direction of the rounding error rather than its size.

The closest published work is Preserving Updates in Low-Bit Recurrent State Quantization (Li, Dong, Xiao, Xing, 2026). It explains state rounding damage through updates that rounding loses or reverses, and proposes fixes (LRB, BDR). Several results here confirm that paper in a different setting rather than add to it; each finding says which.

Setup

• Models: AntonV/mamba2-130m-hf and AntonV/mamba2-370m-hf (Hugging Face ports of state-spaces Mamba-2), fp32.
• Data: WikiText-2 test set, non-overlapping windows of 2048 tokens. 6 windows unless noted; newer runs use 24 windows and report paired 95% bootstrap CIs over windows (experiments/analyze_paired.py).
• Hardware: one RTX 3070 8 GB. A sequential per-token scan (src/quant/seq_scan.py) replaces the library scan so the state can be rounded at every step. It matches the library scan to a max logit difference of 1.1e-3.
• Each head decays its state by lambda = exp(dt * A) per step. Heads are grouped by lambda: at most 0.9, above 0.9 up to 0.99 ("mid band"), and above 0.99.
• Quantizer: one scale per row of the state (per head and head_dim, over the 128 state entries), refit every step. Two grids:
  • symmetric (default): [-amax, amax], step amax / (2^(b-1) - 1).
  • asymmetric (--asym): [min, max] with an integer zero point, step (max - min) / (2^b - 1). This matches the INT4 baseline of arXiv 2610.15709 (App. A.1), except that we read the output from the quantized state and they read it before quantizing.

Findings

Perplexity is WikiText-2 on 130M (control 20.32 on 6 windows, 21.57 on 24 windows) unless noted. Full tables, pre-registered predictions and caveats are in experiments/LOG.md.

1. The grid decides how much 8 bits hurt (expK). With the symmetric grid, 8-bit nearest rounding (RTN) raises perplexity to 38.7 (all heads) and 38.3 (mid band only). With the asymmetric grid it costs +0.2 (mid band, 20.53) and +0.8 (all heads, 21.12). The large 8-bit damage reported in earlier versions of this README is mostly a property of the symmetric grid. At 4 bits the asymmetric grid fails too: rounding only the mid band gives 267.6 (all heads: 296.2), so at 4 bits nearly all of the damage comes from the mid band alone. Symmetric damage is not monotone in bits: mid band +18.0 at 8 bits but +5.6 at 4 bits (25.9); all heads 38.7, 114.5 and 73.9 at 8, 6 and 4 bits. The asymmetric grid is monotone.
2. Damage sits in the mid band (0.9 < lambda <= 0.99), not the slowest heads. Symmetric 8-bit, mid band only: +14.0 ppl [+9.8, +19.3] on 24 windows; 30.4 against 14.3 on 370M. The lambda > 0.99 group costs +1.1 (130M) and +0.5 (370M); the fastest group costs nothing. Every group is still critical: zeroing any one group's output sends perplexity above 10^4.
3. Error size is not the cause; direction is. Keeping RTN's error magnitudes with random signs gives 20.4 (130M) and 14.29 (370M); iid Gaussian noise of the same per-row RMS gives 20.4. Over seeds 0-4 on 24 windows these random modes cost +0.07 to +0.15 (expJ). Already shown differently in 2610.15709, whose fix enlarges per-step error and improves the model.
4. Reversals are necessary but not sufficient (expL, symmetric 8-bit, mid band). A reversal is a step where rounding moves a state entry against the new input (2610.15709). Correcting only reversed entries leaves 1 to 2% of the damage (no_reversal +0.35 [+0.23, +0.49]; BDR +0.22 [+0.14, +0.31]). Injecting only the reversal error into an exact recurrence gives +0.15 [+0.07, +0.24], under 1%. The necessity half confirms 2610.15709 at 8 bits; the sufficiency failure is new here and not yet explained.
5. Reversal rate does not place the damage; reversal error seen by the readout is a candidate correlate (expL diag, exploratory, all heads). Reversal rate by band is 4.8% (lambda <= 0.9), 19.5% (mid), 37.7% (lambda > 0.99): the slowest heads reverse most yet are barely hurt. The reversal error as seen by the readout, relative to the readout signal, is 0.0024, 0.0048 and 0.0024: twice as large in the damaged band. These diag numbers come from quantizing all heads, while the damage figures come from quantizing one band at a time. It does not separate the other two bands, which tie at 0.0024 yet differ in damage (about 0 against +1 in expE and expE2). The total readout-visible error (all of it, not only reversals) is highest in the harmless fast band (0.0079). One model, one diagnostic pass, no confidence interval; not a causal test.
6. Removing a small part of the error makes things worse (24 windows). Removing the part of the RTN error along the state gives +4.8 [+2.5, +7.5] over plain RTN; removing the part along the current readout C_t gives +13.4 [+7.9, +21.3]. Keeping only either part is nearly harmless (+0.003, +0.06). Unexplained; one candidate is that removal moves the state off the grid and changes which entries reverse.
7. Error energy does not predict damage. The lambda > 0.99 heads have the largest RTN output error but cost little perplexity. RTN errors are weakly correlated across time (lag-1 0.02 in the mid band), strongest in the lambda > 0.99 group. These are per-element lag correlations at 8 bits, not the summed coherent gain of 2610.15709, so the two are not directly comparable.
8. Error growth with memory length. With a per-head symmetric scale, relative state error grows like 1/sqrt(1 - lambda^2) at 8 bits for lambda < 0.99 (measured over predicted 0.95). It fails at 4 bits, where errors saturate.
9. Stochastic rounding (SR) fixes 8 bits and fails at 4 bits. SR stays within 0.4 ppl of control at 8 bits under both grids. At 4 bits it is far worse than RTN: symmetric all heads 538.7 against 73.9; asymmetric mid band 2.1e10 against 267.6. 2610.15709 does not test SR on the state.

Open questions

• Why reversals need the full RTN loop to do damage (finding 4), and whether the readout-visible reversal error (finding 5) predicts damage per head better than lambda does. The per-head test of this (expM) is pre-registered in experiments/LOG.md; results are not logged yet.
• Why the two grids rank differently at 8 and 4 bits: symmetric is far worse at 8 bits, asymmetric far worse at 4 bits, and symmetric damage peaks at 6 bits. Step size alone cannot explain both directions. Symmetric 9-bit (same step as asymmetric 8-bit) and a 5/6/7-bit sweep of both grids would separate step size from grid shape.
• Finding 6.

Limits

• Two small models (130M, 370M), one dataset (WikiText-2), up to 24 windows of 2048 tokens. The 2610.15709 results are on 1.3B models at up to 8K tokens.
• CIs exist only for the 24-window rebaseline (expJ) and the newer 6-window runs (expK, expL); older results are single numbers.
• Random-mode results use one seed except expJ (seeds 0-4) and the SR runs in expE (two seeds).
• The 370M group thresholds reuse a lambda file computed once per model. 370M has not been rerun with the asymmetric grid or the reversal tests.
• Prior-work search (see docs/gap_statement.md) is a few search passes, and some entries were read only as snippets. A NeurIPS 2026 paper on a scale-state feedback in recurrent-state quantization ("When Dynamic Scaling Fails") may overlap with finding 4 and has not been read.

Reproduce

python -m venv .venv
.venv\Scripts\pip install -r requirements.txt
.venv\Scripts\python experiments\expG_noise_structure.py --tag repro --windows 24 --seeds 0 1 2 3 4
.venv\Scripts\python experiments\analyze_paired.py repro

requirements.txt pins the CUDA 12.8 PyTorch build. Each script writes progress to experiments/results/ and skips runs already recorded, so an interrupted run resumes. Weights and data download from Hugging Face on first use. On Linux or macOS, use .venv/bin/ in place of .venv\Scripts\.

Layout

• src/quant/ has the sequential scan and the state quantizer (symmetric and asymmetric grids).
• src/harness/ has the perplexity and LAMBADA evaluations.
• experiments/ has one script per experiment, analyze_paired.py, the results, and LOG.md with pre-registrations and results.
• docs/explainer.md explains the formulas at a technical and a plain level. docs/gap_statement.md covers prior work. docs/plan_open_questions.md is the working plan.
• paper/notes.md is an older summary written before expG; the README supersedes it.

License

Code and text in this repository: MIT. See LICENSE.

Third-party material is not redistributed here; scripts download it at run time under its own license:

• Model weights: AntonV/mamba2-130m-hf and AntonV/mamba2-370m-hf (MIT, per their Hugging Face cards), ports of state-spaces/mamba2-130m and state-spaces/mamba2-370m (Apache-2.0). experiments/results/lambda_*.npy hold per-head decay statistics computed from these models.
• Data: WikiText-2 via Salesforce/wikitext (CC BY-SA 4.0). Result files hold only aggregate numbers, no dataset text.
• arXiv 2610.15709 (CC BY-NC-ND 4.0) is cited and paraphrased in experiments/LOG.md and docs/; no text or figures from it are reproduced.