GHOST: Unmasking Phantom States in Mamba2 via Grouped Hidden-state Output-aware Selection & Truncation

AI summary

Overview

Research area: Efficient inference and structured pruning for state-space models (SSMs), specifically the Mamba2 architecture — combining ideas from control theory (balanced truncation) with neural network compression.

Technical level: Intermediate. The paper assumes familiarity with SSM recurrences, grouped-query-style parameter sharing, and standard pruning terminology, but its core idea is explained from first principles.

Scope: The paper introduces GHOST, a gradient-free structured pruning framework that removes hidden-state channels from Mamba2 by scoring each channel's controllability and observability from forward-pass statistics alone.

What This Paper Is About

Mamba2 expanded the SSM state dimension from Mamba1's 16 to 128, which for a 1.3B model grows the recurrent state from roughly 12 MB to roughly 100 MB — an 8x expansion that saturates memory bandwidth and hurts cache locality during autoregressive generation. Existing pruning methods do not fix this: unstructured sparsity leaves activations dense, magnitude pruning cannot distinguish a state's static weight size from its actual runtime activity, and gradient-based structured pruning is memory-prohibitive and brittle. GHOST's goal is to compress the state dimension directly, using only forward passes, while matching the fidelity of expensive gradient-based methods.

Key Contributions

1. GHOST framework: A data-driven structured pruning method for Mamba2 that prunes the SSM state dimension N using inter-group thresholding, scoring each state channel by measured controllability (hidden-state covariance) multiplied by observability (a local output-energy Hessian).

2. Diagnosis of "proxy failure": The paper formally characterises two failure modes of magnitude pruning — corporeal states (structurally large weights but dynamically inert) and phantom states (small weights but high runtime activity) — and shows GHOST prunes the former while retaining the latter.

3. Empirical gains at zero gradient cost: GHOST outperforms Magnitude and Random baselines and rivals Taylor pruning while requiring no gradient computation and a fraction of the memory.

4. Robustness evidence across axes: Experiments covering sparsity levels, sequence-length generalisation and robustness, model scale from 130M to 2.7B parameters, zero-shot downstream tasks, out-of-distribution tasks, and calibration-data efficiency.

Main Findings

• Proxy failure is real and measurable. For Mamba2-1.3B, the paper analyses 6,144 hidden states (48 layers with 128 states each) and finds no positive correlation between a state's static score derived from projection weights and its dynamic energy. Phantom states sit top-left (highly active, low weight norms) and corporeal states bottom-right (large weights, low utilisation).

• 50% compression at roughly 1 perplexity point. The headline result: across models from 130M to 2.7B parameters, GHOST achieves 50% state-dimension reduction with approximately 1 perplexity point increase on WikiText-2.

• Sparsity sweep (Table 1, WikiText-2 perplexity, dense baseline 13.17). At 50% sparsity GHOST reaches 14.23, versus 13.94 for Taylor, 13.25 for SparseGPT, 29.34 for Magnitude and 17.77 for Random. At 70% sparsity GHOST yields 16.16 while Taylor collapses to 4255. At 90%, GHOST reaches 25.07, Taylor 4690, Magnitude 69.49 and Random 64.21. SparseGPT stays strongest throughout (15.51 at 90%).

• Length generalisation favours GHOST. Calibrated on L_cal = 128 and evaluated up to L_eval = 2048, GHOST's perplexity falls from 23.99 to 14.27 as context grows, tracking the dense trend (23.04 to 13.18). Taylor goes the other direction, rising to 1613 at L_eval = 2048 and reaching 317.3 even at the calibration length.

• Length robustness anomaly. GHOST tracks SparseGPT and Dense consistently across the sweep; Taylor shows a sharp phase transition in the L ∈ [16, 512] regime, underperforming even Magnitude and Random, before converging back outside that window.

