Convergence Rate Analysis of SOAP with Arbitrary Orthogonal Projection Matrices

被廣泛用於訓練大語言模型的矩陣優化器 SOAP 問世超過一年，卻始終沒有嚴格的收斂性理論保證。北京大學 Huan Li 與 Zhouchen Lin 在這篇 arXiv 短注中首次建立了 SOAP 的收斂速率，並將框架推廣至任意正交投影矩陣——唯一要求是投影矩陣在每次迭代時與當前隨機梯度條件獨立，這是絕大多數常見構造方式都能自然滿足的條件。

矩陣優化器的崛起：從 Adam 家族到 SOAP

過去十年，Adam 家族（AdaGrad、RMSProp、Adam、AdamW）是訓練深度神經網路的事實標準。這些方法把全部參數拉直為一個高維向量，對每個元素分別維護動量 M 和自適應二階矩 V，再逐元素更新。設計簡單通用，但代價是完全忽略了神經網路層本身的矩陣幾何結構——全連接層和 attention 投影層的參數本質上是一個 m×n 矩陣，向量化等於主動丟棄了這個低秩結構的資訊。

近年興起的矩陣優化器（matrix-based optimizer）直接在矩陣空間操作，在多項大模型預訓練實驗中展現了超越 Adam 家族的潛力。代表作包括 Shampoo（以 Kronecker 積近似完整曲率）、Muon（通過 Newton-Schulz 迭代做梯度正交化）、SOAP（Vyas et al., 2025，在預條件矩陣特徵向量張成的子空間中跑 Adam）、Splus（Frans et al., 2025）以及 ARO（Gong et al., 2026）。其中 SOAP、Splus、ARO 統稱「基於投影的優化器」（projection-based optimizer），共同思路是先把梯度投影到低秩子空間，再在投影空間內做 Adam 風格的自適應更新。

SOAP 的理論缺口：SVD 投影與 Adam 機制耦合難分析

向量類優化器的收斂性已有成熟文獻，矩陣類方法的理論基礎相對薄弱。Muon 因結構較簡單（純梯度正交化）已被充分分析；Shampoo 的 AdamW 風格實作的收斂保證則由 Li et al.（2026）建立。SOAP 的分析困難源於兩個機制的深度耦合：結構化的 SVD 正交投影，加上投影空間內部的 Adam 自適應步長——這兩者的糾纏讓標準隨機優化工具難以直接套用。

最接近的前置工作是 He et al.（2025）對 GaLore（Zhao et al., 2024）的一個變體進行分析，但他們替換了兩個關鍵設計：把 SVD 投影換成隨機投影，並以 MSGD（動量隨機梯度下降）取代 Adam 的自適應機制。論文作者直接指出，He et al. 的技術工具「遠遠不足以建立完整 SOAP 算法的收斂性」，SOAP 的理論保證是當時的開放問題。

首個 SOAP 收斂定理：核范數以 K 四次方根衰減

Theorem 1 建立了廣義 SOAP 的兩個收斂界。第一個（公式 3）以逐元素 ℓ₁ 范數度量投影後梯度，在 ‖X‖₁ = Θ(√(mn))‖X‖_F 的理想條件下可與隨機非凸優化的理論下界匹配（Arjevani et al., 2023）。第二個（公式 4）更直觀，以核范數（nuclear norm，即矩陣所有奇異值之和，可視為矩陣版的梯度大小）直接衡量收斂：

(1/K) ∑ E[‖∇f(X_k)‖_*] ≤ 10√(σ̂_F² L Δ / K σ_op²) + 4√(mn) ⁴√(σ̂_F² L Δ / K)

其中 Δ = f(X₁) - f 為初始函數值與最優值之差；m×n 是參數矩陣維度；σ_F²、σ_op² 分別為梯度 Frobenius 范數與算子范數（最大奇異值）的噪聲方差上界；L 為梯度 Lipschitz 常數。右側主導項以 K 的四次方根*衰減，與 Adam 系列在非凸隨機優化中的典型速率形式一致。證明框架沿用 Li et al.（2025a）的 AdamW 分析技術，核心鏈條是 Lipschitz 展開 → 動量誤差遞推 → Hölder 不等式收斂求和。

條件獨立假設統一覆蓋 Splus 和 ARO

本文最大的靈活性在於投影矩陣幾乎沒有構造限制——唯一的要求是條件獨立：在給定第 k-1 步之前所有資訊的條件下，投影矩陣 P_{k-1}、Q_{k-1} 與第 k 步的隨機梯度估計 G_k 條件獨立。形式上，這讓推導中出現的關鍵期望可以因子化：E_k[P^T G_k Q | F_{k-1}] = P^T · E_k[G_k | F_{k-1}] · Q，從而使梯度估計的隨機性能被乾淨地控制。

滿足這個條件的常見構造方式包括：原始 SOAP 中延遲一步更新的預條件矩陣特徵向量；前一步動量矩陣 M_{k-1} 的左右奇異向量；或前一步梯度 G_{k-1} 的奇異向量。這些方法都只使用到第 k-1 步之前的資訊，天然滿足條件獨立。因此 Theorem 1 不只適用於 SOAP 本身，也直接覆蓋 Splus 和 ARO，為整類基於投影的矩陣優化器提供了統一的收斂框架。

速率差距：核范數界比 Shampoo 慢 √min{m,n} 倍

定理也明確指出 SOAP 的理論局限。公式 4 的核范數收斂速率比 AdamW 風格 Shampoo（Li et al., 2026）慢 √min{m,n} 倍，比理論最優下界至少慢 √max{m,n} 倍。這個差距的數學根源是：核范數是么正不變范數，可以消去投影矩陣；但 ℓ₁ 范數不是么正不變的，投影矩陣留在界中帶入了 √(mn) 這個維度因子。這究竟是分析工具的局限，還是 SOAP 算法本身的計算瓶頸，目前仍是開放問題——作者並未聲稱這個差距不可逾越。

SOAP 首獲收斂保證，條件獨立投影統一覆蓋三類算法；但核范數速率比 Shampoo 慢 √min{m,n} 倍，理論收緊是下一步課題。