常用的optimizer优化器和优化函数::[MUON-＞MuonClip-＞Gluon-＞AdaMuon ]

Kimi.ai 的 Moonlight 模型（一个 16B 参数的 MoE 模型，使用 5.7T tokens 训练）成功采用 Muon 作为主要优化器，证明了其在工业级规模应用的可行性。

1 muon【MomentUm Orthogonalized by Newton-schulz】

苏剑林-科学空间：https://spaces.ac.cn/archives/10592
从向量到矩阵的本质跨越

1.1 Muon 是什么？

• 核心概念
  Muon 是一个专门为神经网络隐藏层的二维参数设计的优化器。与传统优化器将所有参数视为扁平化向量不同，Muon 直接利用权重矩阵的矩阵结构，采用基于谱范数的更新规则。

简单来说，Muon 使用矩阵正交化来指导参数更新，确保权重变化以稳定训练的方式对齐。这是一个从传统的基于动量、逐元素更新的 Adam 和 AdamW 的重大转变。

• 技术原理
  Muon 是一个专门为神经网络隐藏层的二维参数设计的优化器。与传统优化器将所有参数视为扁平化向量不同，Muon 直接利用权重矩阵的矩阵结构，采用基于谱范数的更新规则。

在 Muon 中，我们利用的是作为 RMS 范数向量空间之间线性算子的矩阵几何。

Muon 是一个用于优化神经网络隐藏层二维参数的优化器。
它的定义如下：

1.2 实际应用案例

• CIFAR-10 速度记录
  Muon 将 CIFAR-10 训练到 94% 准确率的记录从 3.3 A100-秒降低到 2.6 A100-秒。
• GPT-2 训练
  用于以 175 美元的计算成本将 Transformer 训练到 GPT-2 (XL) 性能水平。
• Moonlight 模型
  Kimi.ai 的 Moonlight 是一个 3B/16B 参数的混合专家（MoE）模型，使用 Muon 进行了 5.7T tokens 的训练。这个模型改进了当前的帕累托前沿，与之前的模型相比，在更少的训练 FLOPs 下实现了更好的性能。
• Kimi K2 - 万亿参数级突破 ⭐
  Kimi K2 是 Moonshot AI 开发的最新一代大语言模型，也是 MuonClip 优化器的首次万亿参数级应用。这个模型代表了 Muon 优化器从理论到工业级应用的重大里程碑。
  模型规模
  总参数: 1.04 万亿（1T）
  激活参数: 320 亿（32B）每次前向传播
  训练数据: 15.5 万亿（15.5T）tokens
  训练稳定性: 零损失尖峰 - 在整个训练过程中没有出现一次训练崩溃

2 MuonClip 优化器：Muon 的进化版

在将 Muon 扩展到万亿参数级别时，Moonshot AI 团队发现了一个严重的问题：注意力 logits 爆炸。在中等规模(9B 激活/53B 总参数)的 MoE 模型训练中，注意力 logits 快速超过 1000 的量级，可能导致数值不稳定甚至训练发散。

为了解决这个问题，团队提出了 MuonClip 优化器，它在 Muon 的基础上集成了三个关键改进：
权重衰减 (Weight Decay)
RMS 一致性匹配 (Consistent RMS Matching)
QK-Clip：核心创新

2.1 三个关键改进

2.1.1 权重衰减 (Weight Decay)

与 Moonlight 模型类似，在大规模训练中权重和层输出的 RMS 会持续增长，可能超出 BF16 的精度范围。权重衰减机制可以有效控制权重的增长。

MuonClip 中的权重衰减实现
W_t = (1 - λ) * W_{t-1} + learning_rate * Update_t
其中 λ 是权重衰减率，典型值为 0.1

2.1.2 RMS 一致性匹配 (Consistent RMS Matching)

Vanilla Muon 的更新幅度会根据权重矩阵的形状自然变化，导致某些层更新过大，其他层更新过小。
MuonClip 通过显式缩放确保更新规模在所有层之间保持一致：
更新矩阵 O_t 按形状缩放
scaling_factor = sqrt(max(A, B)) * 0.2
其中 A, B 是矩阵的维度
常数 0.2 用于匹配 AdamW 的典型更新 RMS
这使得 MuonClip 可以直接复用 AdamW 的超参数，大大降低了调优成本。

2.1.3 QK-Clip：核心创新 🎯

QK-Clip 是 MuonClip 最关键的创新，它直接解决了注意力 logits 爆炸的根本原因。

问题根源：在 Transformer 中，注意力分数由 Query (Q) 和 Key (K) 矩阵的点积决定：
attention_score = Q @ K^T
当 Q 和 K 权重矩阵的谱范数（spectral norm）过大时，点积结果会爆炸式增长，导致训练不稳定。

QK-Clip 工作原理：
算法流程

