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

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

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

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

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

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

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

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

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

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

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

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

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

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

根據(jù)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回路能夠完成上面的例子=，那換成：

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

動態(tài)組合多頭注意力長啥樣？

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

如下圖所示，得到DCMHA：

△圖1. DCMHA總體結構

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

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

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

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

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

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

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

△圖3. Compose的計算

效果如何？

規(guī)模擴展

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

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

△圖4. Transformer和DCFormer的規(guī)模擴展效果

△圖5. Pythia和DCPythia的規(guī)模擴展效果

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

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

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

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

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

下游任務評測

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

△表1. DCFormer 和 Pythia 在下游任務中的表現(xiàn)

從表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%，而且隨著模型參數(shù)的增長，額外的計算開銷會逐漸降低。

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

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

消融實驗

結果如下：

△表3. DCMHA的消融實驗

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

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

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

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

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

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

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

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

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