改進Transformer核心機制注意力，讓小模型能打兩倍大的模型！

ICML 2024高分論文，彩云科技團隊構建DCFormer框架，替換Transformer核心組件多頭注意力模塊（MHA），提出可動態組合的多頭注意力（DCMHA）。

DCMHA解除了MHA注意力頭的查找選擇回路和變換回路的固定綁定，讓它們可以根據輸入動態組合，從根本上提升了模型的表達能力。

可以近似理解為，原來每層有固定的H個注意力頭，現在用幾乎同樣的參數量和算力，可按需動態組合出多至HxH個注意力頭。

DCMHA即插即用，可在任何Transformer架構中替換MHA，得到通用、高效和可擴展的新架構DCFormer。

這項工作由來自北京郵電大學、AI創業公司彩云科技的研究人員共同完成。

研究人員用在DCFormer基礎上打造的模型DCPythia-6.9B，在預訓練困惑度和下游任務評估上都優于開源Pythia-12B。

DCFormer模型在性能上與那些計算量是其1.7-2倍的Transformer模型相當。

多頭注意力模塊有何局限？

大模型的scaling law告訴我們，隨著算力的提升，模型更大、數據更多，模型效果會越來越好。雖然還沒有人能明確說明這條路的天花板有多高，能否達到AGI，但這確實是目前大家最普遍的做法。

但除此以外，另一個問題同樣值得思考：目前絕大多數大模型都基于Transformer，它們都是用一個一個Transformer塊像搭積木一樣搭起來的，那作為積木塊的Transformer本身，還有多大的改進提升空間？

這是模型結構研究要回答的基本問題，也正是彩云科技和北京郵電大學聯合完成的DCFormer這項工作的出發點。

在Transformer的多頭注意力模塊（MHA）中，各個注意力頭彼此完全獨立的工作。

這個設計因其簡單易實現的優點已在實踐中大獲成功，但同時也帶來注意力分數矩陣的低秩化削弱了表達能力、注意力頭功能的重復冗余浪費了參數和計算資源等一些弊端。基于此，近年來有一些研究工作試圖引入某種形式的注意力頭間的交互。

根據Transformer回路理論，在MHA中 ，每個注意力頭的行為由W^Q、W^K、W^V、W^O四個權重矩陣刻畫（其中W^O由MHA的輸出投影矩陣切分得到）。

其中，W^QW^K叫做QK回路（或叫查找選擇回路），決定從當前token關注上下文中的哪個（些）token，例如：

W^OW^V叫做OV回路（或叫投影變換回路），決定從關注到的token取回什么信息（或投影什么屬性）寫入當前位置的殘差流，進而預測下一個token。例如：

研究人員注意到，查找（從哪拿）和變換（拿什么）本來是獨立的兩件事，理應可以分別指定并按需自由組合（就像在SQL查詢中，WHERE后的選擇條件和SELECT后的屬性投影是分開寫的一樣），MHA硬把它們放到一個注意力頭的QKOV里“捆綁銷售”，限制了靈活性和表達能力。

例如，假設有個模型存在注意力頭A、B、C其QK和OV回路能夠完成上面的例子=，那換成：

需要交叉組合現有注意力頭的QK和OV回路，模型就可能“轉不過彎兒”了（經研究人員系統構造的合成測試集驗證，<=6B的中小尺寸模型在這類看似簡單的任務上確實表現不佳）。

動態組合多頭注意力長啥樣？

以此為出發點，本文研究團隊在MHA中引入compose操作：

如下圖所示，得到DCMHA：

△圖1. DCMHA總體結構

將QW^Q和KW^K算出的注意力分數矩陣A^S和注意力權重矩陣A^W，與VW^V相乘之前，對其在num_heads維上做線性映射得到新的矩陣A’，通過不同的線性映射矩陣（composition map），以實現各種注意力頭組合的效果。

例如圖2（c）中將head 3和7的QK回路與head 1的OV回路組合在一起，形成一個“新的”注意力頭。

△圖2. 8個注意力頭的簡化的典型composition map的功能，淺色表示大值

為了最大限度的增強表達能力，研究人員希望映射矩陣由輸入動態生成，即動態決定注意力頭怎樣組合。

但他們要生成的映射矩陣不是一個，而是對序列中每對源位置的query Q_i和目的位置的key K_j，都要生成這樣一個矩陣，計算開銷和顯存占用都將難以接受。

為此，他們進一步將映射矩陣分解為一個輸入無關的靜態矩陣W_b、一個低秩矩陣w₁w₂和一個對角矩陣Diag(w_g)之和，分別負責基礎組合、注意力頭間的有限方式（即秩R<=2）的動態組合和頭自身的動態門控（見圖2（d）和圖3（b））。其中后兩個矩陣由Q矩陣和K矩陣動態生成。

