ICML2024高分！魔改注意力，讓小模型能打出兩倍大的模型

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

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

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

原文意思是，每層都有固定的H個注意力頭，現在用幾乎理解為，原來每層都有固定的H個注意力頭，現在用幾乎相同的參數量和計算力，可以動態組合出多至HxH個注意力頭。 微調後的內容可以更明確地表達原文的意思，如下： 原來的模型每層都包含固定數量的H個注意力頭，現在我們可以利用

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迴路，模型就可能「轉不過彎兒」了（經研究人員系統構造的合成測試集驗證，

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

以此為出發點，本文研究團隊在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（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。

以上是ICML2024高分！魔改注意力，讓小模型能打出兩倍大的模型的詳細內容。更多資訊請關注PHP中文網其他相關文章！