新火種
2024-12-08 
1行代碼改進(jìn)大模型訓(xùn)練，Llama訓(xùn)練速度提升至1.47倍，全華人團(tuán)隊(duì)出品
只要改一行代碼，就能讓大模型訓(xùn)練效率提升至1.47倍。
擁有得州大學(xué)奧斯汀分校背景四名華人學(xué)者，提出了大模型訓(xùn)練優(yōu)化器Cautious Optimizers。
在提速的同時(shí)，Cautious能夠保證訓(xùn)練效果不出現(xiàn)損失，而且語(yǔ)言和視覺(jué)模型都適用。
該優(yōu)化器以哈密頓量和下降動(dòng)力學(xué)為理論基礎(chǔ)，在加速的同時(shí)不影響收斂特性。
作者在600M到1B不同參數(shù)規(guī)模的Llama模型上進(jìn)行了試驗(yàn)，獲得了最高47%的加速率。
該研究相關(guān)代碼已經(jīng)開(kāi)源，在GitHub上有使用方法的詳細(xì)講解。
一行代碼改進(jìn)大模型訓(xùn)練
Cautious Optimizers在PyTorch當(dāng)中增加的一行代碼，核心思路是引入實(shí)現(xiàn)一種掩蔽機(jī)制，從而避免參數(shù)更新的方向與當(dāng)前梯度方向相悖。
因?yàn)檫@兩個(gè)方向一旦不一致，就有可能導(dǎo)致?lián)p失函數(shù)暫時(shí)增加，造成收斂速度的減緩。
不過(guò)作者并未在方向不一致的來(lái)源問(wèn)題上過(guò)度糾結(jié)，而是引入了一種判斷機(jī)制，在參數(shù)更新之前增加一步計(jì)算，從而過(guò)濾掉方向不一致的情形。
這也正是上面代碼的直接作用。
△
GD：梯度下降，GDM：帶動(dòng)量的梯度下降，C-GDM：本項(xiàng)目
具體來(lái)說(shuō)，加入的兩行代會(huì)對(duì)u和g兩個(gè)向量求內(nèi)積，u向量對(duì)應(yīng)優(yōu)化器給出的參數(shù)更新方向，而g向量對(duì)應(yīng)當(dāng)前時(shí)刻的梯度方向。
作者設(shè)計(jì)了一個(gè)對(duì)齊掩碼函數(shù)?，當(dāng)u和g的內(nèi)積小于0時(shí)（即方向不一致），?的輸出為0向量；當(dāng)內(nèi)積大于等于0時(shí)，?的輸出為全1向量。
而一旦?為零向量時(shí)，w_t計(jì)算式中含u的項(xiàng)也會(huì)變?yōu)榱阆蛄?，?dǎo)致此項(xiàng)更新被跳過(guò)。
這樣就可以判斷參數(shù)更新和梯度方向是否一致，如果不一致則不會(huì)用于參數(shù)更新，避免了訓(xùn)練過(guò)程中損失函數(shù)的回升。
訓(xùn)練效率提升47%
為了評(píng)估Cautious Optimizers的具體效果，作者分別在語(yǔ)言模型Llama和視覺(jué)模型MAE上進(jìn)行了試驗(yàn)。
作者選取了60M、100M、350M和1B四種參數(shù)規(guī)模的Llama模型，在C4語(yǔ)料庫(kù)上進(jìn)行預(yù)訓(xùn)練。
優(yōu)化器選用了AdamW和Lion，以及它們對(duì)應(yīng)的Cautious版本:C-AdamW和C-Lion，每個(gè)實(shí)驗(yàn)中進(jìn)行1萬(wàn)步迭代。
結(jié)果C-AdamW和C-Lion在所有規(guī)模上都表現(xiàn)出明顯的收斂加速效果。
尤其是在1B規(guī)模上，相比原版的AdamW和Lion，它們的樣本效率分別提高了47%和28%，這表明Cautious Optimizer能有效減少訓(xùn)練震蕩，使收斂更平穩(wěn)高效。
并且，Cautious Optimizer在所有情況下都取得了更低的困惑度，印證了其出色的泛化性能。
為了評(píng)估模型的實(shí)際效果，研究者在語(yǔ)句匹配、文本蘊(yùn)含、情感分類(lèi)等6個(gè)GLUE下游任務(wù)上測(cè)試了AdamW和C-AdamW優(yōu)化后1B模型的表現(xiàn),
結(jié)果表明，C-AdamW的平均得分比AdamW高出2%，在大多數(shù)任務(wù)上都取得了進(jìn)步，說(shuō)明Cautious跳過(guò)部分參數(shù)更新的方式不會(huì)引起模型性能下降。
對(duì)于視覺(jué)模型，作者以ViT為骨干網(wǎng)絡(luò)，在ImageNet-1K數(shù)據(jù)集上預(yù)訓(xùn)練了MAE模型。
由于視覺(jué)任務(wù)的特殊性，訓(xùn)練過(guò)程采用了隨機(jī)遮擋圖像塊并重建的范式，因此優(yōu)化目標(biāo)是最小化重建誤差，而非通常的分類(lèi)損失。
作者對(duì)比了AdamW和C-AdamW的表現(xiàn)，即訓(xùn)練50輪后的最終重建誤差，結(jié)果C-AdamW的誤差為0.5926，低于AdamW的0.6085。
一作曾在一周內(nèi)復(fù)刻o1
本項(xiàng)目是由四名華人學(xué)者共同打造的。
第一作者Kaizhao Liang，是AI推理加速服務(wù)商SambaNova公司的一名高級(jí)ML工程師。
在o1模型發(fā)布一周內(nèi)，該公司就推出了一個(gè)類(lèi)似o1模型思考過(guò)程的開(kāi)源平替，主要作者正是Liang。
其他三名作者是得州大學(xué)奧斯汀分校CS助理教授Qiang Liu，以及他的兩名博士生，Lizhang Chen和Bo Liu。
此外，Liang的人工智能碩士學(xué)位也是從該校獲得。
相關(guān)推薦
- 免責(zé)聲明
- 本文所包含的觀點(diǎn)僅代表作者個(gè)人看法，不代表新火種的觀點(diǎn)。在新火種上獲取的所有信息均不應(yīng)被視為投資建議。新火種對(duì)本文可能提及或鏈接的任何項(xiàng)目不表示認(rèn)可。 交易和投資涉及高風(fēng)險(xiǎn)，讀者在采取與本文內(nèi)容相關(guān)的任何行動(dòng)之前，請(qǐng)務(wù)必進(jìn)行充分的盡職調(diào)查。最終的決策應(yīng)該基于您自己的獨(dú)立判斷。新火種不對(duì)因依賴(lài)本文觀點(diǎn)而產(chǎn)生的任何金錢(qián)損失負(fù)任何責(zé)任。
