Ultralytics Muon 优化器与 MuSGD 实现原理Newton-Schulz 正交化在 YOLO26 训练中的应用【免费下载链接】ultralyticsUltralytics YOLO26, YOLO11, YOLOv8 — object detection, instance segmentation, semantic segmentation, image classification, pose estimation, object tracking项目地址: https://gitcode.com/GitHub_Trending/ul/ultralyticsUltralytics 在ultralytics/optim/muon.py中实现了基于Newton-Schulz 正交化Newton-Schulz iteration的 Muon 优化器并将其与标准 SGD 融合成混合优化器MuSGD用于 YOLO26 等大规模视觉模型的稳定训练。本文以 参考文档 为骨架结合 源码实现 与 训练器集成代码系统讲解 Muon 的数学动机、zeropower_via_newtonschulz5正交化过程、muon_update的动量与分桶批处理细节以及MuSGD/Muon两个优化器的 API 与在yolo train命令中的实际用法。读完本文你将掌握在 Ultralytics 训练流程中开启 MuSGD、理解参数分组规则并能在自己的 PyTorch 工程中独立复用这套正交化优化原语。背景为什么要引入 Muon 风格的正交化更新传统基于动量的优化器如 SGD with momentum依赖梯度方向本身而在深度网络尤其是 Transformer、大型卷积模型中权重矩阵的谱结构会对收敛稳定性产生显著影响。Muon 优化器算法详见 arXiv:2502.16982由 Moonshot AI 的 Kimi K2 训练实践推广的核心思想是在应用动量之后对更新方向做矩阵正交化——即把更新矩阵投影到其 SVD 分解中UV^T方向的近似——从而抑制病态谱方向带来的振荡换取更平滑的收敛曲线。从 模型介绍文档 与 FAQ 可以看到Ultralytics 将 Muon 思路从大语言模型训练迁移到计算机视觉YOLO26 系列官方权重在 Objects365 预训练与 COCO 微调阶段均使用 MuSGD 优化器见 YOLO26 训练配方。它不是用 Muon 替换全部优化逻辑而是构造SGD Muon 的混合体只对二维线性权重与四维卷积滤波器施加正交化更新其余参数BatchNorm 权重、偏置等仍走标准 SGD在保留 SGD 良好泛化性的同时获得更平滑、稳定的优化轨迹。源码架构总览ultralytics/optim/muon.py模块按从底层算法原语到高层优化器类的顺序组织共包含四个核心成员也正是 参考文档 所列的四个 APIAPI类型作用zeropower_via_newtonschulz5(G, eps)函数用 5 步五次 Newton-Schulz 迭代近似计算矩阵的零次幂正交化逼近 SVD 的UV^Tmuon_update(grad, momentum, beta, nesterov)函数对梯度做指数滑动平均可选 Nesterov再把多组矩阵分桶、零填充后一次性批量正交化并按维度比缩放Muon优化器面向非分布式场景的纯 Muon 优化器对全部参数执行正交化更新权重衰减以乘法方式施加MuSGD优化器Muon 与 SGD 的混合优化器可通过参数组use_muon开关决定某组参数走混合路径还是纯 SGD模块导出optim/__init__.py仅暴露MuSGD与Muon并列入__all__而trainer.py中以from ultralytics.optim import MuSGD形式导入的是MuSGD即 YOLO26 训练实际使用入口。Newton-Schulz 五次迭代零次幂正交化实现剖析数学动机为什么要计算零次幂设矩阵G的 SVD 为G USV^T。所谓零次幂指用某种数值方法在不显式做 SVD的前提下逼近U V^T相当于把奇异值矩阵S中每个奇异值都替换为 1。UV^T是G在正交群上的投影方向把它作为参数更新方向可以去除奇异值尺度差异这正是 Muon 与自适应梯度类方法思路的分水岭前者不归一化逐参数学习率而是把整个矩阵的更新方向拉到谱空间中的等距点。五次迭代系数与数值流程zeropower_via_newtonschulz5的实现位于 muon.py采用了固定系数的 5 步迭代X G.reshape(-1, G.size(-2), G.size(-1)).bfloat16() X / X.norm(dim(-2, -1), keepdimTrue) eps # 保证最大奇异值 1 if G.size(-2) G.size(-1): X X.transpose(-2, -1) # 行多于列时转置以提升效率 a, b, c 3.4445, -4.7750, 2.0315 for _ in range(5): A X X.transpose(-2, -1) B torch.baddbmm(A, A, A, betab, alphac) # b*A c*AA X torch.baddbmm(X, B, X, betaa) # a*X BX关键点说明系数优化目标迭代系数(a, b, c) (3.4445, -4.7750, 2.0315)是针对在零点的收敛斜率最大化调优得到的使得结果在经验上对优化足够有效docstring 注明输出逼近USV^T其中S的对角元近似服从Uniform(0.5, 1.5)而不追求 SVD 那样精确的UV^T。数值稳定性先按矩阵的(-2, -1)维求范数归一化确保最大奇异值不超过 1再叠加eps默认1e-7防止除零。精度选择内部统一转成bfloat16计算兼顾吞吐与稳定性。输入约束仅接受二维矩阵或三维批量assert G.ndim in {2, 3}三维情况按第 0 维作为 batch 维每个切片独立参与同一轮迭代。形状处理行数多于列数时先转置、迭代完成后再转回减少X X^T的矩阵尺寸对2D输入直接reshape(G.shape)还原。为什么固定 5 步就够参考文档与源码注释明确恰好执行 5 次迭代。这是因为正交化只用于产生一个合理方向的更新而非求解高精度 SVD五次 Newton-Schulz 在奇异值已归一化到[0,1]区间的前提下足以把矩阵推向近似正交进一步迭代带来的增益对优化器而言可以忽略反而增加开销。文档同时提醒该实现不做精确UV^T分解属于经验上对神经网络优化表现良好的松弛方案。muon_update动量 分桶批量正交化 维度缩放在任意一组参数进入正交化之前muon_updatemuon.py先统一完成动量计算再组织高效的批量正交化。函数签名与语义def muon_update(grad, momentum, beta0.95, nesterovTrue):动量更新规则含 Nesterov 变体动量缓冲按指数滑动平均更新in-place修改调用方传入的缓冲momentum beta * momentum (1 - beta) * grad。启用 Nesterov默认Trueupdate beta * momentum (1 - beta) * grad即在平均动量基础上再向前看一步关闭 Nesterov直接取update momentum。在实现层面滑动平均与加法分别用torch._foreach_mul_/torch._foreach_add_完成属于fused foreach 算子避免逐个参数启动 kernel 的开销同时函数支持传入单张量或张量列表通过isinstance(grad, torch.Tensor)判定列表场景可一次处理整组参数。分桶与零填充一次调用正交化尽量多的矩阵正交化前更新张量先做形状归一多维张量展平ndim 2时reshape(len(u), -1)只保留第一个维度作为矩阵的行其余拍平为列——这正是四维卷积滤波器[out, in, kh, kw]能被当作二维矩阵处理的机制转置方向统一行数大于列数时先转置保证进入迭代时行数 列数按键分桶以(行数, scale, device, dtype)为 key 把更新分组其中scale max(1, dim[-2] / dim[-1]) ** 0.5是按参数两维比例计算出的维度补偿系数用于修正矩形矩阵在正交化后引入的尺度偏差同桶零填充合批同一桶内各矩阵列数不同时用torch.nn.functional.pad补齐到最大列数后torch.stack成三维批量一次调用zeropower_via_newtonschulz5。docstring 特别指出零填充列在 Newton-Schulz 迭代中保持为零因此填充不引入伪信息。正交化完成后先mul_(scale)施加维度补偿再把每个切片裁回各自列数、按需转置并reshape回原始形状返回全程最终to回梯度原本的 dtypebfloat16 仅用于计算中间态。MuSGDMuon SGD 混合优化器构造参数与参数组约定MuSGD(optim.Optimizer)muon.py构造签名MuSGD(params, lr1e-3, momentum0.0, weight_decay0.0, nesterovFalse, use_muonFalse, muon0.5, sgd0.5)其中muon与sgd是混合模式下的两个更新分量权重默认各 0.5保存为优化器实例属性。真正决定某组参数走哪条路径的是param_groups中每个组的use_muon开关构造默认False时所有组走纯 SGDuse_muonTrue该组参数同时接收正交化的 Muon 更新与标准 SGD 动量更新use_muonFalse该组只走标准 SGD等价于带动量的 SGD。因为use_muon是 per-parameter-group 配置用户可以按需把卷积/全连接权重与其余参数拆进不同组。参考文档给出的构造示例即展示了这种分组param_groups [ { params: model.conv_params, lr: 0.02, use_muon: True, momentum: 0.95, nesterov: True, weight_decay: 0.01, }, { params: model.other_params, lr: 0.01, use_muon: False, momentum: 0.9, nesterov: False, weight_decay: 0, }, ] optimizer MuSGD(param_groups, muon0.5, sgd0.5)step() 内部流程混合与降权细节step以torch.no_grad()包装若传入closure则在其内部恢复梯度后再求 loss返回该 loss否则返回None。对每个参数组过滤出grad is not None的参数无梯度参数整组跳过首次见到某参数时为其初始化状态缓冲Muon 动量缓冲momentum_buffer所有组都会建若use_muonTrue额外再建一个独立的momentum_buffer_SGD——Muon 动量与 SGD 动量各自独立维护互不串扰Muon 半区对整组梯度与momentum_buffer调用muon_update(betamomentum, nesterovnesterov)随后torch._foreach_add_(params, updates, alpha-(lr * self.muon))施加负学习率——即 Muon 分量只吃到lr * muon的有效步长接着lr * self.sgd把剩余学习率留给 SGD 分量SGD 半区将momentum_buffer_SGD混合组或momentum_buffer纯 SGD 组作为缓冲先对梯度加weight_decay权重衰减torch._foreach_add(grads, params, alphaweight_decay)再走经典动量公式并_foreach_add_(params, updates, alpha-lr)。值得强调的语义与源码 docstring 一致混合模式下 weight decay 只施加于 SGD 分量正交化只作用于二维及以上张量一维偏置/标量类参数天然适合留在纯 SGD 通道。Muon独立使用的非分布式优化器Muonmuon.py是不依赖 SGD 的纯 Muon 版本签名Muon(params, lr0.02, weight_decay0, momentum0.95)要点所有参数统一走动量 Newton-Schulz 正交化经muon_update没有按组区分混合与否梯度同步兜底某参数grad is None时step 中会以p.grad torch.zeros_like(p)强制赋零梯度用于分布式环境下保证状态同步一致性权重衰减为乘法式torch._foreach_mul_(params, 1 - lr * weight_decay)即p * (1 - lr * weight_decay)与 MuSGD 中把 weight decay 揉进梯度加法式的路径不同随后_foreach_add_(params, updates, alpha-lr)施加正交化后的更新docstring 明确其定位为non-distributed settings非分布式训练场景使用。optimizer Muon(model.parameters(), lr0.02, weight_decay0.01, momentum0.95) loss model(data) loss.backward() optimizer.step()在 Ultralytics 训练流程中的集成与使用入口BaseTrainer.build_optimizerMuSGD并非由用户手工构造而是通过训练参数optimizer触发。在 trainer.py 的build_optimizer中可用优化器集合为{Adam, Adamax, AdamW, NAdam, RAdam, RMSprop, SGD, MuSGD, auto}与 默认配置 中optimizer: auto的说明取值SGD, MuSGD, Adam, Adamax, AdamW, NAdam, RAdam, RMSProp, auto一致。关键选择逻辑auto自动选择当训练轮次迭代数大于 10000iterations 10000时自动选用MuSGD学习率取0.01、动量取0.9短训练回退到 AdamWlr_fit round(0.002 * 5 / (4 nc), 6)nc为类别数。这也印证了 训练模式文档 中MuSGD 推荐用于较长 YOLO26 训练与大数据集的表述参数分拣规则use_muon name MuSGD时启用四桶分拣param.ndim in {2, 4}的权重线性层矩阵、卷积滤波器→ 第 3 桶标注use_muon: True, param_group: muon这正是 Muon 仅正交化矩阵与卷积滤波器的落地方式bias→ 第 2 桶无 weight decayBatchNorm 类模块与logit_scale→ 第 1 桶无 weight decaybn通过遍历nn.__dict__中名称含Norm的类收集其余普通权重 → 第 0 桶施加正常weight_decay混合权重MuSGD 启用时muon, sgd (0.2, 1.0)即 Muon 分量取 0.2、SGD 分量取 1.0区别于类默认的 0.5/0.5通过partial(MuSGD, muonmuon, sgdsgd)构造头部高学习率对检测头cv3/one2one_cv3与语义分割头相关参数在微调时拆分出lr * 3的独立参数组源码注释为 finetuning 时对部分参数给更高学习率其目的在 events.py 中也说明为每个组被 MuSGD 拆成两份微调 lr*3 的那一半在前。命令行与 Python 用法训练模式文档给出了最直接的开启方式见 train.md 与 model-training-tips.mdyolo train modelyolo26n.pt datacoco8.yaml optimizerMuSGD等价 Python 写法from ultralytics import YOLO model YOLO(yolo26n.pt) model.train(datacoco8.yaml, optimizerMuSGD, epochs100, lr00.01, momentum0.9)在 Python 训练接口中相关超参映射到 trainer 的argsoptimizer之外lr0、momentum、weight_decay会被传入build_optimizer并进入optim_argsMuSGD 分支中optim_args {lr: lr, momentum: momentum, nesterov: True}即 YOLO 训练里的 MuSGD 默认开启 Nesterov。实用建议来自 训练技巧指南不确定时先用optimizerauto再在自有数据集上与optimizerMuSGD对比MuSGD 适用于追求 SGD 式泛化、又希望比普通 SGD 收敛更平滑的长时间训练。YOLO26 官方配方中的 MuSGD 配置YOLO26 训练配方文档 展示了一个超参文件片段摘要muon_w/sgd_w对应 MuSGD 两个分量权重经进化搜索得到optimizer: MuSGD官方两阶段Objects365 预训练 150 epochs → COCO 微调均在 640×640、batch 128 下使用 MuSGD其中微调阶段部分超参经进化搜索寻得如不同尺寸模型muon_w落在约 0.436–0.528、sgd_w落在约 0.479–0.674 区间详见 训练配方文档。这些数值佐证了 MuSGD 中两个分量权重即muon/sgd因子是可调且需要按模型规模搜索的关键超参数。工程细节与性能权衡把源码中几个容易被忽略的工程点汇总如下便于在自研工程中复用foreach 融合算子动量滑动平均、参数施加全部使用torch._foreach_*系列减少 kernel 启动次数分桶正交化进一步把形状兼容的多个矩阵合并进一次baddbmm链式迭代bfloat16 中间计算正交化数值热区使用 bfloat16结果再回退到原 dtype对支持 bf16 的 GPU 硬件较友好零填充无污染分桶后的列零填充保持为迭代不动点因此不同形状矩阵可以安全合批由 muon.py 的分桶逻辑与pad实现支撑两组动量状态MuSGD 混合组同时维护momentum_bufferMuon 用与momentum_buffer_SGDSGD 用状态字典以参数为 key便于state_dict序列化与断点续训权重衰减路径差异MuSGD 在 SGD 分支把 decay 叠加进梯度等价 L2Muon 类则用乘法(1 - lr * decay)两者语义不同迁移参数时需注意无梯度参数处理差异Muon 类赋零梯度强制同步MuSGD 则直接跳过该组参数——若自定义分组请确保行为符合预期正交化适用范围仅对ndim 2的张量有意义正交化把高维张量按首维保留、其余拍平一维参数应留在纯 SGD 通道这与build_optimizer中只把 2D/4D 权重送入 muon 桶的设计完全自洽。总结ultralytics/optim/muon.py用约 340 行代码完整落地了动量 Newton-Schulz 正交化 高效合批的 Muon 思想并向上提供Muon纯正交化与MuSGDSGD 混合两个优化器。对使用者而言最短路径是在训练命令中指定optimizerMuSGD或依赖optimizerauto在长训练下自动切换剩下的参数分拣、双动量维护、0.2/1.0 分量权重都由 trainer.py 与优化器自身完成对二次开发者而言zeropower_via_newtonschulz5与muon_update是无框架依赖的纯 PyTorch 原语可直接移植到任意自定义训练循环中。【免费下载链接】ultralyticsUltralytics YOLO26, YOLO11, YOLOv8 — object detection, instance segmentation, semantic segmentation, image classification, pose estimation, object tracking项目地址: https://gitcode.com/GitHub_Trending/ul/ultralytics创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
