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