Muon: An optimizer for hidden layers in neural networks

引言

  • 在深度学习领域,优化算法对模型训练效率和性能起着关键作用。从经典的随机梯度下降 (SGD) 及其动量法,到自适应优化方法 Adam/AdamW 等,一系列优化器大大加速了神经网络的收敛。然而,随着模型规模和数据量的爆炸式增长,研究者们不断探索更高效的新型优化策略,以进一步缩短训练时间和提高性能
  • Muon (Momentum Orthogonalized by Newton-Schulz) 便是近期引起广泛关注的创新成果之一,代表一种利用牛顿–舒尔茨迭代对梯度动量进行正交化处理的优化算法。Muon优化器专门针对神经网络隐藏层的二维权重参数设计,其核心思想是在每次参数更新前,对梯度更新矩阵进行近似正交化处理,从而改善优化动力学。凭借这一独特机制,Muon 在实际训练中取得了惊人的加速效果:例如,它被用于 NanoGPT 和 CIFAR-10 等任务的训练,并刷新了这些任务的速度纪录​
  • 一句话总结 Muon 的原理:利用牛顿–舒尔茨迭代高效地逼近梯度矩阵的正交化形式,并将其用于更新神经网络中的二维权重参数,从而加速训练收敛

Muon 相关链接

  • 代码实现:https://github.com/KellerJordan/Muon
  • 博客介绍:
    • https://kellerjordan.github.io/posts/muon/
    • https://jeremybernste.in/writing/deriving-muon

背景

Muon 优化器的设计初衷是针对神经网络中的线性变换层(即权重为矩阵的层,例如全连接层、卷积层等)提供一种专门优化方案。传统优化器(如 SGD 或 Adam)对所有参数一视同仁地应用统一的更新规则,而 Muon 采取了一种模块化的视角:根据层类型的不同采用不同的优化策略,以充分利用每类层结构的特性。这种“为不同网络模块定制优化器”的理念正是深度学习优化最新趋势的一部分。Muon 优化器聚焦于二维权重矩阵参数(例如全连接层的权重矩阵或卷积核张量展平后的矩阵),通过对这些矩阵形式的梯度更新进行特殊处理,以提升优化效果
在这里插入图片描述

Muon 优化器具体实现

Muon 是针对神经网络隐藏层二维参数的优化器。其算法定义如下:

  • 梯度计算与动量累积:在每次参数更新迭代 t t t,首先计算当前参数 θ t − 1 \theta_{t-1} θt1的梯度 G t = ∇ θ L t ( θ t − 1 ) G_t = \nabla_{\theta}\mathcal{L}t(\theta{t-1}) Gt=θLt(θt1)。然后将其累积到动量矩阵 B t B_t Bt中,即 B t ← μ B t − 1 + G t B_t \leftarrow \mu B_{t-1} + G_t BtμBt1+Gt,其中 μ \mu μ是动量超参数(类似于SGD动量法中的动量因子)。初始化时 B 0 = 0 B_0 = 0 B0=0。​
  • 以下伪代码(Algorithm 2 Muon)展示了Muon优化算法的主要循环步骤,其中 B t B_t Bt表示动量梯度累积, μ \mu μ是动量系数, η \eta η是学习率。在每次迭代中,先计算梯度 G t G_t Gt并累积到 B t B_t Bt中,然后对 B t B_t Bt执行Newton-Schulz迭代(即步骤5的 O t ← N e w t o n S c h u l z 5 ( B t ) O_t \leftarrow \mathrm{NewtonSchulz5}(B_t) OtNewtonSchulz5(Bt)),最后用得到的正交化更新 O t O_t Ot对参数进行梯度下降更新(步骤6)
    在这里插入图片描述
    其中,‘NewtonSchulz5’被定义为以下 Newton-Schulz 矩阵迭代
