在訓(xùn)練大型語言模型（LLM）時，Adam(W) 基本上已經(jīng)成為了人們默認使用的優(yōu)化器。

Adam 盡管性能優(yōu)異，但使用成本很高。具體來說，Adam 需要內(nèi)存來保存其優(yōu)化器狀態(tài)：一階動量 m 和二階動量 v^2。這總共需要模型大小至少 2 倍的內(nèi)存。這樣的內(nèi)存消耗已經(jīng)成為了 LLM 訓(xùn)練的一大主要負擔(dān)。

舉個例子，要訓(xùn)練一個 7B 模型，只是 Adam 就需要每張卡有大約 56 GB 來保存 m 和 v；而如果再加上梯度，則總共需要 86 GB。即使使用最先進的 A100-80GB，成本也過高了。

為了支持這樣的高內(nèi)存算法，實踐中必須要使用 CPU 卸載與分片，但這又會增加延遲，減慢訓(xùn)練速度。在訓(xùn)練 PaLM （有 5400 億參數(shù)）這樣的更大型模型時，情況還會更糟。在這種情況下，Adam 自身就要占用超過 50 GB，并且這也是預(yù)訓(xùn)練階段的一大主要開銷。

因此，人們希望設(shè)計出內(nèi)存需求更少又有效的優(yōu)化器。首先，減少內(nèi)存可以減輕 CPU 卸載的負擔(dān)并能減輕對模型參數(shù)執(zhí)行分片的需求。這些都能減少 GPU 和 CPU 之間的通信量，并進一步提升訓(xùn)練過程的吞吐量和速度。其次，這允許實踐者使用更少的 GPU 來訓(xùn)練所需大小的模型，從而極大地節(jié)省成本和能源。第三，這能降低訓(xùn)練 LLM 的門檻并鼓勵更多 GPU 資源有限的研究者參與進來。

但是，修改 Adam 但不影響其性能的難度非常大。其中一個主要原因是我們?nèi)匀徊惶斫?Adam 的 m 和 v 的作用。我們還不清楚 Adam 的哪些組件對其卓越性能而言必不可少，也就更不知道可以重新設(shè)計和改進哪些組件了。

Adafactor 是一個頗受歡迎的嘗試，其能降低在 v 上低秩分解所用的內(nèi)存。但是，很多研究發(fā)現(xiàn) Adafactor 訓(xùn)練 LLM 的性能不佳。可能的原因有兩個。第一，當前 Adam 中的 v 可能對有效性來說非常關(guān)鍵，已經(jīng)沒有縮減空間。第二，v 是有可能縮減的，只是 Adafactor 使用的方法并不是最合適的：矩陣分解是一種可以廣泛應(yīng)用的通用方法，但它沒有利用太多特定于問題的結(jié)構(gòu)，因此它在特定的神經(jīng)網(wǎng)絡(luò)任務(wù)上效果不佳。

近日，香港中文大學(xué)（深圳）、深圳市大數(shù)據(jù)研究院、杜克大學(xué)和斯坦福大學(xué)的一個聯(lián)合研究團隊發(fā)現(xiàn)：可以通過一個簡單技巧來降低 v 的使用量。

目前而言，Adam 的 v 會為每個參數(shù)單獨分配一個學(xué)習(xí)率，即第 i 個參數(shù)獲得的學(xué)習(xí)率為 ，，其中 v_i 是 v 的第 i 個組件。對于十億參數(shù)量的模型，Adam 就要設(shè)計十億個學(xué)習(xí)率。

該團隊認為這些學(xué)習(xí)率資源可以大幅降低，同時還不會影響模型性能，甚至能讓模型獲得更優(yōu)性能。該團隊注意到，Transformer 的 Hessian 有一種接近塊對角線的結(jié)構(gòu)，其由不同大小的密集子塊構(gòu)成。他們發(fā)現(xiàn)，對于每一個這樣的密集子塊，都存在性能優(yōu)于 Adam 的單個高質(zhì)量學(xué)習(xí)率——只要有足夠的資源將其搜索出來。由于密集子塊的數(shù)量遠少于參數(shù)數(shù)量，這個發(fā)現(xiàn)表明使用遠遠更少的學(xué)習(xí)率是有可能取得優(yōu)良性能的。剩下的問題如何高效地找到它們。

為了找到性能足以比肩甚至優(yōu)于 Adam 的優(yōu)良學(xué)習(xí)率，該團隊提出了一種低成本的簡單方法 Adam-mini。另外，該團隊也發(fā)布了 Adam-mini 的一種實現(xiàn)。

