七年前，論文《Attention is all you need》提出了 transformer 架構，顛覆了整個深度學習領域。
如今，各家大模型都以 transformer 架構為基礎，但 transformer 內部運作原理，仍是一個未解之謎。
去年，transformer 論文作者之一 Llion Jones 宣佈創立人工智慧公司 Sakana AI。近期，Sakana AI 發表了一篇題為《Transformer Layers as Painters》的論文，探究了預訓練 transformer 中的資訊流，並針對僅解碼器和僅編碼器凍結 transformer 模型進行了一系列實驗。請注意，該研究沒有對預訓練模型進行任何型別的微調。
論文地址：https://arxiv.org/pdf/2407.09298v1
該研究認為 transformer 的內部機制（特別是中間層）可以類比畫家作畫流水線來理解。
作畫流水線通常是將畫布（輸入）傳遞給一系列畫家。有些畫家擅長畫鳥類，而另一些畫家則擅長畫輪子。每個畫家從其下一級畫家那裡收到畫布，然後其決定是否給畫作新增一些筆畫，或者只是將其傳遞給其上一級畫家（使用剩餘連線）。
這個類比並不是一個嚴格的理論，而是一個思考 transformer 層的工具。受這個類比的啟發，該研究測試驗證了一些假設：
各層是否都在使用相同的表徵空間？
所有層都是必要的嗎？
中間層都執行相同的功能嗎？
層的順序重要嗎？
這些層可以並行執行嗎？
對於某些任務來說，順序是否比其他因素更重要？
迴圈有助於層並行嗎？
哪些變體對模型效能影響最小？
該研究對預訓練 LLM 進行了一系列實驗，其中包括試驗標準 transformer 執行策略的變化，並在僅解碼器 (Llama) 和僅編碼器 (BERT) 模型的各種基準上測量這些變化對模型效能的影響。
各層是否都在使用相同的表徵空間？
為了回答不同層是否使用相同的表徵空間，作者測試了 Transformer 在跳過特定層或切換相鄰層的順序時是否具有穩健性。例如，在 Llama2-7B 中，第 6 層通常期望接收第 5 層的輸出。如果給第 6 層以第 4 層的輸出，它是否會出現「災難性」的行為？
在圖 2 中，我們可以看到，除了第一層和最後幾層之外，Llama2-7B 的各層對跳層或切換層都相當穩健。
該實驗表明，中間層共享一個表徵空間，且與「外圍層」（第一層和最後幾層）擁有不同的表徵空間。為了進一步驗證這一假設，作者效仿之前的研究，測量了基準中模型（Llama2-7B、Llama2-13B 和 BERT-Large）不同層的隱藏狀態啟用之間的平均餘弦相似度。圖 3 顯示了所有中間層之間的一致性。
這表明該模型可能具有「開始」、「中間」和「結束」層的三個不同的表徵空間。回答問題 1：是的，中間層似乎共享一個共同的表徵空間。
所有層都是必要的嗎？
為了進一步測試中間層的重定向空間是否真正共享（除了具有接近的餘弦相似度之外），該研究嘗試了「跳過層」，即將第 N 層的輸出直接傳送到第 N + M 層（其中 M > 1）的輸入中，從而「跳過」M − 1 層，如圖 1a 所示。該實驗是為了看看第 N + M 層是否可以理解第 N 層的啟用，儘管它僅根據從第 N + M − 1 層發來的輸入進行訓練。圖 4 顯示，Llama2-7B 和 BERT-Large 在許多基準測試上效能均出現適度下降。回答問題 2，是否所有層都是必要的：
不，至少可以刪除一些中間層而不會發生災難性故障。
中間層都執行相同的功能嗎？
如果中間層都共享一個共同的表徵空間，這是否意味著除此之外的中間層是多餘的呢？為了測試這一點，研究者們重新執行了前一子節中的「跳過」實驗，他們將中間層的權重替換為中心層的權重，有效地在被替換的每一層上迴圈 T - 2N + 1 次，其中 T 是總層數（Llama2-7B 為 32 層，BERT-Large 為 24 層）。
如圖 5 所示，可以觀察到，隨著被替換層數的增加，模型在基準測試的得分迅速下降。從後文的圖 11 看來，這種替換層的做法比研究者們嘗試的其他方法都更糟糕。因此，研究者得出結論：中間層執行的是不同的功能，讓中間層之間共享權重並不可行。
層的順序重要嗎？
之前的實驗表明，中間層共享一個表示空間，但在該空間中負責不同的功能。下一個需要解決的問題是，這些功能的順序有何意義。為了解決這個問題，研究者們設計了兩組實驗。首先，以與訓練時相反的順序來執行中間層。具體來說，取第 T - N 層的輸出，將其輸入到第 T - N - 1 層，然後將這一層的輸出輸入到第 T - N - 2 層，依此類推，一直到第 N 層，再將這一層的輸出傳送到後面的 T - N 層。在第二組實驗中，研究者採用隨機順序執行中間層，並在 10 個種子值上取平均值。
圖 6 和圖 7 分別顯示了反向和以隨機順序執行中間層的結果，模型在所有基礎測試集中都顯示出了逐漸下降的趨勢。這也表明雖然層的順序對模型來說有一定的重要性，但即使改變了順序，這些層仍然能夠發揮作用。
更有趣的是，隨機打亂層的順序比完全反過來效果更好。這可能是因為，隨機打亂的順序在某些方面保留了層之間的一些原有關係（即層 i 在層 j 之後，其中 i > j），而完全反過來則完全打破了這些關係。
這些層可以並行執行嗎？
為了驗證層本身存在比執行的順序更重要，研究者們設計了一個實驗，並行執行中間層，將它們的平均結果傳送給最終的 N 層。
如圖 8 所示，模型在所有基準測試中的表現均呈現了一種平緩下降趨勢，然而，這種趨勢並不適用於 GSM8K 中的數學應用題。
實驗結果顯示，大部分情況下這種方法都是有效的，只是一些複雜的數學題處理得不太好。這種並行處理方法相比直接跳過一些層，效果更好，但不如按反向順序執行層的效果出色。基於此，研究者得出結論：並行執行層在一般情況下是可行的，但對於需要順序邏輯理解的數學問題，這種方法可能不太適用。
對於某些任務來說，順序是否比其他因素更重要？
對於大多數經過「改造」的模型，在面對抽象推理（ARC）或數學推理（GSM8K）基準測試時，它們往往顯示出最陡峭的下降趨勢。這一現象可能源於逐步推理任務對於模型層級順序的敏感度遠高於那些主要依賴語義理解的常識性任務。與那些僅透過理解語義便能完成的任務不同，推理任務要求模型同時把握結構與含義。這種觀察與模型在單次處理過程中可能進行一定程度的順序依賴性推理的假設相吻合。
研究者使用了一個比喻來說明：如果畫一幅由許多不同元素組成的拼貼畫，那麼畫的順序可能不那麼重要；但如果是要畫一幅精確的建築場景，那麼每一筆的順序就變得非常重要了。據此，研究者得出了結論：數學和推理任務對模型層的順序具有更高的依賴性，而對於那些主要依賴語義理解的任務，順序的影響則相對較小。
迴圈有助於層之間並行嗎？
沿用上一節中畫畫的的比喻，當畫家在畫一幅畫時，不是一開始就畫所有東西，而是先畫一部分，比如車身，然後再根據這部分來新增其他的東西，比如車輪。在 AI 模型中，層就是所謂的畫家，處理資訊就是在畫畫，如果先得到了正確的資訊，也就先畫出了所謂的車身，那麼它們就能更好地完成自己的工作，為畫作新增車輪。
對於 transformer 而言，當給予適當的輸入時，層可能只在前向傳播中做出貢獻，並非透過殘差連線「傳遞」輸入。如果情況確實如此，那麼迭代上一個實驗中的並行層應該比單次執行並行層更能提高模型的效能。基於此，研究者透過將並行層的平均輸出反饋到同一層中進行固定次數的迭代來測試這一點。
圖 9 展示了將並行層迴圈 3 次的結果。迴圈並行 3 次的結果顯著優於單次迭代（並行層）。起始層 N 設定為 15（針對 Llama2-7B 模型）或 11（針對 BERT 模型）時，即處於每種情況的極左端點，僅有單一的層級受到影響。在這種特定情況下，三次迴圈並行的效果等同於單純地將中間層重複三次。與此同時，對於這一點上的並行層而言，其效能與完整模型無異。
研究者們還針對不同的迭代次數重複了相同的實驗。圖 10 展示了 Llama2-7B 的效能隨並行化層數 M 和迭代次數的變化情況。每個 M 的最高效能迭代次數用紅框標出。除了 M=29 和 M=31（幾乎並行化所有層）外，最佳迭代次數大致與並行化層數成線性比例。因此，研究者得出的結論是：最佳迭代次數與並行化層數成正比。
如何調整層，對模型效能的影響最小？
最後，在圖 11 中，研究者們將所有實驗中對 Transformer 的「改造」進行了比較，在一個圖表上顯示了所有基準測試的中位數或平均性 。
中間重複 —— 用相同數量的中間層副本替換中間層 —— 表現最差， 很快就降到了隨機基線的效能。相反，迴圈並行和隨機層順序的影響最小。因此，研究者得出的結論是：重複單一層的影響最嚴重。隨機化層順序和迴圈並行的影響最小。
這些實驗整體上顯示出平緩的效能下降，但研究者仍然不清楚為什麼這些層在大多數擾動下還能保持一定的穩健性，這個問題還需在未來的研究中進一步探討。
更多細節請參見原論文。
參考連結：https://arxiv.org/pdf/2407.09298v1