任何 transformer 變體、任何數(shù)據(jù)集都通用。

在一塊消費級 GPU 上只用一天時間訓(xùn)練，可以得到什么樣的 BERT 模型？

最近一段時間，語言模型再次帶火了 AI 領(lǐng)域。預(yù)訓(xùn)練語言模型的無監(jiān)督訓(xùn)練屬性使其可以在海量樣本基礎(chǔ)上進行訓(xùn)練，并獲得大量語義語法知識，不論分類還是問答，似乎沒有 AI 解決不了的問題。

然而，大模型既帶來了技術(shù)突破，也對算力提出了無窮無盡的需求。

最近，來自馬里蘭大學(xué)的 Jonas Geiping、Tom Goldstein 討論了所有關(guān)于擴大計算規(guī)模的研究，深入探討了縮小計算規(guī)模的改進方向。他們的研究引發(fā)了機器學(xué)習(xí)社區(qū)的關(guān)注。

在新研究中，作者對于單塊消費級 GPU（RTX 2080Ti）能訓(xùn)練出什么樣的語言模型進行了討論，并獲得了令人興奮的結(jié)果。讓我們看看它是如何實現(xiàn)的：

模型規(guī)模的擴展

在自然語言處理（NLP）領(lǐng)域，基于 Transformer 架構(gòu)的預(yù)訓(xùn)練模型已經(jīng)成為主流，并帶來諸多突破性進展。很大程度上，這些模型性能強大的原因是它們的規(guī)模很大。隨著模型參數(shù)量和數(shù)據(jù)量的增長，模型的性能會不斷提高。因此，NLP 領(lǐng)域內(nèi)掀起了一場增大模型規(guī)模的競賽。

然而，很少有研究人員或從業(yè)者認為他們有能力訓(xùn)練大型語言模型（LLM），通常只有行業(yè)內(nèi)的科技巨頭擁有訓(xùn)練 LLM 的資源。

為了扭轉(zhuǎn)這一趨勢，來自馬里蘭大學(xué)的研究者進行了一番探索。

論文《Cramming: Training a Language Model on a Single GPU in One Day》：

論文鏈接：https://arxiv.org/abs/2212.14034

這個問題對于大多數(shù)研究人員和從業(yè)者來說具有重要意義，因為這將成為模型訓(xùn)練成本的參考，并有望打破 LLM 訓(xùn)練成本超高的瓶頸。該研究的論文迅速在推特上引發(fā)關(guān)注和討論。

IBM 的 NLP 研究專家 Leshem Choshen 在推特上評價道：「這篇論文總結(jié)了所有你能想到的大模型訓(xùn)練 trick?！?/p>

馬里蘭大學(xué)的研究者認為：如果按比例縮小的模型預(yù)訓(xùn)練是大型預(yù)訓(xùn)練的可行模擬，那么這將開啟一系列目前難以實現(xiàn)的大規(guī)模模型的進一步學(xué)術(shù)研究。

此外，該研究嘗試對過去幾年 NLP 領(lǐng)域的整體進展進行基準測試，而不僅僅局限于模型規(guī)模的影響。

該研究創(chuàng)建了一項稱為「Cramming」的挑戰(zhàn) —— 在測試前一天學(xué)習(xí)整個語言模型。研究者首先分析了訓(xùn)練 pipeline 的方方面面，以了解哪些修改可以實際提高小規(guī)模模擬模型的性能。并且，該研究表明，即使在這種受限環(huán)境中，模型性能也嚴格遵循在大型計算環(huán)境中觀察到的擴展定律。

雖然較小的模型架構(gòu)可以加快梯度計算，但隨著時間的推移，模型改進的總體速度幾乎保持不變。該研究嘗試利用擴展定律在不影響模型大小的情況下通過提高梯度計算的有效率獲得性能提升。最后，該研究成功訓(xùn)練出性能可觀的模型 —— 在 GLUE 任務(wù)上接近甚至超過 BERT—— 而且訓(xùn)練成本很低。