• 論文標題：Adam-mini: Use Fewer Learning Rates To Gain More
• 論文地址：https://arxiv.org/pdf/2406.16793
• 實現(xiàn)代碼：https://github.com/zyushun/Adam-mini

在實驗中，Adam-mini 的表現(xiàn)非常優(yōu)秀。如圖 1 所示， 在預(yù)訓(xùn)練 Llama2-7B 時，Adam-mini 的內(nèi)存占用可以大幅降低，同時吞吐量也有明顯提升。在保證性能與 Adam 相當?shù)耐瑫r，其能帶來 33% 的速度提升。

圖 2 展示了 Adam-mini 的大致圖示和學(xué)習(xí)率使用情況。可以看到，Adam-mini 是根據(jù) Hessian 結(jié)構(gòu)來分配學(xué)習(xí)率，因此其使用的學(xué)習(xí)率比 SGD 多，但遠少于 Adam。

本研究的通訊作者，香港中文大學(xué)（深圳）數(shù)據(jù)科學(xué)學(xué)院副教授孫若愚（Ruoyu Sun）表示：「如果你喜歡 Adam(W)，那 Adam-mini 就是你的不二之選！」

方法

動機和觀察

為了應(yīng)對特征值的異構(gòu)性，Transformer 的每一參數(shù)塊都需要不同的學(xué)習(xí)率。Adam 的 v 可以提供這一點。

論文《Why Transformers Need Adam: A Hessian Perspective》發(fā)現(xiàn)：必須為每個塊使用不同的學(xué)習(xí)率。但是，Adam 所做的卻遠不止此：其并不只是為每個塊分配不同的學(xué)習(xí)率，而是為每個參數(shù)都分配不同的學(xué)習(xí)率。而參數(shù)量（可能超過數(shù)十億）要遠遠多于塊的數(shù)量（通常數(shù)百計）。那問題就來了：

有必要為每個參數(shù)都使用單獨的學(xué)習(xí)率嗎？如果沒必要，又能節(jié)省多少？

該團隊進行了一番探索，得到了如下發(fā)現(xiàn)：

• 如圖 4 (a) 和 (b) 所示，Adam 的表現(xiàn)優(yōu)于最優(yōu)的單學(xué)習(xí)率方法。這符合預(yù)期。
• 如圖 4 (c) 和 (d) 所示，這里研究了其 Hessian 是 (a) 的一個密集子塊的新問題。該團隊探究了針對這個問題的最優(yōu)單學(xué)習(xí)率方法，結(jié)果發(fā)現(xiàn)其表現(xiàn)優(yōu)于 Adam，即使 Adam 分配的學(xué)習(xí)率要多得多。(a) 的所有三個子塊都有類似的現(xiàn)象。
• 如果收集 (2) 中的那些最優(yōu)學(xué)習(xí)率，并將其用于梯度下降的一個「逐塊」版本，那么其在原始問題上的速度會比 Adam 快，見圖 4(b) 中的綠線。

也就是說，對于帶有塊對角 Hessian 的一般問題而言，許多學(xué)習(xí)率并不一定會帶來額外收益。

該團隊也在 Transformer 中觀察到了類似的現(xiàn)象。

總結(jié)起來，該團隊發(fā)現(xiàn)：對于 Transformer，有可能使用比 Adam 遠遠更少的學(xué)習(xí)率實現(xiàn)與之相當或更好的性能。剩下的問題就是在不使用網(wǎng)格搜索的情況下如何找到這些學(xué)習(xí)率。

新方法：Adam-mini

基于上述討論，該團隊提出了 Adam-mini，見算法 1。Adam-mini 的目標是降低 Adam 中的學(xué)習(xí)率資源，而無需費力地對學(xué)習(xí)率執(zhí)行網(wǎng)格搜索。Adam-mini 包含兩個步驟。步驟 1 只在初始化時進行。

步驟 1-1：將模型參數(shù)切分為塊。對于 Transformer，該團隊選擇的策略是「Partition for Transformers」，即根據(jù)頭來切分所有的 Query 和 Key，并使用 PyTorch 來切分其余部分。對于其它網(wǎng)絡(luò)，則使用默認的 PyTorch 切分即可，他們將其稱之為「Partition for non-Transformers」。

