An exact instrument for measuring how a trained selective state-space model (Mamba-1, Falcon-Mamba, Mamba-2) uses its per-channel state, and the input-driven mode migration it reveals.
Because a selective SSM layer has a diagonal, fixed state matrix, each channel's output is an exact sum of per-mode contributions. Accumulating the outer products of those contributions over a window gives a per-(layer, channel, window) Gram tensor, from which the exact output error of pruning any subset of modes follows in closed form, offline, at any budget. This repository is the code behind that instrument: the audit, the mechanism counterfactuals, the end-to-end pruning experiments, and the baselines.
audit/
audit_run.py exact per-mode decomposition (modal_scan) + validation gate
gram_utils.py Gram packing and nested tail sums
prune_forward.py deployed pruned forward path + instrument cross-check + e2e
cross_sweep.py swept deployment cross-check (independent reference side)
e2e_baselines.py static / modal-HSV / LAST baselines (empirical Gramians)
bootstrap_ci.py paired and block bootstrap confidence intervals
mamba2_check.py depth-localized migration in Mamba-2 (SSD)
mamba2_prune.py Mamba-2 end-to-end static vs scheduled pruning
mamba2_mechanism.py Mamba-2 frozen-signal mechanism
normmatch_freezeB.py energy-matched freeze-B control
churn_gap.py churn-vs-gap relation
churn_gap_theory.py energy-weighted churn vs migration gap
timescales.py pole-timescale structure of migration
ltv_gramians.py time-varying Hankel singular values sigma_i(t)
probe.py, scheduler.py, learned_scheduler.py, half_transfer.py
switched_signal.py switched-system control with ground-truth regimes
downstream.py downstream tasks under mode pruning (lm-eval-harness)
tests/ synthetic-mixer unit tests (no downloads)
ghost_comparison/ native comparison against the released GHOST code
pip install -r requirements.txt
The audit runs the reference (fp32, non-fused) SSM path so the exactness gate
is meaningful; on machines with the fused Mamba/causal-conv1d kernels
installed, the scripts disable the fast path where required.
The pipeline consumes public released checkpoints (Mamba-1 130M-2.8B, Falcon-Mamba-7B, Mamba-2 130M/780M). A typical run:
python audit/audit_run.py --model 130m # per-mode Gram audit + gate
python audit/prune_forward.py --cross-check # instrument == deployment
python audit/cross_sweep.py --model 130m # swept validation
python audit/e2e_baselines.py --model 130m --v2 # baselines at equal budget
Each script prints its validation gate first and writes summary JSON; the
tables and figures in the paper are produced from those summaries. The
ghost_comparison/ directory documents how to run the released GHOST code
head-to-head, including the environment pins it requires.
pytest audit/tests
If you use this code, please cite the accompanying paper (see the repository description for the current reference).