Learning Rate Transfer for Hybrid Transformer-SSM Architectures¶

会议: NeurIPS 2026
arXiv: 2610.01172
代码: https://github.com/cml-snu/hybrid-mup
领域: 优化/理论
关键词: 学习率迁移、最大更新参数化、混合架构、状态空间模型、AdamW

一句话总结¶

本文系统探究了 Transformer 与 SSM 混合架构在实际有限宽度下的学习率扩展规律，证实无需针对 SSM 设计复杂的专用规则，基于 AdamW 的朴素最大更新参数化（Naïve μP）即可在多种串并行混合结构及生产级模型（Nemotron-H）上实现近乎零间隙的学习率跨宽度与深度迁移。

研究背景与动机¶

结合注意力机制与状态空间模型（SSM，如 Mamba）的混合架构，凭借全局依赖建模能力与亚二次方推理复杂度的双重优势，已成为新一代大语言模型的重要发展方向。在数十亿参数规模下，如 Mamba-2-Hybrid 等混合模型不仅在多项基准上达到或超越纯 Transformer，还能带来数倍的推理加速，并在 Nemotron-H、Jamba 等前沿模型中得到实际应用。然而，混合模型在大规模训练中面临着一个极其昂贵的工程痛点：随着网络宽度与深度的缩放，最优超参数（尤其是学习率）会剧烈漂移，必须在每个规模下重新进行大范围超参数搜索，其网格调优开销往往远超单次全量训练成本。

在纯 Transformer 领域，最大更新参数化（μP）通过在无限宽度极限下对权重初始化与各层学习率进行协同缩放，成功实现了最优学习率在不同宽度间的零样本无缝迁移。然而，后续理论工作指出，SSM 中的多项关键算子并不满足 Tensor Programs 框架的表示性条件；为此，学者们推导了专门针对纯 SSM 的分组件规则（μP-SSM），并指出 SSM 在无限宽度极限下无法保持特征学习动态。由此产生了三个关键割裂：其一，现有 μP-SSM 理论仅针对纯 SSM 和无限状态维度（$N_x \to \infty$），未考虑混合架构中 Transformer 与 SSM 异构模块间的动态交互；其二，理论推导假设了真实零阶保持（ZOH）离散化与 SGD 优化器，而工业界标准 Mamba 普遍采用固定状态维度（$N_x=16$）与简化版 ZOH；其三，基于坐标检查（coordinate check）的传统 μP 诊断在混合模型上全面失效，使得实践者缺乏可落地的超参迁移指导。

本文切入的角度是：重新审视工业实践体制与理论渐近假设之间的断层，评估通用 μP 结合主流自适应优化器在有限宽度混合模型中的实际表现。核心 idea：在实用简化 ZOH 与固定状态维度体制下，AdamW 的逐参数二阶矩归一化能自动平衡 Transformer 与 SSM 的跨模块梯度异构性，使得完全不含 SSM 专用缩放规则的朴素 μP 即可解耦全局更新比率与宽度依赖，在混合架构上实现 0.0% 的跨宽度与深度学习率迁移。

方法详解¶

整体框架¶

论文所研究的学习率迁移范式旨在通过小规模代理模型（低宽度、浅层数）上的低成本粗网格搜索确定最优基准学习率 $\eta^*$，并将其直接外推至数十亿参数的工业级目标模型，免去目标尺度上的超参数调优。方法整体覆盖了纯架构基准（纯 Mamba、纯 Transformer）、串行混合（如 MTMT、TMTM）、并行双分支混合（如 MHMH、HMHM、HHHH）以及工业级非对称混合模型（Nemotron-H、GLA 混合结构）。

整个参数化与优化协同工作流如下：输入序列经嵌入层后进入交替或并行的混合模块堆叠；在前向传播阶段，朴素 μP 对隐藏层权重采用 $\mathcal{O}(1/d_{\text{model}})$ 方差初始化，并将隐藏层更新步长显式缩放 $\eta_{\text{base}}/d_{\text{model}}$；在反向传播阶段，由于 SSM 内部状态更新梯度与自注意力机制存在巨大的模长不一致，AdamW 的逐参数二阶矩归一化自适应地抵消模块间梯度异构性，实现组件更新平衡；最终使得全局更新-权重比在不同模型尺度下保持恒定。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入序列与嵌入层<br/>嵌入权重采用 N(0, 1) 初始化"] --> Prop["异构混合层堆叠<br/>串行交替 MTMT / 并行分支 HHHH / 生产级 Nemotron-H"]
    Prop --> InitScale["全局更新-权重比不变性<br/>隐藏层权重按 1/d 缩放方差，学习率按 1/d 缩放"]
    InitScale --> AdamNorm["局部组件更新自适应平衡<br/>AdamW 逐参数梯度二阶矩归一化 g/sqrt(v)"]
    AdamNorm --> OutTransfer["无缝学习率零样本迁移<br/>代理规模最优学习率直接泛化至目标尺度 (间隙 0.0%)"]

