眾所周知，參數量是機器學習算法的關鍵。在大規模參數量和數據集的支持下，簡單的體系結構將遠遠超過複雜的算法。
在自然語言領域，被稱為史上最強NLP的GPT-3擁有1750億參數。近日，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）”來選擇為任何給定數據要諮詢哪些專家。
自90年代初首次被提出以來，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層（淺藍色）。該層對序列中的令牌進行獨立操作。他們兩個標記令牌（X1=More 和 X2=Parameters ）在四個FFN專家之間進行匹配（實線），其中路由器獨立地傳輸每個令牌。Switch FFN層返回輸出的FFN並乘以路由器門值（虛線）
簡化稀疏通道研究人員提出了Switch Routing簡化策略：針對不同的輸入，匹配最適合的一個專家。實驗證明，這種簡化可以保留模型質量，減少路由計算並提高性能。具體優勢如下：
1、通道計算量大大降低，因每個樣本僅需要一個專家通道參與計算。
2、每個專家通道所計算樣本的batchsize被大大縮小(至少減半)。
3、每個專家通道的實現複雜度減小了，專家間的通信開銷降低。
研究人員使用Mesh-Tensorflow（MTF）庫—與Tensorflow 類似的語義和API ，能夠能夠高效支持分佈式的數據與模型結構。它將物理內核集抽象為處理器的邏輯網格。然後利用張量和計算按指定的維度進行切分，便於跨維度輕鬆劃分模型。
與MoE Transformer的比較第一個測試是在“Colossal Clean Crawled Corpus”數據集上進行的。該數據集也被成為C4，含750GB的文本語料，全部來自於Reddit、Wikipedia和其他網絡資源中。研究人員使用了32個TPU核在C4數據集上預先訓練了幾種不同的Switch Transformer模型。
模型的任務是預測段落中15%被遮住的單詞，以及通過檢索文本回答一些有難度的問題。
從以上對比結果中，可以發現：
1、在速度-質量的指標上，Switch-Transformer超過了精密微調的稠密模型，與MoE Transformer。在有限的計算量與工作時間下，Switch Transformer都得到了最優的結果。
2、Switch Transformer有更小的計算足跡（computational footprint）。
3、在更小的capactiy factor下，Switch Transformer工作的效果更佳。
提升訓練與微調技術稀疏專家模型相比普通Transformer更難訓練，這是由於切換通道帶來的不穩定引起的，同時，低精度計算會導致惡化softmax計算中的問題。
研究人員通過在局部引入高精度的訓練，而非全局的方法，達到了理想的穩定性，測試效果如圖：
同時作為不穩定性的另一種補救措施，研究人員減小了默認的Transformer初始化比例，將 S=1.0降低了10倍。這既提高了質量，又降低了不穩定訓練的可能性。如圖：
論文中，研究人員對預訓練階段 Switch Transformer的預訓練可擴展性進行了研究。為了避免數據受限，他們採用了大型 C4 數據庫，並以固定的計算成本測試了基於步數和時間的可伸縮性。
下圖為多個模型在訓練步數恆定、專家數量增加時表現出的可擴展性提升情況。可以看到：在保持每個 token 的 FLOPS 不變時，擁有更多的參數（專家）可以提高訓練速度。
此外還可以發現，專家數量的增加會導致更有效的樣本模型。Switch-Base 64專家模型在60k和450k步態達到了與T5-Base模型相同的性能，這是7.5倍的加速。
時間的可擴展性從基於步數的設置中，可以觀察到：樣本效率提升未必能轉化為時間上的模型質量提升。那麼，在訓練時間和計算成本都固定的前提下，是應該訓練一個稠密模型還是稀疏模型？下圖解決了這個問題。
圖中展示了預訓練模型質量隨訓練時間增加所產生的變化。在訓練時間和計算成本都固定的情況下，Switch Transformer 的速度優勢非常明顯。在這種設置下，如果要達到相似的困惑度，Switch-Base 64 專家模型的訓練時間僅為 T5-Base 模型的 1/7。
微調針對 T5-Base 和 T5-Large 模型，研究人員設計了具備更多參數的 FLOP-matched Switch Transformer。在 SuperGLUE基準上， FLOP-matched Switch Transformer 相比 T5-Base 和 T5-Large 的性能分別提升了 4.4% 和 2%。同時，該模型在多項推理和知識任務中也帶來了顯著提升。
這説明該模型架構不只對預訓練有用，還可以通過微調將質量改進遷移至下游任務中。
為了解決超大規模神經網絡的部署問題，研究人員提出了多種蒸餾技術：將大型稀疏模型蒸餾到小型密集模型中。
使用表 7 中最優的蒸餾技術，可以將具備 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層之外的專家層，初步的證據表明，這同樣可以改善模型質量。
目前的工作僅考慮了語言任務，未來可嘗試將模型稀疏性類似地應用到其他模態（例如圖像）或多模態中。
引用鏈接：
https://arxiv.org/pdf/2101.03961.pdf
https://venturebeat.com/2021/01/12/google-trained-a-trillion-parameter-ai-language-model/
https://arxiv.org/search/cs?searchtype=author&query=Fedus%2C+W
雷鋒網雷鋒網雷鋒網