![]()
如果你在過去一年關注過大模型訓練的技術,大概率聽過 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 計算單步更新量;最后,將該更新量旋轉回原來的坐標。這個過程可以寫成:
![]()
![]()
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 過訓練。結果顯示(圖 3),ARO 對 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 則指向一個新的可能:從 “矩陣優化” 走向 “全模型耦合優化”—— 優化器的設計,也許應該和架構綁得更緊一些。
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.