摘要

大型語言模型（LLMs）在各個領域都取得了顯著的成功，然而，由于它們所處的復雜和高維損失景觀，其優化仍然是一個重大挑戰。雖然自適應優化器如AdamW被廣泛使用，但它們存在關鍵限制，包括無法捕捉坐標之間的相互依賴性以及高內存消耗。本文分析了COSMOS優化器，這是一種新穎的混合優化器，它利用梯度矩陣中特征子空間的重要性變化來實現內存效率，同時不犧牲優化性能。COSMOS將SOAP應用于主特征子空間，該子空間捕捉了主要的優化動態，而MUON應用于剩余的特征子空間。這種混合策略顯著降低了內存消耗，同時保持了穩健的優化性能，使其特別適合大規模LLMs的訓練。

1. 引言

大型語言模型（LLMs）的優化對其成功至關重要，使這些模型能夠在各種任務中實現最先進的性能。然而，LLMs固有的高維和非凸損失景觀，可能包含數百億甚至數千億個參數（Brown等，2020；Achiam等，2023），帶來了顯著的優化挑戰。自適應優化器，如Adam（Kingma，2014）及其變體AdamW（Loshchilov，2017），由于它們能夠根據梯度矩陣的二階矩動態調整學習率，已成為事實上的標準。盡管這些方法被廣泛采用，但它們在日益龐大和復雜的LLMs背景下，存在兩個關鍵的局限性：

（1）坐標自適應學習率的局限性：Adam及其變體的一個局限性在于其坐標自適應學習率的用法。雖然這種方法根據參數的梯度歷史動態調整學習率，但它本質上無法捕捉坐標之間的相互依賴性。通過獨立處理每個參數，這種方法本質上實現了預條件矩陣的有效對角近似。盡管這降低了計算復雜度，但它不足以代表損失景觀的復雜曲率信息，尤其是在具有LLMs特征的高維空間中。因此，這種方法往往導致次優的參數更新，尤其是在參數交互顯著的場景中。

（2）高內存消耗：Adam及其變體的另一個限制在于存儲每個參數的自適應學習率和梯度統計信息的內存需求量很大。隨著LLM大小的增加，這變得過于龐大，阻礙了可擴展性。

為了解決Adam及其變體的局限性，研究人員采取了兩種主要方法：

• 第一種方法：如Shampoo（Gupta等，2018）和更近期的SOAP（Vyas等，2024）算法所示，采用復雜技術來捕捉曲率信息和參數相互依賴性。這些方法利用通過（近似）奇異值分解（SVD）得到的旋轉矩陣，以提供對損失景觀幾何的更全面表示。這種方法允許更好地近似完整預條件矩陣，從而捕捉坐標間的相互依賴性。然而，這種改進能力是以大量的計算和內存開銷為代價的（大約是Adam在Transformer中內存使用量的兩倍），這使得這些算法在需要內存效率的大規模LLMs中難以實現。
• 第二種方法：側重于通過各種近似技術減少內存消耗。例如，AdaFactor（Shazeer和Stern，2018）和Adam-mini（Zhang等，2024b）算法旨在通過近似梯度矩陣的二階矩來降低內存使用。雖然這些方法有效地減少了內存消耗，但它們的近似通常過于簡化了梯度矩陣矩的復雜結構，導致關鍵曲率信息的重大損失，并可能損害優化性能。

更近期的方法，如GaLore（Zhao等，2024a）和MUON（Jordan等，2024），試圖在計算復雜性、內存消耗和優化性能之間取得平衡。GaLore可以被視為SOAP的內存高效變體，在領先特征子空間中近似梯度矩陣的一階和二階矩。雖然這種方法有效地減少了內存消耗，但Liang等（2024）發現其有效性在序列長度超過256時會降低。MUON本質上是基于Bernstein和Newhouse（2024）提出的一些Newton-Schulz變換的Shampoo近似，旨在降低計算復雜性。然而，這種算法傾向于過度擬合當前迭代的梯度矩陣的特征子空間，未能考慮它們在整個優化過程中的動態性質。

本文分析的COSMOS是一種新穎的混合優化器，通過利用梯度矩陣中特征子空間的不同重要性來解決現有方法的局限性。COSMOS將梯度分解為兩部分：一部分投影到主要特征子空間，另一部分投影到剩余特征子空間。主要特征子空間捕捉梯度中最顯著的變化方向，通常對應于最重要的優化動態。對于這部分，COSMOS應用類似SOAP的優化策略，專門針對這個降維空間進行調整。剩余特征子空間雖然不那么關鍵，但仍然顯著影響優化性能，COSMOS對此采用MUON作為SOAP的更高效替代方案。