def qk_clip_step(W_q, W_k, tau=100, alpha=0.5):
    """
    Args:
        W_q: Query 权重矩阵
        W_k: Key 权重矩阵  
        tau: 最大允许的 QK 分数阈值
        alpha: Q 和 K 之间的缩放平衡因子
    """
    # 1. 执行标准 Muon 更新
    W_q, W_k = muon_update(W_q, W_k)
    
    # 2. 后更新检查：计算最大 QK 分数
    max_qk_score = compute_max_qk_score(W_q, W_k)
    
    # 3. 如果超过阈值，重新缩放权重矩阵
    if max_qk_score > tau:
        eta = tau / max_qk_score  # eta < 1
        W_q = W_q * (eta ** alpha)
        W_k = W_k * (eta ** (1 - alpha))
    
    return W_q, W_k

关键特性：

• 权重级别干预：在权重层面直接控制，而不是在输出 logits 上做限制
• 后更新安全网：每次优化器步骤后检查，确保不会超出安全范围
  保持更新方向：通过缩放而非截断，保持了优化方向的一致性
  平衡调整：通过参数 α (通常为 0.5) 在 Q 和 K 之间平衡缩放
  与传统方法的对比：

方法	作用位置	问题
Logit Soft-capping	Softmax 之后	扭曲注意力分布，影响学习；未解决根本问题（权重过大）
Query-Key Normalization	向量归一化	未解决根本问题（权重过大）
QK-Clip	权重矩阵级别	从源头预防，不影响注意力分布

2.2 MuonClip 完整算法

def MuonClip(params, lr=2e-4, momentum=0.95, weight_decay=0.1,
             tau=100, alpha=0.5, ns_steps=5):
    """
    Args:
        params: 模型参数
        lr: 学习率
        momentum: 动量系数
        weight_decay: 权重衰减率
        tau: QK-Clip 阈值
        alpha: Q/K 缩放平衡因子
        ns_steps: Newton-Schulz 迭代次数
    """
    for W in params:
        # 1. 应用权重衰减
        W = (1 - weight_decay) * W
        
        # 2. 计算梯度更新（带动量）
        G = compute_gradient_with_momentum(W, momentum)
        
        # 3. Newton-Schulz 正交化
        G_orth = newton_schulz_orthogonalize(G, ns_steps)
        
        # 4. RMS 一致性缩放
        scaling = sqrt(max(W.shape)) * 0.2
        Update = lr * scaling * G_orth
        
        # 5. 应用更新
        W = W + Update
        
        # 6. QK-Clip（仅用于 Q/K 权重矩阵）
        if is_attention_weight(W):
            W = apply_qk_clip(W, tau, alpha)
    
    return W

2.3 训练配置与成果

2.3.1 超参数设置：

• 预训练阶段（前 10T tokens）
  learning_rate = 2e-4
  weight_decay = 0.1
  batch_size = 67M tokens
  qk_clip_tau = 100
  context_length = 4096

• 退火阶段（5.5T tokens）
  learning_rate = 2e-4 → 2e-5 (cosine decay)
  weight_decay = 0.1
  batch_size = 67M tokens

• 长上下文激活阶段
• 400B tokens @ 4k sequence length
• 60B tokens @ 32k sequence length
• 使用 YaRN 方法扩展到 128k

2.3.2 训练表现：

✅ 15.5T tokens 训练过程中零损失尖峰
✅ 损失曲线完全平滑稳定（见 Figure 3）
✅ 最大注意力 logits 在训练初期被限制在 100，约 30% 训练步骤后自然衰减到稳定范围
✅ 相比 AdamW 实现约 2 倍计算效率

2.3.3 MuonClip 的意义

Kimi K2 和 MuonClip 的成功证明了：

• 可扩展性突破：Muon 系列优化器可以稳定地扩展到万亿参数级别
• 问题导向创新：通过 QK-Clip 机制优雅地解决了大规模训练中的注意力爆炸问题
• 工程实用性：在真实的工业级训练中实现了零损失尖峰，这在大规模 LLM 训练中极为罕见
• 开源贡献：Kimi K2 模型和训练方法的开源为社区提供了宝贵的实践参考
  这标志着 Muon 优化器从学术研究走向大规模工业应用的成功转型，也为未来的超大规模模型训练提供了新的技术路径。

2.4 与传统优化器的对比

vs. SGD
SGD：基础的梯度下降，需要仔细调整学习率
Muon：通过矩阵正交化提供更稳定的更新方向

vs. Adam/AdamW
Adam/AdamW：逐元素自适应学习率，需要存储每个参数的一阶和二阶动量
Muon：矩阵结构感知，内存效率更高，收敛速度更快

vs. Shampoo
Shampoo：使用完整的二阶信息，计算成本高
Muon：使用 Newton-Schulz 迭代简化计算，在保持性能的同时大幅降低计算开销

2.5 Moun 扩展与改进

2.5.1 AdaMuon

AdaMuon 通过集成逐元素二阶矩估计来增强 Muon，实现更新级别的自适应性。它在 Muon 的结构化更新之上添加了动态学习率调制，既保持了全局几何保持性，又实现了细粒度的自适应性。