• Small models favour GHOST, larger models favour Taylor. On 130M and 370M models, Taylor explodes (1E16 and 12647 perplexity) while GHOST gives 29.06 and 19.96 against dense values of 25.85 and 18.13. GHOST is the only method the paper reports as providing consistent, usable post-pruning performance across the whole 130M–2.7B spectrum. Taylor recovers at 1.3B (13.94) and 2.7B (11.50) versus GHOST's 14.23 and 12.14.

• Zero-shot competitiveness. GHOST outperforms Taylor on PIQA at 73.12% accuracy, within 0.22% of the dense baseline. Taylor leads among compressed methods on Lambada (59.23%), ARC-e (63.85%) and ARC-c (29.69%). Magnitude and Random collapse on Lambada to 17.50% and 31.21% respectively.

• Out-of-distribution robustness. Calibrated only on WikiText-2, GHOST matches SparseGPT on HumanEval code (5.49) and on MMLU Elementary Mathematics (25.66), ahead of Taylor (4.88 code, 22.49 math). The paper attributes this to identifying channels fundamental to SSM operation rather than channels useful for English next-word prediction.

• Low calibration cost. All seven methods show equivalent sample efficiency, with GHOST moving only from 14.28 at 8 calibration samples to 14.23 at 128.

• Joint saliency beats either component alone. Ablating the score, controllability-only (P) gives 13.26/13.64/14.56/16.71/2158 and observability-only (Q) gives 13.37/14.23/16.16/19.83/33.12 from 10% to 90% sparsity, versus GHOST's 13.24/13.51/14.23/16.16/25.07. The joint measure is reported as consistently best and most stable.

• Cost profile. GHOST's calibration on Mamba2-1.3B without custom kernels takes 8 minutes (2 forward passes over 128 samples), compared to under 1 minute for Magnitude/Random and 8 minutes for Taylor, which additionally requires more than 45 GB VRAM. For a 1.3B model, Taylor's gradient requirement is 45 GB VRAM versus GHOST's 15 GB.

Methodology in Plain English

GHOST starts from a control-theory idea called balanced truncation: in a dynamical system, the most valuable internal states are those that are both easy to drive with input (controllable) and influential on output (observable). Mamba2 has input-dependent dynamics terms and a non-linear SiLU activation, so the classical algebraic machinery for balanced truncation does not apply. GHOST therefore estimates these two quantities empirically, purely from forward passes.

Controllability is measured as the hidden-state covariance — the squared magnitude of each state channel's activations, since the hidden state is the accumulation of processed inputs.

Observability is measured as the Hessian of the output energy with respect to the hidden state. Because the immediate output is a linear projection of the state through C, this Hessian works out to C-transpose times C. The authors deliberately use only the instantaneous term rather than summing over all future timesteps, justified by the exponential decay in Mamba2's discretised transition dynamics.

Saliency for a channel is the product of the two terms, then aggregated over the sequence length, the calibration set, the K heads in a group and the P head channels, and finally square-rooted. The square root aligns the metric with the definition of Hankel singular values and restores it to the linear scale of input features.

Pruning is inter-group: scores from all G groups in a layer are pooled, sorted, and a global threshold is chosen for a target sparsity. This lets the model allocate state capacity adaptively — groups modelling complex dynamics keep more channels (r > N(1−κ)) while redundant groups are pruned harder. Pruning is soft: projection weights and convolution filters for the selected channels are zeroed, and layers are processed sequentially so each layer calibrates against the already-sparsified activations of the previous one. The paper also shows the score minimises the expected increase in local mean-squared reconstruction error, connecting the control-theoretic derivation to a loss-preservation argument.

The whole procedure needs two forward passes per layer. No gradients, no Hessians stored, no quadratic memory. Time complexity per layer is O(|D_cal| · L · G · K · P · N) — the cost of a standard inference pass — and space complexity is O(G · N), avoiding the O(N²) storage of full-covariance or Hessian-based methods.

Experimental setup: random seeds fixed to 42; Mamba2-1.3B with HuggingFace default precision (float32) and sequence length (2048); calibration on 128 WikiText-2 samples at batch size 1; compression to 50% sparsity by default; evaluation via EleutherAI's LM Evaluation Harness; all runs on a single H100 80 GB GPU.

