在训练大模型时很多人的默认优化器是 AdamW。它足够稳定几乎成了“标配”。但最近几个月Muon 优化器打破了这种惯性Mistral 用它在部分场景中取得了比 AdamW 更好的收敛效果和最终指标也让越来越多的研究者开始重新思考大模型训练不一定非要 Adam矩阵结构信息可能比逐元素缩放更重要。不过Muon 也有一个明显短板——它虽然对梯度做了矩阵正交化处理但并没有显式利用损失曲面的曲率信息。简单说它对所有方向用了近乎一致的“旋转缩放”策略而真实损失面在不同方向上的弯曲程度差异很大。MALT 的出现恰好是想解决这个问题在不引入二次方内存开销的前提下通过对角预条件Diagonal Preconditioning给 Muon 增加曲率感知Curvature-Aware能力。这篇文章我会从 Muon 的核心思想讲起分析 MALT 为什么选择对角预条件给出可运行的 PyTorch 演示实现并讨论“这个优化器到底适合谁、迁移时需要验证什么”。如果你正在调大模型训练、对比优化器或者在思考“AdamW 之后下一个主流优化器长什么样”这篇文章值得读完并收藏。1. 为什么 Muon 能挑战 AdamW 的地位要理解 MALT先要理解它继承的对象——Muon。Muon 由 Mistral AI 的研究者提出核心思路非常直接把神经网络的权重矩阵视为一个线性变换而不是一堆“互不相干”的标量参数。基于这个视角它不再像 AdamW 那样对每个参数做逐元素归一化而是先对梯度做正交化处理再结合动量更新。这里的关键技术是 Newton-Schulz 迭代。它是一种快速逼近矩阵极分解的方法能够把梯度矩阵逐步“拉”向正交矩阵。每次迭代的复杂度可控通常只需要 5 到 10 次迭代就能得到足够好的正交化结果。Muon 在动量之后施加这个正交化步骤相当于给更新方向加了一个“形状约束”。对比一下两者的差异维度AdamWMuon对参数的处理方式逐元素独立归一化对权重矩阵整体做正交化主要数学工具一阶矩、二阶矩估计动量 Newton-Schulz 迭代内存开销需要额外保存两个动量变量额外保存一个动量变量整体略低适合场景参数规模大的通用训练权重矩阵结构强的模型如 Transformer理论视角对角自适应学习率利用矩阵谱结构的预条件从实践结果看Muon 在 Transformer 类模型上通常能显著减少训练步数或者在相同步数下获得更低的 loss。它的成功让很多人意识到AdamW 的逐元素缩放并不一定是最优的矩阵层面的结构约束可能更符合神经网络权重的本质。2. Muon 的局限正交化不等于曲率感知Muon 的效果令人兴奋但它并不是优化器的终点。一个容易忽略的问题是正交化本质上是一种“不考虑损失面几何信息”的预条件策略。损失函数的曲率决定了最优步长。在曲率大的方向步子应该小在曲率小的方向步子可以大。理想情况下优化器应该根据每个方向的曲率调整步长这正是二阶优化方法的目标。但完整的二阶方法例如自然梯度法需要维护 Fisher 信息矩阵代价是参数量平方级别的内存和计算大模型根本扛不住。Muon 的做法是对所有方向统一做正交化相当于只调整了更新方向的“形状”却没有针对不同方向的曲率差异做区分。可以类比成你在不同坡度的山路上都用同样大小的步幅走路方向走对了但步伐并不总是合适的。AdamW 至少通过二阶矩估计对每个参数方向做了缩放虽然是一个很粗糙的曲率近似但确实具备方向差异性。Muon 在这方面的信息反而比 AdamW 少。MALT 的切入点就在这里它希望在保留 Muon 矩阵结构优势的同时补充方向级别的曲率信息而且必须足够轻量不能把二阶方法的成本带回来。3. 为什么是对角预条件而不是更复杂的近似既然要引入曲率感知为什么不直接上 K-FAC 这类更精细的近似原因很现实复杂度。K-FAC 利用 Kronecker 分解来近似 Fisher 矩阵理论上能在保持一定精度的同时降低开销但实现复杂度非常高需要针对不同层结构做定制分布式场景下同步和更新也有额外成本。对于大多数研究团队和工程师来说这属于“知道但用不起来”的范畴。对角预条件是一个务实的折中。它假设参数方向之间没有相关性只估计每个方向的方差或曲率。虽然是一个粗糙的近似但它在实践中已经被证明有效——AdamW 的二阶矩估计本质上就是一种对角预条件。MALT 的思路是把这个机制对接进 Muon 的正交化框架中让正交化之后的更新方向再按对角曲率信息做缩放。从概率图模型和信息几何的角度看Fisher 矩阵对角线上的元素确实包含了各参数方向的局部敏感性信息。虽然忽略相关性会丢失一些信息但考虑到神经网络参数规模通常在千万到百亿级别保留相关性在计算上根本不现实。对角预条件在“信息量”和“可实现性”之间找到了一个合理的平衡点。4. MALT 的算法框架与数学直觉MALT 的全称是 Lightweight Curvature-Aware Muon via Diagonal Preconditioning。从命名就能看出它的三个关键设计决策轻量、曲率感知、基于 Muon。它的大致工作流程可以拆解为几步第一计算梯度并按照 Muon 的方式维护动量。动量项不仅提供惯性也是后续预条件的基础。第二维护一个对角曲率估计。最常用的方式是类似 Adam 的二阶矩滑动平均也就是对梯度的平方做指数加权平均。这个统计量可以看作是参数方向曲率的粗略代理代表“这个方向最近梯度的波动幅度”。第三对动量执行正交化处理得到结构化的更新方向。第四用对角曲率统计量缩放更新方向。直观理解是梯度变化剧烈的方向说明损失面弯曲厉害步长要保守梯度变化平缓的方向说明比较平坦步长可以激进一些。从数学直觉上看MALT 相当于在做“自然梯度”的轻量级近似。自然梯度通过 Fisher 矩阵逆来调整更新方向保证参数更新在分布空间中移动的是最短距离。而 MALT 只保留 Fisher 矩阵的对角部分再用正交化来捕捉矩阵整体结构——两部分叠加等于把“结构调整”和“方向缩放”解耦了。这个解耦很关键。它意味着 MALT 可以让正交化专注于改善梯度方向的几何形态而让对角预条件专注于步长调整两者各司其职不会互相干扰。5. 从 AdamW 迁移到 MALT 的 PyTorch 演示实现下面我用 PyTorch 实现一个演示版本的 MALT 优化器。需要说明的是这是一个为了讲清原理而简化的版本具体实现细节要以官方版本为准本文重点演示通用思路和关键步骤。# 文件路径malt_optimizer.py import torch from torch.optim import Optimizer def newton_schulz_(g, iterations5): 对梯度矩阵做 Newton-Schulz 正交化近似。 输入 g 的形状为 [..., M, N]输出近似正交的矩阵。 # 初始缩放帮助迭代稳定 with torch.no_grad(): a, b, c 1.0, 1.0, 1.0 for _ in range(iterations): g a * g b * torch.bmm(g, torch.bmm(g.transpose(-1, -2), g)) \ c * torch.bmm(g, torch.bmm(g, torch.bmm(g.transpose(-1, -2), torch.bmm(g, g)))) return g class MALT(Optimizer): 轻量级曲率感知 Muon 优化器演示版。 核心思想Muon 正交化 对角预条件。 参数: params: 模型参数 lr: 学习率 momentum: 动量系数 beta: 对角曲率估计的 EMA 系数 eps: 数值稳定项 ns_iter: Newton-Schulz 迭代次数 def __init__(self, params, lr0.01, momentum0.9, beta0.999, eps1e-8, ns_iter5): defaults dict(lrlr, momentummomentum, betabeta, epseps, ns_iterns_iter) super().__init__(params, defaults) def step(self, closureNone): loss None if closure is not None: loss closure() for group in self.param_groups: lr group[lr] momentum group[momentum] beta group[beta] eps group[eps] ns_iter group[ns_iter] for p in group[params]: if p.grad is None: continue grad p.grad.data if grad.dim() 2: # 对向量参数退化为 Adam 风格更新 state self.state[p] if momentum_buffer not in state: state[momentum_buffer] torch.zeros_like(p.data) state[variance] torch.zeros_like(p.data) buf state[momentum_buffer] var state[variance] buf.mul_(momentum).add_(grad) var.mul_(beta).addcmul_(grad, grad, value1 - beta) denom var.sqrt().add_(eps) p.data.addcdiv_(buf, denom, value-lr) else: state self.state[p] if momentum_buffer not in state: state[momentum_buffer] torch.zeros_like(p.data) state[variance] torch.zeros_like(p.data) buf state[momentum_buffer] var state[variance] # 1. 更新动量 buf.mul_(momentum).add_(grad) # 2. Newton-Schulz 正交化动量 orth_buf newton_schulz_(buf.clone().unsqueeze(0), ns_iter).squeeze(0) # 3. 更新对角曲率估计基于原始梯度平方 var.mul_(beta).addcmul_(grad, grad, value1 - beta) # 4. 对角预条件缩放 denom var.sqrt().add_(eps) scaled orth_buf / denom p.data.add_(scaled, alpha-lr) return loss这段代码有几个关键点第一对二维以上的权重矩阵使用完整的 MALT 逻辑动量、正交化、曲率缩放。第二对向量参数退化为 Adam 风格更新因为正交化主要针对矩阵设计向量参数强行正交化没有意义。第三曲率估计用的是梯度平方的滑动平均不是动量平方这样可以让曲率信息更贴近即时梯度变化。下面是一个完整的使用示例用一个小型 Transformer 训练任务来演示如何替换 AdamW# 文件路径train_with_malt.py import torch import torch.nn as nn from malt_optimizer import MALT class SimpleTransformer(nn.Module): def __init__(self, vocab_size1000, d_model128, nhead4, num_layers2): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.pos_encoding nn.Parameter(torch.randn(1, 128, d_model)) layer nn.TransformerEncoderLayer( d_modeld_model, nheadnhead, dim_feedforward512, batch_firstTrue ) self.encoder nn.TransformerEncoder(layer, num_layersnum_layers) self.output nn.Linear(d_model, vocab_size) def forward(self, x): seq_len x.shape[1] x self.embedding(x) self.pos_encoding[:, :seq_len, :] x self.encoder(x) return self.output(x) def generate_batch(batch_size8, seq_len32, vocab_size1000): data torch.randint(1, vocab_size, (batch_size, seq_len)) return data def train(): device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleTransformer().to(device) # 用 MALT 替换原来的 AdamW optimizer MALT(model.parameters(), lr1e-3, momentum0.9, beta0.999) criterion nn.CrossEntropyLoss() model.train() for step in range(100): src generate_batch().to(device) tgt src[:, 1:].contiguous() src src[:, :-1] logits model(src) loss criterion(logits.reshape(-1, logits.size(-1)), tgt.reshape(-1)) optimizer.zero_grad() loss.backward() optimizer.step() if step % 10 0: print(fstep {step}, loss {loss.item():.4f}) if __name__ __main__: train()运行这段代码你会看到 loss 在各 step 间平稳下降。如果 loss 不降或振荡严重优先检查学习率和 Newton-Schulz 迭代次数。6. 如何验证迁移是否成功替换优化器只是第一步真正重要的是验证收益。很多人在这一步翻车直接在完整大模型上跑成本高、周期长出了问题又很难定位。更稳妥的做法是先在中小规模任务上做对照实验。建议从三个维度验证第一个维度是收敛曲线。固定训练步数分别记录 AdamW、Muon、MALT 的 loss 曲线。重点看两个指标达到相同 loss 所需的步数、相同步数下的最终 loss。如果 MALT 在前 20% 步数内就能拉开差距说明它的优势是结构性的不是噪声造成的。第二个维度是最终指标。对于语言模型可以看验证集困惑度或下游任务准确率。这个维度最可靠但成本也最高。建议在模型已经能跑通的基础上下调学习率再做一次对比。第三个维度是训练稳定性。有些优化器前期 loss 下降快后期却发散或停滞。观察 loss 尾部波动、梯度范数变化、更新范数变化能判断优化器是否存在稳定性隐患。这里真正容易踩坑的地方是学习率。Muon 和 MALT 的更新尺度与 AdamW 完全不同直接沿用 AdamW 的学习率很可能导致震荡或发散。稳妥的做法是从比 AdamW 高一个量级的学习率开始做粗搜索再逐步精调。如果你在跑参数规模较大的模型建议先用小规模实验确定学习率区间再应用到正式训练。7. MALT 的适用场景与不适用场景结合 Muon 的经验和 MALT 的设计思路可以大致画出适用边界。MALT 最可能有优势的场景是权重矩阵较多的模型尤其是 Transformer 架构训练步长较长、需要精细调整方向的任务以及希望降低优化器状态内存开销的团队。由于 MALT 不需要保存二阶矩之外额外的大量状态相比某些二阶近似方法内存占用更可控。MALT 不太适合的场景也很明显。一个是对参数量很小、结构简单的模型例如逻辑回归或小型 MLP。这类模型几乎不会因为矩阵正交化而受益对角预条件用 Adam 就够了引入 Newton-Schulz 迭代反而是额外开销。另一个是 Embedding 层和 LayerNorm 这类非矩阵参数占主导的网络MALT 的优势同样难以发挥。另外如果你的训练流程已经针对 AdamW 做了大量超参数搜索替换优化器的收益未必能覆盖重新调参的成本。从工程角度看“当前最优方案”和“理论上更优方案”之间还需要考虑团队维护成本。8. 常见问题与排查思路问题现象可能原因排查方式解决方案loss 不降学习率过低或过高分别尝试 lr 提升/降低 10 倍做学习率粗搜确定合理区间loss 发散Newton-Schulz 迭代不稳定观察更新范数是否异常增大降低学习率或减少 ns_iter 到 3~5训练速度比 AdamW 慢正交化引入额外计算profile 前向与优化器 step 耗时减少 ns_iter或只在特定层启用 MALT显存峰值高于预期曲率估计状态与动量同时占用显存检查模型状态和优化器状态内存确认没有额外副本必要时混合精度训练最终指标不如 AdamW超参数不匹配对比相同步数下 loss 与下游指标精调 lr 和 beta确认任务是否适合 MALT向量参数表现异常对非矩阵参数错误应用正交化检查参数维度分支向量参数应退化为 Adam 风格更新排查时不要只看 loss。重点观察梯度范数和更新范数的比例关系。如果更新范数一直偏大可能是正交化或曲率缩放失效如果梯度范数在训练后期异常升高则需要检查模型结构和损失设计。9. 最佳实践与工程建议关于优化器选择我的核心建议是不要因为某个优化器在 benchmark 上表现好就盲目切换你的数据分布、模型结构、训练策略都会影响结果。从工程角度以下几点值得关注第一给优化器模块预留切换接口。把优化器初始化放在独立函数里用配置项控制 optimizer 类型。这样可以在不改动训练主逻辑的情况下快速做 A/B 对比。第二固化超参数搜索流程。每个优化器都有不同的最优学习率区间建议在换优化器时先跑一个 3~5 组的小规模 grid search确定基础学习率再做完整训练成本远低于直接踩坑。第三关注内存与吞吐。MALT 虽然是对角预条件但 Newton-Schulz 迭代需要额外的临时张量。在显存紧张时可以考虑梯度累积或混合精度也可以对模型中的部分层使用 MALT、其余层使用 AdamW作为一种折中方案。第四记录训练曲线的标准格式。统一记录 step、loss、token 数和吞吐方便跨优化器对比。不要只用 loss 曲线还要保存下游指标、更新范数、梯度范数这些是定位问题的关键线索。第五谨慎看待“论文效果”。论文中的结果通常基于特定模型、数据和调参力度复现时往往有差距。如果初步实验没有看到明显收益不一定代表方法无效也可能是超参数或任务特征不匹配。10. 后续可以深入的方向如果你对优化器方向感兴趣可以沿着三条线继续深入研究。第一条线是了解自然梯度和 Fisher 矩阵的更多细节。MALT 用对角预条件近似曲率而自然梯度理论给出了更完整的框架。理解 Fisher 矩阵的统计含义会让你对预条件方法有更本质的认识。第二条线是关注 K-FAC 和其他结构化近似方法。虽然实现复杂度高但它们在很多任务中的收益是对角方法无法达到的。理解 Kronecker 分解如何降低二阶信息的内存开销是进阶优化器研究的重要一环。第三条线是关注与你任务接近的训练实践报告。优化器领域论文很多但真正有参考价值的是大规模训练实验中的工程经验。看看其他人如何在超大模型上调整超参数、处理数值稳定性问题远比空谈优化器理论更有帮助。MALT 这个名字代表了一个清晰方向把矩阵结构先验和曲率信息结合起来同时保持工程上的轻量。它不一定是最终答案但它的设计思路——既不放弃结构信息也不忽视几何信息——很可能影响下一代优化器的发展方向。如果你正在做大规模模型训练值得花一个周末把它跑起来用自己的任务验证一下。 MALT优化器：轻量级曲率感知Muon，挑战AdamW的大模型训练新选择
在训练大模型时很多人的默认优化器是 AdamW。它足够稳定几乎成了“标配”。但最近几个月Muon 优化器打破了这种惯性Mistral 用它在部分场景中取得了比 AdamW 更好的收敛效果和最终指标也让越来越多的研究者开始重新思考大模型训练不一定非要 Adam矩阵结构信息可能比逐元素缩放更重要。不过Muon 也有一个明显短板——它虽然对梯度做了矩阵正交化处理但并没有显式利用损失曲面的曲率信息。简单说它对所有方向用了近乎一致的“旋转缩放”策略而真实损失面在不同方向上的弯曲程度差异很大。MALT 的出现恰好是想解决这个问题在不引入二次方内存开销的前提下通过对角预条件Diagonal Preconditioning给 Muon 增加曲率感知Curvature-Aware能力。这篇文章我会从 Muon 的核心思想讲起分析 MALT 为什么选择对角预条件给出可运行的 PyTorch 演示实现并讨论“这个优化器到底适合谁、迁移时需要验证什么”。如果你正在调大模型训练、对比优化器或者在思考“AdamW 之后下一个主流优化器长什么样”这篇文章值得读完并收藏。1. 为什么 Muon 能挑战 AdamW 的地位要理解 MALT先要理解它继承的对象——Muon。Muon 由 Mistral AI 的研究者提出核心思路非常直接把神经网络的权重矩阵视为一个线性变换而不是一堆“互不相干”的标量参数。基于这个视角它不再像 AdamW 那样对每个参数做逐元素归一化而是先对梯度做正交化处理再结合动量更新。这里的关键技术是 Newton-Schulz 迭代。它是一种快速逼近矩阵极分解的方法能够把梯度矩阵逐步“拉”向正交矩阵。每次迭代的复杂度可控通常只需要 5 到 10 次迭代就能得到足够好的正交化结果。Muon 在动量之后施加这个正交化步骤相当于给更新方向加了一个“形状约束”。对比一下两者的差异维度AdamWMuon对参数的处理方式逐元素独立归一化对权重矩阵整体做正交化主要数学工具一阶矩、二阶矩估计动量 Newton-Schulz 迭代内存开销需要额外保存两个动量变量额外保存一个动量变量整体略低适合场景参数规模大的通用训练权重矩阵结构强的模型如 Transformer理论视角对角自适应学习率利用矩阵谱结构的预条件从实践结果看Muon 在 Transformer 类模型上通常能显著减少训练步数或者在相同步数下获得更低的 loss。它的成功让很多人意识到AdamW 的逐元素缩放并不一定是最优的矩阵层面的结构约束可能更符合神经网络权重的本质。2. Muon 的局限正交化不等于曲率感知Muon 的效果令人兴奋但它并不是优化器的终点。一个容易忽略的问题是正交化本质上是一种“不考虑损失面几何信息”的预条件策略。损失函数的曲率决定了最优步长。在曲率大的方向步子应该小在曲率小的方向步子可以大。理想情况下优化器应该根据每个方向的曲率调整步长这正是二阶优化方法的目标。但完整的二阶方法例如自然梯度法需要维护 Fisher 信息矩阵代价是参数量平方级别的内存和计算大模型根本扛不住。Muon 的做法是对所有方向统一做正交化相当于只调整了更新方向的“形状”却没有针对不同方向的曲率差异做区分。可以类比成你在不同坡度的山路上都用同样大小的步幅走路方向走对了但步伐并不总是合适的。AdamW 至少通过二阶矩估计对每个参数方向做了缩放虽然是一个很粗糙的曲率近似但确实具备方向差异性。Muon 在这方面的信息反而比 AdamW 少。MALT 的切入点就在这里它希望在保留 Muon 矩阵结构优势的同时补充方向级别的曲率信息而且必须足够轻量不能把二阶方法的成本带回来。3. 为什么是对角预条件而不是更复杂的近似既然要引入曲率感知为什么不直接上 K-FAC 这类更精细的近似原因很现实复杂度。K-FAC 利用 Kronecker 分解来近似 Fisher 矩阵理论上能在保持一定精度的同时降低开销但实现复杂度非常高需要针对不同层结构做定制分布式场景下同步和更新也有额外成本。对于大多数研究团队和工程师来说这属于“知道但用不起来”的范畴。对角预条件是一个务实的折中。它假设参数方向之间没有相关性只估计每个方向的方差或曲率。虽然是一个粗糙的近似但它在实践中已经被证明有效——AdamW 的二阶矩估计本质上就是一种对角预条件。MALT 的思路是把这个机制对接进 Muon 的正交化框架中让正交化之后的更新方向再按对角曲率信息做缩放。从概率图模型和信息几何的角度看Fisher 矩阵对角线上的元素确实包含了各参数方向的局部敏感性信息。虽然忽略相关性会丢失一些信息但考虑到神经网络参数规模通常在千万到百亿级别保留相关性在计算上根本不现实。对角预条件在“信息量”和“可实现性”之间找到了一个合理的平衡点。4. MALT 的算法框架与数学直觉MALT 的全称是 Lightweight Curvature-Aware Muon via Diagonal Preconditioning。从命名就能看出它的三个关键设计决策轻量、曲率感知、基于 Muon。它的大致工作流程可以拆解为几步第一计算梯度并按照 Muon 的方式维护动量。动量项不仅提供惯性也是后续预条件的基础。第二维护一个对角曲率估计。最常用的方式是类似 Adam 的二阶矩滑动平均也就是对梯度的平方做指数加权平均。这个统计量可以看作是参数方向曲率的粗略代理代表“这个方向最近梯度的波动幅度”。第三对动量执行正交化处理得到结构化的更新方向。第四用对角曲率统计量缩放更新方向。直观理解是梯度变化剧烈的方向说明损失面弯曲厉害步长要保守梯度变化平缓的方向说明比较平坦步长可以激进一些。从数学直觉上看MALT 相当于在做“自然梯度”的轻量级近似。自然梯度通过 Fisher 矩阵逆来调整更新方向保证参数更新在分布空间中移动的是最短距离。而 MALT 只保留 Fisher 矩阵的对角部分再用正交化来捕捉矩阵整体结构——两部分叠加等于把“结构调整”和“方向缩放”解耦了。这个解耦很关键。它意味着 MALT 可以让正交化专注于改善梯度方向的几何形态而让对角预条件专注于步长调整两者各司其职不会互相干扰。5. 从 AdamW 迁移到 MALT 的 PyTorch 演示实现下面我用 PyTorch 实现一个演示版本的 MALT 优化器。需要说明的是这是一个为了讲清原理而简化的版本具体实现细节要以官方版本为准本文重点演示通用思路和关键步骤。# 文件路径malt_optimizer.py import torch from torch.optim import Optimizer def newton_schulz_(g, iterations5): 对梯度矩阵做 Newton-Schulz 正交化近似。 输入 g 的形状为 [..., M, N]输出近似正交的矩阵。 # 初始缩放帮助迭代稳定 with torch.no_grad(): a, b, c 1.0, 1.0, 1.0 for _ in range(iterations): g a * g b * torch.bmm(g, torch.bmm(g.transpose(-1, -2), g)) \ c * torch.bmm(g, torch.bmm(g, torch.bmm(g.transpose(-1, -2), torch.bmm(g, g)))) return g class MALT(Optimizer): 轻量级曲率感知 Muon 优化器演示版。 核心思想Muon 正交化 对角预条件。 参数: params: 模型参数 lr: 学习率 momentum: 动量系数 beta: 对角曲率估计的 EMA 系数 eps: 数值稳定项 ns_iter: Newton-Schulz 迭代次数 def __init__(self, params, lr0.01, momentum0.9, beta0.999, eps1e-8, ns_iter5): defaults dict(lrlr, momentummomentum, betabeta, epseps, ns_iterns_iter) super().__init__(params, defaults) def step(self, closureNone): loss None if closure is not None: loss closure() for group in self.param_groups: lr group[lr] momentum group[momentum] beta group[beta] eps group[eps] ns_iter group[ns_iter] for p in group[params]: if p.grad is None: continue grad p.grad.data if grad.dim() 2: # 对向量参数退化为 Adam 风格更新 state self.state[p] if momentum_buffer not in state: state[momentum_buffer] torch.zeros_like(p.data) state[variance] torch.zeros_like(p.data) buf state[momentum_buffer] var state[variance] buf.mul_(momentum).add_(grad) var.mul_(beta).addcmul_(grad, grad, value1 - beta) denom var.sqrt().add_(eps) p.data.addcdiv_(buf, denom, value-lr) else: state self.state[p] if momentum_buffer not in state: state[momentum_buffer] torch.zeros_like(p.data) state[variance] torch.zeros_like(p.data) buf state[momentum_buffer] var state[variance] # 1. 更新动量 buf.mul_(momentum).add_(grad) # 2. Newton-Schulz 正交化动量 orth_buf newton_schulz_(buf.clone().unsqueeze(0), ns_iter).squeeze(0) # 3. 更新对角曲率估计基于原始梯度平方 var.mul_(beta).addcmul_(grad, grad, value1 - beta) # 4. 对角预条件缩放 denom var.sqrt().add_(eps) scaled orth_buf / denom p.data.add_(scaled, alpha-lr) return loss这段代码有几个关键点第一对二维以上的权重矩阵使用完整的 MALT 逻辑动量、正交化、曲率缩放。第二对向量参数退化为 Adam 风格更新因为正交化主要针对矩阵设计向量参数强行正交化没有意义。第三曲率估计用的是梯度平方的滑动平均不是动量平方这样可以让曲率信息更贴近即时梯度变化。下面是一个完整的使用示例用一个小型 Transformer 训练任务来演示如何替换 AdamW# 文件路径train_with_malt.py import torch import torch.nn as nn from malt_optimizer import MALT class SimpleTransformer(nn.Module): def __init__(self, vocab_size1000, d_model128, nhead4, num_layers2): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.pos_encoding nn.Parameter(torch.randn(1, 128, d_model)) layer nn.TransformerEncoderLayer( d_modeld_model, nheadnhead, dim_feedforward512, batch_firstTrue ) self.encoder nn.TransformerEncoder(layer, num_layersnum_layers) self.output nn.Linear(d_model, vocab_size) def forward(self, x): seq_len x.shape[1] x self.embedding(x) self.pos_encoding[:, :seq_len, :] x self.encoder(x) return self.output(x) def generate_batch(batch_size8, seq_len32, vocab_size1000): data torch.randint(1, vocab_size, (batch_size, seq_len)) return data def train(): device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleTransformer().to(device) # 用 MALT 替换原来的 AdamW optimizer MALT(model.parameters(), lr1e-3, momentum0.9, beta0.999) criterion nn.CrossEntropyLoss() model.train() for step in range(100): src generate_batch().to(device) tgt src[:, 1:].contiguous() src src[:, :-1] logits model(src) loss criterion(logits.reshape(-1, logits.size(-1)), tgt.reshape(-1)) optimizer.zero_grad() loss.backward() optimizer.step() if step % 10 0: print(fstep {step}, loss {loss.item():.4f}) if __name__ __main__: train()运行这段代码你会看到 loss 在各 step 间平稳下降。如果 loss 不降或振荡严重优先检查学习率和 Newton-Schulz 迭代次数。6. 如何验证迁移是否成功替换优化器只是第一步真正重要的是验证收益。很多人在这一步翻车直接在完整大模型上跑成本高、周期长出了问题又很难定位。更稳妥的做法是先在中小规模任务上做对照实验。建议从三个维度验证第一个维度是收敛曲线。固定训练步数分别记录 AdamW、Muon、MALT 的 loss 曲线。重点看两个指标达到相同 loss 所需的步数、相同步数下的最终 loss。如果 MALT 在前 20% 步数内就能拉开差距说明它的优势是结构性的不是噪声造成的。第二个维度是最终指标。对于语言模型可以看验证集困惑度或下游任务准确率。这个维度最可靠但成本也最高。建议在模型已经能跑通的基础上下调学习率再做一次对比。第三个维度是训练稳定性。有些优化器前期 loss 下降快后期却发散或停滞。观察 loss 尾部波动、梯度范数变化、更新范数变化能判断优化器是否存在稳定性隐患。这里真正容易踩坑的地方是学习率。Muon 和 MALT 的更新尺度与 AdamW 完全不同直接沿用 AdamW 的学习率很可能导致震荡或发散。稳妥的做法是从比 AdamW 高一个量级的学习率开始做粗搜索再逐步精调。如果你在跑参数规模较大的模型建议先用小规模实验确定学习率区间再应用到正式训练。7. MALT 的适用场景与不适用场景结合 Muon 的经验和 MALT 的设计思路可以大致画出适用边界。MALT 最可能有优势的场景是权重矩阵较多的模型尤其是 Transformer 架构训练步长较长、需要精细调整方向的任务以及希望降低优化器状态内存开销的团队。由于 MALT 不需要保存二阶矩之外额外的大量状态相比某些二阶近似方法内存占用更可控。MALT 不太适合的场景也很明显。一个是对参数量很小、结构简单的模型例如逻辑回归或小型 MLP。这类模型几乎不会因为矩阵正交化而受益对角预条件用 Adam 就够了引入 Newton-Schulz 迭代反而是额外开销。另一个是 Embedding 层和 LayerNorm 这类非矩阵参数占主导的网络MALT 的优势同样难以发挥。另外如果你的训练流程已经针对 AdamW 做了大量超参数搜索替换优化器的收益未必能覆盖重新调参的成本。从工程角度看“当前最优方案”和“理论上更优方案”之间还需要考虑团队维护成本。8. 常见问题与排查思路问题现象可能原因排查方式解决方案loss 不降学习率过低或过高分别尝试 lr 提升/降低 10 倍做学习率粗搜确定合理区间loss 发散Newton-Schulz 迭代不稳定观察更新范数是否异常增大降低学习率或减少 ns_iter 到 3~5训练速度比 AdamW 慢正交化引入额外计算profile 前向与优化器 step 耗时减少 ns_iter或只在特定层启用 MALT显存峰值高于预期曲率估计状态与动量同时占用显存检查模型状态和优化器状态内存确认没有额外副本必要时混合精度训练最终指标不如 AdamW超参数不匹配对比相同步数下 loss 与下游指标精调 lr 和 beta确认任务是否适合 MALT向量参数表现异常对非矩阵参数错误应用正交化检查参数维度分支向量参数应退化为 Adam 风格更新排查时不要只看 loss。重点观察梯度范数和更新范数的比例关系。如果更新范数一直偏大可能是正交化或曲率缩放失效如果梯度范数在训练后期异常升高则需要检查模型结构和损失设计。9. 最佳实践与工程建议关于优化器选择我的核心建议是不要因为某个优化器在 benchmark 上表现好就盲目切换你的数据分布、模型结构、训练策略都会影响结果。从工程角度以下几点值得关注第一给优化器模块预留切换接口。把优化器初始化放在独立函数里用配置项控制 optimizer 类型。这样可以在不改动训练主逻辑的情况下快速做 A/B 对比。第二固化超参数搜索流程。每个优化器都有不同的最优学习率区间建议在换优化器时先跑一个 3~5 组的小规模 grid search确定基础学习率再做完整训练成本远低于直接踩坑。第三关注内存与吞吐。MALT 虽然是对角预条件但 Newton-Schulz 迭代需要额外的临时张量。在显存紧张时可以考虑梯度累积或混合精度也可以对模型中的部分层使用 MALT、其余层使用 AdamW作为一种折中方案。第四记录训练曲线的标准格式。统一记录 step、loss、token 数和吞吐方便跨优化器对比。不要只用 loss 曲线还要保存下游指标、更新范数、梯度范数这些是定位问题的关键线索。第五谨慎看待“论文效果”。论文中的结果通常基于特定模型、数据和调参力度复现时往往有差距。如果初步实验没有看到明显收益不一定代表方法无效也可能是超参数或任务特征不匹配。10. 后续可以深入的方向如果你对优化器方向感兴趣可以沿着三条线继续深入研究。第一条线是了解自然梯度和 Fisher 矩阵的更多细节。MALT 用对角预条件近似曲率而自然梯度理论给出了更完整的框架。理解 Fisher 矩阵的统计含义会让你对预条件方法有更本质的认识。第二条线是关注 K-FAC 和其他结构化近似方法。虽然实现复杂度高但它们在很多任务中的收益是对角方法无法达到的。理解 Kronecker 分解如何降低二阶信息的内存开销是进阶优化器研究的重要一环。第三条线是关注与你任务接近的训练实践报告。优化器领域论文很多但真正有参考价值的是大规模训练实验中的工程经验。看看其他人如何在超大模型上调整超参数、处理数值稳定性问题远比空谈优化器理论更有帮助。MALT 这个名字代表了一个清晰方向把矩阵结构先验和曲率信息结合起来同时保持工程上的轻量。它不一定是最终答案但它的设计思路——既不放弃结构信息也不忽视几何信息——很可能影响下一代优化器的发展方向。如果你正在做大规模模型训练值得花一个周末把它跑起来用自己的任务验证一下。