微軟等提出ARO優(yōu)化器：訓練提速1/3，揭示矩陣優(yōu)化新「藍海」

如果你在過去一年關注過大模型訓練的技術，大概率聽過 Muon 這個名字 —— 這個在月之暗面 K2 模型的相關討論中走紅的優(yōu)化器，被視為是可能挑戰(zhàn) Adam 的新秀。它的思路很直接：對動量矩陣進行正交化，讓各個奇異方向上的更新速率一致，提升訓練效率。

Muon 帶動了基于正交化算法的改進熱潮，但一個根本問題始終較少被討論：正交化方法，究竟是通往高效訓練的必經(jīng)之路，還是某個更深層原則的一個特例？我們是否能跳出 “正交化” 這個框，找到矩陣優(yōu)化算法的新 “藍海”？

微軟研究院聯(lián)合港中文（深圳）、威斯康星大學麥迪遜分校最新放出的長篇論文，從方法論創(chuàng)新、工程驗證到理論詮釋，給出了肯定的答案。

• 團隊首先將現(xiàn)有常用矩陣優(yōu)化器統(tǒng)一到基于旋轉的視角 ——在旋轉后的坐標系中最速下降
• 論文把 “梯度旋轉” 作為第一原則，讓旋轉策略動態(tài)地提升最速下降的速率，推導出一類新的優(yōu)化器：ARO（自適應旋轉優(yōu)化，Adaptively Rotated Optimization）。Muon 可被視為 ARO 的一個特例。
• 通過嚴格控制的大規(guī)模訓練， ARO 將大語言模型的訓練效率相對 AdamW 提升了約 1/3（額外時間開銷壓在 3% 以內），比 Muon 還要高效 10%～15%，且在最多 80 億參數(shù)、多倍過訓練的壓力測試下，未出現(xiàn)收益遞減跡象
• 最后，論文還進一步探究更深層問題：為什么旋轉是本質的？首先，他們通過理論分析，提出了對稱性假設—— 即旋轉 / 乃至矩陣優(yōu)化，本質上可能是利用了大模型架構豐富的對稱性；而 ARO 的旋轉策略則進一步利用了這種 “紅利”，在收斂效率與魯棒性之間取得了更好的權衡。作者將對稱性觀點反饋在 ARO 完善上，進一步開發(fā)跨層耦合等新特性，取得良好效果。

• 論文標題：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
• 機構：微軟研究院，香港中文大學（深圳），威斯康星大學麥迪遜分校

旋轉：更一般的優(yōu)化框架

論文指出，如果把 Muon、SOAP、SPlus、Galore 等常見矩陣優(yōu)化方法進行簡化和抽象，它們本質上都是在一個被旋轉后的坐標系中，使用 Adam 或者變體進行模型優(yōu)化。它們先找到一個旋轉矩陣 R，把梯度 G 旋轉到新的坐標系下；用某個基座優(yōu)化器 f 計算單步更新量；最后，將該更新量旋轉回原來的坐標。這個過程可以寫成：

ARO 優(yōu)化器：將梯度旋轉作為第一原則

論文提出將旋轉最速下降提升到設計優(yōu)化器的新原則， 從而可以考慮更一般的旋轉 R，和更廣泛的基座優(yōu)化器 f。能不能讓這兩個部分有機地聯(lián)動起來，去優(yōu)化一個具體的訓練效率指標，例如模型訓練損失的下降速度？

論文提出：給定一般的基座優(yōu)化器 f，我們可以近似地求解旋轉 R，使旋轉更新下的訓練損失下降速率得以提升。于是，我們推導出 ARO 的更新規(guī)則：

（其中M是動量）

它的直覺很簡單：ARO 是在拿上一輪旋轉后的基座優(yōu)化器更新量，尋找新的旋轉去大概 “對齊” 原始的梯度動量。換句話說，ARO 在主動地去尋找一個能讓當前優(yōu)化器 f 發(fā)揮得更好的旋轉角度。實驗發(fā)現(xiàn)（圖 1），用 ARO 的更新方向，比基于傳統(tǒng)特征旋轉的更新，能帶來更優(yōu)的瞬時損失下降率 —— 該優(yōu)勢在整個訓練過程中持續(xù)存在。

嚴格控制的實驗準則：為了結論的可靠性，論文給自己加了道檻

優(yōu)化器評估常面臨一個痛點：在研究級場景下所得出的結論，很難遷移到實際場景。原因可能在于基準設置中的一些實驗準則未與真實環(huán)境對齊，導致指導性有限。

對此，論文規(guī)定了一套實驗準則：從混合精度選取、學習率衰減、非隱層優(yōu)化器統(tǒng)一、到學習率遷移策略等環(huán)節(jié)都進行去偏控制；并盡可能采用大的 batch size（最高 1400 萬）、長的序列長度（最高 4K），足夠大的模型規(guī)模（最高 80 億）和訓練預算（最高 8 倍過訓練），盡可能貼近真實訓練場景；在可行的情況下對 AdamW 基線進行端到端調參，避免用外推法估算超參。在這種規(guī)范下，論文得出的加速率較為溫和，但在跨尺度測試下卻更加一致、更可遷移。

大規(guī)模實驗：顯著、穩(wěn)定、一致的效率提升

在以上原則下，評估分為兩部分。

• 小規(guī)模驗證（1 億 - 15 億參數(shù) GPT）中，ARO 的旋轉策略在多種基座優(yōu)化器下均展現(xiàn)出普適性提升。橫向對比無旋轉和傳統(tǒng)特征旋轉版本，以及橫向對比 AdamW 和正交化方法，全部表現(xiàn)更優(yōu)（圖 2）。這也側面說明，梯度旋轉是一個非常關鍵的設計維度。

