混合 Transformer-SSM 架构的学习率迁移

arXiv:2610.01172v1 公告类型：新论文 摘要：我们研究了结合 Transformer 和状态空间模型（SSM）模块的混合架构中的学习率（LR）缩放问题，这类架构已被近期的多种生产级语言模型采用。具体而言，我们关注以下差距：一方面，是在无限宽度且状态规模不断增大的条件下，针对采用零阶保持（ZOH）离散化的 SSM 所推导出的理论缩放规则；另一方面，是业界标准的实际实现——在固定状态规模下使用简化 ZOH 的 Mamba。令人惊讶的是，在这一实际场景中，对于宽度 256–2048、深度 4–32 的架构，直至十亿参数规模，仅使用原始的 $\mu$P 方案，混合架构就能实现近乎为零的学习率迁移差距；尽管 SSM 运算不满足 Tensor Programs 的可表示性条件，并且我们测试的每一种参数化都未通过 $\mu$P 正确性的标准坐标检查诊断。我们将这一现象归因于混合架构中学习率迁移的双条件分解：其一是全局的更新到权重不变性，由 $\mu$P 的初始化和学习率缩放共同保证；其二是局部的逐组件平衡，由 AdamW 的逐参数归一化提供。我们的观察表明，最优学习率在宽度扩大至 8 倍的范围内保持不变；这种宽度不变性还跨越不同的深度、序列长度、批大小以及 Transformer 与 SSM 的比例而成立，并且能够迁移到 Nemotron-H——一个不属于我们自定义架构集合的生产级混合模型。我们希望这些发现能够弥合理论缩放规则与实际混合架构实现之间的差距，并推动进一步研究以实现二者的衔接。