An Exact Instrument for State Usage in Selective State-Space Models, and the Input-Driven Migration It Reveals

Noosaga Placements

• Machine Learningsubfield90%
  The paper develops an exact instrument for measuring state usage in selective state-space models, which is a core machine learning problem of model analysis and pruning. It introduces a novel method for quantifying per-mode importance and demonstrates its effectiveness across multiple model scales.
  We give an exact instrument for measuring how a trained model uses these modes.Because the state matrix is diagonal, each channel's output decomposes exactly into per-mode contributions, and a per-(layer, channel, window) Gram tensor yields the exact output error of dropping any subset of modes, offline, at any budget.Input-scheduled mode pruning on this measurement outperforms static, Hankel-based, and layer-adaptive rankings at every scale from 130M to the deployed 7B Falcon-Mamba, and at half the state budget it matches the unpruned model.
• Deep Learningframework90%
  The paper is situated within the deep learning framework, specifically focusing on selective state-space models (Mamba, Falcon-Mamba, Mamba-2) which are deep learning architectures. The analysis and pruning methods are applied to these models.
  Selective state-space models such as Mamba route information through a bank of first-order modes whose input coupling is set by a learned selection mechanism.Applying the instrument across the Mamba-1 family (130M–2.8B), the deployed 7B Falcon-Mamba, and Mamba-2Input-scheduled mode pruning on this measurement outperforms static, Hankel-based, and layer-adaptive rankings at every scale from 130M to the deployed 7B Falcon-Mamba
• Deep Learningsubfield80%
  The paper focuses on selective state-space models (Mamba, Falcon-Mamba, Mamba-2), which are deep learning architectures. The analysis of mode migration and pruning is directly relevant to understanding and optimizing these deep learning models.
  Selective state-space models such as Mamba route information through a bank of first-order modes whose input coupling is set by a learned selection mechanism.Applying the instrument across the Mamba-1 family (130M–2.8B), the deployed 7B Falcon-Mamba, and Mamba-2, we find that trained models re-allocate their state space with the inputInput-scheduled mode pruning on this measurement outperforms static, Hankel-based, and layer-adaptive rankings at every scale from 130M to the deployed 7B Falcon-Mamba
• Interpretabilityframework70%
  The paper's instrument provides a direct measurement of how trained models use their internal modes, which is an interpretability tool. It reveals input-driven mode migration, offering insights into the internal workings of SSMs.
  For interpretability, the modes are the layer's internal degrees of freedom, and their usage is a direct measure of what the layer stores.We give an exact instrument for measuring how a trained model uses these modes.Applying the instrument across the Mamba-1 family (130M–2.8B), the deployed 7B Falcon-Mamba, and Mamba-2, we find that trained models re-allocate their state space with the input
• Natural Language Processingsubfield50%
  The paper evaluates its method on language models (Mamba, Falcon-Mamba) using text datasets (English prose, code, technical text) and perplexity benchmarks. However, the core contribution is not NLP-specific but applies to SSMs in general.
  Windows are 1024 tokens, 32 per domain from three domains (English prose, code, technical text), with a disjoint evaluation set of 24 windows per domainOn downstream tasks that use long-range state it recovers most of static pruning's loss (Section 4.3).on LAMBADA, static pruning to r = 8 collapses 790M accuracy from 0.63 to 0.41 and the two-pass oracle mask (read from the scored passage) recovers 86% of the drop
• Statistical Learningframework50%
  The paper uses statistical concepts like Gram matrices and energy-based measures to quantify mode importance and pruning error. The analysis of mode migration and the development of input-scheduled pruning rely on statistical learning principles.
  Accumulating the outer products of these contributions over a window yields a small Gram tensor, per (layer, channel, window), from which the exact output error of pruning any subset of modes follows in closed form, offline, at any budget.From the Gram tensor and the per-mode energies eh,i = Gh,ii we read off, per layer: Migration gap.The effective order PRh = (P i eh,i)2 / P i e2 h,i, the number of modes that carry the channel's energy

Abstract

Selective state-space models such as Mamba route information through a bank of first-order modes whose input coupling is set by a learned selection mechanism. We give an exact instrument for measuring how a trained model uses these modes. Because the state matrix is diagonal, each channel's output decomposes exactly into per-mode contributions, and a per-(layer, channel, window) Gram tensor yields the exact output error of dropping any subset of modes, offline, at any budget. Validated against the reference implementation to a relative error of $2.3\times10^{-7}$ on the Mamba-1 family where it is exact, the instrument predicts a layer's deployed pruning error to a median relative deviation of $5\times10^{-7}$ over $4{,}464$ configurations, its floor set by the reconstruction. Applying the instrument across the Mamba-1 family (130M--2.8B), the deployed 7B Falcon-Mamba, and Mamba-2, we find that trained models re-allocate their state space with the input: which modes carry the signal migrates across contexts, and at the most affected layers a per-input oracle roughly halves the output error of a fixed mode set. Frozen-signal counterfactuals attribute the migration primarily to the input-dependent write map $B_t$; the timestep usually identified with selectivity carries almost none of it. Input-scheduled mode pruning on this measurement outperforms static, Hankel-based, and layer-adaptive rankings at every scale from 130M to the deployed 7B Falcon-Mamba, and at half the state budget it matches the unpruned model. Because the scheduler reads each window's mode usage from a first pass, this demonstrates realizable headroom; we claim no deployed compute or memory saving.

Paper Context