作者：九羽，公眾號：煉丹筆記

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

Google Brain科學家Barret Zoph表示，他們設計了一個名叫「Switch Transformer」的簡化稀疏架構，可以將語言模型的參數量擴展至 1.6 萬億。萬萬沒想到，模型規模的演進如此之快，沒幾個月的時間，就從千億走向了萬億，當我們還在研究BERT的各種迭代時，世界上那批頂尖的人已經開啟了另一扇“暴力美學”的大門。而這，才是真正的深度領域的“軍備競賽“。

對于長文沒有閱讀習慣的朋友，可以直接讀一下本文摘要。

1）Switch Transformer在網絡結構上最大的改進是Sparse routing的稀疏結構，相比于OpenAI在GPT-3里所使用的Sparse Attention，需要用到稀疏算子而很難發揮GPU、TPU硬件性能的問題。Switch Transformer不需要稀疏算子，可以更好的適應GPU、TPU等硬件

2）Switch Transformer雖然有1.6萬億參數，但通過Sparse routing的改進，每輪迭代只會觸發部分Expert的計算，而每個token也只會路由給一個Expert，所以對算力的需求并沒有隨著參數量的增加而大幅增長，使得這個模型更加容易訓練。

3）數據并行、模型并行、Expert并行的并行策略設計，在MoE網絡結構上能夠獲得更低的通信開銷，提高并行的效率。

在深度學習中，模型通常對所有輸入重復使用相同的參數。而MoE模型則是為每個例子選擇不同的參數。于是一個稀疏激活的模型（參數數量驚人但計算成本不變）誕生了。然而，盡管取得了一些顯著的成功，但由于復雜性、通信成本和訓練的不穩定性，模型廣泛采用仍需優化。

我們用Switch Transformer來解決這些問題。同時，我們簡化了MoE路由算法，設計了直觀的改進模型，降低了通信和計算成本。我們提出的訓練方法減輕了不穩定性，并且我們首次展示了用較低精度（bfloat16）格式訓練大型稀疏模型的可能性。

同時，基于T5 Base和T5 Large（Raffel et al.，2019）設計模型，以在相同計算資源的情況下獲得高達7倍的預訓練速度。這些改進擴展到多語言設置中，我們在所有101種語言中測量mT5基本版本的增益。最后，通過在“巨大的干凈的爬蟲語料庫”上預訓練多達萬億個參數的模型，提高了當前語言模型的規模，并實現了比T5-XXL模型4倍的加速。

大規模訓練是實現靈活和強大的神經語言模型的有效途徑。雖然有效，但計算量也非常大（Strubell等人，2019年）。為了提高計算效率，我們提出了一種稀疏激活模型：Switch Transformer。在我們的例子中，稀疏性來自于為每個傳入的例子激活一個子集的神經網絡權重。

Switch Transformer在Mix of Expert的基礎上，采用sparsely activated方法，只使用了模型權重的子集，轉換模型內輸入數據的參數達成相同的效果。

MoE(Mix of Expert)是一種神經網絡，也屬于一種combine的模型，上個世紀90年代被提出。適用于數據集中的數據產生方式不同。不同于一般的神經網絡的是它根據數據進行分離訓練多個模型，各個模型被稱為專家，而門控模塊用于選擇使用哪個專家，模型的實際輸出為各個模型的輸出與門控模型的權重組合。各個專家模型可采用不同的函數（各種線性或非線性函數）。混合專家系統就是將多個模型整合到一個單獨的任務中。

在分布式訓練設置中，模型將不同的權重分配到不同的設備上，雖然權重會隨著設備數量的增加而增加，但每個設備可以保持內存和計算足跡的自我管理。

Switch Transformer在許多任務上的效果有提升。

（1）在使用相同數量的計算資源的情況下，它可以使預訓練的速度提高了7倍以上。

（2）大型稀疏模型可以用來創建更小、更稠密的模型，這些模型可以對任務進行微調，其質量增益只有大型模型的30% 。

（3）Switch Transformer 模型在100多種不同的語言之間進行翻譯，研究人員觀察到其中101種語言都得到提升 ，而其中91% 超過基線模型4倍以上的速度。

Shazeer（2018）和Lepikhin（2020）通過將MoE層添加到Transformer的密集前饋網絡（FFN）計算中，設計了MoE變壓器（Shazeer et al.，2017）。同樣，我們的工作也替換了變壓器中的FFN層，但在此簡要探討了另一種設計。我們將開關層添加到Transformer自我注意層中。為此，我們將生成查詢、鍵和值的可訓練權重矩陣替換為交換層。

由于TPU加速器的限制，我們的張量的形狀必須是靜態的。因此，每個expert都有處理token表示的有限且固定的能力。然而，這為我們的模型提出了一個問題，該模型在運行時動態路由token，這可能導致在exper上的不均勻分布。

如果發送給exper的token數小于exper容量，那么計算可能只是簡單地進行填充——這是對硬件的低效使用，但在數學上是正確的。但是，當發送給exper的令牌數大于其容量（exper溢出）時，需要一個協議來處理這個問題。Lepikhin等人（2020年）采用了exper模型的混合模型，并通過將其表示傳遞到下一層來解決exper溢出問題，而無需通過我們也遵循的剩余連接進行處理。

不能保證一個模型在訓練前目標上的效果會轉化為下游任務的結果。下圖顯示了上游模型質量的相關性，包括稠密模型和非稠密模型和Switch模型，在C4預訓練任務上使用兩個下游任務度量：平均SuperGLUE性能和TriviaQA分數。我們選擇這兩個任務作為一個探索模型的推理和其他事實知識。

總結來說，Switch Transformers模型有兩個創新：

（1）基于Transformer MoE網絡結構，簡化了MoE的routing機制，降低了計算量；

（2）進一步通過數據并行、模型并行、Expert并行的方式降低了訓練通信量，提升訓練性能。

參考文獻

1. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity：https://arxiv.org/pdf/2101.03961.pdf

更多干貨，請關注微信公眾號：煉丹筆記

傳送門：Google Brain：從不廢話，直接扔大

總結

以上是生活随笔為你收集整理的NLP炼丹笔记：Switch Transformers 朴实无华 大招秒杀的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。