在学习人工智能时,最难的主题之一可能是用于训练它们的优化算法,因为它基本上处于微积分、几何学和统计学的交汇点。这就是为什么人工智能的先修课程看起来如此可怕的原因所在。你要有矩阵运算、多元微积分、概率论,等等,不胜枚举。然而,训练人工智能的想法听起来却愚蠢地简单:你运行模型,衡量它的错误程度,利用这些误差来更新权重,然后重复这个过程。

当然,听起来简单的部分通常是困难所在的地方,因为“利用箭头更新权重”这句话实际上承担了极其繁重的工作。为了快速分解其工作原理,首先你需要通过反向传播来确定每个权重对误差的贡献程度。这涉及所有微积分内容,从而给你一个称为梯度的东西,它告诉你哪个方向会使误差最陡峭地增加。因此,从逻辑上讲,你可以朝相反的方向移动权重以减少误差。

然而,它并没有告诉你需要在那个方向走多远才能减少误差。你可以迈出一小步,几乎不减少任何误差;或者跨度过大,反而增加误差。你可以想象在黑暗中下山徒步旅行。因此,当前的坡度告诉你哪个方向是下坡路,但它并没有告诉你是该迈出一英寸还是十米。你可能身处山谷中并意外地再次走上坡路,对吧?

因此,这个决策是优化器的任务。它接收来自反向传播的梯度并将其转化为实际的权重更新。至于如何应用这些权重更新,这很大程度上取决于你选择的优化算法,其中最流行的是 Adam(自适应矩估计)的缩写。该算法本身看起来相当可怕,但核心理念却非常简单。Adam 基本上试图同时解决两个问题:首先,梯度是有噪声的,因为出于效率原因,我们不会只在一个数据样本或整个数据集上一次性进行训练。

相反,我们是以小批量(mini-batches)进行的。因此,Adam 保持过去梯度的移动平均值,这给了优化器一种记忆,让它知道它曾经远离过哪些方向。其次,不同的权重表现各异,有些权重的梯度始终很大,而另一些则几乎不动。因此,Adam 还保持平方梯度的移动平均值,以此告诉它每个参数在历史上有多活跃。

随后出现了 AdamW,这基本上就是 Adam 加上一个虽小但重要的修正,即如何应用权重衰减(weight decay)。权重衰减背后的理念很简单:在训练过程中,我们稍微抑制模型过度依赖非常大的权重,因为较大的权重可能使模型过于激进地拟合训练数据,几乎像是在记忆局部模式而不是学习更广泛的模式。

AdamW 中的重要区别在于,这种惩罚是独立于 Adam 通常的梯度更新来应用的。虽然这听起来可能微不足道,但对于像 Adam 这样的自适应优化器来说,它使权重衰减表现得更加干净和可预测。这就是为什么 AdamW 变得如此受欢迎的原因,因为它稳定、自适应,并且与其他方法相比通常更容易训练。

但在此时,我有一种感觉,并没有多少人真正跟上我的解释,因为我刚才说的肯定听起来像是一堆术语。这也可能是为什么当 Frontier Open Research 决定转向 Muon 时,我们几乎没有任何反应,尽管其影响可能像新的架构设计一样巨大。更不用说,这是第一种成功挑战了长达 12 年现状(即 AdamW)的新方法。

因此,在今天的视频中,让我用尽可能少的数学术语为你拆解这个名为 Muon 的新优化器,因为它具有大多数人都未能理解的惊人未来前景。但在深入探讨之前,我用今天的赞助商 Solid 运行了一个小实验。Solid 为你提供拥有自己的设备、账户和预算的 AI 智能体(AI agents)。你可以给它们分配任务,它们会完成任务而无需你不断回来监督。

Solid 的酷之处在于,你可以给它一个大得多的目标,让它自行处理所有设置和执行工作。诸如构建整个应用程序、创建 3D 模型或运行机器学习训练实验等事情都是如此。所以我基本上告诉 Solid,使用两种不同的优化器 AdamW 和 Muon 训练两个小型基于 Transformer 的语言模型,并构建一个原生 Mac 应用程序来比较结果。

正如我所请求的那样,它为我构建了这样一个页面,展示了它的所做所为:包括进行消融实验、构建应用程序,总共耗时超过一小时。这是它构建的应用程序。它是一个真正的 Mac OS 应用。正如你所见,它甚至制作了结果的交互式视图。不仅如此,它是完全自主完成的,没有访问我的计算机。

