©原創作者 | LJ

GLaM: Efficient Scaling of Language Models with Mixture-of-Experts

https://arxiv.org/pdf/2112.06905.pdf

01 摘要

這是上個月谷歌剛剛在arxiv釋出的論文，證明了一種能scale GPT-3但又比較節省耗能的架構。

GPT-3自問世以來在多項自然語言處理的任務上都有超強的表現。但是訓練GPT-3這樣龐大的模型非常耗費能源。

在這篇論文中，作者開發了以Mixture of Experts為基礎的GlaM (Generalist Language Model)。它雖然引數量有GPT-3的7倍之多，但訓練起來只需GPT-3三分之一的能耗，而且在NLP任務的表現上相比GPT-3持平甚至更優。

02 什麼是Mixture of Experts Model (MoE）

MoE這個概念其實已經提出很久了。這個概念本身非常容易理解，有點類似ensemble：與其訓練一個模型，我們訓練數十個獨立的“專家模型”(expert model)。

與簡單的ensemble不同的是，在做訓練或推斷(inference)的時候，我們用一個gating network來“挑選專家” — 在幾十個專家模型中挑選出幾個適合的專家模型用來計算。通俗的講，這些專家“術業有專攻”，根據所長而分工。

那麼，為什麼MoE可以省能耗呢？因為無論是訓練或者推算的時候，每次真正的計算只有幾個專家被啟用。所以，雖然引數量很大，但每次用到的引數只是很小的一部分。

這個團隊在2017年在一篇ICLR的論文[1]裡已經把MoE的概念應用在了當時NLP state-of-the-art的RNN model上，並且超越了當時的state of the art。

這次，因為GPT-3的釋出，作者又將MoE的概念應用在GPT-3這樣以transformer為基礎的模型上。

這次也同樣得到了很好的結果。這篇論文其實是2017年那篇論文的自然延伸，方法上大同小異，所以很多的精華在2017年的那篇論文裡。我們今天一併總結。

03 實驗細節

3.1 困難和挑戰

MoE這個概念既然由來已久，為什麼從前沒有得到很好的實驗結果呢？這是因為在實際應用中MoE有幾個難點之前一直沒被解決，而Google團隊逐一找到了解決方案。

難點一：batch size太小

MoE訓練的時候會把一個batch的訓練資料分佈到幾十個expert model上。一分下來，分到每個model頭上的資料量就很少了。一般GPU和CPU需要很大的batch size保證運算的高效，不然固定成本(overhead cost)太大了。但如果一開始的batch太大的話，memory又不夠。作者巧妙的把模型的某些部分做分散式計算，通過data parallelism和model parallelism來彌補這個問題。

難點二：專家模型分工不均勻

Gating network挑選專家的結果經常會是同樣的幾個專家。這樣就會引起惡性迴圈：好的專家模型能分到的訓練資料越來越多，模型也變得越來越好；不好的專家一直不被呼叫，就一直不好。為此，作者在loss function里加了一項權重用來load balance，鼓勵gating network分配平衡。

難點三：需要在大資料集上才能體現優勢

已有的專家模型的論文所採用的的資料集只有60萬個資料點。專家模型的引數量龐大。所以需要更大的資料集才能真正體現他的優勢。作者採用和GPT-3類似的訓練集，很自然地解決了這個問題。

3.2 實際訓練的注意事項

由於這個模型有幾個兆的引數量，訓練一次就是一個龐大的工程，所以容錯率非常低：如果模型訓練時出現bug了，重頭再訓練非常昂貴，也基本沒什麼空間來優化超引數。

另外，像這一類極稀疏引數被啟用的模型很容易碰到分母為零之類的計算不穩定性。基於此，在理論的基礎上，作者也分享了實際訓練也有一些注意事項。

注意事項一： 先訓練一個小模型。在這個過程中，能夠暴露很多潛在的bug。在保證這個迷你模型能夠收斂之後，再真正訓練大的模型。

注意事項二：如果一個batch中出現導數為零或者無限大的情況時，直接跳過這個batch的訓練資料。

注意事項三：及時儲存checkpoint。一旦出錯，自動把之前的checkpoint調出來，從那裡開始繼續訓練，這樣就不用從頭開始。由於每個batch是隨機抽取的，即使上一次隨機的資料導致了分母為零等不穩定性，新一批資料是重新隨機選擇的，有可能能順利通過之前報錯的地方。

注意事項四：雖然是老生常談，但訓練資料的質量真的很重要。作者花了很多心思來提高資料的質量，並印證這樣做的重要性。

原始的資料集包含了1.6 x 10^12的單詞，包括了書籍和網際網路上各種各樣的網頁。有的網頁是專業人員寫的，文字質量高；有的網頁是各種論壇上零零總總的使用者評論，文字質量低。

和GPT-3團隊的思路一樣，作者也自己訓練了一個線性模型，用以判斷文字質量。

他們將維基百科、出版的書籍和一些其他信得過的網站作為正面樣本，把其他的網站作為負面樣本，讓模型去鑑別。通過這個過程，他們過濾掉了很多低質量網頁，把他們從訓練集中刪除。

這樣做的效果是非常顯著的。在下圖中可以看到綠色的是過濾後的訓練集，灰色的虛線是過濾前的原始訓練集。橫軸是訓練集的大小。之前很多的研究都在看訓練集大小的影響。

這個結果告訴我們，訓練集不能一味只追求數量，而是要追求數量。同樣的原始訓練集，數量擴大三倍，還遠不如在同樣的數量下把質量提上去。另外，相對於自然語言理解（Natural language understanding， NLU)的任務，資料質量在自然語言生產（Natural language generation, NLG)的任務上效果更為顯著。

04 結果

4.1 MoE和GPT-3的比較

GLaM在29個自然語言的任務上總的來說相比GPT-3略勝一籌。最重要的是訓練的總能耗僅僅是GPT-3的三分之一。

以TriviaQA任務舉例，上圖中Dense指的是類似GPT-3架構的單個模型。Dense和MoE model的準確率都會因為引數量增長而增長。但是在相同的運算量下（橫軸），MoE總是表現得更好。

4.2. 需要多少專家模型

相比於dense model，MoE如果想scale的話不僅可以將模型變得更寬更深，還可以增加專家的數量。只要每次被啟用的專家數量不變，增加專家並不增加prediction時的運算量。

上圖中，作者控制每次prediction的計算量不變，調整不同的專家數量。一般看來，專家越多，模型表現越好。

05 總結

總的來說，作者開發了以MoE為基礎的GLaM模型。雖然模型引數量很多，但通過啟用少量的專家，這類的模型訓練和推算的能耗更低，而且結果比GPT-3更好。