深度学习优化的核心挑战与直观类比

训练人工智能的优化算法处于微积分、几何学和统计学的交汇点,其前置知识要求极高。尽管训练过程听起来简单——运行模型、测量误差、更新权重并重复,但利用梯度箭头更新权重的过程实际上承担了巨大的计算负荷。首先需要通过反向传播确定每个权重对误差的贡献,这涉及微积分中的梯度概念,它指示了使误差增加最陡峭的方向。因此,逻辑上应朝相反方向移动权重以减少误差。

然而,梯度仅告诉方向而非步长。如果步子太小,几乎无法减少误差;如果步子太大,反而可能增加误差。这就像在黑暗中下山:当前的坡度告诉你哪边是下坡路,但无法决定下一步该迈1英寸还是10米。你可能身处山谷却误入上坡路段。这一决策任务由优化器承担,它接收反向传播的梯度并将其转化为实际的权重更新。目前最流行的算法是Adam(自适应矩估计),其核心思想看似复杂,实则旨在解决两个关键问题。

Adam优化器的机制与AdamW的改进

梯度噪声是训练中的主要问题之一。为了提高效率,我们通常使用小批量数据而非单个样本或全量数据集进行训练。因此,Adam通过维护过去梯度的移动平均值来赋予优化器“记忆”,从而识别出它曾远离的方向。此外,不同权重的行为差异显著:有些权重始终拥有巨大的梯度,而另一些则几乎不动。为此,Adam还维护了梯度的平方的移动平均值,以此反映每个参数在历史上的活跃程度。

随后出现的AdamW是对Adam的重要改进。其背后的权重衰减(Weight Decay)理念很简单:在训练过程中轻微抑制模型过度依赖大权重的倾向,因为较大的权重可能导致模型过于激进地拟合训练数据,类似于记忆局部模式而非学习更广泛的规律。AdamW的关键区别在于将权重衰减与常规的梯度更新分开应用。虽然这听起来微不足道,但对于自适应优化器而言,这使得权重衰减的行为更加干净和可预测。正因如此,AdamW因其稳定性、自适应性和相对容易训练的特点而变得极其流行。

Muon的提出背景与Solid赞助实验展示

尽管AdamW统治了12年,但前沿开源研究最近转向了一种名为Muon的新优化器。这是首个成功挑战12年现状(AdamW)的新方法,其影响力可能堪比新的架构设计。Muon由独立研究员Keller Jordan于2024年12月首次提出,全称为“Momentum Orthogonalized by Newton Schultz”(牛顿舒尔茨正交化的动量)。为了直观展示其潜力,视频作者使用赞助商Solid进行了实验。

“我告诉Solid用两种不同的优化器Atom W和Muon训练两个小型Transformer语言模型,并构建一个原生Mac应用来比较结果。”

Solid是一个拥有独立设备、账户和预算的AI代理平台。作者下达任务后,Solid在云端自动设置Mac环境、安装工具、运行实验并构建原生应用。整个过程耗时超过一小时,最终生成了一个包含交互式结果视图的Mac OS应用程序。该实验完全由AI自主完成,无需作者干预计算机操作。这一案例展示了Solid如何为代理提供执行复杂任务所需的底层基础设施。

Muon的核心理念:从标量到矩阵视角的转变

Muon的核心思想其实相当直接,但需要回到Adam的数学本质来理解。从数学上讲,Adam将模型视为一系列独立的数字列表。尽管Adam能优雅地处理权重的梯度历史,但它仍将每个权重视为相互独立的个体。这种方法在处理单个数字时完全可行,但忽略了神经网络内部的本质结构。

“这些计算不仅仅是处理孤立的数字,而是矩阵。整个Transformer基本上是由矩阵乘法构建的。”

注意力机制中的QKV投影、MLP层等均为矩阵。在数学上,矩阵并非随机独立数字的集合,而是一种变换(Transformation)。它接收输入并将它们混合,产生一组新的特征。矩阵内的数字通过其产生的变换相互连接:它们共同决定信息如何被旋转、拉伸、压缩或混合。理解这一矩阵本质是掌握Muon突破性的关键,因为它不再孤立地看待参数,而是考虑了权重之间的几何关系。

Adam的局限与矩阵更新的“方向失衡”

传统的Adam优化器在处理权重更新时,往往将矩阵内部的数字视为独立变量进行单独调整。然而,实际神经网络中的权重并非孤立存在,矩阵内的数值之间存在复杂的关联性。Adam风格优化器忽略了这种相关性,导致其无法捕捉模型几何结构的全貌。为了形象地说明这一问题,可以将优化过程比作100个人在黑暗中下山。Adam的做法是向每个人单独发出指令,虽然看似有100个独立源在同时行动,但如果其中90个人都朝着几乎相同的方向移动,那么这种并行性就是低效的。大家重叠在相似的方向上,浪费了大量的容量,因为并没有充分利用这100个独立源去探索不同的路径。相比之下,Muon的核心差异在于它试图正交化更新过程。通过迫使矩阵使用更多独立的方向,Muon防止了少数几个方向主导整个优化过程,从而更充分地利用矩阵的容量。

