MESH：用于混合专家训练的内存高效Sinkhorn优化

# MESH: 面向混合专家训练的内存高效 Sinkhorn 优化
Source: https://arxiv.org/html/2608.04407
###### 摘要

内存高效的矩阵优化器（如 Sinkhorn 梯度下降）为稠密 Transformer 矩阵移除了大部分 AdamW 优化器状态，但直接应用于混合专家（MoE）训练并不可靠。我们在一个受控的 110M 参数 nanowhale DeepSeek 风格 MoE 预训练设置中研究了这一失败。SAGE/Sinkhorn 混合方法将优化器状态从 0.883GB 降至 0.331GB，但评估损失恶化到 3.8265，远高于相同设置下观察到的 AdamW 基线（在我们研究的种子中为 3.58–3.64）。我们表明，路由 MoE 专家矩阵是主要的失败点：其梯度是有条件的、随时间变化的，并且无状态 Sinkhorn 归一化难以处理。我们提出 MESH，一种用于 MoE 专家的隐藏动量 Sinkhorn 更新。MESH 通过梯度缓冲区生命周期恢复时间一阶矩信号，而无需将专家一阶矩存储为优化器状态。MESH-B 是一个可选的块预条件变体，添加了粗略的神经元/块逆 RMS 乘数。在所有消融实验中，矩阵归一化之前的时间平滑是主要因果成分；块/神经元预条件可以改善内存-质量边界，但尚未被证明是普遍必需的。在两个额外的种子中，MESH 和 MESH-B 相对于 AdamW 将优化器状态内存减少 62.5%，峰值 PyTorch CUDA 分配减少约 12.6%，评估损失差距不大。全状态诊断变体在消融中恢复了类似 AdamW 的性能，支持以下结论：MoE 专家需要时间平滑，但不一定需要完整的逐坐标 AdamW 状态。

关键词——混合专家、内存高效优化、Sinkhorn 优化、隐藏动量、语言模型预训练

## 1 引言

