Research map / Architectures and efficiency / RNNs and normalization
On Feature Learning in Structured State Space Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that µP and spectral scaling fail to give feature learning in structured SSMs (Mamba/S4) because they are not Tensor Programs, and derives a corrected scaling (µP-SSM) for initialization and learning rates that yields stable feature learning and hyperparameter transfer.
Problem
Feature learning and hyperparameter transfer in state space models (SSMs such as S4 and Mamba) at infinite width, where standard parameterization, µP, and spectral scaling conditions are known to work for transformers and MLPs but are not validated for SSMs.
Approach
Proves structured SSMs (HiPPO-based S4/Mamba) are not representable as NE⊗OR⊤ Tensor Programs, so µP does not apply; derives via forward and backward signal propagation the correct init scales (σB ∈ Θ(√(Nx/Nu)), σC ∈ Θ(1/√(NxNu))) and learning rates (ηB ∈ Θ(Nx/√Nu), ηC ∈ Θ(1/(Nx√Nu)), ηa ∈ Θ(Nu)) giving µP-SSM.
Key results
- Under SP, Mamba latent states and outputs explode at initialization; under spectral scaling with ZOH the latent states and their updates vanish with width, while µP-SSM gives Θ(√Nx) latent-state norms and non-trivial updates for both ZOH and Euler discretization.
- On WikiText-103 Mamba, µP-SSM improves test loss over SP and gives stable learning-rate transfer and monotonic improvement with width, whereas SP shows non-monotonic behavior.
- Proves Mamba/S4 with HiPPO-based A are not NE⊗OR⊤ Tensor Programs, so µP does not apply to them.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use µP-SSM scaling instead of standard parameterization | Mamba · WikiText-103 | test loss | Helped |
| use µP-SSM scaling instead of heuristic µP scaling | Mamba · WikiText-103 | training stability at large learning rates | Helped |
| use µP-SSM scaling instead of standard parameterization | Mamba · increasing model widths | test loss | Helped |
| use standard parameterization | Mamba · increasing model widths | test loss | Mixed |
| use spectral scaling instead of µP-SSM scaling | Mamba · ZOH discretization, increasing width | latent-state update norm | Hurt |
| replace ZOH with Euler discretization for B (ablation) | Mamba with spectral scaling · increasing width | output-signal stability | Hurt |
| replace ZOH with Euler discretization for B (ablation) | Mamba with µP-SSM scaling · increasing width | feature-signal scaling | No clear effect |
| use heuristic µP scaling instead of µP-SSM scaling | Mamba · large learning rates | training stability | Hurt |
| use spectral scaling instead of µP-SSM scaling | Mamba · ZOH discretization, increasing width | latent-state norm at initialization | Hurt |
| use standard parameterization instead of µP-SSM scaling | Mamba · increasing width | signal stability | Hurt |
Surprising findings
- Spectral scaling fails for SSMs: under ZOH, SSM hidden states and their updates vanish with width, leaving the first SSM block in the lazy regime, even though outputs update non-trivially.
- Optimal learning rate transfers across widths under spectral scaling despite the SSM layer not being in the feature learning regime.
Experimental setup
Datasets: WikiText-103 WikiText
Benchmarks: WikiText-103
Metrics: test loss stability of latent state and output norms update scale
Baselines: Standard Parameterization µP (heuristic) spectral scaling parameterization
Builds on: Tensor Programs µP (Maximal Update Parameterization) spectral scaling Mamba S4
Closest papers by meaning
- How Feature Learning Can Improve Neural Scaling Laws (ICLR 2025 · similarity 0.88)
- Sharp feature-learning transitions and Bayes-optimal neural scaling laws in extensive-width networks (NeurIPS 2026 · similarity 0.88)
- Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $\mu$ Parametrization (ICML 2025 · similarity 0.88)
- Provable Benefits of Complex Parameterizations for Structured State Space Models (NeurIPS 2024 · similarity 0.88)
- $\boldsymbol{\mu}\mathbf{P^2}$: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling (NeurIPS 2024 · similarity 0.88)
- On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.