微軟清華改進Transformer:用降噪耳機原理升級注意力,一作在線答疑
Transformer自問世后就大放異彩,但有個小毛病一直沒解決:
總愛把注意力放在不相關的內容上,也就是信噪比低。
現在微軟亞研院、清華團隊出手,提出全新改進版Differential Transformer,專治這個老毛病,引起熱議。
論文中介紹,整體思路類似差分放大電路或降噪耳機,用兩個信號的差值來濾除共模噪聲。
具體到在語言模型中,如果句子很長,只有少數token會真正影響當前token的含義。而注意力機制允許每兩個詞之間產生交互,其中就包含大量噪聲了。
團隊提出的方法是在注意力層中增加一個Softmax,然后兩個Softmax做減法。
這一減,噪音信息就被大幅抵消,讓注意力更集中在相關內容上。
語言建模任務上的一系列實驗結果顯示,僅需約65%的模型大小或訓練tokens,DIFF Transformer就能達到與傳統Transformer相當的性能。
新架構在長上下文建模、關鍵信息檢索、減少幻覺、提高上下文學習能力以及減少激活異常值等各項指標中,普遍優于Transformer架構。
論文上傳到arXiv平臺后,有不少學者到彈幕版alphaXiv劃線提問。一作Tianzhu Ye正絕贊在線答疑中。
差分Transformer
與傳統Tranformer相比,DIFF Transformer保持宏觀架構不變,主要區別在于用差分注意力替換傳統softmax注意力。
此外還采用了LLaMA系列中的一些改進,如pre-RMSNorm歸一化和SwiGLU激活函數。
在差分注意力模塊中,需要先給Q和K分成兩個組,然后分別計算softmax。
第二組乘了一個標量λ,是可學習的參數,在同一層的注意力頭之間共享。
λ的引入是為了在差分操作中平衡兩組注意力的貢獻,使得差分注意力機制能夠更好地適應不同的任務需求和數據分布。
接下來是一系列實驗結果。
語言建模評估
在1T tokens上訓練3B大小的DIFF Transformer,遵循 StableLM-3B-4E1T的配方,在各種下游任務中與以前訓練良好的Transformer模型相比表現出優勢。
可擴展性評估
只需約65%的模型參數或訓練tokens來匹配Transformer的性能。
長上下文能力評估
在額外1.5B tokens上訓練3B大小的DIFF Transformer,擴展上下文長度至64k。
隨著上下文長度增加,累計平均負對數似然(NLL)持續降低,并且比傳統Transformer的NLL值更低。
表明DIFF Transformer可以有效利用不斷增加的上下文。
關鍵信息檢索能力評估
也就是多個“針”的大海撈針試驗,設置不同的上下文長度(4K和64K)來模擬不同復雜程度的信息檢索場景。
在4K上下文長度下,隨著插入 “針” 數量和查詢數量的增加,DIFF Transformer的準確率保持穩定,而Transformer 的準確率顯著下降。
在64K上下文長度下,DIFF Transformer在不同答案針深度(即關鍵信息在長上下文中的位置)和上下文長度下都能保持穩定性能,且在關鍵信息位于上下文前半部分時優勢明顯。
特別是當關鍵信息位于25%深度時,DIFF Transformer比Transformer的準確率提高了 76%。
上下文學習能力評估
分為兩個角度來評估,分別是多樣本分類和上下文學習的穩健性。
多樣本分類任務,同樣使用64K上下文長度的3B參數模型,DIFF Transformer的準確率始終高于Transformer,提升幅度從5.2%到21.6%不等
上下文學習穩健性采用排列順序任務,DIFF Transformer的結果方差遠小于傳統Transformer。
上下文幻覺評估
主要關注輸入中包含正確事實,但模型仍然無法產生準確輸出的情況。
將模型輸出與ground-truth一起發給GPT-4o,讓GPT-4o來判斷是否存在幻覺,此前試驗表明GPT-4o與人類評判結果一致率較高,相對可靠。
在不同數據集上DIFF Transformer的準確率更高,幻覺更少。
激活異常值分析
Transformer中的激活異常值,導致模型在訓練和推理過程中難以量化。
試驗比較了注意力logits和隱藏狀態兩種激活類型下的最大激活值,DIFF Transformer都表現出更低的頂部激活值,即產生更少的激活異常值。
在對注意力logits進行量化實驗時,DIFF Transformer在降低比特寬度量化時仍能保持較高性能,而Transformer在6-bi 量化時準確性顯著下降。
4-bit的DIFF Transformer能達到與6-bit的Transformer相當的準確性,且比4-bit的Transformer準確率提高約 25%。
代碼已開源, 降噪耳機類比引熱議
對于目前讀者的疑問,作者已做出幾點答復
問題1:Diff Transformer與每個注意力頭溫度可學習的方法有什么不同?與門控注意力對比如何?
作者回應在實驗中,可學習溫度效果不大。而本文方法是門控注意力的改進。
問題2:差分注意力是否意味著將標準注意力矩陣參數翻倍?
作者澄清,單個注意力頭維度翻倍,但是注意力頭數量減半,總體在參數和FLOPS上都是對齊的。
問題3:第二組Softmax乘可學習參數lambda的研究思路。
作者也做出詳細回應。
DIFF Transformer在純學術圈之外也引起非常多的討論,有不少人困惑論文中將方法與降噪耳機的類比。
降噪耳機采集環境噪聲并生成相反的信號,在這種情況下哪些信號屬于噪聲是已知的,但差分注意力中并不能事先確定哪些是噪聲。
一種解釋是,低注意力分數的噪聲也有很低的梯度,因此模型其實已知哪些是噪聲,只是單個Softmax無法輸出0,所以噪聲很難去除。
也有人提出,比起降噪耳機,其實專業音頻中“平衡線”,或者USB、網卡等傳輸方式更適合一些。
使用兩條信號線傳輸正負信號,接收器只對比他們之間的差異,由于在空間中離得很近,受到的干擾是相同的。
有用的正負信號相減會被放大,相同的噪聲相減卻被抵消,大大增強抗干擾能力。
總之,DIFF Transformer代碼已開源在微軟unilm項目下,其中還包含魔改版支持差分注意力的FlashAttention-2的代碼。
感興趣的可以試起來了。
論文:https://arxiv.org/abs/2410.05258
代碼:https://aka.ms/Diff-Transformer