μ子优化器介绍与概述 | AI生成和翻译

Home 2025.09 English 中文

---

Muon优化器全面解析

Muon优化器是一种专为神经网络线性层（如全连接层或嵌入层）训练设计的二阶启发式优化算法，亦可扩展至其他层结构。该算法最初由Keller Jordan、Jeremy Bernstein等研究人员于2024年底提出，其理论根源可追溯至几何优化技术中的极坐标初始化和模对偶框架[1][2]。月之暗面（Moonshot AI）创始人兼Kimi AI首席执行官杨植麟在讨论其千亿参数大模型Kimi K2训练时特别指出，Muon通过适应损失函数几何形态的高秩更新机制，成为模型高效训练的核心支柱[3][4]。然而基础版本存在稳定性缺陷（如长时训练中的损失值突变），促使月之暗面开发了增强版本MuonClip，通过引入注意力层QK裁剪等稳定化机制提升训练鲁棒性[3][2]。

Muon的突出优势在于令牌效率：相较于AdamW等一阶优化器，它以更少的训练令牌量即可达到相当性能，这对LLM预训练等计算密集型任务极具价值。该算法旨在以低于传统二阶方法（如牛顿法）的计算成本实现近似效果，重点通过高秩矩阵更新实现特征值自适应。在大规模模型的梯度噪声环境中，Muon基于自然梯度和矩阵平方根思想的预条件技术展现出独特优势。

核心原理与推导

• 基本理念：Muon植根于几何优化理论，使参数更新适应损失函数的”能量景观”。它采用基于费雪信息矩阵（或其近似）的预条件器对梯度进行缩放，思路类似AdaGrad或Shampoo但针对稠密线性层优化[1][2]。
• 算法步骤：
  1. 梯度计算：计算线性层权重(W)的标准梯度(\nabla W)
  2. 预条件处理：通过Newton-Schulz迭代近似预条件器的矩阵平方根，实现无需完整特征分解的秩自适应
  3. 更新规则：应用能更有效缩放高秩分量的更新策略，常结合动量或裁剪机制增强稳定性
• 数学洞察：若(G)为梯度矩阵，Muon近似执行(W \leftarrow W - \eta \cdot \sqrt{P}^{-1} G)的更新，其中(\sqrt{P})采用迭代矩阵平方根[2][5]。这与AdamW的对角化或动量缩放形成对比，使Muon能更好捕捉参数间关联性
• 效率提升：在NanoGPT等基准测试中，Muon可将训练步数减少20-50%[1]

优势与局限

• 优势：
  • 线性层收敛优势：在LLM典型的高维稠密空间中表现卓越，以更少令牌实现更低损失[4][6]
  • 资源高效：所需梯度计算量减少，单轮训练速度更快
  • 开源可扩展：存在多个实现版本，包括专用于GPU加速的Flash-Muon[4][7]
• 局限：
  • 稳定性问题：在深层网络或稀疏层中易发散，MuonClip通过训练过程中裁剪注意力分数（如查询-键乘积）解决此问题[3][2]
  • 层结构特异性：不适用于卷积或循环层，偏重线性/MoE架构。Keras官方文档明确不建议在非线性层使用[8]
  • 超参数敏感：需精细调整学习率((\eta))和正交化操作，跨模型尺寸迁移需重新调参[2]
• MuonClip变体（Kimi专用）：该演进版本集成QK裁剪机制，在15.5万亿令牌预训练中防止失稳。它成功稳定了Kimi K2模型的320亿激活参数，实现零损失突变的训练过程并在Tau2-Bench等基准取得优异表现（如66.1分）[3][8]。目前未公开代码，属专有实现但基于开源Muon构建

Muon已影响AI优化领域发展，出现在Scion基准测试和Reddit/X平台讨论中，常因其”几何直观性”受赞誉。完整推导请参阅Jeremy Bernstein的技术博客[2]。下面我们来看具体实现方案。

代码实例：PyTorch实现Muon优化器

以下是根据官方仓库（https://github.com/KellerJordan/Muon）改编的基础Muon优化器PyTorch实现，这是针对稠密线性层的简化版本，包含预条件器的Newton-Schulz迭代。

import torch
import torch.nn as nn

class Muon(torch.optim.Optimizer):
    """
    线性层专用Muon优化器
    改编自：https://github.com/KellerJordan/Muon
    """
    def __init__(self, params, lr=1e-3, lr_b=2e-3, b2=0.95, wd=0.0):
        defaults = dict(lr=lr, lr_b=lr_b, b2=b2, wd=wd)
        super().__init__(params, defaults)

    def step(self):
        for group in self.param_groups:
            lr = group['lr']
            lr_b = group['lr_b']
            b2 = group['b2']
            wd = group['wd']

            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad.data.float()
                state = self.state[p]
                if 'momentum' not in state:
                    state['momentum'] = torch.zeros_like(grad)

                # 动量更新
                state['momentum'].mul_(b2).add_(grad)

                # 权重衰减
                if wd != 0:
                    p.data.mul_(1 - lr * wd)

                # Muon正交化（秩自适应）
                grad_vec = state['momentum'].view(-1, grad.shape[-1])
                p_vec = p.data.view(-1, p.shape[-1])

                # Newton-Schulz矩阵平方根近似（简化版）
                G = grad_vec @ grad_vec.t() / grad_vec.shape[0]
                # 完整实现为迭代过程，此处采用幂级数近似
                sqrt_G = torch.sqrt(G + 1e-6 * torch.eye(G.shape[0], device=G.device))

                # 参数更新
                update = grad_vec.t() @ sqrt_G @ grad_vec / sqrt_G.shape[0]
                p.data.sub_(lr_b * update.view(p.shape))

# 使用示例
model = nn.Linear(768, 768)  # 稠密层
optimizer = Muon(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()
data = torch.randn(32, 768)
target = torch.randn(32, 768)

for epoch in range(10):
    optimizer.zero_grad()
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

代码说明：

• 此为精简实现，完整GitHub仓库包含高效CUDA内核（如加速矩阵乘法的Flash-Muon）[3][4]
• Keras/TensorFlow用户可使用from keras.optimizers import Muon; opt = Muon(learning_rate=0.01)[8]
• 调优超参数（如lr_b=2e-3）至关重要，建议根据模型实验调整。Kimi K2使用的MuonClip未公开代码，但其在注意力前向传播中增加了裁剪逻辑[2][5]

如需扩展实现或基准测试资料，欢迎进一步交流！

参考文献： [1] https://kellerjordan.github.io/posts/muon/ [2] https://jeremybernste.in/writing/deriving-muon [3] https://github.com/KellerJordan/Muon [4] https://github.com/nil0x9/flash-muon [5] https://www.lakernewhouse.com/writing/muon-2 [6] https://medium.com/@kyeg/building-the-muon-optimizer-in-pytorch-a-geometric-approach-to-neural-network-optimization-17f4601be548 [7] https://discuss.huggingface.co/t/tutorial-understanding-and-implementing-the-muon-optimizer/167717 [8] https://keras.io/api/optimizers/muon/

---

Back

x-ai/grok-code-fast-1

Donate