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

1. Problem Statement

The transition from Mamba1 to Mamba2 significantly increased the State Space Model (SSM) state dimension ($N$) from 16 to 128 to enhance temporal modeling. While beneficial for performance, this 8× expansion creates a severe inference bottleneck:

• Memory Bandwidth Saturation: The recurrent state size swells (e.g., from ~12 MB to ~100 MB for a 1.3B model), saturating memory bandwidth during autoregressive generation.
• Failure of Existing Pruning Methods:
  • Unstructured Sparsity: Methods like SparseGPT or Wanda sparsify weights but leave activations dense, failing to reduce memory bandwidth.
  • Static Magnitude Pruning: Relying on fixed weight norms ignores runtime dynamics. The paper identifies two critical failure modes:
    • Corporeal States: Channels with large weights but low actual utilization (inert).
    • Phantom States: Channels with low weight norms but high dynamic activity (critical). Standard pruning mistakenly removes these.
  • Gradient-Based Methods: Techniques like Taylor pruning offer high fidelity but require computing Hessians/gradients for the full graph, leading to prohibitive VRAM usage (e.g., >45 GB for a 1.3B model) and distribution shift issues.

2. Methodology: GHOST

The authors propose GHOST (Grouped Hidden-state Output-aware Selection and Truncation), a structured pruning framework that approximates Balanced Truncation (a control-theoretic concept) using only forward-pass statistics.

Core Concepts

GHOST identifies "living" states by jointly measuring:

1. Controllability: How effectively inputs drive a specific state channel.
2. Observability: How strongly a state channel influences future outputs.

Technical Implementation

• Empirical Estimation: Instead of solving Lyapunov equations (which require backward passes), GHOST uses instantaneous empirical Gramians derived from forward-pass data:
  • Controllability Metric ($P$): Estimated via the Hidden State Covariance (squared magnitude of the hidden state $H_t$).
  • Observability Metric ($Q$): Estimated via the Output Energy Hessian, approximated locally as the squared magnitude of the projection vector $C'_t$ (since $H_t$ contributes to output $y$ via $C'_t$).
• Saliency Scoring: The importance of a state channel $i$ is defined as the product of these metrics: $S_i \approx \sqrt{P_i \cdot Q_i}$. This mirrors the Hankel singular values used in balanced truncation.
• Grouped Processing: GHOST operates within Mamba2's Grouped Query Attention (GQA) structure. It pools saliency scores across groups within a layer, allowing for adaptive capacity allocation. Complex dynamic groups retain more states, while redundant groups are aggressively pruned.
• Sequential Pruning: The method processes layers sequentially. After pruning layer $j$, the activations are updated and used to calibrate layer $j+1$, mitigating distribution shift without requiring backpropagation.

Algorithm Complexity

• Time: Requires only two forward passes over a calibration dataset. Complexity is linear with respect to the state dimension ($O(N)$), matching standard inference costs.
• Space: Requires $O(G \cdot N)$ storage for statistics, avoiding the $O(N^2)$ storage of full covariance methods or the massive VRAM overhead of gradient-based methods.

3. Key Contributions

1. Novel Framework: Introduction of GHOST, the first structured pruning method for Mamba2 that approximates balanced truncation without gradients.
2. Solving "Proxy Failure": GHOST successfully distinguishes between "Corporeal" (high weight, low use) and "Phantom" (low weight, high use) states, preventing the catastrophic removal of critical dynamic channels.
3. Efficiency: Achieves high-fidelity pruning with zero gradient computation and a fraction of the memory overhead of second-order methods.
4. Empirical Validation: Demonstrates robustness across model scales (130M to 2.7B), sequence lengths, and out-of-distribution tasks.

4. Experimental Results

Experiments were conducted on Mamba2 models (130M to 2.7B parameters) using WikiText-2 for calibration and various downstream tasks.

• Sparsity vs. Perplexity:
  • At 50% state-dimension reduction, GHOST achieves only a ~1 perplexity point increase on WikiText-2 (14.23 vs. 13.17 baseline).
  • At 70% sparsity, GHOST maintains a perplexity of 16.16, whereas Taylor pruning collapses (perplexity > 4000) and Magnitude pruning fails significantly.
• Sequence Length Generalization:
  • GHOST generalizes well to longer contexts (up to 2048 tokens) when calibrated on short contexts.
  • In contrast, Taylor pruning shows a "phase transition" failure, with perplexity exploding as sequence length increases, indicating that short-context gradients do not transfer to long-range dependencies.
• Scaling Laws:
  • GHOST remains stable across all model sizes, including small models (130M) where gradient-based methods (Taylor) fail catastrophically due to noisy gradient signals.
• Zero-Shot & OOD Robustness:
  • GHOST outperforms baselines on Out-of-Distribution (OOD) tasks (HumanEval code generation, MMLU math), suggesting it captures fundamental system dynamics rather than overfitting to English text patterns.
  • On PIQA, GHOST achieves 73.12% accuracy, nearly matching the dense baseline and outperforming Taylor pruning.

5. Significance

• Democratizing Mamba2: By reducing the recurrent state dimension by 50% with minimal accuracy loss, GHOST significantly lowers memory bandwidth requirements, making large-scale SSMs deployable on consumer-grade hardware (e.g., fitting a 1.3B model's state into standard 40GB A100 VRAM).
• Bridging Control Theory and Deep Learning: The work successfully adapts classical control-theoretic concepts (balanced truncation) to modern deep learning architectures using only forward-pass statistics, offering a new paradigm for model compression.
• Practical Viability: Unlike gradient-based methods that are often too expensive for practical use, GHOST offers a computationally cheap, scalable, and robust solution for post-training compression of SSMs.

In conclusion, GHOST provides a critical solution to the inference bottleneck of Mamba2 by intelligently identifying and pruning redundant state dynamics while preserving the "phantom" states essential for temporal modeling, all without the computational cost of backpropagation.