关键设计¶

1. 全局更新-权重比不变性：消除宽度对特征学习更新尺度的依赖

标准参数化（SP）将权重矩阵按 $W \sim \mathcal{N}(0, 1/\sqrt{d_{\text{model}}})$ 初始化并对所有层采用同一学习率 $\eta$。随着宽度 $d_{\text{model}}$ 的增大，矩阵前向激活值与反向梯度的方差迅速积累，导致相对更新步长 $\rho(W) = \|\Delta W\|_F / \|W\|_F = \mathcal{O}(\eta \sqrt{d})$ 随着模型加宽而膨胀，迫使最优学习率向更小值漂移。朴素 μP（Naïve μP）通过将隐藏层初始化方差设为 $\mathcal{O}(1/d_{\text{model}})$，并将隐层学习率缩放为 $\eta_{\text{base}} / d_{\text{model}}$，确保无论网络宽度如何扩展，全局参数矩阵的更新-权重比始终稳定在 $\mathcal{O}(\eta)$，即： $$\rho(W) = \frac{\|\Delta W\|_F}{\|W\|_F} = \mathcal{O}(1) \quad (d_{\text{model}} \to \infty)$$ 在混合架构中，该设计完全沿用纯 Transformer 的矩阵缩放原则，不为 SSM 内部的分立权重（如连续状态转换投影参数）引入任何定制化幂次调整，保持了参数化实现的通用性与轻量性。

2. 局部组件更新自适应平衡：AdamW 二阶矩归一化对冲异构梯度

在包含 Transformer 与 SSM 的异构网络中，单纯保证各层全局更新比率 $\mathcal{O}(1)$ 仍不足以确保学习率迁移，还必须满足任意两个异构组件 $W_X, W_Y$ 之间的相对更新强度 $\beta_{X/Y} = \|\Delta W_X\|_F / \|\Delta W_Y\|_F$ 在网络缩放时保持守恒。在传统 SGD 优化器下，该比例直接受制于梯度的均方根比值 $\beta_{X/Y} \propto (\eta_X / \eta_Y) \cdot (\mathrm{RMS}(\nabla_{W_X}\mathcal{L}) / \mathrm{RMS}(\nabla_{W_Y}\mathcal{L}))$，由于 SSM 状态计算对宽度的敏感性不同于自注意力，必须施加人工设计的层级学习率补偿（如 μP-SSM 中规定的 $\eta_A \propto N_u$）。而在现代预训练所标配的 AdamW 优化器中，有效单步更新为： $$\Delta W \propto \frac{g}{\sqrt{v} + \epsilon} \approx \operatorname{sign}(g)$$ AdamW 的逐参数二阶矩估计 $v$ 自动抵消了反向传播中跨组件的模长级差，使得模块间相对更新比与固有梯度模长彻底解耦。相反，若在 AdamW 基础之上机械套用理论推导的 $\mu\text{P-SSM}$ 规则，显式补偿将与 AdamW 内部归一化发生破坏性“重叠缩放”，反而造成更新严重失衡并引发训练发散。

3. 传统坐标检查失效分析与实用诊断解耦

传统 Tensor Programs 体系将“各层预激活值在初始化及前几步的模长对宽度保持不变”（即坐标检查 Coordinate Check）视为 μP 正确实现的金标准。但在混合模型中，由于 Mamba 的卷积与循环扫描机制包含特征维度的局部累加，实验检测显示所有参数化方法（包括 SP、朴素 μP、近似分块 μP 及忠实分块 μP）均未通过坐标检查，其预激活值随着宽度增加呈现出显著波动。作者指出，坐标检查是无限宽度连续极限下的充分条件而非必要条件；在实用的离散化混合架构中，只要同时满足“全局更新-权重比不变”与“局部更新平衡”两个条件，即使微观激活尺度出现局部漂移，宏观优化轨迹的最优学习率依然能够实现零间隙精确迁移。

损失函数 / 训练策略¶

