Learning Rate Transfer for Hybrid Transformer-SSM Architectures

📄 arXiv: 2610.01172v1 📥 PDF

作者: 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

分类: cs.LG

发布日期: 2026-10-01

备注: Accepted at NeurIPS 2026. 42 pages, 14 figures

💡 一句话要点

提出学习率转移方法以优化混合Transformer-SSM架构

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 学习率转移 混合架构 Transformer 状态空间模型 深度学习 模型优化 自然语言处理

📋 核心要点

  1. 现有的混合Transformer-SSM架构在学习率缩放方面存在理论与实践之间的差距，影响模型性能。
  2. 论文提出了一种新的学习率转移方法，强调全局和局部平衡的结合，以优化混合架构的学习率设置。
  3. 实验结果显示，在不同宽度和深度下，学习率转移差距接近零，且最优学习率在多个条件下保持不变。

📝 摘要（中文）

本研究探讨了混合架构中学习率（LR）缩放的问题，这些架构结合了Transformer和状态空间模型（SSM）模块，广泛应用于最新的语言模型。我们重点关注在无限宽度和增长状态大小下，SSM的理论缩放规则与实际应用中使用的简化ZOH Mamba之间的差距。令人惊讶的是，在实际情况下，混合架构在256到2048的宽度和4到32的深度范围内，使用原始的μP处方实现了近乎零的学习率转移差距。我们认为这归因于混合架构中学习率转移的两个条件分解：由μP的初始化和学习率缩放强制的全局更新到权重的不变性，以及由AdamW的每参数归一化提供的局部平衡。我们的观察表明，最优学习率在宽度上保持不变，并且这种不变性在深度、序列长度、批量大小和Transformer与SSM的比例之间也成立。

🔬 方法详解

问题定义：本论文旨在解决混合Transformer-SSM架构中学习率缩放的理论与实践之间的差距，现有方法在不同宽度和深度下的学习率转移效果不佳。

核心思路：论文提出了一种新的学习率转移机制，强调全局更新到权重的不变性和局部每参数归一化的结合，以实现更有效的学习率设置。

技术框架：整体架构包括混合Transformer-SSM模块，采用μP初始化和AdamW优化器，确保学习率在不同条件下的有效转移。

关键创新：最重要的创新点在于提出了学习率转移的双重条件分解，突破了SSM操作在μP正确性条件下的限制，实现了在实际应用中的有效性。

关键设计：关键设计包括使用μP的初始化策略和AdamW的每参数归一化技术，确保学习率在宽度、深度等多个维度上的不变性。通过这些设计，模型在大规模参数设置下依然能够保持良好的学习效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在宽度256到2048和深度4到32的范围内，混合架构实现了近乎零的学习率转移差距，且最优学习率在不同条件下保持不变，显示出显著的性能稳定性和可迁移性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和对话系统等，能够提升混合模型在实际生产环境中的性能和稳定性。未来，这一研究可能推动更多高效的学习率调整策略的开发，进一步优化深度学习模型的训练过程。

📄 摘要（原文）

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.