AdamW 仍然是语言模型预训练的默认优化器，因为它将时间一阶矩平滑与逐坐标自适应缩放结合在一起[1 (https://arxiv.org/html/2608.04407#bib.bib1),2 (https://arxiv.org/html/2608.04407#bib.bib2)]。然而，它的内存成本很高：AdamW 为每个可训练参数存储两个完整的状态张量。这促使优化器减少或移除矩状态。Sinkhorn 风格的矩阵优化器通过交替进行行和列操作来归一化梯度矩阵，从而为稠密二维权重生成无状态更新[4 (https://arxiv.org/html/2608.04407#bib.bib4)]。SAGE 风格的混合方法进一步认识到，嵌入对于纯无状态方法来说是困难的，因此为词表张量分配轻量状态更新，同时将稠密矩阵留给 Sinkhorn[5 (https://arxiv.org/html/2608.04407#bib.bib5)]。

这些角色分配主要针对稠密 Transformer 设置而开发。我们询问它们是否适用于混合专家（MoE）语言模型。答案是否定的：路由专家矩阵是二维的，但其梯度与稠密 MLP 梯度并不是同一对象。每个专家接收由路由器选择的 token 子集，并乘以学习到的门控分数，而这种条件分布会在训练过程中变化。在一个 nanowhale 110M DeepSeek 风格 MoE 模型中，直接采用 SAGE/Sinkhorn 混合方法可将优化器状态降至 0.331GB，但评估损失仅达到 3.8265，远差于 AdamW。将 SAGE 应用于专家更糟；而降低 Sinkhorn 缩放、融合专家张量或使用路由置信度阻尼都无法弥合这一差距。

因此，关键问题不在于 Sinkhorn 是否内存高效，而在于它对于路由 MoE 专家缺少什么。我们的消融实验表明，主要缺失的要素是矩阵归一化之前的时间平滑。无状态 Sinkhorn 根据瞬时专家梯度计算 `Sinkhorn\(G_t\)`。在路由专家中，`G_t` 是专家特定 token 分布的高方差条件样本。由于 Sinkhorn 是非线性的，`Sinkhorn\(\mathbb{E}[G_t]\)` 和 `\mathbb{E}[\mathrm{Sinkhorn}(G_t)]` 的表现可能不同；后者会放大逐步路由噪声。优化器失败是一个顺序问题：先归一化再平均与先平均再归一化表现不同。MESH 将更新转换为 `Sinkhorn\(H_t\)`，其中 `H_t` 在矩阵归一化之前对专家的梯度历史进行平均。

我们提出 MESH，一种 MoE 感知的隐藏动量 Sinkhorn 优化器。MESH 保留 SAGE/Sinkhorn 混合方法用于词表和非专家矩阵，但用隐藏动量专家 Sinkhorn 替换无状态专家 Sinkhorn。其专家一阶矩通过梯度缓冲区生命周期维护，而不是作为优化器状态张量。我们还评估了 MESH-B，它添加了粗略的块/神经元逆 RMS 乘数。主要的实验信息刻意保持保守：MoE 专家主要需要在矩阵归一化之前进行时间平滑；块级自适应性可能有所帮助，尤其是在 seed-42 诊断中，但它不是唯一的因果成分。

我们的贡献是：

- •我们确定路由 MoE 专家是在 nanowhale 预训练中直接采用 SAGE/Sinkhorn 优化失败的关键点。
- •我们引入 MESH，一种隐藏动量 Sinkhorn 专家更新，它在不添加显式专家一阶矩优化器状态的情况下恢复一阶矩信号。
- •我们评估 MESH-B，一种可选的块预条件扩展，并表明无块和块预条件的隐藏动量变体在多个种子上给出相似的内存-质量权衡。
- •我们提供了否定性消融实验，表明仅二阶矩、仅符号/动量、仅置信度、仅缩放和融合 Sinkhorn 的干预措施都不能解释这一差距。

## 2 相关工作

### 2.1 自适应与内存高效优化

Adam 和 AdamW 为每个参数维护一阶和二阶矩估计[1 (https://arxiv.org/html/2608.04407#bib.bib1),2 (https://arxiv.org/html/2608.04407#bib.bib2)]。这种鲁棒性以每个参数两个状态张量为代价。内存高效的 Adam 变体以不同方式降低这一成本：Adam-mini 通过利用块结构使用更少的学习率资源[7 (https://arxiv.org/html/2608.04407#bib.bib7)]，而 HMAdam/HMAdamW 通过将动量类量存储在梯度缓冲区中，移除显式一阶矩张量[6 (https://arxiv.org/html/2608.04407#bib.bib6)]。MESH 仅将后一种思想用于 MoE 专家一阶矩，并将其与 Sinkhorn 风格矩阵归一化结合。

### 2.2 矩阵归一化与 SAGE 混合方法

Sinkhorn 和梯度多归一化优化器通过行和列归一化计算矩阵更新[3 (https://arxiv.org/html/2608.04407#bib.bib3),4 (https://arxiv.org/html/2608.04407#bib.bib4)]。它们的吸引力在于，可以在不存储 AdamW 两个完整状态张量的情况下更新大型稠密矩阵。SAGE 通过将嵌入视为单独困难类别来扩展这一策略[5 (https://arxiv.org/html/2608.04407#bib.bib5)]。我们的工作识别出第二个困难类别：路由 MoE 专家矩阵。

### 2.3 混合专家模型

稀疏 MoE 层将 token 路由到专家子集，从而在不于每个 token 上激活每个参数的情况下增加模型容量[8 (https://arxiv.org/html/2608.04407#bib.bib8),9 (https://arxiv.org/html/2608.04407#bib.bib9),10 (https://arxiv.org/html/2608.04407#bib.bib10)]。这种路由改变优化器问题。专家梯度以路由器决策和门控权重为条件，因此即使路由专家矩阵与稠密 MLP 矩阵形状相同，二者也不等价。MESH 正是针对这种专家特定的梯度结构而设计。

## 3 背景与诊断

### 3.1 稠密矩阵的 Sinkhorn 更新

设 `W∈Rm×nW\in\mathbb{R}^{m\times n}` 为一个矩阵参数，`G_t=\nabla_W\mathcal{L}_t` 为其在步骤 `t` 处的梯度。Sinkhorn 风格矩阵优化器通过重复归一化行范数和列范数来形成更新方向，

`U_t=\mathrm{Sinkhorn}(G_t;K,\gamma)`，(1)

其中 `K` 是归一化轮数，`γ` 是标量更新缩放。这种更新的吸引力在于它是无状态的：它使用当前矩阵梯度，并且不为稠密矩阵存储一阶或二阶矩。

### 3.2 为什么路由专家不是普通稠密矩阵

路由 MoE 专家权重也是一个矩阵，但其梯度具有不同的采样结构。在 top-k MoE 层中，专家 `e` 仅接收由路由器选择的 token：

`G_{e,t}=\sum_{i:e\in\mathrm{topk}(x_i,t)}a_{i,e,t}\,g_{i,e,t}`，(2)

其中 `a_{i,e,t}` 是门控权重。因此 `G_{e,t}` 是在条件分布 `p(x\mid e\in\mathrm{topk}(x,t))` 下的梯度，而不是在完整小批量分布下的梯度。随着训练过程中路由器与隐藏状态的变化，这种条件分布也在变化。

这一点很重要，因为 Sinkhorn 是非线性的。稠密训练通常可以容忍 `Sinkhorn\(G_t\)`，因为 `G_t` 是对全批量稠密矩阵梯度的合理稳定估计。然而，在路由专家中，序列 `G_{e,t}` 携带由路由引起的时间噪声。在平滑之前应用 Sinkhorn 会对每个有噪声的专家样本计算 `Sinkhorn\(G_{e,t}\)`。而动量则近似于对时间平均的专家信号应用 Sinkhorn。这一区别是本文的核心诊断。

### 3.3 我们检验的假设

我们考虑了专家失败的几种解释。原始路由计数稀疏性可能降低有效批量大小；门控权重偏斜可能在每个专家都被选择时仍产生有效稀疏性；语义路由可能使专家梯度非独立同分布；Sinkhorn 可能对专家张量划分敏感；或者无状态 Sinkhorn 可能只是缺乏时间平滑。我们的实验排除了前四种作为完整解释。剩余得到支持的假设是，路由专家需要在矩阵归一化之前获得一阶矩信号。粗略的块/神经元预条件可以改进这种更新，但多种子结果表明，一阶矩是主要因果成分。

## 4 MESH

### 4.1 问题：MoE 专家梯度不是稠密矩阵梯度

设 `W_e` 为一个专家矩阵，`G_{e,t}=\nabla_{W_e}\mathcal{L}_t` 为其在步骤 `t` 处的梯度。稠密 MLP 矩阵从小批量中几乎所有 token 接收梯度贡献。而路由专家接收的是

`G_{e,t}=\sum_{i\in\mathcal{R}_{e,t}}a_{i,e,t}\,g_i(W_e)`，(3)

其中 `\mathcal{R}_{e,t}` 是为专家 `e` 选择的 token 子集，`a_{i,e,t}` 是门控权重。子集和权重都随路由器和隐藏状态而变化。因此，专家梯度是随时间变化的条件估计。将无状态 Sinkhorn 直接应用于 `G_{e,t}` 使用

`U_{e,t}=\mathrm{Sinkhorn}(G_{e,t};K,\tau)`，(4)

这会归一化瞬时矩阵，而不记住专家信号是持续性的还是一步路由波动。

MESH-B 还会将 Sinkhorn 方向乘以一个粗略的均值归一化块/神经元逆 RMS 因子。

图 1：MESH 核心机制。失败是一个顺序问题：归一化每个瞬时路由专家梯度并不等同于归一化时间平均的专家信号。
### 4.2 MESH：在 Sinkhorn 之前使用隐藏动量

MESH 只改变 MoE 专家角色。主算法 MESH 维护一个未归一化的专家动量累积器

`H_{e,t}=\beta_1 H_{e,t-1}+G_{e,t}`。(5)

这与传统 Adam 一阶矩只相差一个常数缩放。我们直接使用 `H_{e,t}`，因为 Sinkhorn 归一化和均值归一化块乘数对全局缩放在显式更新缩放范围内不变。然后专家矩阵方向为

`U_{e,t}=\mathrm{Sinkhorn}(H_{e,t};K,\tau)`，(6)

其中 `K` 是交替行/列归一化的轮数，`τ` 是 Sinkhorn 更新缩放。

隐藏动量实现将 `H_{e,t}` 存储在优化器状态字典之外。在优化器步骤之后，专家梯度缓冲区被暂存。在下一次反向传播之前，缓冲区被重新附加为 `p.grad\leftarrow\beta_1 H_{e,t}`。反向传播累积会加上当前梯度，从而得到 `H_{e,t+1}`。这种“反向传播前准备”的生命周期在 Trainer/Accelerate/DDP 风格循环中是必要的，因为在这些循环中，朴素的 `zero_grad` 桥接可能被包装模型的梯度清除绕过。我们最终实现中的诊断报告大约每次优化器步骤一次 prepare 调用，缓冲区影子余弦接近 1，相对误差约为 `10^{-8}`。

### 4.3 MESH-B：可选的块/神经元缩放

MESH 可以单独用作隐藏动量 Sinkhorn。我们还评估了一种块缩放变体，记为 MESH-B，它对 Sinkhorn 方向应用粗略的逆 RMS 乘数。对于专家 `w_1,w_3` 矩阵，我们使用行统计量；对于 `w_2`，我们使用列统计量。设 `R(H_{e,t})` 表示这个块 RMS。默认 MESH-B 统计量从平滑源 `H_{e,t}` 计算，而不是来自单独的原始梯度流：

`V_{e,t}=\beta_2 V_{e,t-1}+(1-\beta_2)R(H_{e,t})^2`。(7)

我们使用 `H` 作为该统计量，因为最好的 seed-42 诊断使用了 `V(H)`，而且它使隐藏动量实现保持简单，无需存储额外的原始梯度流。原始梯度块统计量是一种自然替代方案，但我们将系统性比较留待未来工作。

乘数被裁剪并均值归一化，

`Q_{e,t}=\mathrm{clip}(\frac{(V_{e,t}+\epsilon)^{-p/2}}{\mathrm{mean}((V_{e,t}+\epsilon)^{-p/2})},q_{\min},q_{\max})`，(8)

专家更新为

`W_{e,t+1}=W_{e,t}-\eta\,Q_{e,t}\odot U_{e,t}-\eta\lambda W_{e,t}`，(9)

其中 `Q_{e,t}` 在适当的块维度上广播。块乘数刻意保持粗略。它不是完整的逐坐标二阶矩，也不是本文的核心因果主张。我们的结果表明，时间平滑是主要需求，而块/神经元缩放可以在某些设置中改善边界。

图 2：MESH 保留 SAGE/Sinkhorn 混合方法用于非专家参数，并且只替换 MoE 专家更新。MESH 是隐藏动量专家 Sinkhorn；MESH-B 添加可选的块预条件。全状态变体是诊断对照。
### 4.4 全状态诊断对照

我们仅将全状态变体用于诊断。MESH-Full 在优化器状态中显式存储相同的未归一化专家动量 `H_{e,t}`，而不是放在梯度缓冲区中。MESH-B-Full 额外存储 MESH-B 使用的块统计量。比较隐藏变体和全状态变体，可以将隐藏动量的系统效应与时间平滑的优化器效应区分开来。在我们的实验中，全状态变体通常是更强的质量诊断，而 MESH 和 MESH-B 是内存高效的算法。

## 5 实验

### 5.1 设置

我们使用 HuggingFace nanowhale 训练代码库，它实现了一个 1