作者｜AI Box

?來源｜RUC AI Box

本文深入解讀了由 Google Brain 設(shè)計的名叫「Switch Transformer」的簡化稀疏架構(gòu)，可以將語言模型的參數(shù)量擴展至 1.6 萬億（GPT-3 是 1750 億）。在計算資源相同的情況下，Switch Transformer 的訓練速度可以達到 T5 模型的 4-7 倍。本文將從「為什么選擇MoE」、「如何設(shè)計高效的網(wǎng)絡(luò)結(jié)構(gòu)」、「訓練技巧」和「一些重要的討論」4 個方面進行解讀。

論文標題：

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

論文鏈接：

https://arxiv.org/abs/2101.03961

代碼鏈接：

https://github.com/tensorflow/mesh

Why is MoE

在現(xiàn)有的深度神經(jīng)網(wǎng)絡(luò)方法中，針對模型的輸入，所有的參數(shù)都會參與計算。在預訓練模型參數(shù)量變的越來越大的情況下，計算資源的需求也會變得巨大。而Mixture of Experts（MoE）改變了這種情況。MoE可以為不同的輸入選擇性地激活模型中的一部分參數(shù)參與計算，這樣在增大模型參數(shù)量的同時，計算量可以維持相對不變。

一種典型的MoE框架由一個門控子網(wǎng)絡(luò)（Gating network）和多個專家子網(wǎng)絡(luò)（Expert odel）構(gòu)成，門控網(wǎng)絡(luò)為輸入x計算各個專家網(wǎng)絡(luò)輸出所占的比重，然后采取加權(quán)求和的方式得到最終的輸出。

另有使用門控子網(wǎng)絡(luò)對輸入進行路由選擇，即根據(jù)各個專家網(wǎng)絡(luò)對應(yīng)的門控值（gating value），選擇出Top-K個專家子網(wǎng)絡(luò)參與當前輸入的實際計算，這樣可以顯著降低計算量。

本文就基于MoE的思想，將Transformer中的前饋全連接子層（Feed-Forward Network，FFN）視為Expert，使用多個FFN代替原來單一的FFN，并且使用了最簡單的路由選擇策略，將K設(shè)置為1，即不同的輸入只會選擇一個FFN進行計算。這樣相比較于原來的結(jié)構(gòu)，計算量只增加了路由選擇的計算量，而新增的計算量相比較于原來的計算量而言可以忽略，這樣就實現(xiàn)了增大模型參數(shù)的同時維持相對不變的計算量。

How?Efficient

1. 模型結(jié)構(gòu)——簡化稀疏路由

混合專家模型（MoE）：x表示每一個token的輸入，則通過路由權(quán)重計算得到的logits為h(x)，門控值（gated value）通過所有專家的logits使用softmax計算得到。輸出y來自不同的專家的加權(quán)和，權(quán)重即路由門控值得到。?

Switch Routing:?

本文中使用了一個創(chuàng)新的路由策略，即每次只發(fā)給一個專家，這樣可以簡化路由的計算量的同時保證模型的性能。這樣做的優(yōu)勢：

（1）路由計算量減少，只有一個expert激活；

（2）expert中的batch_size（專家容量）至少減半；

（3）簡化路由的實現(xiàn)，減少傳統(tǒng)MOE方法中通信的代價。

2. 模型結(jié)構(gòu)——高效稀疏路由

分布式Switch實現(xiàn)：?

（1）問題：模型編譯是靜態(tài)確定的，計算是動態(tài)的，如何確定每一個expert維度；

（2）方法：使用capacity factor擴展，太多的容量導致計算和內(nèi)存消耗增加，太小的容量導致token被放棄計算傳到下一層，結(jié)論是經(jīng)驗性的根據(jù)“model quality and speed”的權(quán)衡來選擇較低比例“dropped tokens”。

Load Balancing Loss:

為了促使每個expert都可以拿到近似均勻分布的樣本，這里引入負載均衡損失。當f_i= p_i= 1/N的時候損失是最小的。

實驗：

本文采用的實驗是T5模型的基礎(chǔ)上應(yīng)用switch transformer和MOE，下面是一些結(jié)論：

（1）Switch Transformer比MoE和Dense模型都要好；

（2）Switch Transformer在capacity比較小的時候效果更好；

（3）Switch Transformer比MoE要少一點計算量，如果讓它們計算量相等的話，那么Switch Transformer還可以有更多提升(Switch-base-Expand)。

?

3. 為什么高效——數(shù)據(jù)、模型及專家并行