# Pytorch 代码
def newtonschulz5(G, steps=5, eps=1e-7):assert G.ndim == 2a, b, c = (3.4445, -4.7750, 2.0315)X = G.bfloat16()X /= (X.norm() + eps)if G.size(0) > G.size(1):X = X.Tfor _ in range(steps):A = X @ X.TB = b * A + c * A @ AX = a * X + B @ Xif G.size(0) > G.size(1):X = X.Treturn X

牛顿–舒尔茨迭代介绍

这一步是 Muon 的核心创新,即对累积梯度 B t B_t Bt应用 Newton-Schulz 迭代以计算其近似正交化矩阵 O t O_t Ot。Newton-Schulz 迭代是一种用于近似计算矩阵函数(例如矩阵平方根或逆)的算法。这里它被用于逼近 B t B_t Bt的逆平方根(或相关变换),以达到正交化的效果​。Muon使用固定步数的Newton-Schulz迭代来处理 B t B_t Bt,典型地迭代 5 次就能取得足够精度​。在实现中,作者选择在半精度(bfloat16)下运行该迭代以提高效率,并发现这种迭代在低精度下仍然数值稳定。相比直接进行奇异值分解(SVD),Newton-Schulz迭代可以充分利用GPU的并行计算能力,大幅降低计算开销。例如,作者最初实现的Newton-Schulz迭代每次更新对形状为 n × m n\times m n×m的参数需约 6 n m 2 6n m^2 6nm2次浮点运算,随后经过改进将复杂度降低到 4 n m 2 + 2 m 3 4n m^2 + 2 m^3 4nm2+2m3​(在实际中 m ≤ n m \le n mn,因此复杂度进一步受限于参数矩阵较小的一维),显著提升了运算效率。通过这样处理,Muon在每步更新中增加的计算和内存开销相对可控,却换来了更有效的梯度方向调整。

Muon 原理分析

Muon 在每次迭代时首先使用常规方法计算梯度并累积动量,然后对动量梯度矩阵进行**正交化(orthogonalization)**处理,再用于更新权重。所谓对梯度矩阵正交化,是指找到一个与原梯度矩阵最接近的“半正交矩阵”作为更新,其中“半正交”意味着矩阵的行向量或列向量是正交的(满足 O T O = I O^T O = I OTO=I O O T = I O O^T = I OOT=I )。等价地,这相当于将梯度矩阵进行奇异值分解 G = U S V T G=USV^T G=USVT ,然后用 U V T UV^T UVT (即将奇异值矩阵 S S S 置为单位矩阵)替换原梯度——这一结果被称为“正交化梯度”。Muon 的更新步骤正是用这样的正交化梯度替代原始梯度动量进行参数调整。

那么,引入正交化的直观意义何在? 简单来说,它可以丰富更新的方向性。研究者通过对 Transformer 等网络的梯度观察发现,传统优化器产生的梯度更新矩阵往往条件数很高,接近于秩亏(也就是被少数几个主导方向所支配)。换言之,许多神经元对应的更新方向非常相似,梯度矩阵接近低秩。这意味着一些“罕见方向”(对应较小奇异值的方向)在更新中作用很弱,但这些方向可能对进一步降低损失仍然重要。通过正交化处理,Muon 实质上放大了这些罕见方向在更新中的作用,因为正交化将所有奇异值归一,无论原本大小,从而赋予每个独立方向以均等的尺度权重​。这种调整有助于避免训练过程中某些重要方向被忽略,提高了搜索参数空间的效率。 从经验上看,正是这种对梯度更新方向分量的重新均衡,使得 Muon 相较于传统优化方法展现出更快的收敛和更高的效率。

Muon 实践注意事项

