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

Adam 盡管性能優異，但使用成本很高。具體來說，Adam 需要內存來保存其優化器狀態：一階動量 m 和二階動量 v^2。這總共需要模型大小至少 2 倍的內存。這樣的內存消耗已經成為了 LLM 訓練的一大主要負擔。

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

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

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

但是，修改 Adam 但不影響其性能的難度非常大。其中一個主要原因是我們仍然不太理解 Adam 的 m 和 v 的作用。我們還不清楚 Adam 的哪些組件對其卓越性能而言必不可少，也就更不知道可以重新設計和改進哪些組件了。

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

近日，香港中文大學（深圳）、深圳市大數據研究院、杜克大學和斯坦福大學的一個聯合研究團隊發現：可以通過一個簡單技巧來降低 v 的使用量。

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

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

為了找到性能足以比肩甚至優于 Adam 的優良學習率，該團隊提出了一種低成本的簡單方法 Adam-mini。另外，該團隊也發布了 Adam-mini 的一種實現。

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

在實驗中，Adam-mini 的表現非常優秀。如圖 1 所示， 在預訓練 Llama2-7B 時，Adam-mini 的內存占用可以大幅降低，同時吞吐量也有明顯提升。在保證性能與 Adam 相當的同時，其能帶來 33% 的速度提升。

圖 2 展示了 Adam-mini 的大致圖示和學習率使用情況。可以看到，Adam-mini 是根據 Hessian 結構來分配學習率，因此其使用的學習率比 SGD 多，但遠少于 Adam。

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

方法

動機和觀察

為了應對特征值的異構性，Transformer 的每一參數塊都需要不同的學習率。Adam 的 v 可以提供這一點。

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

有必要為每個參數都使用單獨的學習率嗎？如果沒必要，又能節省多少？

該團隊進行了一番探索，得到了如下發現：

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

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

該團隊也在 Transformer 中觀察到了類似的現象。

總結起來，該團隊發現：對于 Transformer，有可能使用比 Adam 遠遠更少的學習率實現與之相當或更好的性能。剩下的問題就是在不使用網格搜索的情況下如何找到這些學習率。

新方法：Adam-mini

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

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

步驟 1-2：使用 embd_blocks。對于 Transformer，其包含嵌入層和輸出層。對于其它網絡，不選擇任何參數。

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

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

• 對于 Adam：

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

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

分片策略的原則

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

原則 1：應當將參數切分成不同的參數塊，使得每個參數塊都關聯了 Hessian 中最小的密集子塊。

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

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

他們猜測這是因為 PyTorch 切片無法完整地理解 Hessian結構。經過一番探索后，他們發現 Hessian 子模塊分為兩類：

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

第二類：這種 Hessian 子塊有密集的結構，無法進一步分成更小的塊。這一類包含 Value、注意力投射和 MLP 層。請注意 Value 的 Hessian 結構不同于 Query 和 Key 的 Hessian 結構，但它們全都由 4 個頭組成。這是因為 Value 位于自注意力設計的 softmax 算子之外，而 Query 和 Key 卻不是。

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

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

Adam-mini 的一些特點

減少內存用量

根據設計，Adam-mini 可為 Transformer 減少學習率的數量——從參數總數減少到嵌入層的大小、輸出層的大小和其它層中塊的數量的總和。因此，減少的內存比例取決于模型中非嵌入參數的比例。

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

吞吐量更高

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

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

實驗

預訓練

設置

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

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

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

Llama 系列

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

GPT2 系列

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

敏感度分析

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

監督式微調和 RLHF

該團隊評估了 Adam-mini 對下游微調任務的有效性。

具體來說，他們考慮了兩種有代表性的任務：監督式微調（SFT）和根據人類反饋的強化學習（RLHF）。這里使用的預訓練模型是 Llama-2-7B。

結果見表 3，可以看到 Adam-mini 的表現優于 AdamW，即便 Adam-mini 使用了單個學習率并且內存效率更高。

非 LLM 任務

該團隊也在非 LLM 任務評估了 Adam-mini。表 4 給出了在 ImageNet 上訓練 ResNet18、在 CelebA 上訓練擴散模型、在 OGB-arxiv 上訓練圖卷積網絡（GCN）和圖注意力網絡（GAT）的結果。結果發現 Adam-mini 使用更少的內存便能取得與 AdamW 相當或更好的性能。