那么这是如何工作的呢?幕后运作是 Solid 在其云中运行各种电脑和手机。为了运行我们的优化器实验,Solid 自动设置了一台 Mac,安装了所需的工具,运行了实验并构建了原生应用,这简直太疯狂了。你可以看到 Solid 如何为智能体提供基础设施以支持各种行动。那基本上就是我的实验。

所以,如果你有自己的项目并希望 Solid 接手,请现在通过描述中的链接查看它们。感谢 Solid 赞助本期视频。无论如何,Muon 最初由独立研究员 Keller Jordan 在 2024 年 12 月提出。Muon 代表的是“由牛顿舒尔茨正交化的动量”(Momentum Orthogonalized by Newton Schultz)。一个听起来相当可怕的名字,但我断言其理念实际上非常简单。

首先,让我们回到 Adam 一会儿。从数学上讲,Adam 将模型视为一个巨大的独立数字列表。我的意思是,即使 Adam 能够优雅地处理权重的梯度历史,它仍然将每个权重视为相互独立的。这种方法完全没问题,直到你意识到神经网络中的这些计算处理的不仅仅是单个数字,而是矩阵。

整个 Transformer 基本上就是由矩阵乘法构建而成的。注意力机制中的 QKV(查询、键、值)投影都是矩阵,MLP 层也是矩阵。这里的矩阵不仅仅是一堆随机独立的数字集合。它是一种变换(transformation)。它接收一些输入并将它们混合在一起,产生一组新的特征。

因此,矩阵内部的数字通过它们创建的变换而相互连接,因为它们共同决定了信息如何被旋转、拉伸、压缩或混合。

而不是独立行为。因此,矩阵内部的数字实际上是相关的,这正是 Adam 风格优化器所忽视的地方。想象一下,有100个人在黑暗中试图下山。虽然 Adam 给每个人单独发出指令,但它仍然可能忽略大局。比如,如果这100人中有90人都朝着几乎完全相同的方向前进呢?虽然你名义上有100个人,但你并没有真正最优地利用这100个独立源,因为大家都在相似的方向上重叠,浪费了大量的容量。

因此,Muon 的不同之处在于它对更新进行了正交化处理,试图让矩阵更新使用更多独立的方向,以防止少数方向主导一切。更具体地说,矩阵更新可以分为两部分:方向和强度。方向告诉你更新试图移动信息的位置;强度则告诉你要主导每个方向的力度有多大。如果一个方向的强度巨大,那么大部分更新都花在那里了;如果另一个方向强度微小,那么这个方向就几乎没被使用。

所以,即使你的矩阵更新在数值上很大,实际的学习更新可能主要集中在少数几个方向上。这意味着更新本身表现得好像只知道如何通过少数主导路线移动,也就是说你并没有真正利用矩阵的全部容量。Muon 的做法是通过归一化这些方向强度来修复这个问题。它保留方向,但消除它们之间的不平衡,使各个方向的贡献更加均匀,这就是正交化的部分。

这也是为什么 Muon 从根本上与 Adam 不同:Muon 是感知矩阵的,而 Adam 不是。如果你的优化器不理解矩阵结构,它仍然可以工作,但它可能会错过模型几何结构中最重要的部分之一。这引出了 Muon 名称的下一部分:牛顿-舒尔茨方法(Newton-Schulz approach)。

正交化矩阵的干净数学方法是使用类似 SVD(奇异值分解)的东西。它基本上是将一个混乱的矩阵转化为几个干净的垂直方向,每个方向都有一个数字表示该矩阵拉伸该方向的强度。然而,在每个训练步骤中对每个权重矩阵进行完整的 SVD 计算会过于昂贵,因为它必须缓慢地计算出所有这些方向和拉伸值。

这比在训练期间进行实际的向前和向后传播要花费多得多的工作。因此,Muon 使用的是另一种方法:牛顿-舒尔茨迭代(Newton-Schulz iteration)。不深入复杂的数学细节,其核心思想是牛顿-舒尔茨是一种快速的近似正交化方法。因为 Muon 只需要最终的正交化更新,所以它不需要计算 SVD 所需的许多其他内容。