在具体实现中,还有一些实用技巧和细节:首先,Muon 目前仅针对二维权重张量(包括将卷积核展平成矩阵的情况)应用对于标量或向量参数(例如偏置项、LayerNorm 中的缩放参数等),以及输入层和输出层的权重,作者建议仍采用常规优化器 (如 AdamW) 进行训练​。这部分原因在于:嵌入层(输入层)参数的优化动态确实有别于其他层,需要特殊对待;输出层是否需要区别对待则是基于经验的优化选择​。实际实验表明,如果对 Transformer 模型的词嵌入层和最后分类头仍使用 AdamW 优化,而对中间的线性层使用 Muon 优化,能够取得最佳性能​。其次,在动量的具体形式上,Muon 采用Nesterov动量略优于传统动量,因此作者在公开实现中将 Nesterov 作为默认方案​。最后,对于 Transformer 中特有的多头注意力,研究者发现将查询 Q、键 K、值 V 各自的权重矩阵分别应用 Muon 优化效果更佳,而不是将它们合并为一个大矩阵一起正交化​。这一拆分处理使每个子矩阵的正交化更精确,有助于训练稳定。上述这些实现层面的经验总结,体现了 Muon 在不同网络结构中的适配性和灵活性。

实验

  • Muon 取得以下优化结果

    • 将 CIFAR-10 训练到 94% 准确率的速度记录从 3.3 A100-秒提高到 2.6 A100-秒
    • 在 FineWeb 上将训练到 3.28 验证损失(一项被称为 NanoGPT 速度竞赛的竞争任务)的速度记录提高了 1.35 倍。
    • 在扩展到 7.74 亿和 15 亿参数时,持续显示训练速度的改进。
    • 在 10 个 8xH100 小时内,将一个 15 亿参数的 Transformer 训练到 GPT-2 XL 在 HellaSwag 上的性能水平。使用 AdamW 实现相同结果需要 13.3 小时。
  • 下图展示了在NanoGPT语言模型任务中,不同优化器随时间推进的验证集损失下降曲线(横轴为在8×H100 GPU上消耗的实际时间,纵轴为验证损失)比较。可以看到,Muon(紫色曲线)相对于标准的Adam(蓝色)收敛更快。与此同时,图中标注的每步迭代耗时显示,Muon每步约142毫秒,与Adam的139毫秒几乎相当。这意味着Muon在提供更快收敛的同时,并未引入明显的单步时间开销在这里插入图片描述

Muon 后续工作进展

  • kimi 团队对 Muon 做了进一步优化:https://ar5iv.labs.arxiv.org/html/2502.16982
    • 由于Muon需要对每个二维权重矩阵执行Newton-Schulz迭代,如果在分布式环境下一个矩阵分散在不同设备上,如何高效地完成迭代成为一个问题。kimi 这里开发了内存优化且通信高效的分布式Muon实现。
    • 最初的Muon实验主要集中在十亿参数量级以内的模型上(如1.5B参数的Transformer)。一个悬而未决的问题是:Muon能否在数百亿乃至上万亿参数、万亿级别token数据的大规模训练中保持优势并稳定运行?kimi 这里的结果展示了通过在 Muon 中引入权重衰减(weight decay)以及按参数尺度调整更新幅度等改进,可以使 Muon 在无需特殊超参数调优的情况下,直接应用于超大规模模型的训练
      在这里插入图片描述
      在这里插入图片描述

总结

  • Muon 激发了人们对于优化器创新的热情——长久以来,Adam/AdamW 几乎是大多数任务的不二之选,而 Muon 证明了通过深入挖掘梯度结构特性,我们仍能找到显著优于现有方法的新途径。
  • 当然,Muon 目前也并非完美无缺。它对模型结构有特定假设(主要作用于稠密的线性层),在一些场景下的通用性还有待进一步验证。此外,对于更大规模模型、更长训练阶段(如微调、强化学习等)的效果也需要更多实验来检验

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.rhkb.cn/news/35149.html

如若内容造成侵权/违法违规/事实不符,请联系长河编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

数据结构与算法-图论-拓扑排序

