Tilde Research 发布了 Aurora,一个新的「leverage-aware」优化器,修复了 Muon 中一个隐藏的 bug —— 该优化器一直在悄悄地为 modded-nanoGPT speedruns 和越来越多的前沿训练管道提供动力。Bug:在像 MLP 层这样的高矩阵中,Muon 的极因子更新创建行范数各向异性,导致一些神经元接收大规模更新,而其他神经元几乎不接收任何更新。在第 500 个训练步骤时,超过四分之一的神经元实际上已死。Aurora 以大约 6% 的额外计算开销防止这种情况,作为 Muon 的替代品插入,并附带 GitHub 上的开放代码。
诊断比修复更重要。NorMuon(一个先前的中间版本)通过事后归一化到单位范数来纠正行范数各向异性,获得了良好的经验结果,但没有解释为什么存在底层问题。Aurora 的分析:Muon 的极因子更新对方形矩阵做正确的事情,对高矩阵做错误的事情 —— 而「高」描述了大多数具有大扩展因子的 MLP,所以这个 bug 在每个人正在训练的架构中加剧。Aurora 将权重更新重新表述为联合约束:左半正交性和均匀行范数,同时解决而不是事后修补。两种实现交付:黎曼 Aurora(在受约束流形上的梯度投影)和 vanilla Aurora(更简单的实际变体)。Tilde 报告在 1.1B 规模上的开源互联网数据上 100× 数据效率,在 modded-nanoGPT speedrun 上的新 state-of-the-art(超过 NorMuon 之前的 SOTA),并在 HellaSwag 上超过更大的模型。100× 声明需要 harness 公开,然后才能被视为福音 —— 那是一个代际结果,不是渐进的 —— 但 speedrun SOTA 是更可验证的点,因为它有每个人都可以比较的公共参考数字。
Muon 自 2024 年末以来一直在获得采用,作为比 AdamW 更高效的计算替代方案,特别是对于 nanoGPT 风格的 speedruns 和越来越多地用于生产前沿训练运行。Aurora 的诊断意味着每个从 Muon 获得良好结果的人都在第 500 步时悄悄地损失了大约四分之一的 MLP 容量给死神经元 —— 推测训练后期会更多。NorMuon 已经是人们感觉到有什么不对但没有解释的标志。更广泛的模式:优化器研究有一个安静的十年,AdamW 被视为已解决,而最近的浪潮(Lion、Sophia、Muon、NorMuon,现在是 Aurora)正在重新打开这个问题。drop-in 替代框架和 6% 计算开销是使 Aurora 真正可采用而不是研究好奇心的部分 —— 如果它干净地移植到现有训练管道,从 Muon 切换的门槛很低。死神经元数字也是任何人添加到训练运行仪表板的有用诊断,无论他们最终选择哪个优化器。
GitHub 代码在 `tilde-research/aurora-release`,论文在 blog.tilderesearch.com。如果你在 100M 参数规模以上训练 transformers 并使用 Muon,在相信 100× 头条数字之前,Aurora 值得在你的具体工作负载上进行受控 A/B 测试。神经元死亡框架是应该让任何在 Muon 上运行生产训练的人担心的部分 —— 你可能一直在损失你不知道自己在损失的容量。对其他所有人来说,优化器研究浪潮继续表明「训练稳定性」和「训练效率」背后仍有大量未解决的问题,以及解决它们的实验室对堆栈其余部分获得不成比例的杠杆。