跳到主内容
快讯直播
AI智模界
AI 词典

Muon 优化器:把动量矩阵掰成正交方阵

一句话定义

Muon(MomentUm Orthogonalized by Newton–Schulz)是一种专门为神经网络里二维权重矩阵设计的优化器:它不逐个元素地更新参数,而是把累积出来的动量矩阵先"正交化",再拿这个规整过的结果当作更新方向。

它到底做了什么

先看 AdamW。它给每个参数单独维护一阶矩和二阶矩,然后逐元素地归一化:某个参数梯度一直很大,就压一压;某个参数梯度很小,就抬一抬。粒度是"每个数字管自己"。

Muon 换了个视角。它把一整层权重看成一个矩阵,把梯度累积成动量矩阵 M,然后问一个问题:如果只保留 M 的"方向",把各个奇异方向的强度拉平,会得到什么?答案就是 M 的极分解里那个正交因子,记作 UVᵀ。

问题是不显式做奇异值分解(SVD)太贵了。Muon 用 Newton–Schulz 迭代来近似:反复做几次矩阵乘法,就能让 M 一步步逼近那个半正交矩阵。次数固定、只用到 matmul,所以开销可控。

打个比方:AdamW 像调音台上每一路都按自己的音量单独旋钮,谁的信号强就压谁;Muon 则是先把整支乐队站成一个方阵——每个人步幅、朝向统一,再一起往前走。歪的队伍被"掰正"之后,每个方向的步长都一样,少数特别大的元素不会再主导整步。

和相邻概念的区别

维度AdamWMuon
更新粒度逐元素整个矩阵
优化器状态一阶矩 + 二阶矩只有动量
每步主要开销逐元素运算(吃内存带宽)几次矩阵乘法(吃算力)
典型学习率偏小、敏感常可放大
常见用法全参数通用隐藏层二维权重,其余仍配 AdamW

和 Shampoo 这类二阶方法比,Muon 不构造、也不开方二阶矩阵,省得多;和带动量的 SGD 比,它多做了一步"掰正",正是这步换来更稳的收敛。

对从业者的实际意义

第一,优化器状态少了一半——没有二阶矩,显存和通信更省。第二,在中型模型上,它常允许更大的学习率,同一份算力下损失下降更快,或者说达到同样效果需要的步数更少,这才是"省算力"的真正来源:不是每步更便宜,而是步数更少。第三,它的每步开销并不免费,矩阵越大,那几次 Newton–Schulz 迭代越重;在超大规模下收益是否成立仍有讨论。

工程上的常见搭配是:二维隐藏权重交给 Muon,嵌入层、输出层、偏置和归一化参数仍交给 AdamW。学习率要按矩阵形状调,权重衰减策略也和平常不同。具体实现细节和最新结论,以原始论文与官方代码库为准。

AI 生成本文由 AI 基于公开信息自动生成,仅供参考。