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

MARS是什么

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

MARS的主要功能

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

MARS的技術原理

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

MARS的項目地址

• arXiv技術論文：https://arxiv.org/pdf/2411.10438

MARS的應用場景

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

常見問題

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