Learning rate scaling shows stability across hybrid transformer state models
Learning Rate Transfer for Hybrid Transformer-SSM Architectures
Machine Learning
Summary
Training complex language models that mix Transformers and State-Space Models (SSMs) requires choosing the right learning rate to update the model efficiently. The authors studied existing theories and real-world practices, finding that practical models behave more predictably than expected when adjusting learning rates across different model sizes and depths. They explain this consistency by how the learning rate updates relate globally to weights and locally to individual parts, influenced by both initialization and the optimizer used. Their work helps connect theoretical rules with how models are actually trained in production.
What this means in practice
- •For deep learning engineers: Choose learning rates that reliably transfer across model sizes and depths when training hybrid Transformer-SSM language models.
- •For machine learning platform teams: Implement scalable training protocols using the observed learning rate stability in hybrid architectures to improve resource efficiency.
Authors
Jimin Seo, Gyubok Lee, Yeonsik Jo, Kiwoong Yoo, Yeongoon Kim, Minhae Oh, Jin Woo Koo, Suhwan Kim, Nakyung Lee, Minsik Seol, Idris Nechnech, Jaehyeon Kim, Giho Lee, Jungwoo Lee
Abstract
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$\times$, 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.