前置芝士 概念 拓扑排序(Topological Sorting)是对有向无环图(DAG,Directed Acyclic Graph)的顶点进行排序的一种算法。它将图中的所有顶点排成一个线性序列,使得对于图中的任意一条有向边 (u, v)&#x…

市长海报/ Mayor‘s posters

AB 省 Bytetown 的市民无法忍受市长竞选活动的候选人随心所欲地将他们的选举海报贴在各个地方。市议会最终决定建造一面选举墙来放置海报,并引入以下规则: 每个候选人都可以在墙上放置一张海报。所有海报的高度都与墙壁的高度相同;海报的宽度可以是任意整…

LeetCode hot 100—验证二叉搜索树

题目 给你一个二叉树的根节点 root ,判断其是否是一个有效的二叉搜索树。 有效 二叉搜索树定义如下: 节点的左子树只包含 小于 当前节点的数。节点的右子树只包含 大于 当前节点的数。所有左子树和右子树自身必须也是二叉搜索树。 示例 示例 1&#…

ccfcsp3402矩阵重塑(其二)

//矩阵重塑&#xff08;其二&#xff09; #include<iostream> using namespace std; int main(){int n,m,t;cin>>n>>m>>t;int c[10000][10000];int s0,sum0;int d[10000],k[100000];for(int i0;i<n;i){for(int j0;j<m;j){cin>>c[i][j];d[s…

MCP和Function Calling的区别

文章目录 1、什么是MCP1.1、定义和特点1.2、架构和工作原理3.3、MCP 的主要优势 2、什么是Function Calling3、MCP和Function Calling的区别4、总结 &#x1f343;作者介绍&#xff1a;双非本科大四网络工程专业在读&#xff0c;阿里云专家博主&#xff0c;前三年专注于Java领域…

裂缝识别系统 Matlab GUI设计

使用说明 裂缝识别系统 Matlab GUI设计 &#xff0c;运行环境Matlab2023b及以上&#xff1b; 一种基于MATLAB图形用户界面&#xff08;GUI&#xff09;的裂缝自动识别系统&#xff0c;该系统利用数字图像处理技术实现裂缝图像的预处理&#xff0c;集成均衡化、噪声滤波、对比…

【源码分析】Nacos实例注册流程分析-事件驱动框架

【踩坑记录】 本人下载的Nacos 服务端版本是2.3.2&#xff0c;在开始进行源码编译便遇到问题&#xff0c;下面是各个问题记录 源码大量爆红 在最开始用Idea加载Maven项目的时候&#xff0c;发现项目中大量的代码爆红&#xff0c;提示其类或者包不存在&#xff0c;后来结果查…

51单片机指令系统入门

目录 基本概念讲解 一、机器指令​ 二、汇编指令​ &#xff08;一&#xff09;汇编指令的一般格式 &#xff08;二&#xff09;按字节数分类的指令 三、高级指令 总结​ 基本概念讲解 指令是计算机&#xff08;或单片机&#xff09;中 CPU 能够识别并执行的基本操作命令…

mysql5.x和mysql8.x查看和设置隔离级别

MySQL的隔离级别 级别标志值描述读未提交READ-UNCOMMITTED0存在脏读、不可重复读、幻读的问题读已提交READ-COMMITTED1解决脏读的问题&#xff0c;存在不可重复读、幻读的问题可重复读REPEATABLE-READ2mysql 默认级别&#xff0c;解决脏读、不可重复读的问题&#xff0c;存在幻…

【函数式编程】【C#/F#】第四讲:单子与函子 - 抽象的编程模式

在第二讲中我们探讨了一个诚实的函数应该要做到什么事&#xff0c;并运用了一种方法&#xff0c;让我们可以去准确的描述数据。 不过有一种情况让我们始料未及&#xff0c;例如网站需要收集一些信息&#xff0c;但有些信息不是必须的&#xff0c;是可有可无的。如果我们要去准…

