Transformer 訓練成本高昂的問題日益凸顯，不僅需要耗費巨額的資金與大量的計算資源，還對環境產生了不可忽視的影響，最近由北京大學與谷歌聯合發表的一篇論文，為這一棘手難題帶來了全新的曙光。論文中提出的創新方案，有望打破 Transformer 訓練成本居高不下的困境，
一、研究動機
如今，很難想象人工智能領域如果沒有Transformer會怎樣。這些模型是大型語言模型的基礎架構，為人工智能帶來了革命性變化。而且，它們的影響力不僅限于自然語言處理領域。在計算機視覺等其他領域，視覺Transformer（ViT）也發揮著重要作用。隨著技術的發展，模型規模不斷擴大，從頭開始訓練模型的成本越來越高，且這種做法難以持續，還引發了環境方面的擔憂。
訓練Transformer模型成本高昂
二、引入Tokenformer
我們今天要解讀的論文題目是。該論文對Transformer架構提出了一項引人注目的改進，名為Tokenformer。Tokenformer能夠在增加模型規模時，避免從頭開始重新訓練模型，從而大幅降低成本。
三、Transformer與Tokenformer架構對比
讓我們從整體上對比一下傳統Transformer架構和Tokenformer，相關對比圖來自論文。
3.1 Transformer整體流程
左邊展示的是原始Transformer架構的簡化圖。底部是輸入的詞元序列。
- 輸入首先經過線性投影層，計算注意力模塊的輸入，即矩陣Q、K和V。這一階段涉及模型參數和輸入詞元（token）之間的相互作用，通過線性投影計算得出。
- 然后，自注意力組件讓輸入詞元之間相互作用，這一過程通過注意力模塊計算。
- 最后，前饋神經網絡（FFN），為下一層生成輸出，這同樣體現了詞元和參數之間通過線性投影計算的相互作用。
3.2 Tokenformer的核心思路
通常，詞元與參數之間的相互作用是通過固定參數大小的線性投影組件計算的。這就導致在增加模型規模時，必須從頭開始訓練。
Tokenformer的主要思路是構建一個完全基于注意力機制的模型，包括詞元與參數之間的相互作用，從而擁有更靈活的架構，支持逐步增加參數數量。
3.3 Tokenformer整體流程
在上面圖的右側，我們可以看到Tokenformer架構的簡化圖。
- 為了計算自注意力模塊的輸入（Q、K和V矩陣），我們將輸入詞元輸入到一個名為詞元 - 參數注意力的新組件中。在這個組件中，除了輸入詞元，我們還會傳入參數。輸入詞元代表詞元 - 參數注意力模塊中的查詢部分，參數則代表鍵和值部分。
- 接著，我們使用與之前相同的自注意力組件。
- 為了為下一層準備輸出，我們用另一個詞元 - 參數注意力模塊取代前饋神經網絡。在這個模塊中，我們將自注意力模塊的輸出作為查詢，同時為鍵矩陣和值矩陣使用不同的參數。
3.4 Tokenformer架構詳解
-
計算自注意力模塊的輸入：深入研究Tokenformer架構，輸入詞元在經過歸一化處理后，會被送入三個Pattention模塊（Pattention即詞元 - 參數注意力模塊的名稱）。三個Pattention模塊各自擁有權重，分別負責為自注意力模塊創建輸入，即每個Pattention模塊生成一個Q、K或V矩陣。
-
Pattention與注意力模塊的區別：
-
原始注意力模塊計算公式為：（此處原文未給出原始注意力公式）。
-
Pattention模塊計算公式為：
P a t t e n t i o n ( X , K p , V p ) = Θ ( X ? K p T ) ? V p Pattention(X, K_p, V_p)=\Theta(X ? K_p^T ) ? V_p Pattention(X,Kp?,Vp?)=Θ(X?KpT?)?Vp?。 -
輸入詞元作為查詢，鍵矩陣和值矩陣由Pattention模塊的參數給出。與注意力模塊相比，Pattention模塊使用了帶有不同縮放機制的改進softmax函數（即 Θ \Theta Θ函數），以提高優化的穩定性。
-
X X X與 K K K相乘后的改進softmax函數將輸入詞元與鍵參數進行匹配，得到注意力分數。然后，利用V參數矩陣對注意力分數進行加權，生成該模塊的最終輸出。
-
-
用Pattention取代FFN：在標準自注意力模塊處理完三個Pattention模塊的輸出后，我們不再使用前饋神經網絡，而是使用兩個連續的Pattention模塊。從上圖左下角可以看到，在這兩個模塊之間，我們通過殘差連接將其與輸入詞元合并。
-
模型的增量式擴展：在架構圖的右下角，展示了通過添加新參數來逐步增加模型規模的示意圖。我們在每個Pattention模塊的鍵矩陣和值矩陣中追加更多的參數詞元行。已訓練的參數詞元會被保留，無需從頭開始訓練。然后，我們訓練規模更大的模型，并在這個過程中更新所有詞元。從下面的結果可以看出，與從頭開始訓練相比，增量式訓練的模型訓練速度要快得多。
四、實驗結果
上圖展示了Tokenformer相對于從頭開始訓練的Transformer在性能和訓練成本方面的對比。
藍色線代表使用3000億詞元從頭開始訓練的Transformer模型。不同大小的圓圈代表不同的模型規模，中間灰色圓圈列出了具體規模。其他線條代表Tokenformer模型，每種顏色代表訓練增量版本時使用的不同詞元數量。例如，紅線代表模型從1.24億參數擴展到14億參數，在用于從頭訓練的3000億詞元中，僅使用了300億詞元進行訓練。最終版本在性能上與相同規模的Transformer相當，但訓練成本卻大幅降低！黃線表明，使用600億詞元訓練增量版本的模型，性能優于Transformer，且訓練成本更低。
五、參考文獻與鏈接
- 論文頁面
- 代碼
- AI Papers Academy解析
