Adam的稳+Muon的快?华为诺亚开源ROOT破解大模型训练「既要又要」的两难困境

在 LLM 优化领域,有两个响亮的名字:Adam(及其变体 AdamW)和 Muon。

它们一个是久经沙场的「守门员」,凭借动量和自适应学习率统治了深度学习的半壁江山,却在面对十亿级参数的混合精度训练时,常常因数值不稳定性而显得力不从心;一个是横空出世的「破局者」,试图通过将权重矩阵视为整体来重塑训练几何,却因为一刀切(one-size-fits-all approach)的系数设计和对异常值噪声的极度敏感,在鲁棒性上留下了缺口。

当训练规模不断指数级膨胀,我们是否只能在 Adam 的「稳」与 Muon 的「快」之间做单选题?

华为诺亚方舟实验室的最新力作 ROOT (Robust Orthogonalized OpTimizer) 给出了否定的答案。

作为一款直击痛点的鲁棒正交化优化器,ROOT 不仅精准修复了 Muon 在不同矩阵维度上的「精度近视」,更通过巧妙的软阈值机制为梯度噪声装上了「减震器」。它正试图用更快的收敛速度和更强的稳定性,为大模型训练建立一套全新的、兼顾精确与稳健的优化范式。

  • 论文标题:ROOT: Robust Orthogonalized Optimizer for Neural Network Training

  • 论文地址:arxiv.org/abs/2511.20...

  • 开源地址:github.com/huawei-noah...

  • 作者:Wei He, Kai Han, Hang Zhou, Hanting Chen, Zhicheng Liu, Xinghao Chen, Yunhe Wang

  • 机构:华为诺亚方舟实验室

LLM 优化史:从 SGD 到 ROOT

要理解这项工作的重要性,我们需要先了解优化器(Optimizer)在 LLM 训练过程中至关重要的地位。

简单打个比方:在深度学习的浩瀚宇宙中,优化器扮演着飞船「引擎」的角色。

其中,最早的优化器是 SGD(Stochastic Gradient Descent),即随机梯度下降。作为深度学习的基石,它确立了神经网络训练的基本范式:通过计算小批量数据的梯度来迭代更新参数。

SGD 是最经典的一阶优化方法。然而,在面对高维且复杂的损失函数曲面(Loss Landscapes)时,原始的 SGD 往往难以兼顾收敛速度与稳定性。为了帮助模型更高效地穿越复杂的「山谷」找到极小值,研究者们在 SGD 的基础上引入了动量机制,这不仅成为了 SGD 的标准配置,也为后来更复杂的自适应方法奠定了基础。

后来,以 Adam 和 AdamW 为代表的自适应方法崛起,成为训练深度学习模型的「事实标准」。

它们通过引入动量和逐参数(Per-parameter)的自适应学习率,让收敛效率大幅超越 SGD。然而,这类方法的底层逻辑是将模型参数视为独立的「标量」或向量进行更新。当模型参数量突破十亿大关,这种忽略参数矩阵内部结构相关性的处理方式,在混合精度训练中逐渐暴露出了数值不稳定的缺点。

为了突破这一瓶颈,以 Muon 为代表的矩阵感知型优化器应运而生。

Muon 不再仅仅盯着单个参数,而是将权重矩阵视为一个整体。它利用 Newton-Schulz 迭代对动量矩阵进行正交化处理,从而在不增加额外计算复杂度(保持 O (N))的前提下,规范了更新的几何结构。

这种方法在理论上等同于在谱范数下进行最速下降,显著提升了训练效率和显存利用率。

尽管 Muon 开启了新的一页,但研究人员发现它并非完美无缺。

华为诺亚方舟实验室的分析指出,现有的正交化优化器存在两个核心局限:

  • 算法鲁棒性的缺失: 现有的 Newton-Schulz 迭代通常使用一组固定的系数。然而,神经网络不同层的权重矩阵形状各异(从正方形到极度扁平的矩形),固定系数在某些维度下会导致近似误差激增,产生「维度脆弱性」。

  • 对梯度噪声的缺乏防御:在大规模训练中,异常数据往往会产生极大幅度的梯度噪声。现有的自适应优化器对这些噪声异常敏感,不仅会破坏更新方向,还可能导致训练彻底失稳。

随机梯度中异常值噪声的概念可视化。大多数梯度值集中在中心附近,存在一个高幅度异常值的尾部。这些异常值会不成比例地影响优化过程。

正是在这种既要「矩阵感知的快」又要「传统方法的稳」的博弈中,ROOT 应运而生,试图填补这一关键的拼图空缺。