2.5.2 Gluon

Gluon 是另一个基于 Muon 的改进版本，进一步优化了 LMO（线性矩阵优化）的理论和实践之间的桥梁。

2.5.3 MuonClip - 工业级稳定性方案 ⭐

MuonClip 是 Moonshot AI 在 Kimi K2 项目中开发的 Muon 增强版本，专门解决万亿参数级训练中的稳定性问题。

适用场景对比
优化器 最佳规模 稳定性 计算效率 超参数调优
AdamW 通用 高 基准 简单
Muon <100B 中 2× AdamW 中等
MuonClip >100B 极高 2× AdamW 简单
AdaMuon 中等规模 高 1.5× AdamW 复杂
MuonClip 代表了 Muon 系列在大规模工业应用中的成熟形态，为训练超大规模模型提供了生产级解决方案。

使用注意事项

• 1. 适用性
     ✅ 适合：大型 Transformer 模型、卷积神经网络
     ⚠️ 需要测试：强化学习工作负载、微调任务
• 2. 分布式训练
     开源的分布式 Muon 实现是内存最优和通信高效的。Newton-Schulz 迭代可以在大规模 GPU 集群上正确分布。
• 3. 超参数调整
     扩展 Muon 的两个关键技术是：（1）添加权重衰减和（2）仔细调整每个参数的更新规模。这些技术使 Muon 能够在大规模训练中开箱即用，无需超参数调整。
• 4. 兼容性考虑
     对于监督微调（SFT），如果预训练优化器也是 Muon，则 Muon 效果更好。在切换优化器时需要考虑这一点。

已验证的能力
✅ 超大规模扩展：MuonClip 已成功扩展到 1T 参数和 15.5T tokens 的训练
✅ 零损失尖峰训练：在 Kimi K2 的整个训练过程中实现完全稳定
✅ 分布式效率：开源了内存最优和通信高效的分布式实现
待探索的方向
多模态扩展：Muon/MuonClip 在视觉-语言模型中的表现如何？
强化学习适配：如何优化 Muon 以更好地支持 RLHF 和 RL 训练？
微调策略：预训练用 MuonClip，微调时是否也应该使用？最佳切换策略是什么？
自动化超参数：能否进一步降低超参数调优负担，实现真正的"开箱即用"？
更极端的稀疏性：在更高稀疏度的 MoE 架构中，MuonClip 的表现如何？
社区生态
随着 Kimi K2 的开源和 MuonClip 技术细节的公开，预期会看到：

更多研究机构和公司采用 Muon 系列优化器
针对特定场景的优化变体（如 AdaMuonClip）
与其他训练技术的组合创新（如 curriculum learning、dynamic batching）
更多开源工具和库的支持
技术演进路径
Muon (2024)
├─→ Moonlight 实践 (5.7T tokens, 16B 参数)
├─→ MuonClip (2025)
│ └─→ Kimi K2 (15.5T tokens, 1T 参数) ⭐
├─→ AdaMuon (学术探索)
└─→ 未来变体
├─→ 多模态适配版本
├─→ 强化学习优化版本
└─→ 自动超参数版本
这些问题的答案将随着更多研究者和从业者采用 Muon/MuonClip 而逐渐明朗。Kimi K2 的成功已经证明，Muon 系列不仅是学术研究的有趣课题，更是能够在真实工业场景中产生重大影响的实用技术。

参考资源

• 官方资源
  Muon 原始论文：Keller Jordan 等人的技术报告
  Muon GitHub 仓库：https://github.com/KellerJordan/Muon
  Muon 博客：https://kellerjordan.github.io/posts/muon/

• Kimi K2 & MuonClip
  Kimi K2 论文：Kimi K2: Open Agentic Intelligence
  Kimi K2 GitHub：https://github.com/MoonshotAI/Kimi-K2
  Kimi K2 官方博客：https://moonshotai.github.io/Kimi-K2/
  Hugging Face 模型库：https://huggingface.co/moonshotai/Kimi-K2-Base

• MuonClip 技术解析：Fireworks AI 博客 - https://fireworks.ai/blog/muonclip
  Moonlight 模型
  Kimi.ai 应用案例：Moonlight 模型训练报告

• 相关研究
  Shampoo 优化器：分布式二阶优化方法
  AdaMuon：Muon + 自适应学习率
  Gluon：基于 Muon 的改进版本
  muP (Maximum Update Parametrization)：学习率迁移理论
  深度相对信任 (Deep Relative Trust)：权重空间与函数空间几何
  模块范数与对偶性：优化理论基础

• 技术博客与解析
  Medium 深度解析：
  “MuonClip: The Optimizer That Made Trillion-Parameter Kimi K2 Possible”
  “The Truth About KIMI K2 Pretraining: Muon Optimizer + MoE Unpacked”
  Baseten 技术博客：Kimi K2 架构与部署详解

• 开源社区
  讨论论坛：GitHub Discussions, Reddit r/MachineLearning
  技术社区：Hugging Face Forums, Papers with Code