步驟 1-2：使用 embd_blocks。對于 Transformer，其包含嵌入層和輸出層。對于其它網(wǎng)絡(luò)，不選擇任何參數(shù)。

步驟 2：對于每個位于 embd_blocks 之外的參數(shù)塊，都使用單個學(xué)習(xí)率。為了高效地為每個塊選擇合適的學(xué)習(xí)率，Adam-mini 的做法是直接將原始 Adam 中的 g?g 替換成其均值。Adam-mini 像 Adam 一樣在這些均值上使用了移動平均值。

一個簡單示例。為了說明 Adam-mini 的關(guān)鍵設(shè)計，該團隊在論文中給出了一個簡單的示例。假設(shè)有一個問題，其有 5 個參數(shù) w ∈ ?^5，Adam 和 Adam-mini 都會執(zhí)行 w = w ? u ? m，其中 m 是一階動量，而 u 有不同的形式，如下所示：

• 對于 Adam：

• 對于 Adam-mini：假設(shè)分片方式為 (1, 2, 3) 和 (4, 5)，則

請注意，有效元素 u_mini 的數(shù)量就等于塊的數(shù)量，這遠遠小于 u_Adam，其等于參數(shù)數(shù)量。結(jié)果表明，這能讓 LLM 的 v 中的元素數(shù)量減少 90% 以上。

分片策略的原則

現(xiàn)在討論如何為 Adam-mini 選擇參數(shù)分片。該團隊基于前述分析得出了一個廣義原則：

原則 1：應(yīng)當將參數(shù)切分成不同的參數(shù)塊，使得每個參數(shù)塊都關(guān)聯(lián)了 Hessian 中最小的密集子塊。

由于參數(shù)的塊對角結(jié)構(gòu)，PyTorch 的默認分片方法是一個合理的候選項。實際上，這種分片確實能很好地用于非 Transformer 任務(wù)，比如 ResNet、擴散模型和圖模型。算法 3 Partition for non-Transformers 展示了該策略。

但不幸的是，該團隊發(fā)現(xiàn)，默認的 PyTorch 切片并不總是能很好地應(yīng)對 Transformer。比如他們發(fā)現(xiàn) Adam-mini 在 1B 模型上會出現(xiàn)訓(xùn)練不穩(wěn)定問題（見圖 6(d)）。

他們猜測這是因為 PyTorch 切片無法完整地理解 Hessian結(jié)構(gòu)。經(jīng)過一番探索后，他們發(fā)現(xiàn) Hessian 子模塊分為兩類：

第一類：類似于 Hessian 整體，這種 Hessian 子塊本身也有進一步的塊對角結(jié)構(gòu)，由更小的密集矩陣構(gòu)成。這一類包含 Query 和 Key。他們通過實驗發(fā)現(xiàn)，小型密集子塊的數(shù)量就等于多頭注意力中頭的數(shù)量。

第二類：這種 Hessian 子塊有密集的結(jié)構(gòu)，無法進一步分成更小的塊。這一類包含 Value、注意力投射和 MLP 層。請注意 Value 的 Hessian 結(jié)構(gòu)不同于 Query 和 Key 的 Hessian 結(jié)構(gòu)，但它們?nèi)加?4 個頭組成。這是因為 Value 位于自注意力設(shè)計的 softmax 算子之外，而 Query 和 Key 卻不是。

基于上述發(fā)現(xiàn)，該團隊發(fā)現(xiàn)默認的 PyTorch 分片確實不是最適合 Transformer 的。可以根據(jù)頭將 Query 和 Key 進一步切分成不同的塊。根據(jù)原則 1，不同的頭應(yīng)當屬于不同的塊。根據(jù)直覺就能知道，不同的頭在理解 token 時的作用也不同，因此它們需要不同的學(xué)習(xí)率就很合理了。

由此，該團隊得到了算法 2：Partition for Transformers，其可根據(jù)頭來切分 Query 和 Key。該過程遵循基于最小 Hessian 子塊的原則如圖 6 (d) 所示。該策略確實能穩(wěn)定訓(xùn)練過程并提升性能。

Adam-mini 的一些特點

減少內(nèi)存用量

根據(jù)設(shè)計，Adam-mini 可為 Transformer 減少學(xué)習(xí)率的數(shù)量——從參數(shù)總數(shù)減少到嵌入層的大小、輸出層的大小和其它層中塊的數(shù)量的總和。因此，減少的內(nèi)存比例取決于模型中非嵌入?yún)?shù)的比例。

