Newton-Muon優化器揭示Muon隱含牛頓法本質，並透過修正輸入資料幾何實現更高效LLM訓練

Newton-Muon優化器揭示Muon隱含牛頓法本質，並透過修正輸入資料幾何實現更高效LLM訓練。

Weijie Su團隊發布論文「The Newton-Muon Optimizer」（arXiv:2604.01472），證明熱門的Muon優化器實為隱含牛頓法，卻忽略了輸入激活的各向異性結構，導致效能未盡理想；他們以此洞見設計出Newton-Muon，實測在Modded-NanoGPT配置上僅需6%更少步數、4%更少訓練時間。

Muon優化器的流行與謎團

Muon優化器在「Kimi」和「GLM」等模型中廣受採用，提供約2倍於AdamW的運算效率，但其核心機制——丟棄梯度奇異值的矩陣梯度正交化——設計原理長期不明。Muon更新規則為：令Gk = ∇f(Wk) = U ΣV⊤為SVD分解，則Wk+1 = Wk − γk U V⊤，實務中加入動量Mk = β Mk−1 + (1 − β) Gk。此方法將梯度投影至半正交空間，顯著提升Transformer等矩陣主導模型的訓練速度，卻未解釋為何此舉如此有效。

三元二次代理模型的突破

團隊引入「triplet quadratic surrogate」代理模型，僅用三元素近似損失：梯度G、輸出空間曲率H，以及層輸入激活堆疊資料矩陣Z。最小化此二次代理（模擬牛頓法精神），曲率H神奇完全抵消，剩下依賴位移項的更新；在最簡各向同性位移假設下，得出封閉形式：W ← W − η · msgn(G (ZZᵀ)⁻¹)，其中msgn(X)為X=USVᵀ的緊湊SVD下UVᵀ。此即Newton-Muon，精準捕捉輸入第二矩的右預條件作用。

標準Muon的隱含缺陷

關鍵洞見：標準Muon等同於ZZᵀ = c*I的Newton-Muon，隱含假設輸入激活Z為各向同性；但真實LLM訓練中，ZZᵀ極度各向異性，導致Muon忽略資料幾何，無法最佳化譜結構。Newton-Muon修正此盲點，尊重輸入分佈，讓更新更貼合Transformer的矩陣本質（如注意力QKV與MLP投影）。

實證效能優勢

在最早公開的Modded-NanoGPT Muon配置（Record #4，用於GPT-2預訓練）上，Newton-Muon達相同驗證損失需6%更少迭代步數、約4%更少壁鐘時間，每步成本僅增加1.8%。此結果證明修正資料幾何的微調，即帶來顯著加速，無需大幅犧牲效率。

建基前期各向同性曲率模型

Newton-Muon源自前期工作「Isotropic Curvature Model for Understanding Deep Learning Optimization」（arXiv:2511.00674），該模型假設損失曲率（二階Hessian及高階項）在所有擾動方向各向同性，轉化為凸優化程式，便於分析單步更新對總損失的影響。應用於Muon，證明最佳更新應使梯度矩陣奇異值譜更均勻化（spectrum homogenization），改善條件數；正交化僅在曲率成長有「相變 kink」時才最優，否則過度，Muon方向正確但不嚴格最優。

幻燈片解析的設計哲學

Weijie Su的幻燈片（https://www.weijie-su.com/files/DL_Opt_muon.pdf）強調優化器對LLM訓練至關重要，前沿基模預訓練成本逾數億美元，瓶頸可能來自資料、GPU或架構。AdamW為標準（更新含動量mk、變異數vk、偏差修正及權重衰減），但忽略矩陣結構；Muon雖優，團隊策略為建簡單一步代理模型，先理解再設計。幻燈片詳述各向同性曲率假設：輸入zi近似球面分佈，解耦一階與高階資訊，代理為min Q − Tr(Q G⊤) + E z∼sphere H(‖Q z‖)，顯露矩陣結構。譜均勻化優於牛頓-舒爾茨迭代，正交化需曲率「kink」才最優；Muon「方向對但走太遠」，真原理為譜均勻化。

對LLM優化趨勢的啟示

簡而言之，Muon是「忘記資料幾何的牛頓法」，Newton-Muon「記得」並修正，預示未來優化器須更精準捕捉輸入Z的異質性，而非粗糙正交化。此不僅解釋Muon成功，更開啟代理模型導向設計新途徑，潛在解決LLM縮放瓶頸。相較AdamW的元素級更新，矩陣級洞見凸顯Transformer時代的轉變，預料將影響後續SOTA模型訓練。