因此,它反复应用一个小矩阵公式,使矩阵逐渐接近正交状态。虽然经过几次迭代后可能不会完全正交化,但相比于初始状态,它已经平衡得多。此外,这种迭代方法让 Muon 仅通过重复矩阵乘法就能实现正交化,这非常适合 GPU 操作。

所以,这个新优化器不仅匹配 LLMs(大型语言模型)的结构,而且也非常适合 GPU 操作。嗯,并不完全是完全如此。如果你看过我的 DeepMind V4 基础设施视频,你会记得我简要提到过在那里使用 Muon 的一些困难。因此,Muon 在实际应用中面临的主要挑战实际上来自它最大的吸引力:作为一个感知矩阵的优化器。

Adam 易于扩展,因为每个参数的更新主要是独立的。因此,当将模型拆分到数千个 GPU 上时,每个 GPU 可以持有部分参数并计算局部的 Adam 更新然后继续。但 Muon 不同。Muon 需要知道整个矩阵的样子,以便对矩阵更新进行正交化。

这在 LLMs 的分布式训练中造成了巨大的问题,因为单个矩阵通常被拆分到多个 GPU 上:一个 GPU 持有矩阵的一个切片,另一个 GPU 持有其余部分。更糟糕的是,矩阵越大,需要的切片就越多。因此,要正确应用 Muon,训练系统需要在 GPU 之间进行额外的通信,并等待所有数据收集完毕。

所以,分布式训练本身就存在通信问题,而 Muon 还会增加另一层难度,阻止其被使用。尽管如此,像 Kimi(月之暗面)和 DeepSeek 这样的公司仍然决定采用它。Kimi 实际上很早就开始对其进行实验,他们是第一个发布重大里程碑报告的公司,证明 Muon 可以扩展到超越玩具模型。

因为即使 Kaiming He(何恺明)原始的 Muon 结果令人兴奋,但在大型 LLM 训练中是否有效仍然是一个巨大的问号。然后在这篇名为《Muon is Scalable for LLM Training》(Muon 可扩展用于大语言模型训练)的论文中,他们确定了使 Muon 更加实用的三个缺失环节。第一个是权重衰减(Weight Decay)。

原始的 Muon 在训练早期效果很好,但当他们扩大规模时,注意到某些权重随时间推移变得过大。这是一个问题,因为如果权重增长过于激进,模型在训练后期会变得不稳定。因此,他们添加了类似 AdamW 风格的权重衰减来控制权重的增长。第二个修复是更新缩放(Update Scaling)。

Muon 的更新大小严重依赖于矩阵的形状。这意味着如果矩阵非常宽或非常高,原始的 Muon 更新可能会变得太小;而如果矩阵太小,更新则可能变得太大。这在 Transformer 中是个问题,因为并非每个矩阵都具有相同的形状。注意力机制矩阵、MLP(多层感知机)矩阵、专家网络和投影矩阵都可以具有不同的维度。

Kimi 团队随后调整了每个参数的更新比例,以使 Muon 的更新大小在不同矩阵形状之间保持一致。这样它就不会仅仅因为矩阵形状不同而意外地过度激进或过于保守地更新某一层。第三个问题是分布式训练问题。

正如我之前提到的,Muon 需要完整的矩阵来对更新进行正交化。因此,他们的解决方案是基于 ZeRO(Zero Redundancy Optimizer)风格的分布式 Muon 实现。这个 ZeRO 技巧基本上不会永远保留完整的优化器副本,因为那样会破坏内存效率。相反,它临时收集重建矩阵所需的片段,运行计算过程。

牛顿-舒尔茨迭代保留更新的局部部分并丢弃其余部分。基本上是一种收集、计算、再丢弃的策略。在他们的缩放定律实验中,Moonshot 报告称,对于拥有 160 亿参数的 MoE 模型和 5.7 万亿个 token,Muon 仅使用约 52% 的训练算力即可达到与 AdamW 相当的性能。这相当于计算效率提升了约两倍,这一优势极具吸引力,尤其是对于严重受限于算力的实验室而言。

