深入解读首个万亿级语言模型Switch Transformer
作者｜AI Box
?來源｜RUC AI Box
本文深入解讀了由 Google Brain 設計的名叫「Switch Transformer」的簡化稀疏架構，可以將語言模型的參數量擴展至 1.6 萬億（GPT-3 是 1750 億）。在計算資源相同的情況下，Switch Transformer 的訓練速度可以達到 T5 模型的 4-7 倍。本文將從「為什么選擇MoE」、「如何設計高效的網絡結構」、「訓練技巧」和「一些重要的討論」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
在現有的深度神經網絡方法中，針對模型的輸入，所有的參數都會參與計算。在預訓練模型參數量變的越來越大的情況下，計算資源的需求也會變得巨大。而Mixture of Experts（MoE）改變了這種情況。MoE可以為不同的輸入選擇性地激活模型中的一部分參數參與計算，這樣在增大模型參數量的同時，計算量可以維持相對不變。
一種典型的MoE框架由一個門控子網絡（Gating network）和多個專家子網絡（Expert odel）構成，門控網絡為輸入x計算各個專家網絡輸出所占的比重，然后采取加權求和的方式得到最終的輸出。
另有使用門控子網絡對輸入進行路由選擇，即根據各個專家網絡對應的門控值（gating value），選擇出Top-K個專家子網絡參與當前輸入的實際計算，這樣可以顯著降低計算量。
本文就基于MoE的思想，將Transformer中的前饋全連接子層（Feed-Forward Network，FFN）視為Expert，使用多個FFN代替原來單一的FFN，并且使用了最簡單的路由選擇策略，將K設置為1，即不同的輸入只會選擇一個FFN進行計算。這樣相比較于原來的結構，計算量只增加了路由選擇的計算量，而新增的計算量相比較于原來的計算量而言可以忽略，這樣就實現了增大模型參數的同時維持相對不變的計算量。
How?Efficient
1. 模型結構——簡化稀疏路由
混合專家模型（MoE）：x表示每一個token的輸入，則通過路由權重計算得到的logits為h(x)，門控值（gated value）通過所有專家的logits使用softmax計算得到。輸出y來自不同的專家的加權和，權重即路由門控值得到。?
Switch Routing:?
本文中使用了一個創新的路由策略，即每次只發給一個專家，這樣可以簡化路由的計算量的同時保證模型的性能。這樣做的優勢：
（1）路由計算量減少，只有一個expert激活；
（2）expert中的batch_size（專家容量）至少減半；
（3）簡化路由的實現，減少傳統MOE方法中通信的代價。
2. 模型結構——高效稀疏路由
分布式Switch實現：?
（1）問題：模型編譯是靜態確定的，計算是動態的，如何確定每一個expert維度；
（2）方法：使用capacity factor擴展，太多的容量導致計算和內存消耗增加，太小的容量導致token被放棄計算傳到下一層，結論是經驗性的根據“model quality and speed”的權衡來選擇較低比例“dropped tokens”。
Load Balancing Loss:
為了促使每個expert都可以拿到近似均勻分布的樣本，這里引入負載均衡損失。當f_i= p_i= 1/N的時候損失是最小的。
實驗：
本文采用的實驗是T5模型的基礎上應用switch transformer和MOE，下面是一些結論：
（1）Switch Transformer比MoE和Dense模型都要好；
（2）Switch Transformer在capacity比較小的時候效果更好；
（3）Switch Transformer比MoE要少一點計算量，如果讓它們計算量相等的話，那么Switch Transformer還可以有更多提升(Switch-base-Expand)。
?
3. 為什么高效——數據、模型及專家并行
本章屬于高效執行這小部中最重要的部分，是關于并行部分的解釋，這里補充一下關于各種并行的方法的解釋。標準的數據并行的定義是一個batch的數據在不同的device上并行處理，這時每一個device上都保存了模型的一份完整拷貝，前向計算完進行梯度匯總和更新。模型并行表示模型不同的參數（層、組件）分配到不同的device上，處理一個batch的數據。
本文中下圖中上面一行整體表示權重的分配方式，下面一行表示數據的分配方式，一種顏色表示一個矩陣（a unique weight matrix）。其中每一個方格表示一個core。?
?
數據并行：?
第一列表示數據并行，模型權重拷貝16份，16個同一種顏色矩陣分別表示一個完整的模型，數據側則是一個完整的矩陣，這里可以理解為16個模型計算完成后由于存在梯度匯總再更新的步驟，所以整體更新的是一個batch，因此這里數據側是一個唯一的矩陣。簡單來說就是模型復制，數據并行。?
模型并行：
模型并行部分從模型側看出來，16個cores維護的是一個整體的模型，但是每一個core只分配到其中非常高通信代價，同一個batch在所有的core上計算，由于1個core中分布了不同的模型權重，每次計算完都需要和其他的core進行通信。
專家并行：
原本路由機制只會分配給當前core中的不同的expert，現在則有可能會分配到其他的core下的expert，范圍更大。
實驗結論：
（1）switch-C表示和T5有相同的ppl下速度是T5的4倍，并且隨著訓練進行還會繼續擴大優勢；
（2）switch-xxl是和T5有相同的FLOPs/seq下，效果好于T5；
（3）從最后兩列看出，switch的兩個模型的樣本利用率高于T5，T5在500k steps下才能達到switch效果。
?
Training Tricks
1. 大型稀疏模型中使用隨機精度
Switch Transformer參數量達到了一萬多億，為了保證計算效率，降低到每次過只過一個expert，這樣相當于關閉了模型某些部分的硬切換機制，就引入了稀疏性，而稀疏性會導致模型訓練不穩定，換句話說，就是稀疏模型可能會對隨機種子比較敏感。
當使用bfloat16精度的時候，模型的不穩定性會影響訓練性能。這個bfloat16是谷歌的一個格式，全稱叫Google Brain Floating Point。MoE Transformer中使用了float32精度來訓練，這樣會帶來一定的通信成本。所以這里作者使用了selectively casting，只在模型的local部分使用float32的精度，這樣可以實現穩定性。
具體來說，在router的輸入端使用float32精度，并且只在router函數中使用，也就是說在當前設備的局部進行計算。在函數結束進行廣播通信的時候，使用bfloat16精度，把這個API暴露給網絡其余部分，因此float32的值永遠不會離開本地設備，而且設備間通信仍然保持低精度。這樣既可以從float32精度中獲得收益，也不會帶來額外的通信成本。
下面這個表說明了這種方法的好處，可以看到，作者提出的這種方法可以保證和bfloat16一樣的訓練速度，但是獲得了媲美float32精度訓練的穩定性。
?
2. 使用更小的參數初始化來保證穩定性
其motivation還是為了保證模型的穩定性。作者觀察到，在Switch Transformer中，合適的初始化方法也是成功訓練的一個重要因素。他們的做法是，用均值μ=0，標準差σ=√s / n的截斷正態分布來對權重矩陣進行初始化，其中s是放縮超參，n是權重向量輸入的數量。這里作為減小路由數量不穩定性的一個補救，作者把transformer默認的初始化s從1.0縮小10倍。他們發現這樣的話既可以提高質量，又可以減少實驗中訓練不穩定的可能性。
下面的表三測量了訓練初期，模型質量的改善還有方差的降低。作者發現，用這個negative log perp度量的平均模型質量得到了一個比較大的改善，而且運行多次的方差也減少了很多，從0.68到0.01。他們使用這個方法，把200多兆參數量的baseline穩定地訓練成了超過一萬億個參數的超大模型。
?
3. 對Switch Tranformer這種大型稀疏模型做正則
因為這篇論文是在一個大的語料庫上做預訓練，然后在比較小的下游任務做微調，一個比較自然的問題就是過擬合問題，因為許多微調任務中比較缺乏數據。Switch Transformer比一些dense model的參數要多很多，這樣可能會導致在這種比較小的下游任務上更容易過擬合。
之前有一些工作使用dropout來防止過擬合。這篇文章是提出了一種比較簡單的在微調時候減輕過擬合問題的方法，就是增加expert內部的dropout，他們叫expert dropout。在微調的時候，只在每個expert層的過渡的feed-forward層計算的時候大幅增加dropout的概率。
下面的表四做了這部分的實驗，我們可以發現，只是簡單對所有層都增加dropout之后，會得到一個比較差的結果，當給非expert層設置一個比較小的dropout，也就是0.1，給expert層設置一個比較大的dropout rate會在四個下游任務上得到一定性能的提升。
?
4. No-Token-Left-Behind機制
因為TPU是有限制的，張量的shape必須是靜態的，所以每個expert處理token表示的能力是有限而且固定的，但是模型在運行時是動態路由token的，通過softmax概率來進行路由，這樣可能會導致在expert上的不均勻分布。如果發送給expert的token數小于expert的實際容量，這樣是對硬件的低效使用。所以作者用了一個方法解決這個問題。這里作者構建了No-Token-Left-Behind機制，它重點在于反復地把第一次路由的所有溢出的token重新進行路由，下圖是這個機制的說明。
可以看到，第一輪路由溢出的token，在第二輪被重新路由到了概率第二高的expert中，這樣正好達到了飽和。如果第二輪還有溢出，就還會繼續迭代下去。這里需要注意的是，作者在實驗中發現，這種trick雖然保證了資源的利用率，但是并沒有帶來什么性能上的提升。這里他們猜測，一旦網絡學習到了token和expert之間的聯系，如果使用這種機制，向第二甚至第三的expert發送了token，相當于對token和expert的聯系進行了更改，這樣可能也會影響性能。
?
Some Discussions
Q1:?Switch Transformer效果更好，是否是因為更大的參數量？
A1:?是的，并且是設計成這樣的。大型模型已被廣泛顯示出具有更好的性能 [1]。我們的模型在使用相同的計算資源的情況下，效率更高，速度更快。
?
Q2: 我沒有supercomputer, 這篇工作對我還有用嗎？
A2: 盡管這項工作集中在非常大的模型上。但是只有兩個exports就能夠提升性能，并且可輕松地適用常用的GPU或TPU的內存限制。因此，這項技術在小規模環境中仍然有用。
?
?
Q3: Sparse Model是否在pareto曲線上優于Dense Model?
A3: 是的。實驗表明，在固定的計算量和時間上，Sparse Model的性能都優于Dense Model。
?
Q4: 我無法部署一個萬億參數的模型，能夠縮小使用這些模型嗎？
A4:?無法保證縮小后的模型的質量。但是以10到100倍的壓縮率將Sparse Model蒸餾為Dense Model，可以獲得Export Model 30%的質量增益。
?
Q5: 為什么使用Switch Transformer來代替模型并行方式的Dense Model？
A5: 從時間的角度看，Switch Transformer要比Dense Model高效得多。另外，這兩者并不沖突，可以在Switch Transformer中使用模型并行來增加每個token的Flops，但這會導致傳統模型并行性的降低。
?
Q6:?為什么Sparse Model尚未廣泛使用？
A6:?嘗試Sparse Model的動機受到了Dense Model的巨大成功的阻礙（其成功是由與深度學習硬件的共同適應驅動的 [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.
????
現在，在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關注」訂閱我們的專欄吧
關于PaperWeekly
PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。
總結
以上是生活随笔為你收集整理的深入解读首个万亿级语言模型Switch Transformer的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 如何有效识别航空旅客未申报的危险品
- 下一篇： 克什米尔在哪个地图上有？