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

作者 | 發布日期 2021 年 01 月 18 日 8:45 | 分類 AI 人工智慧 , Google , 網路 Telegram share ! follow us in feedly


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

自然語言領域,稱為史上最強 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 簡化策略:針對不同輸入,匹配一個最適合的專家。實驗證明,這簡化可保留模型品質,減少路由計算並提高性能。具體優勢如下:

  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 倍。這既提高品質,又降低不穩定訓練的可能性。如圖:

「預訓練+微調」的 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. 進一步提高大型模型訓練的穩定性。目前正進行的方法有:正則化函數、適度梯度裁剪等。
  2. 一般來說預訓練越好,下游任務的效果也就越好,但一些任務發現,1.6T 參數的 Switch-C 會較小模型結果更低。
  3. 目前方法使用的同類專家,未來可嘗試透過更靈活的基礎結構支持異構專家。
  4. 調查 FFN 層之外的專家層,初步證據表明,這同樣可以改善模型品質。
  5. 目前工作僅考慮語言任務,未來可嘗試將模型稀疏性類似應用到其他模態(如圖像)或多模態。

(本文由 雷鋒網 授權轉載;首圖來源:Unsplash