正交化之外是什么?微软等提出ARO优化器:训练提速1/3,揭示矩阵优化新「蓝海」

📅 发布时间:2026/7/6 1:57:52 👁️ 浏览次数:
正交化之外是什么?微软等提出ARO优化器:训练提速1/3,揭示矩阵优化新「蓝海」
微软研究院联合港中文深圳、威斯康星大学麦迪逊分校最新放出的长篇论文从方法论创新、工程验证到理论诠释给出了肯定的答案。如果你在过去一年关注过大模型训练的技术大概率听过 Muon 这个名字 —— 这个在月之暗面 K2 模型的相关讨论中走红的优化器被视为是可能挑战 Adam 的新秀。它的思路很直接对动量矩阵进行正交化让各个奇异方向上的更新速率一致提升训练效率。Muon 带动了基于正交化算法的改进热潮但一个根本问题始终较少被讨论正交化方法究竟是通往高效训练的必经之路还是某个更深层原则的一个特例我们是否能跳出 “正交化” 这个框找到矩阵优化算法的新 “蓝海”微软研究院联合港中文深圳、威斯康星大学麦迪逊分校最新放出的长篇论文从方法论创新、工程验证到理论诠释给出了肯定的答案。团队首先将现有常用矩阵优化器统一到基于旋转的视角 —— 在旋转后的坐标系中最速下降。论文把 “梯度旋转” 作为第一原则让旋转策略动态地提升最速下降的速率推导出一类新的优化器ARO自适应旋转优化Adaptively Rotated Optimization。Muon 可被视为 ARO 的一个特例。通过严格控制的大规模训练 ARO 将大语言模型的训练效率相对 AdamW 提升了约 1/3额外时间开销压在 3% 以内比 Muon 还要高效 10%15%且在最多 80 亿参数、多倍过训练的压力测试下未出现收益递减迹象。最后论文还进一步探究更深层问题为什么旋转是本质的首先他们通过理论分析提出了对称性假设 —— 即旋转 / 乃至矩阵优化本质上可能是利用了大模型架构丰富的对称性而 ARO 的旋转策略则进一步利用了这种 “红利”在收敛效率与鲁棒性之间取得了更好的权衡。作者将对称性观点反馈在 ARO 完善上进一步开发跨层耦合等新特性取得良好效果。论文标题ARO: A New Lens On Matrix Optimization For Large Models论文地址https://arxiv.org/abs/2602.09006作者Wenbo Gong, Javier Zazo, Qijun Luo, Puqian Wang, James Hensman, Chao Ma机构微软研究院香港中文大学深圳威斯康星大学麦迪逊分校旋转更一般的优化框架论文指出如果把 Muon、SOAP、SPlus、Galore 等常见矩阵优化方法进行简化和抽象它们本质上都是在一个被旋转后的坐标系中使用 Adam 或者变体进行模型优化。它们先找到一个旋转矩阵 R把梯度 G 旋转到新的坐标系下用某个基座优化器 f 计算单步更新量最后将该更新量旋转回原来的坐标。这个过程可以写成对于旋转 R上述方法无一例外取为梯度内积矩阵的特征向量后文简称为特征旋转。同时它们将基座优化器设定为 Adam 或其变体。而 Muon 的正交化则是使用特定 Adam 变体的一个特例。这表明梯度旋转有潜力成为比正交化更加一般的优化框架。ARO 优化器将梯度旋转作为第一原则论文提出将旋转最速下降提升到设计优化器的新原则 从而可以考虑更一般的旋转 R和更广泛的基座优化器 f。能不能让这两个部分有机地联动起来去优化一个具体的训练效率指标例如模型训练损失的下降速度论文提出给定一般的基座优化器 f我们可以近似地求解旋转 R使旋转更新下的训练损失下降速率得以提升。于是我们推导出 ARO 的更新规则其中 M 是动量它的直觉很简单ARO 是在拿上一轮旋转后的基座优化器更新量寻找新的旋转去大概 “对齐” 原始的梯度动量。换句话说ARO 在主动地去寻找一个能让当前优化器 f 发挥得更好的旋转角度。实验发现图 1用 ARO 的更新方向比基于传统特征旋转的更新能带来更优的瞬时损失下降率 —— 该优势在整个训练过程中持续存在。严格控制的实验准则为了结论的可靠性论文给自己加了道槛优化器评估常面临一个痛点在研究级场景下所得出的结论很难迁移到实际场景。原因可能在于基准设置中的一些实验准则未与真实环境对齐导致指导性有限。对此论文规定了一套实验准则从混合精度选取、学习率衰减、非隐层优化器统一、到学习率迁移策略等环节都进行去偏控制并尽可能采用大的 batch size最高 1400 万、长的序列长度最高 4K足够大的模型规模最高 80 亿和训练预算最高 8 倍过训练尽可能贴近真实训练场景在可行的情况下对 AdamW 基线进行端到端调参避免用外推法估算超参。在这种规范下论文得出的加速率较为温和但在跨尺度测试下却更加一致、更可迁移。大规模实验显著、稳定、一致的效率提升在以上原则下评估分为两部分。小规模验证1 亿 - 15 亿参数 GPT中ARO 的旋转策略在多种基座优化器下均展现出普适性提升。横向对比无旋转和传统特征旋转版本以及横向对比 AdamW 和正交化方法全部表现更优图 2。这也侧面说明梯度旋转是一个非常关键的设计维度。规模化实验将 ARO 推向更大场景架构覆盖稠密和 MoE规模从 3 亿延伸至 80 亿激活参数训练预算拉到 1-8 倍 Chinchilla 过训练。结果显示图 3ARO 对 AdamW 保持约 1.3-1.35 倍加速对 Muon 等正交化方法保持约 1.1-1.15 倍加速且加速比在更大规模、更长周期下未见衰减。同时作者通过工程优化使得 ARO 在大规模分布训练下的额外开销相比 AdamW 控制在 3% 以内。一个有趣的 “副产物”全模型优化在主流的矩阵优化器实践策略中它们通常只用在隐藏层上 ——embedding 和 LM head 等参数还得靠 AdamW 来管。这被称为 “混合 / 分治模式”。其中一个原因是当其被直接用到上述参数上可能会导致训练显著变差甚至不收敛。而 ARO 路线下一个新的 “副产物” 是它可以在全模型参数上跑通。论文在 Sigma-MoE-2B 里对比了几种设置混合ARO 只优化隐藏层、全模型ARO 优化所有矩阵参数。结果表明图 4全模型模式的 ARO 在训练后期3 倍 - 4 倍过训练之后反而比混合模式效果更好。这意味着 ARO 原则上能够从旋转的角度统一地处理全模型的矩阵参数 —— 这也一定程度上挑战了当前矩阵优化器较为流行的 “分而治之” 的设计理念。为什么旋转是本质的一个更底层的视角对称性接下来论文进一步探究更深层的问题为什么旋转原则 “恰好” 隐藏在诸多矩阵优化器的设计中论文的拓展讨论指向了一个概念神经网络的参数对称性。微软团队在此前的工作SliceGPT中提出过一个定理Transformer 存在丰富的残差流对称性 —— 在特定约束下将参数同时旋转模型的输出不变。这意味着参数空间中存在连续区域其中所有点对应同一函数。与传统优化器相比ARO 在这片区域里多了一个可操作的自由度论文证明ARO 理论上等价于非欧几何下的对称瞬移Symmetry Teleportation—— 一类利用对称性信息加速收敛的经典算法。即在不改变损失的前提下ARO 将参数 “瞬移” 到群轨道中另一个更利于优化的位置再迈出下一步。论文进一步分析了 ARO 是如何利用这种自由度的。传统对称瞬移追求瞬时收敛速率的最大化但这在实际当中并不总能取得实际收益。对此论文主要理论证明了两个结论1. 随机梯度下大幅提高瞬时速率可能会导致损失下降不稳定2. 而 Muon/SOAP 等使用的特征旋转则是另一个极端最大化稳定性但同时会削弱下降速率取向于保守。ARO 的实现则采取了一种温和的部分提升策略在提升下降率的同时维持稳定性在收敛效率与鲁棒性之间取得了更好的权衡。这个视角下ARO 不再是单纯的矩阵运算技巧而是利用架构固有对称性的自然产物。论文将这一观察一般化为 “对称性假设”已知的矩阵优化器之所以有效可能是无意中利用了损失景观中的对称性。通过对称性视角进一步解锁优化 “新姿势”对称性视角不仅是对于优化的新诠释也进一步为 ARO 解锁了 “新姿势”。例如残差流对称性自然地包含了 embedding 和 lm head—— 二者在对称性的语义下与隐含层并无本质不同。因此在对称性视角下ARO 可用于全模型优化上这与大规模实验中的观测吻合。对称性关系揭示了跨层、跨模块之间的耦合约束。例如受同一段残差流支配的矩阵如某一层的 QKV 和上一层的输出投影理当绑定同一个旋转。这提供了一种经济利用跨层相关性的途径 —— 不是通过暴力计算全局二阶矩而是通过架构自身的耦合关系绑定旋转。在小规模模型上初步验证跨层绑定旋转不仅能降低计算开销还显著提升了优化性能。写在最后回过头看ARO 的贡献可以分为三部分把 “旋转” 从既有优化器的隐含设计里提炼为第一原则通过严格的规模化实验证明其有效性用架构本身的全局性质为矩阵优化提供新的诠释并衍生出新的耦合设计。如果说 Muon 优化器是从 “向量到矩阵的本质跨越”那么 ARO 则指向一个新的可能从 “矩阵优化” 走向 “全模型耦合优化”—— 优化器的设计也许应该和架构绑得更紧一些。