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

我们研究结合 Transformer 和状态空间模型 (SSM) 块的混合架构的学习率 (LR) 缩放，这是最近几种生产语言模型采用的一类。特别是，我们关注在无限宽度且状态尺寸不断增长的零阶保持（ZOH）离散化下为 SSM 导出的理论缩放规则与在固定状态尺寸下使用简化 ZOH Mamba 的现场标准实际实现之间的差距。令人惊讶的是，在这种实际情况下，混合架构仅使用原始的 $μ$P 处方即可在宽度 256-2048 和深度 4-32 至十亿参数规模上实现接近零的 LR 传输间隙，即使 SSM 操作超出了其张量程序可表示性条件，并且我们测试的每个参数化都未能通过 $μ$P 正确性的标准坐标检查诊断。我们将其归因于混合架构中 LR 传输的两个条件分解：全局更新权重不变性，由 $μ$P 的初始化和 LR 缩放强制执行；以及由 AdamW 的每个参数标准化提供的本地每个组件平衡。我们的观察表明，最佳 LR 对于高达 8$\times$ 的宽度是不变的，这种宽度不变性在深度、序列长度、批量大小和 Transformer-to-SSM 比率上保持不变，并且它转移到 Nemotron-H，这是我们定制架构集之外的生产混合体。我们希望这些发现能够填补理论扩展规则和实际混合实现之间的差距，并激发进一步的研究以弥补这一差距。