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

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

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

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

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

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

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

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

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

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

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

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

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

圖 2 展示了 Adam-mini 的大致圖示和學(xué)習(xí)率使用情況?？梢钥吹剑珹dam-mini 是根據(jù) Hessian 結(jié)構(gòu)來(lái)分配學(xué)習(xí)率，因此其使用的學(xué)習(xí)率比 SGD 多，但遠(yuǎn)少于 Adam。

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

方法

動(dòng)機(jī)和觀察

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

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

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

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

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

也就是說(shuō)，對(duì)于帶有塊對(duì)角 Hessian 的一般問(wèn)題而言，許多學(xué)習(xí)率并不一定會(huì)帶來(lái)額外收益。

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

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

新方法：Adam-mini

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

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

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

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

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

• 對(duì)于 Adam：

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

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

分片策略的原則

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

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

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

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

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

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

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

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

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

Adam-mini 的一些特點(diǎn)

減少內(nèi)存用量

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

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

吞吐量更高

Adam-mini 可取得比 AdamW 更高的吞吐量，尤其是當(dāng)硬件資源有限時(shí)。

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

實(shí)驗(yàn)

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

設(shè)置

該團(tuán)隊(duì)預(yù)訓(xùn)練了一些開(kāi)源 LLM，包括 GPT2 系列和 LLM 系列。他們?cè)谥髁饔⒄Z(yǔ)語(yǔ)料庫(kù)上從頭開(kāi)始訓(xùn)練了這些模型。

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

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

Llama 系列

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

GPT2 系列

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

敏感度分析

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

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

該團(tuán)隊(duì)評(píng)估了 Adam-mini 對(duì)下游微調(diào)任務(wù)的有效性。

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

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

非 LLM 任務(wù)

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