Learning Rate Transfer for Hybrid Transformer-SSM Architectures
Not on the map. Why?
Abstract (arXiv)
We study learning rate (LR) scaling for hybrid architectures combining Transformer and State-Space Model (SSM) blocks, a class adopted by several recent production language models. In particular, we focus on the gap between the theoretical scaling rules derived for SSMs under zero-order-hold (ZOH) discretization at infinite width with growing state size, and the field-standard practical implementations using simplified-ZOH Mamba at fixed state size. Surprisingly, in this practical regime hybrid architectures achieve a near-zero LR transfer gap across widths 256-2048 and depths 4-32 up to billion-parameter scale using only the original μP prescription, even though SSM operations fall outside its Tensor Programs representability conditions and every parameterization we test fails the standard coordinate-check diagnostic of μP correctness. We attribute this to a two-condition decomposition of LR transfer in hybrid architectures: a global update-to-weight invariance, enforced by μP's initialization and LR scaling; and a local per-component balance, provided by AdamW's per-parameter normalization. Our observations show that the optimal LR is invariant to width up to 8×, that this width invariance holds across depth, sequence length, batch size, and Transformer-to-SSM ratio, and that it transfers to Nemotron-H, a production hybrid outside our custom architecture set. We hope these findings fill the gap between theoretical scaling rules and practical hybrid implementations, and stimulate further research toward bridging it.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 39 of its 48 references that are in the map. First posted on arXiv 1 Oct 2026.
Builds on
39 of 48 references.
References · 39
- 2017Attention Is All You Needcited by 59,932
- 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 12,865
- 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
- 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
- 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,581
- 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
- 2016Pointer Sentinel Mixture Modelscited by 2,992
- 2019PIQA: Reasoning about Physical Commonsense in Natural Languagecited by 2,818
- 2021Efficiently Modeling Long Sequences with Structured State Spacescited by 2,424
- 2023Pythia: A Suite for Analyzing Large Language Models Across Training and Scalingcited by 1,734
- 2020GLU Variants Improve Transformercited by 1,724
- 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
27 more
- 2024Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Dualitycited by 1,283
- 2024The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scalecited by 1,104
- 2016The LAMBADA dataset: Word prediction requiring a broad discourse contextcited by 960
- 2020HiPPO: Recurrent Memory with Optimal Polynomial Projectionscited by 600
- 2023Gated Linear Attention Transformers with Hardware-Efficient Trainingcited by 490
- 2022Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfercited by 272
- 2023Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-trainingcited by 260
- 2024An Empirical Study of Mamba-based Language Modelscited by 177
- 2017Dissecting Adam: The Sign, Magnitude and Variance of Stochastic Gradientscited by 136
- 2024Why Transformers Need Adam: A Hessian Perspectivecited by 123
- 2024Hymba: A Hybrid-head Architecture for Small Language Modelscited by 95
- 2025Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Modelscited by 88
- 2026Mamba-3: Improved Sequence Modeling using State Space Principlescited by 76
- 2024Scaling Exponents Across Parameterizations and Optimizerscited by 71
- 2025Falcon-H1: A Family of Hybrid-Head Language Models Redefining Efficiency and Performancecited by 52
- 2023Tensor Programs IVb: Adaptive Optimization in the Infinite-Width Limitcited by 35
- 2024Super Consistency of Neural Network Landscapes and Learning Rate Transfercited by 23
- 2025Completed Hyperparameter Transfer across Modules, Width, Depth, Batch and Durationcited by 22
- 2025Weight Decay may matter more than muP for Learning Rate Transfer in Practicecited by 20
- 2024Sparse maximal update parameterization: A holistic approach to sparse training dynamicscited by 16
- 2025Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scalescited by 15
- 2024μP2: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scalingcited by 10
- 2024An Empirical Study of μP Learning Rate Transfercited by 9
- 2025Gradient Heterogeneity Complements Hessian Heterogeneity in Transformer Optimizationcited by 6
- 2025Scaling Diffusion Transformers Efficiently via μPcited by 6
- 2025On the Surprising Effectiveness of Large Learning Rates under Standard Width Scalingcited by 6
- 2026Expressivity-Efficiency Tradeoffs for Hybrid Sequence Modelscited by 3
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026A KL Lens on Quantization: Fast, Forward-Only Sensitivity for Mixed-Precision SSM-Transformer Models
- 2026Attention to Mamba: A Recipe for Cross-Architecture Distillation
- 2024Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models
- 2025TransXSSM: A Hybrid Transformer State Space Model with Unified Rotary Position Embedding
- 2026Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
- 2026Complete-muE: Optimal Hyperparameter Transfer and Scaling for MoE Models
- 2025TransMamba: A Sequence-Level Hybrid Transformer-Mamba Language Model
- 2026Priming: Hybrid State Space Models From Pre-trained Transformers