微軟清華改進Transformer：用降噪耳機原理升級注意力，一作在線答疑

魚羊 2024-11-03 15:56:26 來源：量子位

夢晨 發自 凹非寺

量子位 | 公眾號 QbitAI

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

參考鏈接：
[1]https://news.ycombinator.com/item?id=41776324

— 完 —