Why This Matters

Impact on research. The paper reframes SSM compression as a system-identification problem rather than a weight-importance problem, giving a concrete, cheap approximation of balanced truncation for a non-linear, time-varying architecture. It also documents a clear failure mode of one-shot gradient-based structured pruning — catastrophic collapse under distribution shift and on smaller models — which is a useful negative result for the wider compression community.

Real-world applications (note: these are implications of the paper's motivation; the paper itself does not run deployment benchmarks):

• Long-context document or transcript processing, where the recurrent state's memory-bandwidth cost currently dominates autoregressive decoding.
• On-device or edge deployment, since GHOST needs only 15 GB rather than 45 GB for calibration on a 1.3B model and produces a smaller recurrent state.
• Streaming and real-time inference pipelines that are bandwidth-bound rather than FLOP-bound.
• Domain-specific assistants trained on non-English or code-heavy data, where the paper's OOD results suggest GHOST preserves task-relevant dynamics better than curvature-based pruning calibrated on English prose.

Industry relevance. Memory bandwidth, not arithmetic, is the binding constraint in autoregressive generation of large SSMs. A method that shrinks the recurrent state with about one perplexity point of damage at 50% sparsity, and that beats gradient-based pruning in the small-model regime where edge deployment actually happens, is directly relevant to teams serving SSM backbones. The paper is transparent that this is quality preservation, not yet wall-clock speedup: converting soft pruning into real latency gains requires sparse kernel support or physically removing channels and adjusting downstream layers, which the authors leave to future work and defer to prior empirical evidence.

Future Directions

1. Physical speedup. The paper explicitly leaves open converting soft pruning into actual runtime gains, either through sparse kernel support or by physically removing pruned channels and adjusting downstream layers.

2. Architecture transfer. GHOST is designed around Mamba2's scalar-identity A matrix and grouped dynamics. Extending the Gramian approximations to other SSM variants (S4, H3, Mamba1) is flagged as open; the authors argue the core principle should transfer but the instantiation depends on architectural details.

3. Extreme sparsity and small-model regimes. Perplexity degrades to 25 beyond 90% sparsity, and the 130M model suffers a sharper increase (29.06 versus 25.85 dense) than larger ones due to lower representational redundancy. Making the method viable at these extremes is unresolved.

4. Kernel-level integration. The paper notes GHOST adds modest overhead by extracting hidden states during the forward pass, and that benchmarked timings exclude custom kernels — an obvious next step is fusing the scoring pass into the model's existing inference kernels.

Target Audience

Researchers and engineers working on efficient inference for state-space models and sequence architectures; practitioners doing post-training compression of Mamba or hybrid SSM models who need a structured, state-dimension-level method rather than weight sparsification; and control-theory-adjacent readers interested in seeing balanced truncation approximated for non-linear, time-varying neural systems. Readers should already be comfortable with SSM recurrences and standard pruning baselines (magnitude, SparseGPT, Taylor) to get the most from the comparison tables.

Authors’ abstract

While Mamba2's expanded state dimension enhances temporal modeling, it incurs substantial inference overhead that saturates bandwidth during autoregressive generation. Standard pruning methods fail to address this bottleneck: unstructured sparsity leaves activations dense, magnitude-based selection ignores runtime dynamics, and gradient-based methods impose prohibitive costs. We introduce GHOST (Grouped Hidden-state Output-aware Selection and Truncation), a structured pruning framework that approximates control-theoretic balanced truncation using only forward-pass statistics. By jointly measuring controllability and observability, GHOST rivals the fidelity of gradient-based methods without requiring backpropagation. As a highlight, on models ranging from 130M to 2.7B parameters, our approach achieves a 50\% state-dimension reduction with approximately 1 perplexity point increase on WikiText-2. Code is available at https://anonymous.4open.science/r/mamba2_ghost-7BCB/.