清華微軟最新力作:用物理學革新Transformer注意力,「大海撈針」精度暴漲30%!
隨著近些年來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×2?d},模塊中還加入了可學習標量λ。
具體來說,給定輸入序列X ∈ ?^{N×d_model},首先將其投影為Q、K、V矩陣Q_1,Q_2,K_1,K_2 ∈ ?^{N×d} , V ∈ ?^{N×2?d},然后是差分注意力算子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)兩種任務上的幻覺評估??梢园l現,與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是清華大學的一年級博士生,導師是王建勇。同時,他也在微軟亞洲研究院實習,由董力指導。
他的研究興趣是大語言模型的骨干網絡、長序列的建模和推理,以及大語言模型在其他領域的應用。