• 規(guī)模化實驗將 ARO 推向更大場景：架構覆蓋稠密和 MoE，規(guī)模從 3 億延伸至 80 億激活參數(shù)，訓練預算拉到 1-8 倍 Chinchilla 過訓練。結果顯示（圖 3），ARO 對 AdamW 保持約 1.3-1.35 倍加速，對 Muon 等正交化方法保持約 1.1-1.15 倍加速，且加速比在更大規(guī)模、更長周期下未見衰減。同時，作者通過工程優(yōu)化使得 ARO 在大規(guī)模分布訓練下的額外開銷相比 AdamW 控制在 3% 以內。

一個有趣的 “副產物”：全模型優(yōu)化

在主流的矩陣優(yōu)化器實踐策略中，它們通常只用在隱藏層上 ——embedding 和 LM head 等參數(shù)還得靠 AdamW 來管。這被稱為 “混合 / 分治模式”。其中一個原因是當其被直接用到上述參數(shù)上，可能會導致訓練顯著變差，甚至不收斂。而 ARO 路線下一個新的 “副產物” 是：它可以在全模型參數(shù)上跑通。

論文在 Sigma-MoE-2B 里對比了幾種設置：混合（ARO 只優(yōu)化隱藏層）、全模型（ARO 優(yōu)化所有矩陣參數(shù)）。結果表明（圖 4），全模型模式的 ARO 在訓練后期（3 倍 - 4 倍過訓練之后）反而比混合模式效果更好。

這意味著 ARO 原則上能夠從旋轉的角度，統(tǒng)一地處理全模型的矩陣參數(shù)—— 這也一定程度上挑戰(zhàn)了當前矩陣優(yōu)化器較為流行的 “分而治之” 的設計理念。

為什么旋轉是本質的？一個更底層的視角：對稱性

接下來，論文進一步探究更深層的問題：為什么旋轉原則 “恰好” 隱藏在諸多矩陣優(yōu)化器的設計中？論文的拓展討論指向了一個概念：神經(jīng)網(wǎng)絡的參數(shù)對稱性。

微軟團隊在此前的工作（SliceGPT）中提出過一個定理：Transformer 存在豐富的殘差流對稱性—— 在特定約束下將參數(shù)同時旋轉，模型的輸出不變。這意味著參數(shù)空間中存在連續(xù)區(qū)域，其中所有點對應同一函數(shù)。

與傳統(tǒng)優(yōu)化器相比，ARO 在這片區(qū)域里多了一個可操作的自由度：論文證明，ARO 理論上等價于非歐幾何下的對稱瞬移（Symmetry Teleportation）—— 一類利用對稱性信息加速收斂的經(jīng)典算法。即，在不改變損失的前提下，ARO 將參數(shù) “瞬移” 到群軌道中另一個更利于優(yōu)化的位置，再邁出下一步。

論文進一步分析了 ARO 是如何利用這種自由度的。傳統(tǒng)對稱瞬移追求瞬時收斂速率的最大化，但這在實際當中并不總能取得實際收益。對此，論文主要理論證明了兩個結論：1. 隨機梯度下大幅提高瞬時速率可能會導致?lián)p失下降不穩(wěn)定；2. 而 Muon/SOAP 等使用的特征旋轉則是另一個極端，最大化穩(wěn)定性但同時會削弱下降速率，取向于保守。ARO 的實現(xiàn)則采取了一種溫和的部分提升策略，在提升下降率的同時維持穩(wěn)定性，在收斂效率與魯棒性之間取得了更好的權衡。

這個視角下，ARO 不再是單純的矩陣運算技巧，而是利用架構固有對稱性的自然產物。論文將這一觀察一般化為“對稱性假設”：已知的矩陣優(yōu)化器之所以有效，可能是無意中利用了損失景觀中的對稱性。

通過對稱性視角，進一步解鎖優(yōu)化 “新姿勢”

對稱性視角不僅是對于優(yōu)化的新詮釋，也進一步為 ARO 解鎖了 “新姿勢”。例如：

• 殘差流對稱性自然地包含了 embedding 和 lm head—— 二者在對稱性的語義下與隱含層并無本質不同。因此，在對稱性視角下，ARO 可用于全模型優(yōu)化上，這與大規(guī)模實驗中的觀測吻合。
• 對稱性關系揭示了跨層、跨模塊之間的耦合約束。例如，受同一段殘差流支配的矩陣（如某一層的 QKV 和上一層的輸出投影）理當綁定同一個旋轉。這提供了一種經(jīng)濟利用跨層相關性的途徑 —— 不是通過暴力計算全局二階矩，而是通過架構自身的耦合關系綁定旋轉。在小規(guī)模模型上初步驗證：跨層綁定旋轉不僅能降低計算開銷，還顯著提升了優(yōu)化性能。

回過頭看，ARO 的貢獻可以分為三部分：把 “旋轉” 從既有優(yōu)化器的隱含設計里提煉為第一原則；通過嚴格的規(guī)模化實驗證明其有效性；用架構本身的全局性質為矩陣優(yōu)化提供新的詮釋，并衍生出新的耦合設計。如果說 Muon 優(yōu)化器是從 “向量到矩陣的本質跨越”，那么 ARO 則指向一個新的可能：從 “矩陣優(yōu)化” 走向 “全模型耦合優(yōu)化”—— 優(yōu)化器的設計，也許應該和架構綁得更緊一些。