蘋果發布高效雙EMA梯度優化方法，適配Transformer、Mamba模型
在訓練深度學習模型時，優化復雜的非凸損失函數是一個非常難的挑戰。目前廣泛使用的優化器是Adam、AdamW等，主要依賴于EMA（指數移動平均）來累積梯度信息，但這種方法存在一些局限性。
例如，在訓練超大規模數據集和復雜架構模型時，EMA 可能無法有效地捕捉到梯度的變化趨勢，可能會忽略掉一些重要的局部信息，導致模型在訓練過程中陷入局部優解，而無法找到全局最優解。
為了解決這一難題，蘋果和瑞士洛桑聯邦理工學院的研究人員提出了AdEMAMix 優化器，通過混合了兩個EMA一個對近期梯度敏感，另一個能整合更久遠梯度的信息實現局部和全局的優化。
論文地址：https://arxiv.org/abs/2409.03137
AdEMAMix優化器的創新在于對動量估計的重新思考。在傳統的Adam優化器中，動量是通過指數移動平均（EMA）來實現的，這種方法雖然在實踐中被證明是有效的，但它在處理歷史梯度信息時存在很大局限性。
這是因為單一的EMA在給予近期梯度較高權重的同時，很難為早期的梯度保留足夠的影響力，在一定程度上限制了優化器利用歷史信息的能力，尤其是在那些需要長期依賴歷史數據的任務中。
而AdEMAMix優化器引入了兩種不同速率的EMA。第一種是快速變化的EMA，它對近期的梯度變化反應靈敏，能夠快速調整優化方向以適應損失景觀的局部變化。
第二種則是慢速變化的EMA，它對歷史梯度給予更高的權重，使得優化器能夠在長時間內保持對早期信息的記憶。通過這種雙EMA的結構，使AdEMAMix優化器能夠在保持對近期變化敏感的同時，有效地利用長期積累的歷史信息。
平衡這兩種EMA的權重也是一個關鍵點，AdEMAMix引入了一個參數，用于調節慢速EMA在最終更新中的貢獻。通過精心設計的調度策略，使參數的值會隨著訓練的進行而動態調整，從而在訓練初期避免過大的更新，同時在后期逐漸增加對歷史信息的利用。
此外，AdEMAMix還對傳統的動量更新規則進行了改進。在每次迭代中，優化器會計算一個結合了兩種EMA的更新向量，并通過這個向量來調整模型參數。
這樣不僅考慮了梯度的方向和大小，還考慮了梯度的歷史信息，使得優化器能夠在復雜的損失景觀中更加有效地尋找最優解。
為了評估AdEMAMix優化器的性能，研究人員對Transformer、Mamba和VIT三種不同架構的大模型進行了綜合測試。
在 Transformer 模型的測試中，研究人員分別對 110M、335M 和 1.3B 三種不同規模的模型進行了實驗。例如，在相同的訓練條件下，對于 110M 參數的模型，傳統的優化器可能需要更多的迭代次數和訓練時間才能達到一定性能，而AdEMAMix 優化器的時間大大減少性能也獲得提升。
在 Mamba 模型的測試中，研究人員使用了 168M 參數的 Mamba 模型和 FineWeb 數據集。AdEMAMix同樣獲得了非常好的成績，也證明了該方法可以擴展到不同類型的模型中。
在 VIT 模型的測試中，研究人員使用了 ImageNet 的兩個子集進行多次 epoch 的訓練。例如，在訓練 24M 參數模型在 11M 圖像上進行 37 個epoch 時，AdEMAMix 能夠很容易地找到優于傳統優化器的參數設置，從而降低訓練損失。當模型參數增加到 86M時，AdEMAMix 仍能較容易地找到優于基線的參數。
本文轉自 AIGC開放社區 ，作者：AIGC開放社區
