论文
从 SGD 到 Muon:通过 Schatten-p 范数进行自适应优化
From SGD to Muon: Adaptive Optimization via Schatten-p Norms
摘要
现代优化器,如 Muon,对其更新施加矩阵式几何约束。这些矩阵式约束可以在线性最小化 Oracle (LMO) 理论下统一。然而,当前的所有方法都对更新规则强加了固定的 LMO 几何形状,这些几何形状是通过设计或经验选择的,根据问题的几何形状不一定是最佳的。我们引入了一种新颖的高效数据驱动标准,用于在各个深度神经网络层上动态选择代理最优更新 LMO 几何结构。我们的标准使用单步随机特征回归代理模型从梯度和激活统计数据以封闭形式导出,导航从 SGD 插值到 Muon 更新的设计空间。此外,集成参数化预处理允许我们的框架将 SGD、Muon、Adam 和 MuAdam 恢复为特定极值。为了使这种自适应方法具有可扩展性,我们将其与高效的计算策略配对,在高度优化的基线上仅实现 $\sim$ 3% 的运行时开销。作为概念证明,我们证明了这种数据驱动的优化器在三个不同的训练场景中击败了 Muon 和 AdamW 之间表现最佳的优化器,或者保持了竞争力。最终,这项工作提供了证据,证明 LMO 几何结构可以成功、高效地根据运行时数据进行调整,为超越静态几何结构的优化器设计开辟了一条新途径。