1行代碼改進大模型訓練，Llama訓練速度升至1.47倍，華人團隊出品

克雷西 2024-11-27 21:08:46 來源：量子位

克雷西 發自 凹非寺

量子位 | 公眾號 QbitAI

只要改一行代碼，就能讓大模型訓練效率提升至1.47倍。

擁有得州大學奧斯汀分校背景四名華人學者，提出了大模型訓練優化器Cautious Optimizers。

在提速的同時，Cautious能夠保證訓練效果不出現損失，而且語言和視覺模型都適用。

該優化器以哈密頓量和下降動力學為理論基礎，在加速的同時不影響收斂特性。

作者在600M到1B不同參數規模的Llama模型上進行了試驗，獲得了最高47%的加速率。

該研究相關代碼已經開源，在GitHub上有使用方法的詳細講解。

一行代碼改進大模型訓練

Cautious Optimizers在PyTorch當中增加的一行代碼，核心思路是引入實現一種掩蔽機制，從而避免參數更新的方向與當前梯度方向相悖。

因為這兩個方向一旦不一致，就有可能導致損失函數暫時增加，造成收斂速度的減緩。

不過作者并未在方向不一致的來源問題上過度糾結，而是引入了一種判斷機制，在參數更新之前增加一步計算，從而過濾掉方向不一致的情形。

這也正是上面代碼的直接作用。

△GD：梯度下降，GDM：帶動量的梯度下降，C-GDM：本項目

具體來說，加入的兩行代會對u和g兩個向量求內積，u向量對應優化器給出的參數更新方向，而g向量對應當前時刻的梯度方向。

作者設計了一個對齊掩碼函數?，當u和g的內積小于0時（即方向不一致），?的輸出為0向量；當內積大于等于0時，?的輸出為全1向量。

而一旦?為零向量時，w_t計算式中含u的項也會變為零向量，導致此項更新被跳過。

這樣就可以判斷參數更新和梯度方向是否一致，如果不一致則不會用于參數更新，避免了訓練過程中損失函數的回升。

訓練效率提升47%

為了評估Cautious Optimizers的具體效果，作者分別在語言模型Llama和視覺模型MAE上進行了試驗。

作者選取了60M、100M、350M和1B四種參數規模的Llama模型，在C4語料庫上進行預訓練。

優化器選用了AdamW和Lion，以及它們對應的Cautious版本:C-AdamW和C-Lion，每個實驗中進行1萬步迭代。

結果C-AdamW和C-Lion在所有規模上都表現出明顯的收斂加速效果。

尤其是在1B規模上，相比原版的AdamW和Lion，它們的樣本效率分別提高了47%和28%，這表明Cautious Optimizer能有效減少訓練震蕩，使收斂更平穩高效。

并且，Cautious Optimizer在所有情況下都取得了更低的困惑度，印證了其出色的泛化性能。

為了評估模型的實際效果，研究者在語句匹配、文本蘊含、情感分類等6個GLUE下游任務上測試了AdamW和C-AdamW優化后1B模型的表現,

結果表明，C-AdamW的平均得分比AdamW高出2%，在大多數任務上都取得了進步，說明Cautious跳過部分參數更新的方式不會引起模型性能下降。

對于視覺模型，作者以ViT為骨干網絡，在ImageNet-1K數據集上預訓練了MAE模型。

由于視覺任務的特殊性，訓練過程采用了隨機遮擋圖像塊并重建的范式，因此優化目標是最小化重建誤差，而非通常的分類損失。

作者對比了AdamW和C-AdamW的表現，即訓練50輪后的最終重建誤差，結果C-AdamW的誤差為0.5926，低于AdamW的0.6085。

一作曾在一周內復刻o1

本項目是由四名華人學者共同打造的。

第一作者Kaizhao Liang，是AI推理加速服務商SambaNova公司的一名高級ML工程師。

在o1模型發布一周內，該公司就推出了一個類似o1模型思考過程的開源平替，主要作者正是Liang。

其他三名作者是得州大學奧斯汀分校CS助理教授Qiang Liu，以及他的兩名博士生，Lizhang Chen和Bo Liu。

此外，Liang的人工智能碩士學位也是從該校獲得。

論文地址：
https://arxiv.org/abs/2411.16085
GitHub：
https://github.com/kyleliang919/C-Optim