如何進一步提升大規模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
1 為什麼會有Primer？
針對訓練大規模Transformer存在高昂成本的問題，作者試圖回答是否能夠透過
模型架構自動搜尋技術
，找到一個高效Transformer變種，實現以下目標：
給定學習效果，相比標準Transformer，這個新的變種有著更低的訓練代價。
給定學習代價，相比標準Transformer，這個新的變種達到更好的學習效果。
作者給出的答案就是Primer（
PRIM
itives searched transform
ER
）。
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模型訓練代價的可能性。
