首個兆級模型！Google 重量級推出語言模型 Switch Transformers，1.6 兆參數秒殺 GPT-3

眾所皆知，參數量是機器學習算法的關鍵。在大規模參數量和資料庫支援下，簡單的體系結構將遠遠超過複雜的算法。

自然語言領域，稱為史上最強 NLP 的 GPT-3 擁有 1,750 億參數。近日，Google 將參數量直接拉高到 1.6 兆。

1 月 11 日，Google 在 arXiv 發表論文《Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity》，提出最新語言模型 Switch Transformer。研究員介紹，Switch Transformer 擁有 1.6 兆參數，是迄今規模最大的 NLP 模型。

論文指出，Switch Transformer 使用稀疏觸發（Sparsely Activated）技術，只使用神經網路權重子集，或轉換模型內輸入數據的參數。在相同計算資源下，訓練速度比 Google 之前研發的最大模型 T5-XXL 還快 4 倍。

Switch Transformer 由混合專家（Mix of Expert，MoE）AI 模型範例發展而來。MoE 模型是指將多個專家或專門從事不同任務的模型放在較入的模型，並有一個「門控網路」（Gating Network）選擇為任何給定數據要諮詢哪些專家。

自 1990 年代初首次提出，MoE 模型在機器翻譯方面取得顯著成果。但因複雜性高、通訊成本高和訓練不穩定，導致廣泛應用受阻礙。

為了解決這些問題，研究人員簡化 MoE 路由算法，設計出直覺性改進模型，使通訊和計算成本大幅度降低。同時提出的訓練技術消除了訓練的不穩定性，並證明大型稀疏模型也能以低精度（bfloat16）格式訓練。

此外，與 T5-Base 和 T5-Large 模型相比，新模型在相同計算資源下將預訓練速度提高 7 倍之多。這些改進擴展到多語言裝置，可測量 101 種語言在 mT5-Base 版的收益。

最後，透過「Colossal Clean Crawled Corpus」預先訓練多達數兆個參數模型提高語言模型的規模，使速度比 T5-XXL 模型提高 4 倍。

研究者還表示，Switch Transformer 架構不僅在具備超級計算機的環境下有優勢，在只有幾個計算核心的計算機也有效。此外，研究者設計的大型稀疏模型可蒸餾為小而稠密的版本，同時還能保留稀疏模型品質提升的 30%。

Switch Transformer 模型架構

Switch Transformer 的設計原則是以簡單且高效計算的方式最大化 Transformer 模型的參數量。模型需在增加參數量同時保持每個範例的浮點運算（FLOPs）不變。研究人員假設，參數量與執行總計算量無關，是可單獨縮放的重要組件，那麼可透過設計稀疏觸發模型達成。

研究人員有效利用為密集矩陣乘法（Dense Matrix Multiplications，廣泛用於語言模型的數學運算）設計的硬體，如 GPU、Google 的張量處理單元 TPU。分散式訓練設備中，模型將不同的權重分配到不同設備，保證權重隨著設備數增加而增加，同時保證每個設備均有可管理的記憶體和計算足跡。

▲ Switch Transformer 編碼器模組圖示。

如上圖，研究人員將 Transformer 的密集前饋網絡（FFN）層替換為稀疏 Switch FFN 層（淺藍色）。該層對序列的令牌獨立操作。兩個標記令牌（X ₁ =More和X ₂ =Parameters ）在 4 個 FFN 專家之間配對（實線），其中路由器獨立傳輸每個令牌。Switch FFN 層返回輸出的 FFN 並乘以路由器門值（虛線）。

簡化稀疏通道

研究人員提出 Switch Routing 簡化策略：針對不同輸入，匹配一個最適合的專家。實驗證明，這簡化可保留模型品質，減少路由計算並提高性能。具體優勢如下：

通道計算量大大降低，因每個樣本僅需一個專家通道參與計算。
每個專家通道所計算樣本的 batchsize 大大縮小（至少減半）。
每個專家通道的達成複雜度減少，專家通訊開銷降低。

高效稀疏通道

研究人員使用 Mesh-Tensorflow（MTF）資料庫──與 Tensorflow 類似的語義和 API，高效支援分散式數據與模型結構。將物理內核集抽象為處理器的邏輯網格，然後利用張量和計算按指定的維度切分，便於跨維度輕鬆劃分模型。

與 MoE Transformer 比較

第一個測試是在「Colossal Clean Crawled Corpus」資料庫進行，也稱為 C4，含 750GB 文本語料，全部來自 Reddit、Wikipedia 和其他網路資源。研究人員使用 32 個 TPU 核在 C4 資料庫預先訓練幾種不同 Switch Transformer 模型。

