arXiv：cs.LG（机器学习，全量分类）· 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·· 15 小时前AI 评分46
混合 Transformer-SSM 架构的学习率迁移研究
Learning Rate Transfer for Hybrid Transformer-SSM Architectures
AI 导读
研究发现在宽度 256-2048、深度 4-32 的混合 Transformer-SSM 架构中，仅用原始 μP 方案即可实现近零学习率迁移差距，最优学习率对宽度在 8 倍范围内保持不变，且该不变性跨深度、序列长度、批大小和 Transformer 与 SSM 的比例成立。
正文
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 $\mu$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 $\mu$P correctness. We attribute this to a two-condition decomposition of LR transfer in hybrid architectures: a global update-to-weight invariance, enforced by $\mu$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.
| Comments: | Accepted at NeurIPS 2026. 42 pages, 14 figures |
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2610.01172 [cs.LG] |
| (or arXiv:2610.01172v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.01172 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Jimin Seo [view email]
[v1]
Thu, 1 Oct 2026 06:46:03 UTC (4,338 KB)
来源：arXiv：cs.LG（机器学习，全量分类） · arxiv.org