最新的 DeepSeek V4 将 Muon 提升到了新的高度。他们表示使用 Muon 处理大多数模块,因为它能带来更快的收敛速度和更好的训练稳定性。但他们并未在所有地方使用 Muon。在嵌入层、预测头、静态偏置以及 MHC 和 RMSNorm 权重的门控因子中,他们仍然保留了 AdamW。这一设计决策的原因在于 Muon 是矩阵感知的(matrix-aware),而模型中的参数并非都是标准的二维矩阵。

因此,DeepSeek 几乎必须采用混合优化器配方:对于大多数计算密集的矩阵组件使用 Muon,而对于更适合逐元素优化的部分则使用 AdamW。此外,DeepSeek 还增加了一个非常有趣的细节:混合牛顿-舒尔茨迭代(hybrid Newton-Schulz iterations)。这意味着他们不再对每一步使用相同的牛顿-舒尔茨系数,而是将过程分为两个阶段。

在前八次迭代中,他们使用旨在迅速推动奇异值趋向于一的系数;而在最后两次迭代中,则切换到更具稳定性的系数,以保持奇异值精确地围绕在一附近。这样,Muon 起初会迅速正交化矩阵,随后通过在末尾细化接近一的奇异值来避免过度射击或破坏矩阵更新的稳定性。

在基础设施方面,DeepSeek 也必须解决相同的零冲突问题:为了节省内存,必须将参数跨 GPU 分割。但他们没有使用 Kimi 所用的相同技术。DeepSeek 为 Muon 设计了一种混合的 Zero Bucket 分配策略。因此,Zero 不再尽可能多地将每个矩阵分割到多个 GPU 上,而是限制允许共享每个密集矩阵的 GPU 数量,防止每个矩阵被切分成过多的小碎片。

然而,这并不意味着 DeepSeek 消除了收集步骤(gathering step)。Muon 在此处仍然需要完整的矩阵。他们只是通过防止每个矩阵被切片到过多的 GPU 上,使得收集过程变得更加容易。因此,你不需要跨如此多的 GPU 进行数据收集。为了实现这一点,他们使用了一种类似背包问题的分配方式来决定哪些矩阵由哪些 GPU 处理,以实现工作负载的平衡。

因为如果随意分配,一个 GPU 可能会接收几个大矩阵,而另一个排名(rank)的 GPU 则主要接收小矩阵。我们不想让第一个 GPU 成为更新过程中的瓶颈。虽然这是一种稍微复杂的技术,但如果 DeepSeek 和 Kimi 已经能在前沿领域利用 Muon(注:原文误写为 MUA,结合上下文应为 Muon),其他许多 AI 实验室也肯定能受益于这项开放研究。

即使这可能无法反映出像玩具实验中那样的 52% 训练效率提升,但确实仍有空间,尤其是对于更大的训练运行。节省哪怕 10% 的算力也意义重大。然而,除了 Muon(注:原文误写为 mu1)能提供的一些预训练效率优势外,最近浮现出的一个更大图景是:Muon 可能比我们预期的更为重要,因为它从根本上改变了哪些研究想法具有可行性。

再次强调,因为过去十年 AI 研究有一个巨大的假设:使用 AdamW 进行训练。特别是对于模型架构而言,如果有效则表明该架构有前景;如果失败,则该架构可能行不通。然而,随着 Muon 的出现,一些研究人员开始思考:如果某些架构失败并非因为架构本身错误,而是因为 AdamW 不是最佳选择呢?这正是 Tilda Research 的新论文《Parallax》所探讨的核心。

Parallax 基于一个较旧的想法:局部线性注意力(local linear attention)。因此,既然 Softmax 注意力可以被解释为一种局部常数估计器——即对于每个查询(query),注意力会查看附近的键(keys)并平均它们的值,那么局部线性注意力试图通过拟合查询周围的局部线性关系来升级这种简单的预测规则,而不是假设附近区域是平坦的并仅仅进行平均。

因此,理论上局部线性注意力应该比标准的 Softmax 注意力具有更强的表达能力。但问题是原始版本难以扩展。它需要稍微更昂贵的优化过程,造成了数值稳定性问题,并且不能完美融入现代语言模型的预训练流程。因此,尽管该想法在理论上很强,但不足以开发成一种合适的注意力机制。

然而,Parallax 复兴了这一想法。它不再每次求解局部线性系统,而是学习一个额外的类似查询的投影来近似校正项。但令人惊讶的是,论文发现 Parallax 的优势在很大程度上依赖于 Muon(注:原文误写为 Muan)。在 AdamW 下,Parallax 和 Softmax 注意力非常接近;但在 Muon 下,Parallax 显示出更明显的优势。