模型训练采用标准的自回归交叉熵损失。所有参数化对比实验均使用 AdamW 优化器，超参数设定为 $\beta_1=0.9, \beta_2=0.95$，梯度裁剪阈值固定为 1.0，并配合余弦退火学习率调度与线性预热（warmup）。验证覆盖 WikiText-2（短期 2k 步网格搜索与周期遍历）、WikiText-103（20k 步收敛验证，约 1.37 个 epoch）以及高质量网络语料 FineWeb-Edu（500M 至 5B token 规模预训练）。

实验关键数据¶

主实验¶

主实验覆盖了 7 种不同架构拓扑（包含纯基线 MM、TT，串行混合 MTMT，并行混合 MHMH、HMHM、HHHH，以及混合 Transformer 交替 HTHT），在宽度从 256、512 扩展至 1024（8倍宽度扩展实验延伸至 2048，模型参数量达 1.24B~1.66B）下的最优学习率迁移表现。

架构类型	参数化方案	$w=256$ 最优 $\eta^*$	$w=512$ 最优 $\eta^*$	$w=1024$ 最优 $\eta^*$	迁移是否成立	迁移相对损失间隙
MHMH (并行)	Naïve μP	1.0	1.0	1.0	是 (完全对齐)	0.0%
MHMH (并行)	Block-wise μP (忠实)	3.0	3.0	1.0	否 (漂移发散)	> 5.0%
MHMH (并行)	SP (标准参数化)	0.003	0.003	0.001	否 (向小漂移)	1.9% ~ 8.7%
HHHH (全并行)	Naïve μP	1.0	1.0	1.0	是 (完全对齐)	0.0%
HHHH (全并行)	Block-wise μP (忠实)	3.0	3.0	1.0	否 (漂移)	> 5.0%
HHHH (全并行)	SP (标准参数化)	0.003	0.003	0.001	否 (向小漂移)	> 2.0%
MTMT (串行)	Naïve μP	1.0	1.0	1.0	是 (完全对齐)	0.0%
MTMT (串行)	Block-wise μP (忠实)	1.0	3.0	1.0	否 (剧烈漂移)	> 6.0%
MTMT (串行)	SP (标准参数化)	0.010	0.003	0.003	否 (向小漂移)	> 3.0%
TT (纯Transformer)	Naïve μP	1.0	1.0	1.0	是 (基准对齐)	0.0%
MM (纯Mamba)	Naïve μP	1.0	1.0	1.0	是 (基准对齐)	0.0%

在生产级模型 Nemotron-H（Mamba-2 + GQA + Squared-ReLU FFN）上，朴素 μP 在深度 16 与 32、宽度 256 至 1024 上均稳定保持最优学习率 $\eta^*=1.0$，迁移间隙严格为 0.0%；而 SP 方案最优学习率由 0.003 下降至 0.001，直接套用代理学习率将导致 1.9% 的训练性能惩罚。在 FineWeb-Edu 5B token 预训练下游零样本评估中，未经大模型专属网格调优的 MHMH-312M 达到 44.1% 平均准确率，优于同等数据训练的 Pythia-410M（42.4%）与 Mamba-370M（41.9%）。

消融实验¶

消融实验重点解构优化器选择对“局部更新平衡”条件的影响，以及在不同优化器机制下各参数组件在扩展过程中的更新失衡偏差 $\Delta \beta_{X/Y}$ 与损失表现。

优化器	参数化方案	自适应二阶矩 $g/\sqrt{v}$	动量支持	$w=256$ 损失	$w=1024$ 损失	$A_{\log}/\text{FFN}$ 更新偏差 $\Delta \beta$ ($128 \to 512$)
AdamW	Naïve μP	是	是 ($\beta_1=0.9$)	4.392	4.099	21% (平衡)
AdamW	Block-wise μP (忠实)	是	是 ($\beta_1=0.9$)	4.512	4.380	436% (严重失衡)
Adam 无动量	Naïve μP	是	否 ($\beta_1=0$)	4.470	4.465	24% (平衡)
SGD + 动量	Naïve μP	否	是 (0.9)	7.732	7.955	342% (严重失衡)
SGD 纯梯度	Naïve μP	否	否	9.398	10.803	342% (严重失衡)
SGD 纯梯度	Block-wise μP (忠实)	否	否	8.830	8.340	835% (剧烈发散)

关键发现¶

