七年前，論文《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 的「改造」進行了比較，在一個圖表上顯示了所有基準測試的中位數或平均性 。 

中間重復 —— 用相同數量的中間層副本替換中間層 —— 表現最差， 很快就降到了隨機基線的性能。相反，循環并行和隨機層順序的影響最小。因此，研究者得出的結論是：重復單一層的影響最嚴重。隨機化層順序和循環并行的影響最小。 

這些實驗整體上顯示出平緩的性能下降，但研究者仍然不清楚為什么這些層在大多數擾動下還能保持一定的穩健性，這個問題還需在未來的研究中進一步探討。 

更多細節請參見原論文。