这可能是因为 Muon(注:原文误写为 Muan)能提供比 AdamW 更强大且连贯的更新,这引出了一个令人震惊的结论:优化器的选择可以决定一个架构想法看起来是有用还是无用。深度学习中的许多旧思想可能具有良好的理论基础,但它们是在 AdamW 训练体制下被评判的。因此,如果 AdamW 未能激活正确的路径,该想法可能看似已死;但 Muon(注:原文误写为 Muan)改变了更新几何结构,这意味着某些想法在 Muon 下可能突然变得可训练。

然而,Muon(注:原文误写为 Muan)究竟是如何首先解锁这些特性的呢?另一篇最近的 Muon 论文对此进行了深入剖析。该论文认为,Muon(注:原文误写为 Muan)不仅比 Adam 优化得更快,而且能学到更好的特征。更具体地说,论文显示 Muon(注:原文误写为 mu1)训练出的模型学到的特征更加鲁棒且更具可迁移性。

所谓“鲁棒”,是指优化器将帮助模型更好地处理损坏的输入。例如,如果你破坏图像并在文本中引入噪声,Muon(注:原文误写为 mu1)训练出的模型往往比 Adam 或任何随机梯度下降(SGD)模型表现更好。至于“可迁移性”,他们表明在预训练之后,学到的特征对于下游任务更有用。

因此,如果你冻结主干网络并训练一个新的分类器,或者在另一个数据集上微调模型,Muon(注:原文误写为 mu1)预训练出的表征往往能更好地适应。为了从机制上理解这一点,论文将其与两个组件联系起来。第一个是更大的对数几率边际(logit margins)。这里的较大边际意味着模型能够更清晰地区分。

从最强或错误的答案中更清晰地得出正确答案。因此,当输入变得嘈杂时,在预测翻转之前,模型保持正确的空间就更大。其次是更高的有效秩(Effective Rank)。有效秩基本上是一种衡量隐藏状态实际使用了多少个独立方向的方法。如果有效秩较低,表示正在坍缩到少数几个主导方向中;而如果有效秩较高,则模型会将信息分散在更多独立的功能方向上。

这与 Muon 算法完美契合,因为它平衡了整个矩阵更新几何时态。因此,Muon 可能有助于防止学习过程坍缩到少数几个主导方向中,从而产生更丰富的内部表示。这意味着对于下一代 AI 的进步,现在需要开展一类新的消融研究,因为我们对完美优化器 AdamW 的假设刚刚因 Muon 的出现而受到挑战。

如果 Mu1(注:此处原文可能指代前文提到的优化器,结合语境应理解为当前讨论的算法)甚至还不是优化器的巅峰所在,而我们仅仅处于冰山一角呢?好吧,我所知道的是,我们肯定会有更多迷人的 AI 研究分支等待我们去探索。

如果了解 Muon(注:原文误拼为 Muan,根据上下文修正)让你想要更深入地了解优化器或大语言模型(LLMs),我的平台 intuitive.academy 可能正是你开始学习的地方。该平台专注于帮助你建立对 LLMs 背后关键概念的直观理解,从基础一直延伸到混合专家模型(MoE)和模型优化等高级主题。

我们不想用令人望而生畏的数学公式轰炸你,而是确保你先充分理解更高层面的逻辑,这样当你接触到技术细节时,它们才会变得有意义。此外,我们最近刚刚进行了内容改版,回顾并更新了所有早期的基础材料,以提供更好的学习体验。

我们还最近添加了一个新的测验系统,在每个章节后测试你的理解程度,并增加了每周排行榜,让学习过程更具竞争性。因此,如果你在工作中使用 LLMs、正在转型进入该领域,或者想要涉足大模型研究,请通过描述栏中的链接查看。

请使用我们的限时代码“back to school”享受年度会员五折优惠。感谢大家观看,特别鸣谢 Spam Match、Chris Loo、Dan、Robert、Zaviasa、Proof 以及 Enu DX Research Group、Alex Leo、Halawani 等许多通过 Patreon 或 YouTube 支持我的人。

如果你还没有关注我,请在 Twitter 上关注,我们下期再见。