在不犧牲效果的前提下，將計算和參數復雜度降低到幾乎可以忽略的程度（詳見論文中復雜度分析）。再結合JAX和PyTorch實現層面的優化，讓DCFormer可以高效訓練和推理。

△圖3. Compose的計算

效果如何？

規模擴展

評估一個架構的好壞，研究人員關注的最核心指標是算力轉化為智能的效率（或叫性能算力比），即投入單位算力能帶來的模型性能提升——花更少的算力，得到更好的模型。

從圖4和圖5的scaling law曲線（在對數坐標下，每個模型架構的損失隨算力的變化可畫出一條近似直線，損失越低，模型越好）可以看出，DCFormer可以達到1.7~2倍算力的Transformer模型的效果，即算力智能轉化率提升了1.7~2倍。

△圖4. Transformer和DCFormer的規模擴展效果

△圖5. Pythia和DCPythia的規模擴展效果

怎么理解這個提升幅度呢？

自2017年Transformer誕生至今，從改進性能算力比的角度，GLU MLP和旋轉位置編碼RoPE是經大量實踐驗證普適有效且被廣泛采用的為數不多的兩項架構改進。

在原始Transformer中加入這兩項改進的架構也叫Transformer++，Llama、Mistral等最強開源模型均采用該架構。無論Transformer還是Transformer++架構，都可通過DCMHA獲得顯著改進。

在1.4B模型規模下，DCMHA的改進幅度大于Transformer++的兩項改進之和，且擴展性更好（圖4下藍綠線和黑線的對比，DCMHA的改進幅度隨算力增加衰減的更慢，以及圖4和圖5的對比）。

可以說，DCFormer讓Transformer的能力又躍上一個新臺階。

下游任務評測

研究團隊訓練了DCPythia-2.8B和DCPythia-6.9B兩個模型在主流NLP下游任務上進行測評并和同規模的開源模型Pythia進行比較（訓練采用和Pythia完全相同超參數設置）。

△表1. DCFormer 和 Pythia 在下游任務中的表現

從表1中可以看出，DCPythia-2.8B和6.9B不僅在Pile驗證集上的ppl 更低，而且在大部分下游任務上都顯著超過了Pythia，DCPythia6.9B在 ppl 和下游任務上的平均準確率甚至超過了Pythia-12B。

DCFormer++2.8B相對于DCPythia-2.8B有進一步的提升，驗證了DCMHA和Lllama架構結合的有效性。

訓練和推理速度

雖然引入DCMHA會帶來額外的訓練和推理開銷，但是從表2中可以看出DCFormer++的訓練速度是Transformer++的74.5%-89.2%，推理速度則是81.1%-89.7%，而且隨著模型參數的增長，額外的計算開銷會逐漸降低。

△表2. Transformer++和DCFormer++的訓練和推理速度對比

訓練速度是在TPU v3 pod，序列長度為2048，batch_size為1k的情況下對比得到的；推理速度是在A100 80G GPU上進行評測的，輸入長度1024，生成長度128。

消融實驗

結果如下：

△表3. DCMHA的消融實驗

從表3中可以看出以下幾點:

• 雖然加入靜態的組合權重就可以降低ppl，但引入動態的組合權重可以進一步降低ppl，說明了動態組合的必要性。
• 低秩動態組合比動態門控的效果更好。
• 只用query-wise或者key-wise的動態組合得到的ppl相當，與DCFormer++的差距很小。
• 在softmax后做注意力頭組合比在softmax前做更有效，可能是因為softmax后的概率能更直接影響輸出。
• 動態組合權重的秩無需設置過大，也說明了組合權重的低秩性。

此外，研究人員還通過增加局部注意力層的比例和只用query-wise動態組合的方式去進一步減少訓練和推理開銷，詳見論文Table 10。

總的來說，研究團隊有兩點總結。

關于動態權重：近期Mamba，GLA，RWKV6，HGRN等SSM和線性注意力/RNN的工作，通過引入動態（input-dependent）權重的方式，追趕上了Transformer++，但DCFormer用動態組合注意力頭的方式說明了在使用 softmax 注意力的情況下，通過引入動態權重也可以大幅提升Transformer++的效果。

關于模型架構創新：這項工作表明，如果存在一個具有極限算力智能轉化效率的“理想模型架構”，當前的Transformer架構雖已非常強大，但距離這個理想架構很可能還存在很大的差距，仍有廣闊的提升空間。因此，除了堆算力堆數據的大力出奇跡路線，模型架構創新同樣大有可為。

研究團隊還表示，彩云科技會率先在旗下產品彩云天氣、彩云小譯、彩云小夢上應用DCformer。

有關更多研究細節，可參閱原始論文。

ICML2024論文鏈接：https://icml.cc/virtual/2024/poster/34047。
Arxiv 論文鏈接：https://arxiv.org/abs/2405.08553。
代碼鏈接：https://github.com/Caiyun-AI/DCFormer。