正交化的数学原理:从SVD到Newton-Schulz迭代

要理解Muon如何修复这种方向失衡,需要拆解矩阵更新的两个组成部分:方向和强度。方向决定了更新试图移动信息的轨迹,而强度则反映了每个方向的支配程度。如果一个方向的强度巨大,大部分更新资源都集中在此;反之,若强度微小则该方向几乎未被使用。即使矩阵更新的数值总量很大,实际的学习更新可能仅集中在少数主导路线上。Muon通过归一化这些方向强度来修复此问题,保留方向但消除其间的失衡,使各个方向的贡献更加均匀。这种正交化操作使得Muon成为矩阵感知型优化器,而Adam则不是。理论上,使用奇异值分解(SVD)可以干净地将矩阵转化为几个垂直的清洁方向。然而,在每一步训练中对每个权重矩阵执行完整的SVD计算量过大,远超过前向和反向传播本身的工作量。因此,Muon采用了Newton-Schulz迭代法。这是一种快速近似方法,通过反复应用小型矩阵公式将矩阵推近正交状态。这种方法仅依赖重复的矩阵乘法,完美契合GPU的操作特性,既匹配LLM的结构又兼顾了计算效率。

分布式训练的挑战与Kimi团队的三大工程突破

尽管Muon在理论上完美匹配LLM结构,但在实际分布式训练中面临巨大挑战。Adam易于扩展是因为每个参数独立更新,GPU只需持有参数的切片并计算本地更新即可。而Muon需要知道整个矩阵的全貌才能进行正交化,这意味着单个大矩阵被分割到多个GPU上时,系统需要额外的通信并等待数据汇聚。这种开销加剧了分布式训练固有的通信瓶颈。尽管存在困难,Kimi和DeepSeek等公司仍选择采用Muon。Kimi团队在早期就进行了实验,并在论文《mu1 is scalable for LM training》中识别出三个关键缺失环节以使其实用化。首先是权重衰减(Weight Decay), vanilla muon在训练初期有效但后期可能导致权重过度增长从而引发不稳定,因此引入了AdamW风格的衰减来控制。其次是更新缩放(Update Scaling),由于Muon的更新幅度高度依赖矩阵形状,不同维度(如Attention、MLP、专家层)会导致更新大小不一致。Kimi团队调整了每参数的更新缩放比例,确保不同形状的矩阵获得一致的更新强度。最后是分布式实现的优化,采用类似ZeRO的策略,不永久复制完整优化器以节省内存,而是临时汇聚所需片段重构矩阵进行计算。这些改进证明了Muon在大规模LLM训练中的可行性。

Muon的训练效率与DeepSeek V4的混合优化策略

Moonshot在缩放定律实验中报告,对于拥有160亿参数的MoE模型和5.7万亿token的训练数据,Muon仅需约52%的训练FLOPS即可达到与AdamW相当的性能。这一计算效率提升约两倍,对于算力受限的实验室极具吸引力。DeepSeek V4进一步将Muon推向新高度,采用混合优化器配方:在大多数模块中使用Muon以加速收敛并提高稳定性,但在嵌入层、预测头、静态偏置以及MHC和RMSNorm权重中保留AdamW。这一设计决策源于Muon是矩阵感知的(matrix aware),而模型中并非所有参数都是标准的2D矩阵,因此在元素级优化更合适的部分仍需AdamW。

"They say they use muon for most modules because it gives faster convergence and better training stability."

此外,DeepSeek引入了混合牛顿-舒尔茨迭代(hybrid Newton Schulz iterations)。在前8次迭代中,使用旨在快速将奇异值推向1的系数;在最后2次迭代中,切换至更稳定的系数以保持奇异值精确围绕1。这种策略使Muon先快速正交化,再通过细化接近1的奇异值来避免矩阵更新的过冲或不稳定。这种混合设计既利用了Muon在密集矩阵组件上的优势,又规避了其在非标准参数结构中的局限性。

基础设施优化:解决零冲突与负载均衡

在基础设施层面,DeepSeek解决了Muon特有的零冲突(zero conflict)问题。由于需要分割参数以节省显存,但Muon仍需完整矩阵进行计算,直接跨过多GPU切片会导致严重的收集开销。DeepSeek设计了混合Zero桶分配策略(hybrid Zero bucket assignment strategy),限制每个密集矩阵可共享的GPU数量上限,防止矩阵被切分成过多碎片。通过背包式分配(knapsack style assignment),系统决定哪些矩阵由哪些GPU处理,以实现工作负载平衡。若分配不当,某块GPU可能因接收过多大型矩阵而成为更新瓶颈。

"Deepseek designed a hybrid zero bucket assignment strategy for Muon... keeping each matrix from being chopped into too many tiny pieces."