資源有限

為了模擬普通從業(yè)者和研究人員的資源環(huán)境，該研究首先構(gòu)建了一個資源受限的研究環(huán)境：

一個任意大小的基于 transformer 的語言模型，完全從頭開始使用掩碼語言建模（masked-language modeling）進行訓(xùn)練；

pipeline 中不能包含現(xiàn)有的預(yù)訓(xùn)練模型；

任何原始文本（不包括下游數(shù)據(jù)）都可以包含在訓(xùn)練中，這意味著可以通過明智地選擇如何以及何時對數(shù)據(jù)進行采樣來實現(xiàn)加速，前提是采樣機制不需要預(yù)訓(xùn)練模型；

原始數(shù)據(jù)的下載和預(yù)處理不計入總預(yù)算，這里的預(yù)處理包括基于 CPU 的 tokenizer 構(gòu)造、tokenization 和 filtering，但不包括表征學(xué)習(xí)；

訓(xùn)練僅在單塊 GPU 上進行 24 小時；

下游性能在 GLUE 上進行評估，GLUE 上的下游微調(diào)僅限于僅使用下游任務(wù)的訓(xùn)練數(shù)據(jù)進行簡單訓(xùn)練（5 個 epoch 或者更少），并且需要使用為所有 GLUE 任務(wù)設(shè)置的全局超參數(shù)，下游微調(diào)不計算在總預(yù)算中。

該研究與一些經(jīng)典大模型的具體訓(xùn)練設(shè)置比較如下表所示：

改進方法

研究人員實施并測試了已有工作提出的一些修改方向，包括通用實現(xiàn)和初始數(shù)據(jù)設(shè)置，并嘗試了修改架構(gòu)、訓(xùn)練以及改動數(shù)據(jù)集的方法。

實驗在 PyTorch 中進行，不使用特質(zhì)化的實現(xiàn)以盡量公平，所有內(nèi)容都保留在 PyTorch 框架的實現(xiàn)級別上，只允許可應(yīng)用于所有組件的自動運算符融合，另外只有在選擇了最終的架構(gòu)變體之后，才會重新啟用高效注意力內(nèi)核。

圖 1：不同 transformer 架構(gòu)變體 MLM 損失函數(shù)與 token 對比的情況。左：全局視圖。右圖：放大到 10e8 和更多 token 情況下。所有模型都用相同算力成本訓(xùn)練，我們可以看到：通過架構(gòu)重塑實現(xiàn)的改進微乎其微。

有關(guān)提升性能，我們最先想到的方法肯定是修改模型架構(gòu)。從直覺上，較小 / 較低容量的模型似乎在一日一卡式的訓(xùn)練中是最優(yōu)的。然而在研究了模型類型與訓(xùn)練效率之間的關(guān)系后，研究人員發(fā)現(xiàn)縮放法則為縮小規(guī)模設(shè)置了巨大的障礙。每個 token 的訓(xùn)練效率在很大程度上取決于模型大小，而不是 transformer 的類型。

此外，較小的模型學(xué)習(xí)效率較低，這在很大程度上減緩了吞吐量的增加。幸運的是，在相同大小的模型中，訓(xùn)練效率幾乎保持不變這一事實，意味著我們可以在參數(shù)量類似的架構(gòu)中尋找合適的，主要根據(jù)影響單個梯度步驟的計算時間來做出設(shè)計選擇。

圖 2：學(xué)習(xí)率 Schedule。盡管在全局范圍內(nèi)行為相似，但在中間的放大圖里可以看到差異確實存在。

在該工作中，作者研究了訓(xùn)練超參數(shù)對 BERT-base 架構(gòu)的影響?？梢岳斫獾氖?，原始 BERT 訓(xùn)練方法的模型在 Cramming 式訓(xùn)練要求中的表現(xiàn)不佳，因此研究人員重新審視了一些標準選擇。

