Thinking Machines 为 Tinker 添加专用于 LoRA 的优化器 Dimuon

10月，Thinking Machines 在 Tinker 中加入了 Dimuon，这是一个专为 LoRA 构建的优化器。除了在更高 rank 下相比 AdamW 表现出的良好 scaling 曲线，以及“Dimuon”这个名字暗示它是 Muon 的变体之外，他们没有分享任何关于该优化器工作原理的细节。几个月前，PoLoRA（Ghosh、Parshakova 和 Gower）为 LoRA 因子提出了一种预条件、正交化的优化器，再往前还有 LoRA-RITE、LoRA+ 和 Riemannion。卖点总是相似的：LoRA adapter 是乘积 $BA$，而不是两个独立的矩阵，因此把 A 和 B 当作扁平参数向量（AdamW）或两个不相关的矩阵（Muon）来处理的优化器，总会漏掉一些东西。

那么，LoRA 专用优化器真的能击败调好的 Muon 吗？（注意 Dimuon 从未与 Muon 做过对比。）如果能，算法的哪一部分起了作用？

手头有几张 5090、时间有限，我尝试在以下设置上复现类似的研究：

AdamW：逐元素自适应矩，把 A 和 B 当作扁平向量。 Muon：Moonshot 变体，动量按 $0.2\sqrt{\max(m,n)}$ 缩放，以匹配 AdamW 的更新 RMS。 PoLoRA：对每个因子做对角曲率预条件（平方梯度的 EMA），用伙伴因子的 Gram 矩阵做白化，使 A 和 B 耦合，极分解步，步长为 $\text{lr}/(\sigma_A+\sigma_B)$，同时约束两个因子的更新和 $BA$ 的变化量。 anpaure 的 Dimuon 变体：我依赖 X:@anpaure 逆向工程出的 Dimuon 实现 dimuon.py。

与 PoLoRA 设计相比，我注意到的第一件事是缺少像 PoLoRA 中 $\sigma_A, \sigma_B$ 那样的 A-B 分解。另外，在 anpaure 的 Dimuon 变体中，我们依赖一个魔法更新尺寸：$\sqrt{rank}x (rank / 32) ** 0.175$，这是根据比较 anpaure 的 Dimuon 变体和 adamw 得到的梯度步拟合出来的。

一件令人遗憾的事是，我们完全不知道“Di”到底指什么。

每次 1,000 步运行的最小 eval loss，按学习率。灰色菱形标记每个 rank 下的最佳学习率。孤立的 Muon rank-256 点是一次在第 300 步停止的单独运行，不是 sweep。

anpaure 的 Dimuon 变体在每个 rank 上都获胜，且领先幅度随 rank 增长。与调好的 AdamW 相比，它在 rank 16 领先 6e-4（处于边界），rank 64 领先 1.2e-3，rank 256 领先 1.7e-3。PoLoRA 从另一侧呈现相同趋势：rank 16 落后 AdamW 5e-4，rank 256 领先 1.1e-3。无论这些优化器做了什么，当 adapter 有更多可用的方向时，它都更有价值。

在 rank 16 时，Muon 对 LoRA 相比 AdamW 没有任何优势（0.4786 vs 0.4785），rank 64 时也只有处于边界的 4e-4。

如果你只是想看看是否该换优化器，可以在这里停下，开始使用 anpaure 的 Dimuon 变体。

这是我写的 bitsandbytes 集成，支持 anpaure 的 Dimuon 变体（QLoRA、FSDPv2）：theblackcat102/bitsandbytes/tree/feature/muon

为了回答“哪一部分重要”，我写了一个优化器（UnifiedLoRA），为每个设计轴都设有开关，并检查了它的预设能复现原始版本：anpaure 的 Dimuon 变体逐位一致，Muon 和 PoLoRA 在六步内相对误差在 5e-6 以内。这些轴是：

由于资源限制，所有消融都只在 rank 16 下进行。

取普通 Muon，只加入 anpaure 的 Dimuon 变体的特征侧条件器，并把更新范数保持在 Muon 自身的水平，这样只有方向改变。在第 500 步：

这就在没有 anpaure 的 Dimuon 变体的步长规则、动量和 rank 侧条件的情况下，把 anpaure 的 Dimuon 变体复现到了 1e-4。镜像测试结果一致：把 anpaure 的 Dimuon 变体的条件器替换为单位矩阵，在第 100–300 步会损失 +0.003。

两者在各自学习率下都在 3e-4 以内。我还尝试用激活二阶矩来对 A 做条件，这是教科书式的 K-FAC 输入度量，从第一步起就精确，而不是噪声动量估计。结果打平：第 500 步时 0.48160 对 anpaure 的 Dimuon 变体的 0.48168。

随着实验进一步扩大，在第 1000 步时它领先 1.2e-4（0.47845 vs 0.47857），这是噪声。在单 epoch 余弦 schedule 下，衰减在最后完成主要工作；几何决定你多快到达那里。此前我们发现，“Muon + anpaure 的 Dimuon 变体的条件器”在第 500 步时领先 Muon 1.5e-3，大约是噪声底的 15 倍。

条件化也有前期成本。在第 50 步，保持更新范数固定时，条件化运行差 0.025（anpaure 的 Dimuon 变体的步长）到 0.052（Muon 的步长），到第 100 步成本消失。这也不是估计问题：干净的激活统计量在第 50 步的成本（+0.038）比噪声动量估计（+0.025）更大。

Muon 的宽度缩放。Muon 将每个更新乘以 (0.2\sqrt{\max(m,n)})，对于 LoRA 因子来说，这意味着按模型宽度缩放，而不是按 rank 缩放。对于一个 16 行的矩阵来说这似乎是错的。但移除它在 rank 64 时与 Muon 的差距在 4e-4 以内，且完全没有恢复 anpaure 的 Dimuon 变体在第 500 步时 3.5e-3 的领先。

极分解精度与弱方向。这是一个真实存在但无后果的现象。动量的稳定 rank 在每个 adapter rank 下都约为 3，因此随着 rank 增长，大多数方向都很弱：rank 16 时 5.5% 的方向低于 0.1·σ_max，rank 64 时 40%，rank 256 时 84%。我用分半探测检查了这些方向是噪声（两个不相交的数据半集在那里几乎不一致）。极分解步会把它们放大到与强方向相同的大小，听起来很糟。但丢弃它们只让第 500 步 loss 变化 ≤2e-4，而精确 SVD 极分解与 Muon 的 5 步 bf16 近似之间差异 ≤3e-4。

PoLoRA 的 A/B 耦合。把它加入 anpaure 的 Dimuon 变体只让结果变化 3e-5。原因是 PoLoRA 用来白化 A 更新的 B 的 Gram 矩阵条件良好（条件数约 1.4），所以用它的逆平方根白化与乘以一个标量的差异在 5% 以内，而步长归一化无论如何会抵消标量。耦合对于强各向异性的因子会有影响；从零开始的 B 并不是。

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）