这种策略简化了收集过程,避免了跨过多GPU的复杂操作。尽管DeepSeek和Kimi已在前沿应用Muon,其他AI实验室也能从中受益。即使无法完全复现52%的效率增益,节省10%的算力在大型训练运行中仍意义重大。这种基础设施层面的优化使得Muon不仅是一个算法改进,更成为可扩展大规模训练的关键工程突破。

优化器如何重塑架构研究的可行性范式

Muon的影响远超预训练效率,它从根本上改变了哪些研究思路具有可行性的判断标准。过去十年,AI研究普遍假设AdamW是基准:若架构在AdamW下有效则 promising,否则 likely 无效。然而,Muon的出现促使研究者反思:某些架构的失败是否源于AdamW并非最优选择? Tilda Research的新论文Parallax揭示了这一观点。该架构基于局部线性注意力(local linear attention),理论上比标准Softmax注意力更具表达能力,但此前因优化困难、数值不稳定且难以融入现代LM预训练而被视为不切实际。

"The paper finds that the parallax's advantage depends heavily on muon. Under ATMW, Parallax and Softmax attention are much more comparable. But under Muan, Parallax shows a much cleaner advantage."

研究发现,Parallax的优势高度依赖优化器。在AdamW下,两者表现相近;而在Muon下,Parallax展现出更清晰的优势。优化器的选择决定了架构思想看起来是有用还是无用。许多深度学习旧理论虽好,但在AdamW训练模式下未被激活。Muon改变了更新几何结构(update geometry),使得一些此前看似不可行的想法突然变得可训练。这暗示了优化器选择可以决定架构设计的生死。

Muon优化器提升模型鲁棒性与表征丰富度

Muon作为一种新兴的LLM训练优化算法,其核心优势在于显著增强了模型对噪声输入的鲁棒性。当输入数据存在干扰时,更强的预测置信度使得模型在发生错误翻转前拥有更大的容错空间。这种特性意味着即使面对嘈杂环境,模型也能更清晰地辨别正确答案与错误答案之间的界限。此外,Muon还带来了更高的有效秩(Effective Rank)表现。有效秩是衡量隐藏状态实际利用的独立方向数量的关键指标,它直接反映了模型内部表征的质量。

"If the effective rank is high, the model is spreading information across more independent feature directions"(如果有效秩较高,模型会将信息分散到更多的独立特征方向上)

在传统的低有效秩情况下,模型的表征往往会坍缩到少数几个主导方向上,导致信息表达单一。相比之下,高有效秩表明模型能够利用更多样化的独立特征方向来存储和处理信息。Muon通过平衡整个矩阵更新的几何结构,完美契合了这一需求,从而防止学习过程陷入少数主导方向的陷阱。这种机制促进了更丰富的内部表征生成,为下一代AI进步奠定了理论基础。

优化器范式的挑战与未来研究方向

随着Muon的出现,长期以来被视为完美标准的AdamW优化器正面临前所未有的挑战。这一变化迫使研究人员重新审视现有的假设,并意识到我们可能仅仅站在了优化器技术发展的冰山一角。Muon并非优化的终点,而是一个全新的起点。这意味着未来的AI研究需要开展全新类别的消融实验(ablation studies),以深入探究不同优化策略对模型性能的具体影响。

这种范式转移不仅关乎算法本身的优劣,更涉及对整个深度学习训练生态的理解。如果Muon不是优化器的顶峰,那么背后可能隐藏着更多未被发现的潜力和机制。研究人员需要探索这些未知的领域,以应对日益复杂的模型训练需求。新的研究路径要求我们跳出对AdamW的路径依赖,去发现能够更好适应大规模语言模型特性的新型优化方案。

学习资源推荐与社区致谢

对于希望深入理解优化器原理或LLM底层逻辑的学习者,intuitive.academy平台提供了系统化的学习路径。该平台专注于帮助学员建立对LLM关键概念的直观理解,内容涵盖从基础理论到混合专家模型(MoE)及模型优化等高级主题。平台采用先理解高层逻辑、后深入技术细节的教学方式,避免用令人望而生畏的数学公式直接轰炸学员。

近期平台进行了内容改版,更新了所有早期基础材料以提升学习体验,并新增了每章后的测验系统以及每周排行榜机制。使用限时代码'back to school'可享受年度会员50%的折扣。最后,感谢Spam Match、Chris Loo、Dan等Patreon及YouTube支持者对内容的持续贡献,鼓励观众关注Twitter以获取最新动态。

结论

Muon通过矩阵正交化解决了AdamW忽略权重几何结构的问题,显著提升了训练效率与模型鲁棒性。实证数据显示,Muon在MoE训练中仅需约52%的FLOPS即可达到同等性能,且能提升有效秩以丰富内部表征。其核心优势在于加速收敛及增强对噪声的容错能力,但需配合AdamW处理非矩阵参数(如嵌入层),并依赖复杂的分布式工程优化以克服通信瓶颈。Muon不仅重塑了架构评估标准,使如Parallax等此前难以训练的模型显现优势,更标志着优化范式从标量向矩阵视角的根本转变。