文 | 舞風小兔
編 | 智商掉了一地

如何進一步提升大規模Transformer的訓練效果？Primer給出了答案，一起來看看吧！

Transformer是深度學習界的明星模型。由于其具有高度的并行性，十分容易在并行計算系統中擴展至超大規模。自然語言處理任務一路見證了Transformer模型規模的爆炸式增長。

微軟和Nvidia不久前聯合發布的"Megatron-Turning"更是創造了最新記錄：其參數數目已經達到了恐怖的5300億。大規模Transformer通過橫掃各大任務榜單，向所有人展示了“大模型+大數據”這一簡單方法的高度通用性。

在更加有效的深度學習技術出現之前，許多已經正在使用Transformer的任務難免都會期待是否擁有一個更大的模型，就能夠更進一步提升應用任務的效果？但訓練大規模Transformer的高昂成本也同樣令人望而生畏。大規模Transformer，讓人又愛又恨。本文要介紹的Primer就是在該背景下開展的一個工作。

論文標題:
Primer: Searching for Efficient Transformers for Language Modeling

論文鏈接:
https://arxiv.org/abs/2109.08668

論文代碼:
https://github.com/google-research/google-research/tree/master/primer

Arxiv訪問慢的小伙伴也可以在 【夕小瑤的賣萌屋】訂閱號后臺回復關鍵詞 【1029】 下載論文PDF~

1 為什么會有Primer？

針對訓練大規模Transformer存在高昂成本的問題，作者試圖回答是否能夠通過模型架構自動搜索技術，找到一個高效Transformer變種，實現以下目標：

給定學習效果，相比標準Transformer，這個新的變種有著更低的訓練代價。

給定學習代價，相比標準Transformer，這個新的變種達到更好的學習效果。

作者給出的答案就是Primer（PRIMitives searched transformER）。

2 什么是Primer？

Primer 對Transformer的結構給出了兩處修改，在下圖中用紅色圈出。在論文進行的各項實證研究中，這兩項修改最為魯棒。論文作者建議：Transformer語言模型的使用者如果想嘗試改進自己的模型，推薦從這兩項修改開始。

▲Primer對Transformer結構的兩處重要修改

這兩處由遺傳算法自動搜索到的修改分別是：

在Feed Forward 模塊（FFN）部分，將原來的ReLU激活修改為Squared ReLU，也就是在ReLU之后再添加一個Square運算。作者聲稱這個小小的修改最為重要，十分有助于加速模型在更短時間內收斂到同樣的學習效果。

在自注意力機制中、、 映射之后添加 Depthwise 卷積，稱之為Multi-DConv Attention (MDHA)；

上面兩幅圖已十分所見即所得地解釋了論文的結果。對經常與深度學習算法打交道的同學來說，根據這幅圖已經可以在1分鐘之內修改好自己的Tranformer模型，將其變為Primer。然后， 保持所有其他因素不變 去試試能否在自己的任務上復現論文的效果：在更短的時間內，模型收斂到和原來模型同樣的精度。

作者為確定搜索出的模型結構具有廣泛的實用性，做了大量的覆蓋性實驗驗證，驗證變量包括：模型參數規模（20M到1.9B）、訓練時長（到加速器小時）、多個數據集（LM1B，C4，PG19）、多個硬件平臺（TPUv2，TPUv3，TPUv4 和 V100）、將Primer的修改插入多個Transformer 代碼庫（Tensor2Tensor^[1], Lingvo^[2] 和 T5^[3]）中的多個模型（dense Transformer，sparse mixture-of-experts Switch Transformer, Synthesizer）。在大量的試驗中，作者發現只有上面兩個修改具有廣泛的有效性。作者還列舉了一些有效但不總是有效的修改，給出了他們在實驗中的表現：

共享和的 depthwise 表示：共享一部分和映射的映射矩陣，由圖1中的MDHA后加depthwise 卷積計算得到，，實驗發現：大部分時候這個修改對學習效果都是有害的。我們可以看到，這是一個令人類專家看來十分奇怪的模型修改，很像是一個典型的自動搜索產生的修改方案，論文中還給出了類似這樣的奇怪修改，大部分也都沒有能夠改進學習代價。

歸一化層添加位置：標準的Transformer實踐在自注意力模塊（Self-Attention）和FFN（Feed Forward）層之前添加歸一化層，論文作者嘗試在自注意力模塊之前，FFN模塊之后添加歸一化層，這個修改會有幫助，但并不總是有幫助。

自定義歸一化：Primer使用自定義歸一化： 替代，這一修改會有幫助，并不總是有效。

Bottleneck 映射大小的修改：將原始Transformer模型隱藏層大小從512維減少到384，但將FFN模塊的Upwards 映射部分的維度從2048增加到4608，收斂效果在模型參數小于35M時有明顯的改善，但對更大的模型有效性降低。