作者也研究了優(yōu)化數(shù)據(jù)集的思路。擴展法則阻礙了通過架構(gòu)修改取得重大收益的方式（超出計算效率），但縮放定律并不妨礙我們在更好的數(shù)據(jù)上進行訓(xùn)練。如果想在在每秒訓(xùn)練更多的 token，我們應(yīng)該尋求在更好的 token 上訓(xùn)練。

圖 3：在 bookcorpus-wikipedia 數(shù)據(jù)的 Cramming 式訓(xùn)練中訓(xùn)練的模型的詞匯量與 GLUE 分數(shù)和 MNLI 準確性。

在 GLUE 上的表現(xiàn)

研究人員系統(tǒng)地評估了 GLUE 基準的性能和 WNLI，并注意到在前面的部分中只使用了 MNLI (m)，并且沒有根據(jù)完整的 GLUE 分數(shù)調(diào)整超參數(shù)。在新研究中對于 BERT-base 作者微調(diào)了 5 個 epoch 的所有數(shù)據(jù)集，batch size 為 32，學(xué)習(xí)率為 2 × 10-5。對于 Cramming 訓(xùn)練的模型這不是最優(yōu)的，其可以從 16 的 batch size 和 4 × 10?5 的學(xué)習(xí)率以及余弦衰減中獲得微小的改進（此設(shè)置不會改進預(yù)訓(xùn)練的 BERT check point）。

表 3 和表 4 描述了此設(shè)置在 GLUE 下游任務(wù)上的性能。作者比較了原始的 BERT-base check point、在達到算力上限后停止的 BERT 預(yù)訓(xùn)練設(shè)置、Izsak 等人 2021 年研究中描述的設(shè)置和修改后的設(shè)置，為每塊 GPU 設(shè)置訓(xùn)練一天?？傮w而言，性能出奇地好，尤其是對于 MNLI、QQP、QNLI 和 SST-2 等較大的數(shù)據(jù)集，下游微調(diào)可以消除完整 BERT 模型和 Cramming 設(shè)置變體之間的剩余差異。

此外，作者發(fā)現(xiàn)新方法與算力有限的普通 BERT 訓(xùn)練及 Izsak 等人描述的方法相比都有很大改進。對于 Izsak 等人的研究，其描述的方法最初是為一個完整的 8 GPU 刀片服務(wù)器設(shè)計的，并且在新的場景中，將其中的 BERT-large 模型壓縮到較小的 GPU 上是導(dǎo)致大部分性能下降的原因。

表格 3：基線 BERT 與 Cramming 版本模型的 GLUE-dev 性能比較。其中所有任務(wù)的超參數(shù)都是固定的，epoch 限制為 5 個，缺失值為 NaN。是為 8 GPU 刀片服務(wù)器設(shè)計的，而在這里，所有計算被塞進了一塊 GPU。

表格 4：基線 BERT 與填充模型的 GLUE-dev 性能比較。

總體而言，使用論文中的方法，訓(xùn)練結(jié)果已經(jīng)非常接近原版 BERT，但要知道后者使用的總 FLOPS 是新方法 45-136 倍（在 16 塊 TPU 上要花費四天時間）。而當訓(xùn)練時間延長 16 倍時（在 8 塊 GPU 上訓(xùn)練兩天），新方法的性能實際上比原始 BERT 提高了很多，達到了 RoBERTa 的水平。

總結(jié)

在該工作中，人們討論了基于 transformer 的語言模型在計算量非常有限的環(huán)境中可以實現(xiàn)多少性能，值得慶幸的是，幾條修改方向可以讓我們在 GLUE 上獲得不錯的下游性能。研究人員表示，希望這項工作可以為進一步的改進提供一個基線，并進一步給近年來為 transformer 架構(gòu)提出的許多改進和技巧提供理論支撐。

審核編輯 ：李倩