MARS：提升大模型訓(xùn)練效率的字節(jié)優(yōu)化框架解析

MARS是一款由字節(jié)跳動推出的先進優(yōu)化框架，旨在顯著提升大型模型訓(xùn)練的效率。通過結(jié)合預(yù)條件梯度方法和方差減少技術(shù)，MARS基于縮放隨機遞歸動量技術(shù)優(yōu)化梯度估計，提供了一種靈活的訓(xùn)練方案。該框架支持全矩陣和對角Hessian近似，衍生出三種優(yōu)化算法實例：MARS-AdamW、MARS-Lion和MARS-Shampoo。實驗結(jié)果顯示，MARS在訓(xùn)練GPT-2模型時，相較于傳統(tǒng)的AdamW優(yōu)化器，表現(xiàn)出更為卓越的性能。

MARS是什么

MARS（Make vAriance Reduction Shine）是由字節(jié)跳動開發(fā)的一種創(chuàng)新性優(yōu)化框架，旨在提高大型模型的訓(xùn)練效率。該框架將預(yù)條件梯度方法與方差減少技術(shù)相結(jié)合，利用縮放隨機遞歸動量技術(shù)來優(yōu)化梯度估計。MARS的設(shè)計靈活，支持不同的Hessian近似方式，能夠生成基于AdamW、Lion和Shampoo的三種優(yōu)化算法實例。實驗結(jié)果表明，MARS在訓(xùn)練GPT-2模型時，展現(xiàn)了優(yōu)于傳統(tǒng)AdamW優(yōu)化器的性能。

MARS的主要功能

• 提升訓(xùn)練效率：MARS通過結(jié)合預(yù)條件梯度方法和方差減少技術(shù)，有效提升大型模型訓(xùn)練的效率，特別是在深度神經(jīng)網(wǎng)絡(luò)和大型語言模型的訓(xùn)練中。
• 統(tǒng)一的優(yōu)化框架：提供適應(yīng)多種Hessian近似方法的統(tǒng)一框架，包括全矩陣和對角矩陣近似。
• 具體算法實施：在MARS框架下，衍生出三種具體的優(yōu)化算法：MARS-AdamW、MARS-Lion和MARS-Shampoo，分別基于不同的預(yù)條件梯度更新策略。
• 方差減少技術(shù)：利用縮放隨機遞歸動量技術(shù)，有效降低訓(xùn)練過程中的梯度方差，促進模型的快速收斂。

MARS的技術(shù)原理

• 預(yù)條件梯度方法：通過預(yù)條件梯度方法調(diào)整學習率，為每個參數(shù)或參數(shù)組提供量身定制的學習率，以適應(yīng)其局部曲率。
• 方差減少技術(shù)：引入方差減少技術(shù)，例如STORM（Stochastic Recursive Momentum），以減少隨機梯度的方差，從而加快優(yōu)化過程。
• 縮放隨機遞歸動量：在STORM的基礎(chǔ)上，增加縮放參數(shù)，以調(diào)節(jié)方差減少的強度，定義新的梯度估計器。
• 梯度裁剪與指數(shù)移動平均：為提高訓(xùn)練的穩(wěn)定性，MARS在梯度估計器中應(yīng)用了梯度裁剪，并通過指數(shù)移動平均（EMA）計算遞歸動量。

MARS的項目地址

• arXiv技術(shù)論文：https://arxiv.org/pdf/2411.10438

MARS的應(yīng)用場景

• 深度學習模型訓(xùn)練：適用于訓(xùn)練深度神經(jīng)網(wǎng)絡(luò)，特別是參數(shù)眾多的復(fù)雜模型。
• 大規(guī)模語言模型：優(yōu)化大型語言模型的訓(xùn)練過程，如GPT系列，從而提高訓(xùn)練效率和模型性能。
• 計算機視覺任務(wù)：在圖像分類、目標檢測等計算機視覺領(lǐng)域，加速模型訓(xùn)練并提升模型的泛化能力。
• 強化學習算法：在強化學習中，優(yōu)化策略網(wǎng)絡(luò)或價值函數(shù)的參數(shù)，尤其是在面對高方差梯度的情況下。
• 推薦系統(tǒng)模型：在構(gòu)建推薦系統(tǒng)時，優(yōu)化模型參數(shù)，以更好地處理大規(guī)模用戶和物品特征。

常見問題

• MARS適用于哪些類型的模型訓(xùn)練？ MARS特別適合于深度學習和大規(guī)模語言模型的訓(xùn)練。
• 使用MARS能帶來什么優(yōu)勢？ MARS通過減少訓(xùn)練過程中的方差，提高了模型的收斂速度和訓(xùn)練效率。
• 是否需要對現(xiàn)有模型架構(gòu)進行重大改動才能使用MARS？ 不需要，MARS可以靈活集成到現(xiàn)有的訓(xùn)練流程中。
• MARS支持哪些優(yōu)化算法？ MARS支持多種算法實例，包括MARS-AdamW、MARS-Lion和MARS-Shampoo。