清華微軟最新力作：用物理學革新Transformer注意力，「大海撈針」精度暴漲30%！
新智元報道
編輯：喬楊 好睏
【新智元導讀】隨着諾貝爾物理學獎頒給了「機器學習之父」Geoffrey Hinton，另一個借鑑物理學概念的模型架構也橫空出世——微軟清華團隊的最新架構Differential Transformer，從注意力模塊入手，實現了Transformer的核心能力提升。
隨着近些年來NLP領域研究的不斷深入，我們逐漸發現，Transformer架構中出現的幻覺問題，以及各種下游任務中的性能不足，都或多或少與注意力缺陷有關。
雖然上下文窗口可以擴展，但是Transformer還是無法真正關注到有價值的信息。
最近，微軟研究院和清華大學的研究人員共同提出了一種新的模型架構——Differential Transformer，不僅保留了原始Transformer中的可擴展性，也能讓模型更加關注上下文中與任務相關的關鍵信息。
實驗表明，注意力機制的改進，不僅顯著提升了檢索精度，還能緩解LLM的幻覺。
論文地址：https://arxiv.org/abs/2410.05258
Transformer的困境
衆所周知，Transformer的核心是注意力機制，採用softmax函數來衡量序列中各種標記的重要性。然而，最近的研究表明，LLM難以從上下文中準確到檢索關鍵信息。
比如去年斯坦福Percy Liang團隊的一篇論文就指出，雖然語言模型能夠接受較長的上下文作爲輸入，但並不能穩健地利用長輸入上下文中的信息。
論文地址：https://arxiv.org/abs/2307.03172
比如，實驗中發現，僅僅改變關鍵信息在文檔中的出現位置，就可以造成GPT-3.5 Turbo檢索性能的大範圍波動。
此外，本篇論文的實驗結果證明，Transformer經常過度關注不相關的上下文，本文將其稱之爲「注意力噪聲」。
如圖1（左）所示，模型分配給正確答案的注意力分數很低，同時不成比例地關注不相關的上下文，這意味着信噪比很低，最終淹沒了正確答案。
由此看來，我們對於LLM檢索、利用長上下文的過程，知之甚少，其注意力過程也需要更多的改進。
本文所提出的Differential Transformer（DIFF Transformer）正是希望用「差分注意力」（differential attention）機制消除注意力噪聲，促使模型關注上下文中的關鍵信息。
圖1的對比結果可以看出，DIFF Transformer給出的注意力分數的分佈明顯不同於傳統Transformer架構，給予關鍵信息更高的注意力分數，進而顯著提升了檢索能力。
這種能力的提升，對於有效利用LLM的長上下文窗口、緩解幻覺、關鍵信息檢索等方面都有重要的意義。
模型架構
DIFF Transformer也可以用於純Encoder或Encoder-Decoder模型，但本篇論文以純Decoder模型爲例進行描述。
整個模型由L個DIFF Transformer層堆疊而成，每層由一個差分注意力模塊和前饋網絡模塊連接形成。
宏觀佈局類似於傳統Transformer架構，但主要區別在於修改了注意力的softmax過程，並且採用了pre-RMSNorm、SwiGLU等改進。
差分注意力
該模塊的結構示意圖和僞代碼如圖2所示，具體的代碼實現可參考項目GitHub。
代碼地址：https://github.com/microsoft/unilm/tree/master/Diff-Transformer
除了傳統注意力中的權重矩陣W^Q、W^K、W^V ∈ ℝ^{d_model×2d}，模塊中還加入了可學習標量λ。
具體來說，給定輸入序列X ∈ ℝ^{N×d_model}，首先將其投影爲Q、K、V矩陣Q_1,Q_2,K_1,K_2 ∈ ℝ^{N×d} , V ∈ ℝ^{N×2d}，然後是差分注意力算子DiffAttn(·)通過公式（1）計算輸出：
λ被初始化爲常量λ_{init} ∈ (0,1)，並依照公式（2）與其他權重參數同步更新：
其中，λ_ , λ_ , λ_ , λ_ ∈ ℝ^d也都是是可學習向量。
之所以命名爲「差分注意力」，是指兩個softmax函數間的差異可以消除注意力噪音。
這個想法類似於電氣工程中提出的差分放大器（differential amplifiler），將兩個信號之間的差異作爲輸出，從而消除輸入中的共模噪聲；降噪耳機的設計也是基於類似的思路。
DIFF Transformer中也可以使用多頭注意力機制，在同一層的多個head間共享參數λ，將每個head的輸出進行歸一化處理後再拼接、投影，就得到了最終輸出，如公式（3）所示。
公式（3）中的LN（·）是指對每個頭使用RMSNorm，但如圖2（左）所示，也可以使用GroupNorm。
加上前饋網絡模塊，每個DIFF Transformer層就可以描述爲：
實驗
下游任務
首先，研究人員在1T token上訓練3B大小的DIFF Transformer模型，並在各種下游任務上與之前有競爭力的Transformer架構模型進行比較，結果如表1所示。
基線模型大小都爲3B，其中，StableLM-3B-4E1T的1T結果取自技術報告，而OpenLLaMA-v2-3B和StableLM-base-alpha-3B-v2同樣使用1T數據訓練，表中分數爲Eval Harness基準測試上的零樣本準確率。
結果顯示，，與之前經過精心調優的Transformer語言模型相比，DIFF Transformer取得了良好的性能。
尤其是對於長上下文任務，如圖4所示，隨着上下文長度不斷增加，累計平均的負對數似然值（NLL）持續降低，說明Diff Transformer可以更有效地利用不斷增加的上下文。
關鍵信息檢索
「大海撈針」（Needle-In-A-Haystack）測試被廣泛用於評估LLM提取長上下文中的關鍵信息的能力。
本文的實驗遵循LWM和Gemini 1.5的「多針」評估方案，在不同長度的上下文中，N根針被插入不同的深度。每根「針」都由一個簡潔的句子組成，爲特定城市分配一個獨特的魔法數字。
答案針被放置在上下文中的5個不同深度：0%、25%、50%、75%和100%，同時隨機放置其他分散注意力的針。待測LLM的目標，就是是檢索與查詢城市相對應的數字。
4k上下文檢索的可結果如表2所示。雖然兩種模型在N=1或N=2時都取得了良好的準確率，但隨着N的增加，DIFF Transformer的性能保持相對一致，Transformer則顯著下降。
4K長度的平均檢索精度，N代表針數，R表示查詢城市的數量
將上下文長度擴展至64k時，差距就更加明顯，尤其是關鍵信息位於前半部分時（即0%、25% 和 50%深度）。
特別是，將針放置在25%深度時，DIFF Transformer相對於傳統Transformer實現了76%的精度提升。
除了檢索精度，表3進一步分析了兩種模型爲上下文分配的注意力分數。可以看出， DIFF Transformer的確將更多的注意力分配給了有用的信息，並有效地消除注意力噪聲。
值得注意的是，DIFF Transformer在提升檢索精度的同時也緩解了幻覺現象。
實驗包含模型在總結（圖4a）和問答（圖4b）兩種任務上的幻覺評估。可以發現，與Transformer相比，DIFF Transformer的上下文幻覺明顯減輕。
這種性能的提高可能源於，改進後的注意力模塊能更好第關注任務所需信息，而非不相關的上下文。
這與之前研究中的觀察結果一致，即Transformer出現上下文幻覺的一個主要原因是注意力分數的錯誤分配。
對文本摘要和問題回答的幻覺評估。準確度越高表示幻覺越少；評估時採用GPT-4o進行自動化的二元判斷
縮放特性
除了下游任務性能，論文還進行了縮放特性的對比。
擴展模型規模
如圖3a所示，分別使用830M、1.4B、2.8B、6.8B和13.1B參數訓練語言模型，發現DIFF Transformer依舊遵循Scaling Law。
根據擬合曲線，68億參數規模的DIFF Transformer達到了與110億參數規模Transformer相當的驗證損失，但僅需62.2%的參數。
同樣，78億參數的DIFF Transformer匹配了131億參數的Transformer的性能，參數量是後者的59.5%。
擴展訓練Token
如圖3b所示，訓練數據的縮放也遵循類似規律，且擬合曲線表明，使用160B token訓練的DIFF Transformer達到了與使用251B token訓練的Transformer相當的性能，但僅消耗了63.7%的訓練數據。
此外，在HellaSwag上的測試結果還可以發現，Diff Transformer對量化和位寬的穩健性顯著高於Transformer。
作者介紹
本文的4位共同一作都來自微軟研究院，其中兩位是清華大學學生。
Tianzhu Ye
Tianzhu Ye本科畢業於清華大學自動化系，今年剛剛進入本系就讀博士一年級，目前是微軟自然語言計算部門實習生。
Li Dong（董力）
Li Dong從2018年起擔任MSRA自然語言計算組的首席研究員。
他2012年畢業於北京航空航天大學，獲得了計算機科學與工程方向的學士和碩士學位，之後前往愛丁堡大學攻讀信息學博士，曾在微軟Redmond研究院自然語言處理組實習。
Yuqing Xia（夏雨晴）
Yuqing Xia是微軟亞洲研究院（MSRA）系統與網絡研究組的研究員，此前於2019年在北京大學獲得了生物學博士學位
她的研究方向是利用現代硬件技術爲計算密集型任務（如機器學習和深度學習）構建大規模計算系統。
此外，她還對如何運用人工智能來推動自然科學（尤其是生物學）的研究進展抱有濃厚興趣。
Yutao Sun（孫宇濤）
Yutao Sun是清華大學的一年級博士生，導師是王建勇。同時，他也在微軟亞洲研究院實習，由董力指導。
他的研究興趣是大語言模型的骨幹網絡、長序列的建模和推理，以及大語言模型在其他領域的應用。
參考資料：
https://arxiv.org/abs/2410.05258