GHOST: Unmasking Phantom States in Mamba2 via Grouped Hidden-state Output-aware Selection & Truncation
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://github.com/Menezmic21/mamba2_ghost.
Lay Summary
Powerful AI models like Mamba2 suffer from a memory bottleneck that slows them down, and current methods to shrink this memory either accidentally delete important information or cost too much computing power to run. We introduce GHOST, a fast tool that solves this by monitoring the AI in action to safely trim away useless "dead weight" while carefully preserving the memory channels that actively process data and shape the model's final answers. Using GHOST, we successfully cut the internal memory size of large Mamba2 models by 50% with practically no loss in accuracy, making these advanced AIs significantly faster, cheaper, and easier to use in the real world.