AMO: Adaptive Muon Orthogonalization

缪子优化器（Muon）近期已成为大规模预训练任务中一种颇具竞争力的AdamW替代方案，其核心操作是通过牛顿–舒尔茨（Newton-Schulz，NS）迭代实现权重矩阵的正交化。现有各类Muon变体均对所有参数矩阵采用统一的NS迭代调度策略，却忽视了不同矩阵在正交化难度上的潜在差异，以及该差异对模型最终性能的影响。我们通过系统性的实证研究发现，这种“按矩阵而异”的异质性普遍存在，且主要由矩阵自身的几何性质所决定；而该几何性质又会随算子类型、训练阶段及网络深度的不同而动态演化。因此，采用统一的NS调度策略将导致模型各部分的正交化质量参差不齐。受此发现启发，我们提出了自适应缪子正交化方法（Adaptive Muon Orthogonalization, AMO）：这是一种“先观测、后决策”的策略——在训练初期即按算子类型分别测量各权重矩阵的几何特征，并基于这些观测信号，为后续整个训练过程动态分配NS迭代资源。AMO在标准预训练、延长预训练以及持续预训练等多种设定下，均稳定优于采用统一NS调度的Muon基线；在涵盖12项下游任务的平均性能评测中，其相较最强基线分别将Llama3.1-1.4B模型和Qwen3-1.7B模型的得分提升了+0.76和+0.51。

现有Muon优化器对所有参数矩阵采用统一的Newton-Schulz（NS）正交化迭代次数调度，忽视了不同矩阵在几何结构（如条件数、谱分布、秩退化程度）上的显著异质性——这种异质性随算子类型（如Attention QKV、FFN）、训练阶段（初期vs后期）和网络深度动态变化，导致正交化质量不均衡，进而限制预训练收敛性与下游泛化性能。该问题此前未被系统识别和建模，属于优化器内部自适应机制设计的新挑战。

提出Adaptive Muon Orthogonalization（AMO）：一种轻量级'先观测、后决策'范式——在训练早期（仅前1–2%步）按算子类型（如attn.q_proj, ffn.w1等）分组统计权重矩阵的几何指标（如最小奇异值、核范数比、Frobenius/核范数比），据此为每类算子分配差异化NS迭代次数预算，并全程固定执行；避免在线测量开销，兼具理论可解释性与工程实用性。

• 实证揭示了Transformer各层/各类算子权重矩阵的正交化难度存在强几何规律性（如QKV矩阵普遍比FFN更难正交化，浅层比深层更敏感）； • 在Llama3.1-1.4B和Qwen3-1.7B上完成标准、延长（2×tokens）及持续预训练三类设置，平均提升12项下游任务0.76/0.51分（绝对gain），显著优于AdamW、原版Muon及MuP调优基线； • 代码已开源（GitHub: amo-opt），支持HuggingFace Transformers无缝集成； • 关键发现：几何信号在训练极早期（<500 steps）即稳定，使AMO无需任何梯度或loss反馈，仅靠静态权重分析即可实现高效适配； • 值得深入：将几何感知扩展至梯度空间、探索NS迭代次数的课程学习动态调整、以及在MoE架构中对专家矩阵的细粒度正交化调度。

• 'Muon: Orthogonal Optimization for Language Models' (ICLR 2024) —— 首提将Newton-Schulz迭代嵌入优化器的正交化框架； • 'AdaFactor: Adaptive Learning Rates with Sublinear Memory Cost' (ICML 2019) —— 早期自适应二阶方法，但未建模矩阵几何； • 'MuP: Deep Learning with Infinite Width' (NeurIPS 2022) —— 从缩放律角度指导初始化与参数化，启发AMO对算子级差异的关注； • 'Orthogonal Weight Initialization and Training Dynamics' (ICLR 2023) —— 理论分析正交初始化对训练稳定性的影响，但未解决训练中动态正交化调度问题； • 'DyNAMO: Dynamic Adaptive Momentum Optimization' (ACL 2024) —— 针对动量项的自适应，与AMO正交化维度正交互补。