本章屬于高效執(zhí)行這小部中最重要的部分，是關(guān)于并行部分的解釋，這里補充一下關(guān)于各種并行的方法的解釋。標準的數(shù)據(jù)并行的定義是一個batch的數(shù)據(jù)在不同的device上并行處理，這時每一個device上都保存了模型的一份完整拷貝，前向計算完進行梯度匯總和更新。模型并行表示模型不同的參數(shù)（層、組件）分配到不同的device上，處理一個batch的數(shù)據(jù)。

本文中下圖中上面一行整體表示權(quán)重的分配方式，下面一行表示數(shù)據(jù)的分配方式，一種顏色表示一個矩陣（a unique weight matrix）。其中每一個方格表示一個core。?

?

數(shù)據(jù)并行：?

第一列表示數(shù)據(jù)并行，模型權(quán)重拷貝16份，16個同一種顏色矩陣分別表示一個完整的模型，數(shù)據(jù)側(cè)則是一個完整的矩陣，這里可以理解為16個模型計算完成后由于存在梯度匯總再更新的步驟，所以整體更新的是一個batch，因此這里數(shù)據(jù)側(cè)是一個唯一的矩陣。簡單來說就是模型復制，數(shù)據(jù)并行。?

模型并行：

模型并行部分從模型側(cè)看出來，16個cores維護的是一個整體的模型，但是每一個core只分配到其中非常高通信代價，同一個batch在所有的core上計算，由于1個core中分布了不同的模型權(quán)重，每次計算完都需要和其他的core進行通信。

專家并行：

原本路由機制只會分配給當前core中的不同的expert，現(xiàn)在則有可能會分配到其他的core下的expert，范圍更大。

實驗結(jié)論：

（1）switch-C表示和T5有相同的ppl下速度是T5的4倍，并且隨著訓練進行還會繼續(xù)擴大優(yōu)勢；

（2）switch-xxl是和T5有相同的FLOPs/seq下，效果好于T5；

（3）從最后兩列看出，switch的兩個模型的樣本利用率高于T5，T5在500k steps下才能達到switch效果。

?

Training Tricks

1. 大型稀疏模型中使用隨機精度

Switch Transformer參數(shù)量達到了一萬多億，為了保證計算效率，降低到每次過只過一個expert，這樣相當于關(guān)閉了模型某些部分的硬切換機制，就引入了稀疏性，而稀疏性會導致模型訓練不穩(wěn)定，換句話說，就是稀疏模型可能會對隨機種子比較敏感。

當使用bfloat16精度的時候，模型的不穩(wěn)定性會影響訓練性能。這個bfloat16是谷歌的一個格式，全稱叫Google Brain Floating Point。MoE Transformer中使用了float32精度來訓練，這樣會帶來一定的通信成本。所以這里作者使用了selectively casting，只在模型的local部分使用float32的精度，這樣可以實現(xiàn)穩(wěn)定性。

具體來說，在router的輸入端使用float32精度，并且只在router函數(shù)中使用，也就是說在當前設(shè)備的局部進行計算。在函數(shù)結(jié)束進行廣播通信的時候，使用bfloat16精度，把這個API暴露給網(wǎng)絡(luò)其余部分，因此float32的值永遠不會離開本地設(shè)備，而且設(shè)備間通信仍然保持低精度。這樣既可以從float32精度中獲得收益，也不會帶來額外的通信成本。

下面這個表說明了這種方法的好處，可以看到，作者提出的這種方法可以保證和bfloat16一樣的訓練速度，但是獲得了媲美float32精度訓練的穩(wěn)定性。

?

2. 使用更小的參數(shù)初始化來保證穩(wěn)定性

其motivation還是為了保證模型的穩(wěn)定性。作者觀察到，在Switch Transformer中，合適的初始化方法也是成功訓練的一個重要因素。他們的做法是，用均值μ=0，標準差σ=√s / n的截斷正態(tài)分布來對權(quán)重矩陣進行初始化，其中s是放縮超參，n是權(quán)重向量輸入的數(shù)量。這里作為減小路由數(shù)量不穩(wěn)定性的一個補救，作者把transformer默認的初始化s從1.0縮小10倍。他們發(fā)現(xiàn)這樣的話既可以提高質(zhì)量，又可以減少實驗中訓練不穩(wěn)定的可能性。

下面的表三測量了訓練初期，模型質(zhì)量的改善還有方差的降低。作者發(fā)現(xiàn)，用這個negative log perp度量的平均模型質(zhì)量得到了一個比較大的改善，而且運行多次的方差也減少了很多，從0.68到0.01。他們使用這個方法，把200多兆參數(shù)量的baseline穩(wěn)定地訓練成了超過一萬億個參數(shù)的超大模型。

?

3. 對Switch Tranformer這種大型稀疏模型做正則

