轉載從: Tech News 科技新報
眾所皆知，參數量是機器學習算法的關鍵。在大規模參數量和資料庫支援下，簡單的體系結構將遠遠超過複雜的算法。
自然語言領域，稱為史上最強 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 1 =More和X 2 =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 改進訓練程序，並對稀疏模型擴展性研究，但未來此方向還有許多未完成的工作，如：
- 進一步提高大型模型訓練的穩定性。目前正進行的方法有：正則化函數、適度梯度裁剪等。
- 一般來說預訓練越好，下游任務的效果也就越好，但一些任務發現，1.6T 參數的 Switch-C 會較小模型結果更低。
- 目前方法使用的同類專家，未來可嘗試透過更靈活的基礎結構支持異構專家。
- 調查 FFN 層之外的專家層，初步證據表明，這同樣可以改善模型品質。
- 目前工作僅考慮語言任務，未來可嘗試將模型稀疏性類似應用到其他模態（如圖像）或多模態。
- SWITCH TRANSFORMERS: SCALING TO TRILLION PARAMETER MODELS WITH SIMPLE AND EFFICIENT SPARSITY
- Google trained a trillion-parameter AI language model