AdamW 逐参数归一化是关键支柱：消融表明，只要保留 $g/\sqrt{v}$ 机制（即使去除动量），模型在大宽度下的训练损失均显著下降且学习率能够平滑迁移；而去除二阶矩后（SGD 系列），损失大幅恶化（从 4.099 骤增至 7.955 以上），模型间相对更新偏差 $\Delta \beta$ 飙升至 340% 以上。
理论规则与实用架构存在“超量抵消”陷阱：在 AdamW 下强行施加针对纯 SSM 推导的 $\mu\text{P-SSM}$ 规则，由于其显式引入 $\eta_A \propto N_u$ 补偿，与 AdamW 内部的二阶矩消除发生叠加冲突，使更新偏差暴增至 436%，在 20k 步 WikiText-103 收敛中产生了 5.6%~7.7% 的额外损失劣化。
鲁棒性广泛跨越架构维度：朴素 μP 的 0.0% 迁移间隙在序列长度（128 至 2048）、批量大小（8 至 64）、并行 Transformer-SSM 比例（1:3 至 3:1）及梯度裁剪阈值变化下均保持完好；并在 10 亿参数（1.24B/1.66B）及 24.9B token 的 Chinchilla 最优训练预算下成功收敛。

亮点与洞察¶

打破理论壁垒的极简实践：此前学术界普遍认为 SSM 不符合 Tensor Programs 条件因而必须定制非线性规则，本文证实工业标准下的 Mamba 与 Transformer 混合模型直接使用朴素 μP 即可完美迁移，显著降低了混合架构大模型训练的技术门槛。
算力节省极其可观：通过在 $w=256$ 代理模型上进行 3 点粗筛并直接迁移至 $w=1024$ 目标模型，相比在目标尺度上全量扫描 7 个学习率，在 2k 步场景下节省 78% 的 GPU 算力，在 20k 步全收敛训练中算力节省高达 82%（加速比达 5.6×）。
坐标检查不可盲从：明确厘清了“坐标检查”在非纯矩阵乘网络中的局限性，提出以“全局更新-权重比不变”和“局部更新平衡”构成的双条件判定准则，为新型异构神经网络的超参迁移研究提供了全新的理论分析工具。

局限与展望¶

混合专家（MoE）及复数状态模型尚未探索：实验主要覆盖密集型混合架构，尚未验证与稀疏混合专家（MoE）结合时的迁移表现；此外，近期提出的 Mamba-3 采用了复数状态空间，超出了实数域 μP 的描述边界，有待未来构建复数参数化理论。
大状态维度真实 ZOH 场景未闭环：由于工业部署多采用固定 $N_x=16$ 的简化 ZOH，本研究所证伪的 μP-SSM 失效仅针对该主流工程体制；对于理论上状态维度同步扩展（$N_u \propto N_x$）且采用严格指数矩阵求逆的真实 ZOH 设定，其渐近理论预测仍需进一步实证考察。
超长步数与多种子联合验证成本受限：由于算力制约，1.24B 规模模型在 24.9B token 完整训练周期上仅验证了单一学习率与随机种子，超大规模下学习率景观尖锐度（Landscape Sharpness）的长期演化仍值得深入探究。

评分¶

新颖性: ⭐⭐⭐⭐☆ [挑战了 SSM 必须采用复杂参数化缩放的固有认知，提出了优雅的双条件解释框架]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖 7 种混合架构、工业级 Nemotron-H、多达 5500 GPU-hours 及 12 亿参数规模的全面验证]
写作质量: ⭐⭐⭐⭐⭐ [逻辑层层递进，严密分析了理论假设与工程实践之间的断层原因]
价值: ⭐⭐⭐⭐⭐ [为工业级混合大语言模型的大规模预训练提供了立即可用的零成本超参调优方案]

优化器	参数化方案	自适应二阶矩 \(g/\sqrt{v}\)	动量支持	\(w=256\) 损失	\(w=1024\) 损失	\(A_{\log}/\text{FFN}\) 更新偏差 \(\Delta \beta\) (\(128 \to 512\))
AdamW	Naïve μP	是	是 (\(\beta_1=0.9\))	4.392	4.099	21% (平衡)
AdamW	Block-wise μP (忠实)	是	是 (\(\beta_1=0.9\))	4.512	4.380	436% (严重失衡)
Adam 无动量	Naïve μP	是	否 (\(\beta_1=0\))	4.470	4.465	24% (平衡)
SGD + 动量	Naïve μP	否	是 (0.9)	7.732	7.955	342% (严重失衡)
SGD 纯梯度	Naïve μP	否	否	9.398	10.803	342% (严重失衡)
SGD 纯梯度	Block-wise μP (忠实)	否	否	8.830	8.340	835% (剧烈发散)