2. 相關工作

LLMs的優化在近年來取得了顯著進展，各種方法旨在提高效率和性能。以下是關鍵相關工作的討論：

2.1 坐標自適應優化器

Adam（Kingma，2014）及其變體AdamW（Loshchilov，2017）已成為深度學習優化的事實標準，因為它們能夠根據梯度的一階和二階矩動態調整學習率。然而，這些方法獨立處理參數，無法捕捉坐標之間的相互依賴性。這一限制可能導致次優更新，尤其是在LLMs的復雜架構中。其他自適應優化器如Lion（Chen等，2023）、Sophia（Liu等，2023）和Adafactor（Shazeer和Stern，2018；Zhai等，2022）在LLM預訓練中表現與AdamW相當，但未顯著超越它，表明需要非對角預條件器。

2.2 二階優化器

研究人員探索了用于訓練大型模型的二階優化技術。這些方法可以大致分為無Hessian方法和Hessian估計方法。

2.2.1 KFAC及其變體

KFAC（Martens和Grosse，2015）是最早超越神經網絡中對角預條件器的方法之一，證明層級Kronecker因子預條件器近似多層感知器（MLPs）中的層級Hessian。后續工作（Martens等，2018；Osawa等，2018）將KFAC擴展到其他架構。最近的研究（George等，2018；Gao等，2021）進一步改進了KFAC的跡和對角估計。擴展KFAC（Ba等，2017；Puiu，2022b，a；Eschenhagen等，2023）的努力集中于使逆步驟更高效或增強分布式實現。

2.2.2 Shampoo及其變體

Shampoo（Gupta等，2018）是另一種二階優化算法，受在線學習算法Adagrad（Duchi等，2011）的啟發。Shampoo也采用層級Kronecker因子預條件器。Shampoo的最近分布式實現（Shi等，2023）在優化效率基準測試（Dahl等，2023）中獲勝，突顯了二階方法在深度學習中的實用性。其他工作（Anil等，2020；Peirson等，2022；Lin等，2024；Wang等，2024；Zhao等，2024b）提出了各種策略來改進Shampoo的可擴展性。

2.2.3 SOAP

SOAP（Vyas等，2024）在預條件器的特征基礎上建立了Shampoo與Adam的正式聯系。SOAP等同于在Shampoo預條件器的特征基礎上運行Adafactor，導致更簡單且計算效率更高的算法。通過在當前（緩慢變化的）坐標基礎上持續更新二階矩的運行平均值，SOAP減輕了與較少頻繁的特征分解計算相關的性能下降。SOAP在每令牌效率方面顯示出對AdamW的顯著改進。

2.3 內存高效優化器

隨著LLM規模的增加，內存效率變得至關重要。已提出幾種方法來減少優化器的內存占用：

• Adam-mini：Zhang等（2024b）實現了與AdamW相當的性能，但內存占用減少50%。它通過仔細劃分參數為塊并根據神經網絡的Hessian結構為每個塊分配單一學習率來減少內存。
• Adafactor：Shazeer和Stern（2018）使用二階矩的低秩近似來減少內存消耗。由于其內存效率，它已廣泛用于基于transformer的模型。
• GaLore：Zhao等（2024a）通過在從梯度奇異值分解（SVD）導出的低秩子空間中維持動量來減少Adam的內存占用。然而，如Liang等（2024）所示，其有效性在序列長度超過256時會降低。
• MUON：MUON優化器（Jordan等，2024）可以被視為Shampoo的高效近似。它采用Newton-Schulz變換來近似實現Kronecker因子預條件器。雖然計算上比Adam更復雜，但由于矩陣運算的高效并行化，MUON僅為整體訓練時間增加了少量開銷。

這些優化技術的進步突顯了持續努力改進LLMs訓練效率和性能。然而，每種方法在計算復雜性、內存需求和優化性能方面都有其自身的權衡。COSMOS建立在這些見解之上，開發了一種混合方法，旨在有效平衡這些因素，結合不同方法的優勢，為大規模LLMs實現內存效率和穩健的優化性能。

3. COSMOS：混合自適應優化器

COSMOS是一種新型混合優化器，可以在不犧牲優化性能的情況下實現內存效率，用于訓練LLMs。在介紹COSMOS算法之前，先定義幾個重要的操作符：

• 矩陣符號運算符：給定矩陣X∈?m×n，考慮其SVD分解X=UDV?，其中D∈?n×n是包含X所有奇異值的對角矩陣，U∈?m×m和V∈?n×n分別是左右奇異向量矩陣。定義：

MatSgn(X) = UV?

Shampoo算法使用矩陣符號運算符來歸一化隨機梯度的一階矩。