ROOT 优化器:双管齐下

前文我们已经介绍过,现有的正交化优化器(尤其是 Muon)存在的核心缺陷。

ROOT(Robust Orthogonalized OpTimizer)的核心方法,是为正交化优化器做出了针对性的鲁棒性增强,让优化器在快速和稳定「两手抓」。

拒绝「一刀切」

正交化优化器的算法不稳定,核心问题源于正交化系数的「一刀切」。

具体来说,Muon 里 Newton-Schulz 迭代的系数 a、b、c 是固定常数。华为诺亚方舟的研究者们发现,这会引发不同维度矩阵的脆弱性。

正交化误差揭示了固定系数 Newton-Schulz 迭代在维度上的脆弱性。

从上表中的数据能看出,矩阵形状(维度或长宽比)一变,正交化误差会大幅波动。尤其是方阵更吃亏,方阵始终产生最高的 MSE 值,比非方阵配置有显著的差距。

这种维度敏感性在优化过程中造成了固有的脆弱性,因为不同维度的层获得的正交化质量完全不同,损害了梯度更新的一致性和可靠性。

为了解决这种维度脆弱性并构建维度鲁棒的正交化过程,研究者们提出了具有细粒度、特定维度系数的自适应 Newton-Schulz 迭代(AdaNewton)。

ROOT 没有为所有矩阵维度使用全局常数 (a, b, c),而是采用针对网络中每个特定矩阵大小 (m, n) 量身定制的细粒度系数

迭代更新规则在结构上保持相似:

然而,系数针对每个矩阵维度的特定奇异值分布进行了优化。这种方法为提高正交化精度提供了理论保证。

这些系数可以在训练期间与模型参数联合优化,允许正交化过程自动适应每种层类型的属性。这种细粒度的适应代表了一种范式转变:从脆弱的维度敏感正交化转向鲁棒的维度不变正交化,确保了整个网络的更新质量一致。

过滤「异常值」

大模型训练的梯度常出现「重尾现象」:小批量梯度经常被异常值噪声污染,这些噪声包含幅度异常大的梯度分量,这些异常值严重影响到了 Muon 中正交化过程的稳定性。

更糟的是,Newton-Schulz 迭代的多项式性质会放大离群噪声,造成不稳定,甚至可能引发 Transformer 的 attention logits 爆炸的严重问题。

为了解决这一问题,华为诺亚方舟的研究者们的做法很直接干脆:把梯度 Mₜ 分解为「正常部分」和「异常部分」两个分量:

  • 基础分量 Bₜ:包含表现良好的梯度信息。

  • 异常分量 Oₜ:代表异常的大幅度元素。

正交化仅应用于鲁棒分量 Bₜ,而丢弃异常值分量 Oₜ。

通过这个目标函数,可以显式地惩罚大异常值分量的存在,同时确保基础分量保持在合理范围内。

该鲁棒分解的最优解由 L1 范数的近端算子给出,因此 ROOT 采用软阈值算子来提取异常值:

这个函数如果值的幅度高于阈值 ε,则提取超出范围的异常值。

在数学上,软阈值可以被解释为硬裁剪(hard clipping)的一种连续、可微的替代方案。软阈值应用了一种平滑的收缩操作,在抑制极端值的同时保留了梯度幅度的相对排序。

完整的 ROOT 优化器算法

ROOT 的实验表现:真的又稳又快

为了验证 ROOT 是不是真的快速又稳定,华为诺亚方舟实验室训练了一个 1B Transformer 模型。他们的测试非常严苛,涵盖了从预训练 Loss 到下游任务的多项基准,甚至跨越到了视觉任务领域。值得注意的是:「所有模型都是在昇腾 NPU 分布式集群上训练的。」

而最终得到的结果也非常亮眼,证明了 ROOT 优化过程的表现极具竞争力。

首先,在预训练效率上,ROOT 展现了卓越的收敛能力。

使用 10B Token 的训练损失对比

如上图所示,在 10B token 的大规模预训练实验中,两个 ROOT 变体(仅软阈值版与完整版)的训练损失均始终保持在 Muon 的 Loss 曲线下方。最终,ROOT 的训练损失达到 2.5407,比 Muon 基线低 0.01。

而更深入分析显示,Muon 由于采用固定系数,在训练过程中存在较大的近似误差;而 ROOT 凭借自适应系数,始终保持着更接近真实 SVD 的正交化精度 。

相对于真实 SVD 的正交化精度

在多项下游任务基准上,ROOT 也带来了全面的提升:ROOT 取得了 60.12 的平均分,不仅击败了传统霸主 AdamW(59.05),也超越了其直接竞争对手 Muon(59.59)。

