在訓練大型語言模型（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