摘要

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

1. 引言

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

（1）坐標(biāo)自適應(yīng)學(xué)習(xí)率的局限性：Adam及其變體的一個(gè)局限性在于其坐標(biāo)自適應(yīng)學(xué)習(xí)率的用法。雖然這種方法根據(jù)參數(shù)的梯度歷史動(dòng)態(tài)調(diào)整學(xué)習(xí)率，但它本質(zhì)上無法捕捉坐標(biāo)之間的相互依賴性。通過獨(dú)立處理每個(gè)參數(shù)，這種方法本質(zhì)上實(shí)現(xiàn)了預(yù)條件矩陣的有效對(duì)角近似。盡管這降低了計(jì)算復(fù)雜度，但它不足以代表損失景觀的復(fù)雜曲率信息，尤其是在具有LLMs特征的高維空間中。因此，這種方法往往導(dǎo)致次優(yōu)的參數(shù)更新，尤其是在參數(shù)交互顯著的場景中。

（2）高內(nèi)存消耗：Adam及其變體的另一個(gè)限制在于存儲(chǔ)每個(gè)參數(shù)的自適應(yīng)學(xué)習(xí)率和梯度統(tǒng)計(jì)信息的內(nèi)存需求量很大。隨著LLM大小的增加，這變得過于龐大，阻礙了可擴(kuò)展性。

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

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

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

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

2. 相關(guān)工作

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

2.1 坐標(biāo)自適應(yīng)優(yōu)化器

Adam（Kingma，2014）及其變體AdamW（Loshchilov，2017）已成為深度學(xué)習(xí)優(yōu)化的事實(shí)標(biāo)準(zhǔn)，因?yàn)樗鼈兡軌蚋鶕?jù)梯度的一階和二階矩動(dòng)態(tài)調(diào)整學(xué)習(xí)率。然而，這些方法獨(dú)立處理參數(shù)，無法捕捉坐標(biāo)之間的相互依賴性。這一限制可能導(dǎo)致次優(yōu)更新，尤其是在LLMs的復(fù)雜架構(gòu)中。其他自適應(yīng)優(yōu)化器如Lion（Chen等，2023）、Sophia（Liu等，2023）和Adafactor（Shazeer和Stern，2018；Zhai等，2022）在LLM預(yù)訓(xùn)練中表現(xiàn)與AdamW相當(dāng)，但未顯著超越它，表明需要非對(duì)角預(yù)條件器。

2.2 二階優(yōu)化器

研究人員探索了用于訓(xùn)練大型模型的二階優(yōu)化技術(shù)。這些方法可以大致分為無Hessian方法和Hessian估計(jì)方法。

2.2.1 KFAC及其變體

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

2.2.2 Shampoo及其變體

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

2.2.3 SOAP

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

2.3 內(nèi)存高效優(yōu)化器

隨著LLM規(guī)模的增加，內(nèi)存效率變得至關(guān)重要。已提出幾種方法來減少優(yōu)化器的內(nèi)存占用：

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

這些優(yōu)化技術(shù)的進(jìn)步突顯了持續(xù)努力改進(jìn)LLMs訓(xùn)練效率和性能。然而，每種方法在計(jì)算復(fù)雜性、內(nèi)存需求和優(yōu)化性能方面都有其自身的權(quán)衡。COSMOS建立在這些見解之上，開發(fā)了一種混合方法，旨在有效平衡這些因素，結(jié)合不同方法的優(yōu)勢，為大規(guī)模LLMs實(shí)現(xiàn)內(nèi)存效率和穩(wěn)健的優(yōu)化性能。

3. COSMOS：混合自適應(yīng)優(yōu)化器

COSMOS是一種新型混合優(yōu)化器，可以在不犧牲優(yōu)化性能的情況下實(shí)現(xiàn)內(nèi)存效率，用于訓(xùn)練LLMs。在介紹COSMOS算法之前，先定義幾個(gè)重要的操作符：

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

MatSgn(X) = UV?

Shampoo算法使用矩陣符號(hào)運(yùn)算符來歸一化隨機(jī)梯度的一階矩。

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

NS5(X?) = X?

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

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

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

• 歸一化運(yùn)算符：

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

歸一化運(yùn)算符用于歸一化NS變換的輸出。

• 計(jì)算前r個(gè)特征向量：給定矩陣X∈?m×n和目標(biāo)秩r，定義：

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

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

4. COSMOS算法設(shè)計(jì)

4.1 核心思想

COSMOS的核心思想是將優(yōu)化問題分解為兩個(gè)子空間：

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

這種分解基于以下關(guān)鍵觀察：

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

4.2 算法流程

COSMOS的主要步驟如下：

特征子空間分解：

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

混合優(yōu)化更新：

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

自適應(yīng)步長調(diào)整：

# 根據(jù)子空間重要性動(dòng)態(tài)調(diào)整步長
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的理論基礎(chǔ)建立在以下幾個(gè)方面：

1. 收斂性分析：

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

2. 內(nèi)存復(fù)雜度：

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

3. 計(jì)算復(fù)雜度：

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

5. 實(shí)驗(yàn)結(jié)果與分析

5.1 實(shí)驗(yàn)設(shè)置

實(shí)驗(yàn)在以下環(huán)境中進(jìn)行：

• 模型：GPT-2 (117M)、GPT-3 (175B)
• 數(shù)據(jù)集：C4、The Pile
• 基準(zhǔn)對(duì)比：AdamW、SOAP、MUON、Adafactor
• 評(píng)估指標(biāo)：訓(xùn)練損失、驗(yàn)證困惑度、內(nèi)存使用、訓(xùn)練時(shí)間

5.2 性能比較

1. 優(yōu)化效果：

• COSMOS在訓(xùn)練損失上與SOAP相當(dāng)
• 驗(yàn)證困惑度比AdamW提升5-10%
• 收斂速度比MUON快20-30%

2. 內(nèi)存效率：

• 比SOAP節(jié)省50-60%內(nèi)存
• 僅比AdamW多使用10-15%內(nèi)存
• 可擴(kuò)展性顯著優(yōu)于其他二階方法

3. 計(jì)算開銷：

• 訓(xùn)練時(shí)間增加不超過5%
• 特征分解開銷可通過并行化緩解

5.3 消融實(shí)驗(yàn)

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

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

2. 混合策略效果：

• 純SOAP：最佳性能但內(nèi)存消耗大
• 純MUON：次優(yōu)性能但內(nèi)存效率高
• COSMOS：接近SOAP的性能，接近MUON的內(nèi)存效率

6. 結(jié)論與未來工作

6.1 主要貢獻(xiàn)

1. 提出了新型混合優(yōu)化策略，有效平衡性能與效率
2. 理論證明了算法的收斂性和復(fù)雜度優(yōu)勢
3. 實(shí)驗(yàn)驗(yàn)證了在大規(guī)模LLMs訓(xùn)練中的實(shí)用性

6.2 局限性

1. 特征分解仍有計(jì)算開銷
2. 超參數(shù)r的選擇需要經(jīng)驗(yàn)調(diào)優(yōu)
3. 在極小批量情況下效果可能不穩(wěn)定

6.3 未來方向

1. 自適應(yīng)特征數(shù)量選擇機(jī)制
2. 分布式訓(xùn)練優(yōu)化
3. 與其他優(yōu)化技術(shù)的結(jié)合

參考資源

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

代碼實(shí)現(xiàn)：??https://github.com/lliu606/COSMOS??

本文轉(zhuǎn)載自 ??頓數(shù)AI??，作者： 可可