因為這篇論文是在一個大的語料庫上做預訓練，然后在比較小的下游任務(wù)做微調(diào)，一個比較自然的問題就是過擬合問題，因為許多微調(diào)任務(wù)中比較缺乏數(shù)據(jù)。Switch Transformer比一些dense model的參數(shù)要多很多，這樣可能會導致在這種比較小的下游任務(wù)上更容易過擬合。

之前有一些工作使用dropout來防止過擬合。這篇文章是提出了一種比較簡單的在微調(diào)時候減輕過擬合問題的方法，就是增加expert內(nèi)部的dropout，他們叫expert dropout。在微調(diào)的時候，只在每個expert層的過渡的feed-forward層計算的時候大幅增加dropout的概率。

下面的表四做了這部分的實驗，我們可以發(fā)現(xiàn)，只是簡單對所有層都增加dropout之后，會得到一個比較差的結(jié)果，當給非expert層設(shè)置一個比較小的dropout，也就是0.1，給expert層設(shè)置一個比較大的dropout rate會在四個下游任務(wù)上得到一定性能的提升。

?

4. No-Token-Left-Behind機制

因為TPU是有限制的，張量的shape必須是靜態(tài)的，所以每個expert處理token表示的能力是有限而且固定的，但是模型在運行時是動態(tài)路由token的，通過softmax概率來進行路由，這樣可能會導致在expert上的不均勻分布。如果發(fā)送給expert的token數(shù)小于expert的實際容量，這樣是對硬件的低效使用。所以作者用了一個方法解決這個問題。這里作者構(gòu)建了No-Token-Left-Behind機制，它重點在于反復地把第一次路由的所有溢出的token重新進行路由，下圖是這個機制的說明。

可以看到，第一輪路由溢出的token，在第二輪被重新路由到了概率第二高的expert中，這樣正好達到了飽和。如果第二輪還有溢出，就還會繼續(xù)迭代下去。這里需要注意的是，作者在實驗中發(fā)現(xiàn)，這種trick雖然保證了資源的利用率，但是并沒有帶來什么性能上的提升。這里他們猜測，一旦網(wǎng)絡(luò)學習到了token和expert之間的聯(lián)系，如果使用這種機制，向第二甚至第三的expert發(fā)送了token，相當于對token和expert的聯(lián)系進行了更改，這樣可能也會影響性能。

?

Some Discussions

Q1:?Switch Transformer效果更好，是否是因為更大的參數(shù)量？

A1:?是的，并且是設(shè)計成這樣的。大型模型已被廣泛顯示出具有更好的性能 [1]。我們的模型在使用相同的計算資源的情況下，效率更高，速度更快。

?

Q2: 我沒有supercomputer, 這篇工作對我還有用嗎？

A2: 盡管這項工作集中在非常大的模型上。但是只有兩個exports就能夠提升性能，并且可輕松地適用常用的GPU或TPU的內(nèi)存限制。因此，這項技術(shù)在小規(guī)模環(huán)境中仍然有用。

?

?

Q3: Sparse Model是否在pareto曲線上優(yōu)于Dense Model?

A3: 是的。實驗表明，在固定的計算量和時間上，Sparse Model的性能都優(yōu)于Dense Model。

?

Q4: 我無法部署一個萬億參數(shù)的模型，能夠縮小使用這些模型嗎？

A4:?無法保證縮小后的模型的質(zhì)量。但是以10到100倍的壓縮率將Sparse Model蒸餾為Dense Model，可以獲得Export Model 30%的質(zhì)量增益。

?

Q5: 為什么使用Switch Transformer來代替模型并行方式的Dense Model？

A5: 從時間的角度看，Switch Transformer要比Dense Model高效得多。另外，這兩者并不沖突，可以在Switch Transformer中使用模型并行來增加每個token的Flops，但這會導致傳統(tǒng)模型并行性的降低。

?

Q6:?為什么Sparse Model尚未廣泛使用？

A6:?嘗試Sparse Model的動機受到了Dense Model的巨大成功的阻礙（其成功是由與深度學習硬件的共同適應(yīng)驅(qū)動的 [2]）。此外，Sparse Model存在以下幾方面的阻礙：（1）模型復雜性；（2）訓練困難；（3）通信成本。Switch Transformer在緩解這些問題上取得了巨大進步。

?

參考文獻

[1] Kaplan et al. Scaling laws for neural language models. ?arXiv preprint arXiv:2001.08361, 2020.

[2] Sara Hooker. The hardware lottery. arXiv preprint arXiv:2009.06489, 2020.

????

現(xiàn)在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關(guān)注」訂閱我們的專欄吧

關(guān)于PaperWeekly

PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術(shù)平臺。如果你研究或從事 AI 領(lǐng)域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

總結(jié)

以上是生活随笔為你收集整理的深入解读首个万亿级语言模型Switch Transformer的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。