怎么通過增加模型的大小來加速Transformer的訓(xùn)練和推理

模型訓(xùn)練會很慢

在深度學(xué)習(xí)中，使用更多的計算(例如，增加模型大小、數(shù)據(jù)集大小或訓(xùn)練步驟)通常會導(dǎo)致更高的準(zhǔn)確性。考慮到最近像BERT這樣的無監(jiān)督預(yù)訓(xùn)練方法的成功，這一點尤其正確，它可以將訓(xùn)練擴(kuò)展到非常大的模型和數(shù)據(jù)集。不幸的是，大規(guī)模的訓(xùn)練在計算上非常昂貴的，尤其是在沒有大型工業(yè)研究實驗室的硬件資源的情況下。因此，在實踐中，我們的目標(biāo)通常是在不超出硬件預(yù)算和訓(xùn)練時間的情況下獲得較高的準(zhǔn)確性。

對于大多數(shù)訓(xùn)練預(yù)算，非常大的模型似乎不切實際。相反，最大限度提高訓(xùn)練效率的策略是使用隱藏節(jié)點數(shù)較小或?qū)訑?shù)量較少的模型，因為這些模型運行速度更快，使用的內(nèi)存更少。

大模型訓(xùn)練的更快

然而，在我們的最近的論文中，我們表明這種減少模型大小的常見做法實際上與最佳的計算效率訓(xùn)練策略相反。相反，當(dāng)在預(yù)算內(nèi)訓(xùn)練Transformer模型時，你希望大幅度增加模型大小，但是早點停止訓(xùn)練。換句話說，我們重新思考了模型必須被訓(xùn)練直到收斂的隱含假設(shè)，展示了在犧牲收斂性的同時，有機(jī)會增加模型的大小。

這種現(xiàn)象發(fā)生的原因是，與較小的模型相比，較大的模型在較少的梯度更新中可以收斂于較低的測試誤差。此外，這種收斂速度的提高超過了使用更大模型的額外計算成本。因此，在考慮訓(xùn)練時間時，較大的模型可以更快地獲得更高的精度。

我們在下面的兩條訓(xùn)練曲線中展示了這一趨勢。在左側(cè)，我們繪制了預(yù)訓(xùn)練的驗證誤差RoBERTa，這是BERT的一個變體。RoBERTa模型越深，其混亂度就越低(我們的論文表明，對于更寬的模型也是如此)。這一趨勢也適用于機(jī)器翻譯。在右側(cè)，我們繪制了驗證BLEU分?jǐn)?shù)(越高越好)，當(dāng)訓(xùn)練一個英語到法語的Transformer機(jī)器翻譯模型。在相同的訓(xùn)練時間下，深度和寬度模型比小模型獲得更高的BLEU分?jǐn)?shù)。

有趣的是，對于訓(xùn)練前的RoBERTa來說，增加模型的寬度和/或深度都會導(dǎo)致更快的訓(xùn)練。對于機(jī)器翻譯，更寬的模型比更深的模型表現(xiàn)得更好。因此，我們建議在深入之前嘗試增加寬度。

我們還建議增加模型大小，而不是batch size大小。具體地說，我們確認(rèn)一旦batch size接近臨界范圍，增加batch size大小只會在訓(xùn)練時間上提供微小的改進(jìn)。因此，在資源受限的情況下，我們建議在這個關(guān)鍵區(qū)域內(nèi)使用batch size大小，然后使用更大的模型。

測試時間怎么辦呢？

盡管更大的模型具有更高的“訓(xùn)練效率”，但它們也增加了“推理”的計算和內(nèi)存需求。這是有問題的，因為推理的總成本遠(yuǎn)遠(yuǎn)大于大多數(shù)實際應(yīng)用的訓(xùn)練成本。然而，對于RoBERTa來說，我們證明了這種取舍可以與模型壓縮相協(xié)調(diào)。特別是，與小型模型相比，大型模型對模型壓縮技術(shù)更健壯。因此，人們可以通過訓(xùn)練非常大的模型，然后對它們進(jìn)行大量的壓縮，從而達(dá)到兩全其美的效果。

我們使用量化和剪枝的壓縮方法。量化以低精度格式存儲模型權(quán)重，修剪將某些神經(jīng)網(wǎng)絡(luò)的權(quán)值設(shè)置為零。這兩種方法都可以減少推理延遲和存儲模型權(quán)值的內(nèi)存需求。

我們首先在相同的時間內(nèi)預(yù)訓(xùn)練不同尺寸的RoBERTa模型。然后，我們在下游文本分類任務(wù)(MNLI)中對這些模型進(jìn)行微調(diào)，并使用修剪或量化。我們發(fā)現(xiàn)，對于給定的測試時間預(yù)算，最好的模型是那些經(jīng)過大量訓(xùn)練然后經(jīng)過大量壓縮的模型。

例如，考慮最深度模型的修剪結(jié)果(左圖中的橙色曲線)。不需要修剪模型，它達(dá)到了很高的精度，但是使用了大約2億個參數(shù)(因此需要大量的內(nèi)存和計算)。但是，可以對這個模型進(jìn)行大量的修剪(沿著曲線向左移動的點)，而不會嚴(yán)重影響準(zhǔn)確性。這與較小的模型形成了鮮明的對比，如粉紅色顯示的6層模型，其精度在修剪后嚴(yán)重下降。量化也有類似的趨勢(下圖)?？偟膩碚f，對于大多數(shù)測試預(yù)算(在x軸上選擇一個點)來說，最好的模型是非常大但是高度壓縮的模型。

我們已經(jīng)證明了增加Transformer模型的大小可以提高訓(xùn)練和推理的效率，即，應(yīng)該先“大模型訓(xùn)練”，然后再“壓縮”。這一發(fā)現(xiàn)引出了許多其他有趣的問題，比如為什么大的模型收斂得更快，壓縮得更好。