模型任務是預測段落 15% 被遮住的單詞，以及透過檢索文本回答有難度的問題。

從對比結果可發現：

速度─品質指標，Switch-Transformer 超越精密微調的稠密模型與 MoE Transformer。在有限計算量與工作時間下，Switch Transformer 都得到最優結果。
Switch Transformer 有更小的計算足跡（computational footprint）。
更小的 capactiy factor 下，Switch Transformer 工作效果更佳。

提升訓練與微調技術

稀疏專家模型比普通 Transformer 更難訓練，這是由於切換通道的不穩定引起，同時低精確度計算會導致惡化 softmax 計算問題。

研究人員透過局部引入高精準度訓練，而非全局方法，達到理想的穩定性，測試效果如圖：

同時不穩定性的另一種補救措施，研究人員減少預設的 Transformer 初始化比例，將 S=1.0 降低 10 倍。這既提高品質，又降低不穩定訓練的可能性。如圖：

「預訓練＋微調」的 NLP 方法，處理下游任務時容易發生過適現象。對此，研究人員提出「expert dropout」的解決方案──微調過程增加每個 expert 的 dropout 比例緩解過適。

簡單增加所有層的 dropout 率會導致性能下降，但在非專家層設置較小 dropout 率（0.1）和在專家層設置較大的 dropout 率（0.4），可提高 4 個較小下游任務的性能。

性能測試：預訓練、微調和多任務訓練

研究人員研究預訓練階段 Switch Transformer 的預訓練可擴展性，為了避免數據受限，採用大型 C4 資料庫，並以固定的計算成本測試基於步數和時間的可伸縮性。

步態的可擴展性

下圖為多模型在訓練步數恆定、專家數量增加時表現的可擴展性提升情況。可看到保持每個 token 的 FLOPS 不變時，擁有更多參數（專家）可提高訓練速度。

此外還可發現，專家數量增加會導致更有效的樣本模型。Switch-Base 64 專家模型在 60k 和 450k 步態達與 T5-Base 模型相同的性能，這是 7.5 倍加速。

時間的可擴展性

從基於步數的設置，可觀察到：樣本效率提升未必能轉化為時間的模型質量提升。那麼，訓練時間和計算成本都固定的前提下，是該訓練稠密模型還是稀疏模型？下圖解決了這個問題。

圖展示預訓練模型品質隨訓練時間增加產生的變化。訓練時間和計算成本都固定的情況下，Switch Transformer 的速度優勢非常明顯。如果要達到相似困惑度，Switch-Base 64 專家模型的訓練時間僅 T5-Base 模型七分之一。

微調

針對 T5-Base 和 T5-Large 模型，研究人員設計具更多參數的 FLOP-matched Switch Transformer。在 SuperGLUE 基準上， FLOP-matched Switch Transformer 相比 T5-Base 和 T5-Large 性能分別提升 4.4% 和 2%。同時模型在多項推理和知識任務也有顯著提升。

這說明模型架構不只對預訓練有用，還可透過微調將品質改進移至下游任務。

蒸餾技術

為了解決超大規模神經網路的部署問題，研究人員提出多種蒸餾技術：將大型稀疏模型蒸餾到小型密集模型。

使用最優的蒸餾技術，可將具備 11 億參數量的模型壓縮 82%，同時保留 37% 性能提升。最極端的情況下，將模型壓縮 99%，且維持 28% 性能提升。

多語言學習

最後一組下游實驗，研究人員利用模型對 101 種不同語言混合預訓練。

下圖展示 Switch T5 Base 模型與 mT5-Base 在所有語言的質量提升情況（負對數困惑度）。對兩個模型經過 100 萬步預訓練後，Switch Transformer 的最終負對數困惑度相較基線有提升。

此外，透過 mT5-Base 使用 Switch Transformer 模型的加速直方圖，可發現 mT5-Base 的平均速度提高 5 倍，同時 91% 語言速度至少提升 4 倍，表明 Switch Transformer 能有效做到多任務和多語言。

總結及未來研究方向

研究人員表示，Switch Transformer 模型是具可擴展性的高效自然語言處理模型，透過簡化 MoE，形成快速理解、訓練穩定且比同等大小模型更有效的樣本體系結構。實驗證明，這些模型在不同語言任務和訓練模式都表現出出色性能，包括預訓練、微調和多任務訓練。同時，這些結果也讓訓練有上千億、上兆參數量的模型成為可能。

最後研究人員指出，雖然 Switch Transformer 改進訓練程序，並對稀疏模型擴展性研究，但未來此方向還有許多未完成的工作，如：