在 Llama2-7B 上，這個比例是 96.2%。對于主流的語言模型，這個比例通常 ≥ 90%。請注意，如果層數(shù)更多，這個比例也會更加接近 100%；比如對于 Llama3-70B，這個比例為 99.25%。因此，Adam-mini 可將 v 減少至少 90%，由此可為 Adam 帶來 45% 到 50% 的內(nèi)存節(jié)省。

吞吐量更高

Adam-mini 可取得比 AdamW 更高的吞吐量，尤其是當硬件資源有限時。

基于此，Adam-mini 可以減少預(yù)訓(xùn)練的總時間。如表 2 所示。當在 2 臺 A800-80GB 上預(yù)訓(xùn)練 Llama2-7B 時，Adam-mini 的吞吐量比 AdamW 高 49.6%。召回吞吐量 (↑) 指的是每秒處理的 token 數(shù)量，因此在預(yù)訓(xùn)練時，Adam-mini 處理相同數(shù)量的 token 可節(jié)省 33.1% 的時間。

實驗

預(yù)訓(xùn)練

設(shè)置

該團隊預(yù)訓(xùn)練了一些開源 LLM，包括 GPT2 系列和 LLM 系列。他們在主流英語語料庫上從頭開始訓(xùn)練了這些模型。

具體來說，他們在 Openwebtext 上訓(xùn)練了 GPT2 系列，在 CommonCrawl 上訓(xùn)練了 TinyLlama-1B 和 Llama2-7B。他們使用 2B、3B 和 25B token 訓(xùn)練了模型。模型配置（比如上下文長度）是根據(jù)標準協(xié)議選擇的。

該團隊將 Adam-mini 與 AdamW 以及 Adafactor、CAME 和 SM3 等常用內(nèi)存高效型方法進行了比較。對于 Adafactor 和 SM3，為了確保與其它方法進行公平的比較，這里整合了 β_1 = 0.9 的動量。他們在同樣的預(yù)算內(nèi)對所有方法的學(xué)習(xí)率進行了調(diào)節(jié)，并報告了最佳性能。

Llama 系列

圖 7 (a) 展示了預(yù)訓(xùn)練 TinyLlama-1B 的驗證損失曲線。Llama2-7B 的訓(xùn)練曲線已經(jīng)在圖 1(c) 中給出。對于 TinyLlama-1B 和 Llama2-7B，可以發(fā)現(xiàn) Adam-mini 使用更少的內(nèi)存便能達到與 AdamW 相近的效果。

GPT2 系列

圖 7 (b) 展現(xiàn)了 GPT2-125M 的驗證損失曲線，圖 8 展示了大小從 330M 到 1.5B 的 GPT2 的驗證損失曲線。可以看到，Adam-mini 使用更少的內(nèi)存就能取得與 AdamW 相當?shù)膬?yōu)良性能，而其它方法的表現(xiàn)要更差一些。

敏感度分析

該團隊在 GPT2-125M 預(yù)訓(xùn)練任務(wù)上測試了 Adam-mini 對超參數(shù)的敏感度。這里報告了使用 2.5B token 訓(xùn)練之后的驗證損失。如圖 7 所示，Adam-mini 看起來對超參數(shù)并不敏感。

監(jiān)督式微調(diào)和 RLHF

該團隊評估了 Adam-mini 對下游微調(diào)任務(wù)的有效性。

具體來說，他們考慮了兩種有代表性的任務(wù)：監(jiān)督式微調(diào)（SFT）和根據(jù)人類反饋的強化學(xué)習(xí)（RLHF）。這里使用的預(yù)訓(xùn)練模型是 Llama-2-7B。

結(jié)果見表 3，可以看到 Adam-mini 的表現(xiàn)優(yōu)于 AdamW，即便 Adam-mini 使用了單個學(xué)習(xí)率并且內(nèi)存效率更高。

非 LLM 任務(wù)

該團隊也在非 LLM 任務(wù)評估了 Adam-mini。表 4 給出了在 ImageNet 上訓(xùn)練 ResNet18、在 CelebA 上訓(xùn)練擴散模型、在 OGB-arxiv 上訓(xùn)練圖卷積網(wǎng)絡(luò)（GCN）和圖注意力網(wǎng)絡(luò)（GAT）的結(jié)果。結(jié)果發(fā)現(xiàn) Adam-mini 使用更少的內(nèi)存便能取得與 AdamW 相當或更好的性能。