在 9 个标准 LLM 基准上的零样本性能,其中 ROOT 在 6 个基准上领先

同时也能看出 ROOT 具有广泛适用性:无论是在考察常识推理的 PIQA,还是考察科学知识的 SciQ,ROOT 都展现出了极具竞争力的性能。

不仅如此,ROOT 还表现出了非常出色的跨模态泛化能力:在计算机视觉领域(训练 ViT 模型识别 CIFAR-10 数据集)的测试中,ROOT 同样证明了其强大的泛化能力。

在 CIFAR-10 上的 Top-1 测试准确度

特别是在引入软阈值机制后,ROOT 能够有效抑制视觉数据中的梯度噪声,取得了 88.44% 的 Top-1 准确率,显著优于 Muon 的 84.67% 。这表明 ROOT 的「去噪+正交化」范式具有极强的跨领域普适性。

该团队也进行了消融实验,证明了 ROOT 各组件的有效性。

ROOT 或将开启新的优化器时代

在 LLM 训练日益昂贵且复杂的今天,华为诺亚方舟实验室提出的 ROOT 优化器,通过 AdaNewton 和软阈值去噪两大创新,成功在 Muon 的高效基础上补齐了鲁棒性这块短板 。

ROOT 不仅在理论上保证了不同维度矩阵更新的一致性,更在实战中证明了其在抗噪、收敛速度和最终性能上的全面优越性。

ROOT 的代码将会开源,随着更多研究者将其投入到更大规模的万亿级模型训练中,我们有理由相信,它很有可能会开启一个新的优化器时代。

正如这篇论文的结语所言:「这项工作为开发鲁棒的优化框架开辟了有前景的方向,这些框架能够处理未来语言模型日益增加的复杂性和规模,从而可能实现下一代 AI 系统更可靠、更高效的训练。」

凭借此一贡献,华为诺亚方舟实验室展示了其「深潜」的创新特质,秉持理论研究与应用创新并重的理念,致力于推动人工智能领域的技术创新和发展:不随波逐流于表层的应用创新,而是潜入深海,解决最基础、最困难、但影响最深远的优化理论问题。这不仅展示了其强大的科研硬实力,更体现了其作为行业领军者,致力于构建更高效、更鲁棒的下一代 AI 训练范式的战略远见。

团队简介

本文有两位共一作者,他们都是华为诺亚方舟实验室研究员。据公开资料显示,其中韩凯(Kai Han)现为华为诺亚方舟实验室专家研究员,博士毕业于中国科学院软件所,硕士和本科分别毕业于北京大学和浙江大学。其主要研究方向为高效深度学习和 AI 基础模型,已在 AI 领域顶会顶刊发表论文 50 余篇,谷歌学术累计被引 2.1 万余次,其中 GhostNet 和 TNT 入围 PaperDigest 年度最具影响力论文榜单。他还担任 NeurIPS、ICML、ICLR、CVPR、ICCV、AAAI 和 ACMMM 等顶会领域主席,入围斯坦福全球 Top 2% 科学家和爱思唯尔中国高被引学者榜单。

另外,今年 3 月接任华为诺亚方舟实验室主任的王云鹤也是本文的通讯作者。

关于该方法更多信息,请参阅原论文。

相关推荐
可观测性用观测云2 小时前
观测云 MCP Server 接入和使用最佳实践
人工智能
掘金一周2 小时前
大部分人都错了!这才是chrome插件多脚本通信的正确姿势 | 掘金一周 11.27
前端·人工智能·后端
xier_ran2 小时前
深度学习:生成对抗网络(GAN)详解
人工智能·深度学习·机器学习·gan
ModestCoder_2 小时前
ROS Bag与导航数据集技术指南
开发语言·人工智能·自然语言处理·机器人·具身智能
海边夕阳20062 小时前
【每天一个AI小知识】:什么是循环神经网络?
人工智能·经验分享·rnn·深度学习·神经网络·机器学习
2501_918126913 小时前
如何用ai做开发
人工智能
f***a3463 小时前
开源模型应用落地-工具使用篇-Spring AI-高阶用法(九)
人工智能·spring·开源
用户5191495848453 小时前
BBDown:高效便捷的哔哩哔哩视频下载工具
人工智能·aigc
CV实验室3 小时前
CV论文速递:覆盖视频生成与理解、3D视觉与运动迁移、多模态与跨模态智能、专用场景视觉技术等方向 (11.17-11.21)
人工智能·计算机视觉·3d·论文·音视频·视频生成