我们可以无损放大一个Transformer模型吗？
?PaperWeekly 原創(chuàng) ·?作者?|?蘇劍林
單位?|?追一科技
研究方向?|?NLP、神經(jīng)網(wǎng)絡(luò)
看了標題，可能讀者會有疑惑，大家不都想著將大模型縮小嗎？怎么你想著將小模型放大了？其實背景是這樣的：通常來說更大的模型加更多的數(shù)據(jù)確實能起得更好的效果，然而算力有限的情況下，從零預(yù)訓(xùn)練一個大的模型時間成本太大了，如果還要調(diào)試幾次參數(shù)，那么可能幾個月就過去了。
這時候“窮人思維”就冒出來了（土豪可以無視）：能否先訓(xùn)練一個同樣層數(shù)的小模型，然后放大后繼續(xù)訓(xùn)練？ 這樣一來，預(yù)訓(xùn)練后的小模型權(quán)重經(jīng)過放大后，就是大模型一個起點很高的初始化權(quán)重，那么大模型階段的訓(xùn)練步數(shù)就可以減少了，從而縮短整體的訓(xùn)練時間。
那么，小模型可以無損地放大為一個大模型嗎？本文就來從理論上分析這個問題。
含義
有的讀者可能想到：這肯定可以呀，大模型的擬合能力肯定大于小模型呀。的確，從擬合能力角度來看，這件事肯定是可以辦到的，但這還不是本文關(guān)心的“無損放大”的全部。
以 BERT 為例，預(yù)訓(xùn)練階段主要就是一個 MLM 模型，那么“無損放大”的含義就是：是否可以通過某種變換，把一個小模型直接變換成一個大模型，并且輸出完全不改變？
這里的變換，指的是對權(quán)重做一些確定性的變換，而不用通過梯度下降來繼續(xù)訓(xùn)練；輸出完全不改變，指的是對于同一個輸入，小模型和大模型給出的預(yù)測結(jié)果是完全一致的，也就是說它們表面上看起來不一樣，但數(shù)學(xué)上它們是完全一致的函數(shù)，所以稱為“無損放大”。由于是無損放大，我們至少可以保證大模型不差于小模型，所以繼續(xù)預(yù)訓(xùn)練理論上有正的收益。至于先小后大這樣預(yù)訓(xùn)練在效果上能不能比得上一開始就從大訓(xùn)練，這個需要實驗來確定，并不是本文關(guān)心的問題。
直覺來想，這種放大也不困難，比如通過“重復(fù)”、“補零”等操作就可以實現(xiàn)模型權(quán)重的自然放大。事實上嘗試的方向也是如此，但難點在于我們需要仔細分析模型的每一個模塊在被放大之后所產(chǎn)生的后果，以確保最終的結(jié)果是無損的。
嘗試
下面我們以“將一個 BERT 放大為 2 倍”為例子進行分析嘗試，來確定最終的變換形式。這里的“放大”指的是僅僅擴大隱層向量的維度，并不改變模型的層數(shù)，也不改變多頭注意力機制的頭數(shù)。
2.1 Embedding
首先，輸入層是 Embedding 層，因此先要解決的是 Embedding 層的放大問題。這也是其中最簡單的一環(huán)，就是直接將每個 token 的向量維度都放大為 2 倍即可，主要就是“重復(fù)”、“補零”兩種操作：
兩種方案都可以作為候選方案，但直覺上來想，補零這種方式引入了太多的零，會導(dǎo)致過度稀疏和同一個值重復(fù)次數(shù)過多，不利于權(quán)重的多樣性，因此我們還是選擇了重復(fù)這種方案。不過，就算只看重復(fù)，也不指上述一種方式，比如 也是一種方案，但后面關(guān)于 Attention 層的分析表明，后一種方案是不可取的。
除此之外，我們通常還希望變換是正交的，這通常能最大程度上保證模型的穩(wěn)定性，具體來說，正交變換的最基本性質(zhì)是不改變向量的模型，所以我們將最終的重復(fù)變換調(diào)整為：
或者簡記成 ，其中 是上取整運算，我們稱之為“重復(fù)再除以 ”。
2.2 LayerNorm
Embedding 的下一層就是 LayerNorm 了，變換前，LayerNorm 的運算為：
而變換后，我們有：
這也就是說，“減均值除以標準差”這一步自動幫我們消去了 這個因子，其結(jié)果是放大前結(jié)果的直接重復(fù)。如果我們將參數(shù)向量 也按照公式（2）進行變換，那么結(jié)果將是 ，跟 Embedding 層的變換結(jié)果一致，而我們就是要盡量使得每一層“凈變換”都是同樣的一個簡單變換：“重復(fù)再除以 ”。
2.3 FeedForward
按照順序，接下來本來應(yīng)該分析 Attention 層才對，不過 FeedForward 層相對簡單一點，并且 FeedForward 層的分析結(jié)果也對后面理解 Attention 層的變換有所幫助，因此這里先來考慮 FeedForward 層的變換。
FeedForward 層只是兩個全連接層的復(fù)合，所以我們只需要分析單個全連接層：
這里的 是激活函數(shù)。鑒于之前的經(jīng)驗，我們嘗試如下變換：
也就是將 按照式（2）進行變換，而對于 則嘗試使用形式下述變換：
這里的 D 就是輸出維度大小，這里我們假設(shè)模型放大 2 倍后，D 也放大 2 倍。不難看出，該變換其實就是對變換矩陣 行列兩個方向都分別執(zhí)行變換（2）。此時：
這說明變換（6）對于線性變換層來說，能夠滿足我們的理想追求——放大后的結(jié)果就是“重復(fù)再除以 ”。然而，這還不夠，因為全連接層還有個激活函數(shù) ，現(xiàn)在的問題在于 未必等于 ，而如果不等，我們就沒法讓整體的變換等價于“重復(fù)再除以 ”。
事實上，BERT 用的 GeLU 激活函數(shù)就不滿足該恒等式；線性激活函數(shù)（不加激活函數(shù)）顯然是滿足這個等式的，而滿足這個等式一個常見的非線性激活函數(shù)便是?ReLU（也包括 LeakyReLU）函數(shù)，因此一個直接的解決方式就是 FeedForward 層換用 ReLU 激活函數(shù)。事實上，這也已經(jīng)是預(yù)訓(xùn)練模型的一個常見選擇了，百度的 Ernie 和 Google 的 T5 模型，它們的 FeedForward 層激活函數(shù)都是用 ReLU。
那么，像 BERT 這樣的非 ReLU 激活函數(shù)的 FeedForward 層就沒辦法了嗎？那也不至于，因為 FeedForward 層是兩個全連接層的復(fù)合，我們只需要在變換第一個全連接的時候少除以一個 ，變換第二個全連接的時候多除以一個 就行了。具體來說，第一個全連接權(quán)重變?yōu)?#xff1a;
此時就有：
此時結(jié)果就是原結(jié)果的直接重復(fù)，沒有除以 ，既然如此，后面緊接著的全連接層多除以一個 就行了，即后面的全連接層權(quán)重變換為：
這樣整個 FeedForward 層的效果就等價于“重復(fù)再除以 ”了。
2.4 Attention
現(xiàn)在到了最難啃的“硬骨頭”——Attention 層的變換。Attention 層首先通過三個線性層將每個輸入向量變換為 q,k,v：
根據(jù)前面對 FeedForward 層的分析可以得知，如果要想 q,k,v 都達到“重復(fù)再除以 ”的效果，只需要按照變換（6）進行。但 Attention 層不是單純的全連接層，變換完之后，我們要檢查 Attention 矩陣是否不變，我們來算內(nèi)積：
其中 d' 是對應(yīng)的 head_size。這個結(jié)果告訴我們，上述變換保持了內(nèi)積不變，所以應(yīng)該也保持 Attention 矩陣不變。但是，這里有一個陷阱！如果是 T5 這樣的模型，它的內(nèi)積之后是沒有尺度縮放的，所以這樣的確完事了；然而像 BERT 這樣的模型，它是內(nèi)積之后除了個 再做 Softmax 的，，而一旦放大模型后，除以 變成了除以 ，內(nèi)積不變也不能保持 Attention 矩陣不變，而應(yīng)當(dāng)還需要往 q,k 的權(quán)重分別再乘個 ，所以最終的變換應(yīng)該是：
經(jīng)過這樣變換后，Attention 矩陣不變，而 ，所以最終的輸出結(jié)果也是 。
上述內(nèi)容只是針對 Attention 的單個頭進行分析，事實上 Attention 有多個頭，多個頭的輸出結(jié)果還要拼接起來再接一個全連接層。當(dāng)然，由于每個頭都是平等的、獨立的，因此上述結(jié)論基本不變，最后全連接層也只需要按照式（6）進行變換，就可以讓 Attention 的變換效果。但是，多頭帶來的一個效應(yīng)是，我們在重復(fù)的時候，必須局部地進行重復(fù)。
具體來說，我們在實現(xiàn)多頭的時候，并非是真的做了多個全連接運算，而是做了一個大的全連接運算后再 reshape，這樣一來我們可以比較兩種不同的重復(fù)方式的 reshape 結(jié)果：
注意放大前 reshape 結(jié)果是 ，所以對比兩種不同的重復(fù)方式的 reshape 結(jié)果，我們發(fā)現(xiàn)第二種重復(fù)方式 reshape 之后的結(jié)果全亂了，不等價于每個頭分別重復(fù)。因此我們只能選擇前一種重復(fù)方式。
2.5 輸出概率分布
通過以上分析，我們可以使得整個 Encoder 在放大到 2 倍之后，實現(xiàn)“重復(fù)再除以 ”的效果。最后剩下的就是輸出部分，即將 Encoder 的輸出向量轉(zhuǎn)化為 token 的概率分布，這里邊包含幾種情況。
像 GPT、T5 等模型，它們是直接在 Encoder 輸出后面乘以了 Embedding 矩陣的轉(zhuǎn)置來做作為概率分布的 logits（當(dāng)然有可能還有個偏置），由于 Embedding 矩陣本身就包含了“重復(fù)再除以 ”的操作，而 Encoder 的輸出也是“重復(fù)再除以 ”，兩者結(jié)合剛好抵消，所以從概率分布角度看，輸出是完全不變的。
不過 BERT 多了一層全連接，也就是說它先接了一個 GeLU 激活的全連接層，然后才乘以 Embedding 矩陣的轉(zhuǎn)置并加上偏置項作為 logitis。在“FeedForward”那一節(jié)我們已經(jīng)討論了，非 ReLU 激活的全連接層無法實現(xiàn)“重復(fù)再除以 ”的效果，而只能通過變換（9）來實現(xiàn)單純的“重復(fù)”效果，這時候乘以 Embedding 矩陣的轉(zhuǎn)置的話，得到的是原來的 logits 乘以 的效果，輸出會有所改變。
當(dāng)然，由于只是乘以了一個常數(shù)倍，所以分布雖然改變了，但是每個 token 的概率大小順序并沒有改變，這也就意味著，如果只看 MLM 的準確率，那么是完全沒有改變的，所以問題應(yīng)該不大。
當(dāng)然，如果是 ReLU 激活，那么按照式（6）來變換，那么可以實現(xiàn)完全不改變了。此外，如果是像 mT5 那樣，最后轉(zhuǎn)為 logits 的變換矩陣跟 Embedding 層不共享，那么可以同時調(diào)整最后的變換矩陣，也能實現(xiàn)輸出的完全不變。
2.6 RoPE位置編碼
前面的分析都只適用于每個神經(jīng)元都是不相關(guān)的情形，也就是說向量的任意兩個分量 是沒啥關(guān)聯(lián)的。但如果我們在模型中用了“旋轉(zhuǎn)式位置編碼（RoPE）”，那么這個假設(shè)就不成立了，因為 RoPE 是以每兩個分量為一組進行變換的，即 為一組、 為一組，依此類推。
如果還是按照之前式（2）進行重復(fù)變換，那么變換之后就變成了 為一組、 為一組、...，跟原來的分組不一致，所以會帶來很大的偏差。這種情況下，重復(fù)的時候也應(yīng)當(dāng)按照兩個為一組來進行：
當(dāng)然，由于默認的 RoPE 是沒有可訓(xùn)練權(quán)重的，它是按照固定的方式進行漸變的，所以哪怕按照該方式進行重復(fù)，那不能完全保證結(jié)果一致。也就是說，如果使用了 RoPE，那么基本上不能實現(xiàn)無損放大。不過實際測試結(jié)果表明，按照該方式進行重復(fù)放大后，對應(yīng)的 RoFormer 雖然性能有所損失，但不多，可以很快通過繼續(xù)訓(xùn)練恢復(fù)。
結(jié)論
現(xiàn)在我們可以確認，對于 BERT 來說，如果非線性激活函數(shù)用 ReLU，那么 BERT 是可以直接無損放大的，如果非線性激活函數(shù)不是 ReLU，那么可以實現(xiàn) MLM 準確率無損的放大（事實上經(jīng)過更精細的調(diào)整，也可以實現(xiàn)完全無損放大，但每個層的變換有點不統(tǒng)一了，不夠優(yōu)雅）；對于 GPT、T5 等模型來說，不管激活函數(shù)用啥（包括 mT5 用的 GLU 激活，也可以定制適當(dāng)），其實都可以實現(xiàn)無損放大。
其中，將 BERT 權(quán)重進行放大為 2 倍的變換匯總?cè)缦?#xff1a;
如果是其他略有不同的模型，那么就模仿前面的思想進行類似的分析即可。如果是 RoPE，那么將重復(fù)的方案改為式（15）就好；如果是擴大k倍，那么將表格中的多數(shù) 2 換為 k 就好。簡單來說，如果 Attention 沒有尺度縮放（除以 ），以及 FeedForward 的激活函數(shù)是 ReLU（或者 LeakyReLU），那么放大 k 倍的變換就最簡單的，將權(quán)重的每一維都執(zhí)行“重復(fù) k 次并除以 ”就好了。
小結(jié)
本文從數(shù)學(xué)上分析了直接放大 Transformer 模型的可能性，得到了若干可用的變換，在部分情況下可以無損放大 Transformer 模型，另外一些情況則可以將損失降到很小（比如保持 MLM 的準確率完全不變）。而研究 Transformer 模型的無損放大操作，可以為我們實現(xiàn)漸進式地訓(xùn)練大模型提供參考思路。
更多閱讀
#投 稿?通 道#
?讓你的文字被更多人看到?
如何才能讓更多的優(yōu)質(zhì)內(nèi)容以更短路徑到達讀者群體，縮短讀者尋找優(yōu)質(zhì)內(nèi)容的成本呢？答案就是：你不認識的人。
總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學(xué)者和學(xué)術(shù)靈感相互碰撞，迸發(fā)出更多的可能性。?
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優(yōu)質(zhì)內(nèi)容，可以是最新論文解讀，也可以是學(xué)術(shù)熱點剖析、科研心得或競賽經(jīng)驗講解等。我們的目的只有一個，讓知識真正流動起來。
?????稿件基本要求：
? 文章確系個人原創(chuàng)作品，未曾在公開渠道發(fā)表，如為其他平臺已發(fā)表或待發(fā)表的文章，請明確標注?
? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發(fā)送，要求圖片清晰，無版權(quán)問題
? PaperWeekly 尊重原作者署名權(quán)，并將為每篇被采納的原創(chuàng)首發(fā)稿件，提供業(yè)內(nèi)具有競爭力稿酬，具體依據(jù)文章閱讀量和文章質(zhì)量階梯制結(jié)算
?????投稿通道：
? 投稿郵箱：hr@paperweekly.site?
? 來稿請備注即時聯(lián)系方式（微信），以便我們在稿件選用的第一時間聯(lián)系作者
? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿
△長按添加PaperWeekly小編
????
現(xiàn)在，在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關(guān)注」訂閱我們的專欄吧
關(guān)于PaperWeekly
PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學(xué)術(shù)平臺。如果你研究或從事 AI 領(lǐng)域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。
總結(jié)
以上是生活随笔為你收集整理的我们可以无损放大一个Transformer模型吗？的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 祖名豆腐皮自带咸味吗？
- 下一篇： 肥肠家常做法？