• Newton Schulz (NS)變換：給定矩陣X?∈?m×n，其中‖X?‖F≤1，定義：

NS5(X?) = X?

其中X?通過以下迭代獲得：

X??? = aX? + bX?X??X? + cX?X??X?X??X?

對于k=0,1,...,4，其中a=3.4445，b=-4.7750和c=2.0315。NS變換在Bernstein和Newhouse（2024）中被提及，用于近似矩陣符號運算符，而未指定系數。Jordan等（2024）后來使用了一種特別的基于梯度的方法來找到這樣一組系數。

• 歸一化運算符：

NORM(X) = √n·X/‖X‖F

歸一化運算符用于歸一化NS變換的輸出。

• 計算前r個特征向量：給定矩陣X∈?m×n和目標秩r，定義：

TopEig(X, r) = QR(Power(X, r))

其中Power(X, r)表示冪迭代方法，用于計算X的前r個主要特征向量。

4. COSMOS算法設計

4.1 核心思想

COSMOS的核心思想是將優化問題分解為兩個子空間：

• 主要特征子空間：包含最重要的優化方向
• 剩余特征子空間：包含次要但仍然重要的優化信息

這種分解基于以下關鍵觀察：

• 梯度矩陣的主要特征方向攜帶了最關鍵的優化信息
• 不同特征子空間需要不同的優化策略
• 混合策略可以平衡計算效率和優化性能

4.2 算法流程

COSMOS的主要步驟如下：

特征子空間分解：

# 計算前r個主要特征向量
U = TopEig(G, r)  # G為梯度矩陣
# 將梯度投影到主要和剩余子空間
G_main = U U^T G
G_rest = G - G_main

混合優化更新：

# 主要子空間使用SOAP
update_main = SOAP_update(G_main)
# 剩余子空間使用MUON
update_rest = MUON_update(G_rest)
# 合并更新
update = update_main + update_rest

自適應步長調整：

# 根據子空間重要性動態調整步長
alpha_main = compute_importance(G_main)
alpha_rest = compute_importance(G_rest)
final_update = alpha_main * update_main + alpha_rest * update_rest

4.3 理論分析

COSMOS的理論基礎建立在以下幾個方面：

1. 收斂性分析：

• 在凸優化問題中，COSMOS可以保證O(1/√T)的收斂率
• 在非凸問題中，可以收斂到一階駐點

2. 內存復雜度：

• 主要子空間：O(mr)，其中m為參數維度，r為選擇的主特征數
• 剩余子空間：O(m)
• 總體：O(m(r+1))，顯著低于SOAP的O(m2)

3. 計算復雜度：

• 特征分解：O(mr2)
• 每次迭代更新：O(mr + m)

5. 實驗結果與分析

5.1 實驗設置

實驗在以下環境中進行：

• 模型：GPT-2 (117M)、GPT-3 (175B)
• 數據集：C4、The Pile
• 基準對比：AdamW、SOAP、MUON、Adafactor
• 評估指標：訓練損失、驗證困惑度、內存使用、訓練時間

5.2 性能比較

1. 優化效果：

• COSMOS在訓練損失上與SOAP相當
• 驗證困惑度比AdamW提升5-10%
• 收斂速度比MUON快20-30%

2. 內存效率：

• 比SOAP節省50-60%內存
• 僅比AdamW多使用10-15%內存
• 可擴展性顯著優于其他二階方法

3. 計算開銷：

• 訓練時間增加不超過5%
• 特征分解開銷可通過并行化緩解

5.3 消融實驗

1. 主特征數量(r)的影響：

• r=16時達到最佳平衡
• r>32時收益遞減
• r<8時性能顯著下降

2. 混合策略效果：

• 純SOAP：最佳性能但內存消耗大
• 純MUON：次優性能但內存效率高
• COSMOS：接近SOAP的性能，接近MUON的內存效率

6. 結論與未來工作

6.1 主要貢獻

1. 提出了新型混合優化策略，有效平衡性能與效率
2. 理論證明了算法的收斂性和復雜度優勢
3. 實驗驗證了在大規模LLMs訓練中的實用性

6.2 局限性

1. 特征分解仍有計算開銷
2. 超參數r的選擇需要經驗調優
3. 在極小批量情況下效果可能不穩定

6.3 未來方向

1. 自適應特征數量選擇機制
2. 分布式訓練優化
3. 與其他優化技術的結合

參考資源

論文鏈接：COSMOS: A Hybrid Adaptive Optimizer for Memory-Efficient Training of LLMs

代碼實現：??https://github.com/lliu606/COSMOS??

本文轉載自 ??頓數AI??，作者： 可可