3 Scaling Law 對比

由于這篇論文是一篇實驗研究，文章用了長達35頁的篇幅解釋了在TensorFlow中進行模型架構搜的設計、搜索空間設計、諸多無規律的修改。閱讀這篇文章時，研究神經網絡架構搜索的讀者，可以只關注模型搜索空間設計和搜索算法設計部分；研究自然語言處理任務本身的讀者，可以只關注上面兩個簡單的結論。在這里我們只重點地摘要作者如何通過實驗驗證Primer能夠減少大規模Transformer訓練代價這一關鍵論點。

作者對Transformer模型的以下幾組變量進行全排列組合：（1）層數：、（2）隱藏層大小 、（3）FFN 模塊upwards projection 比率 ，產生出模型參數從23M到385M變化的一系列模型。作者在LM1B搜索任務上，使用序列長度64，batch大小4096 token，分別在TPUv2和V100上訓練24小時, 用Tensor2Tensor和T5中的幾種典型Transformer變種作為對比對象，與Primer進行對比。

▲35M模型參數下，幾種Transformer模型與Primer加速因子對比

上圖繪制了35M參數模型的實驗結果，以標準Transformer模型訓練24小時的學習效果為基準，橫軸是幾種Transformer模型變種，縱軸是各個變種達到基準效果的加速比，從中可以看到Primer在所有實驗組中都有1.7倍以上的加速率。

▲幾種Transformer模型與Primer scaling曲線對比

上圖又進一步展示了Primer在所有模型規模下的實驗數字。橫軸是時間，縱軸是驗證集上的損失，不同顏色的曲線代表不同的模型規模。左圖是對所有Transformer變種用MSE實驗去擬合出達到給定精度的最佳模型大小，繪制成二維線圖，也就是scaling law^[4]曲線。可以看到Primer的scaling曲線在最下方，這意味著在所有模型大小情況下，Primer都能更快地達到指定學習精度。

4 小結與討論

盡管這篇論文長達35頁，但是關鍵結論十分簡單，而剩下的篇幅都在闡述如何進行模型搜索和在不同任務上重復驗證Primer能否節約Transformer的訓練代價，對不關心這些細節的實踐者，完全可以跳過作者的長篇大論來直接使用作者的結論。

關于這兩個由遺傳算法自動搜索到的修改，第一個：squared ReLU 在其它學習任務中已經被多次使用，能夠進一步增強ReLU的激活值，或許容易被人類專家注意和想到。第二個：卷積能增加特征向量的局部穩定性，由于，， projection已經足夠簡潔，在和 projection之后添加depthwise卷積，可能是一個連人類專家也不容易主動嘗試的想法，被自動搜索找到了。

讀過這篇論文，Primer依然可能有一定的限制。這里指出值得注意的三點。

盡管Primer的目標是減少大規模Transformer訓練的代價，但是文章實驗的大模型也遠遠小于GPT-3，當模型參數進一步提升時，這兩個修改是否有效，也未經作者的驗證。可能還是在從側面說明，盡管作者想使用自動搜索模型結構這項技術去減少Transformer的訓練代價，但進行實驗本身的代價依然過于昂貴。

作者自己也指出，實驗只在自回歸語言模型上進行，而初步測試表明了這兩項修改應用于其它類型任務時并不總是有效，也就是說這兩項修改有可能只適用于部分任務。

一個小小的遺憾是作者在搜索空間構造時并沒有對自注意力機制這樣的高階模塊的潛在變種進行搜索，畢竟這是Transformer的核心。由于這一步也存在著大量的選擇空間，或許也潛藏著壓縮Transformer模型訓練代價的可能性。

后臺回復關鍵詞【入群】

加入賣萌屋NLP/IR/Rec與求職討論群

后臺回復關鍵詞【頂會】

獲取ACL、CIKM等各大頂會論文集！

?

[1] Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N. Gomez, Stephan Gouws, Llion Jones, ?ukasz Kaiser, Nal Kalchbrenner, Niki Parmar, Ryan Sepassi, Noam Shazeer, and Jakob Uszkoreit. Tensor2tensor for neural machine translation. arXiv preprint arXiv:1803.07416, 2018. (https://arxiv.org/pdf/1803.07416.pdf)

[2] Lingvo: A TensorFlow Framework for Sequence Modeling. (https://blog.tensorflow.org/2019/02/lingvo-tensorflow-framework-for-sequence-modeling.html)

[3] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1–67, 2020. (https://arxiv.org/pdf/1910.10683.pdf)

[4] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. (https://arxiv.org/abs/2001.08361)

總結

以上是生活随笔為你收集整理的如何提升大规模Transformer的训练效果？Primer给出答案的全部內容，希望文章能夠幫你解決所遇到的問題。

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