【vue2 + Cesium】使用Cesium、添加第三方地图、去掉商标、Cesium基础配置、地图放大缩小事件、获取可视区域、层级、高度

参考文章&#xff1a; vue2 使用 cesium 篇【第一篇】 vue2 使用 cesium 【第二篇-相机视角移动添加模型】 vue2 项目模版&#xff1a; vue2-common 安装 cesium npm install cesium --save这个就很简单&#xff0c;只需要一句简简单单的命令就可以实现在 vue 项目中安装 ce…

vllm-openai多服务器集群部署AI模型

服务器配置是两台ubantu系统电脑,每台电脑安装两张4090-48G显存的显卡,共计192G显存。 服务器1 服务器2 准备工作: 1.两台电脑都已经安装了docker 2.两台电脑都已经安装了nvidia驱动 参考vllm官方资料 https://docs.vllm.ai/en/latest/serving/distributed_serving.html…

【电源】斩波电路

文章目录 前言定义概念 缩写降压斩波电路使用步骤总结参考文献 前言 进行大创项目开发的学习 bilibili 定义概念 缩写 斩波电路&#xff1a;分为降压&#xff0c;电荷泵&#xff0c;升压&#xff0c;升降压&#xff0c;Cuk&#xff0c;Speic&#xff0c;Zeta 等等 降压斩…

Hadoop集群组成

&#xff08;一&#xff09;Hadoop的组成 对普通用户来说&#xff0c; Hadoop就是一个东西&#xff0c;一个整体&#xff0c;它能给我们提供无限的磁盘用来保存文件&#xff0c;可以使用提供强大的计算能力。 在Hadoop3.X中&#xff0c;hadoop一共有三个组成部…

c++基础知识-图论进阶

一、拓扑排序 1、基础知识 1&#xff09;什么是拓扑排序 对一个有向无环图G进行拓扑排序&#xff0c;是将G中所有顶点排成一个线性序列&#xff0c;使得图中任意一对顶点u和v&#xff0c;若&#xff0c;则u在线性序列中出现在v之前。 2&#xff09;拓扑排序的操作方法 重复执行…

从Scaling Laws中解析大模型训练的边际递减临界点

前言 当我们拆解GPT-4到DeepSeek的演进路径&#xff0c;会发现一个反直觉的真相&#xff1a;​AI的智能跃迁不依赖参数堆砌&#xff0c;而取决于对"结构-能量-信息"三元关系的精准把控。就像人类大脑在进化中通过皮层折叠而非单纯增大体积来实现智能突破&#xff0c…

Word 小黑第20套

对应大猫21 特定一页设为横向 上下用分页符

【从0到1搞懂大模型】RNN基础(4)

先说几个常用的可以下载数据集的地方 平台&#xff1a;kaggle&#xff08;https://www.kaggle.com/datasets&#xff09; 和鲸社区&#xff08;https://www.heywhale.com/home&#xff09; 阿里天池&#xff08;https://tianchi.aliyun.com/&#xff09; 其他&#xff1a;海量公…

openEuler24.03 LTS下安装MySQL8

前提条件 拥有openEuler24.03 LTS环境&#xff0c;可参考&#xff1a;Vmware下安装openEuler24.03 LTS 步骤 卸载原有mysql及mariadb sudo systemctl stop mysql mysqld 2>/dev/null sudo rpm -qa | grep -i mysql\|mariadb | xargs -n1 sudo rpm -e --nodeps 2>/dev/…

如何在Odoo 18中实现OWL通知服务

如何在Odoo 18中实现OWL通知服务 OWL&#xff08;Odoo Web Library&#xff09;是Odoo的前端框架&#xff0c;用于构建现代化的动态响应式用户界面。在早期版本中&#xff0c;Odoo 前端设计与开发使用的是诸如 QWeb 这类较为老旧的框架&#xff0c;而随着 Odoo 每发布一个新版本…