ICLR2025|差分注意力機制爆火!DIFFTransformer擊穿長文字建模天花板

近年來,Transformer 架構在自然語言處理領域取得了巨大成功,從機器翻譯到文字生成,其強大的建模能力為語言理解與生成帶來了前所未有的突破。
然而,隨著模型規模的不斷擴大和應用場景的日益複雜,傳統 Transformer 架構逐漸暴露出缺陷,尤其是在處理長文字、關鍵資訊檢索以及對抗幻覺等任務時,Transformer 常常因過度關注無關上下文而陷入困境,導致模型表現受限。
為攻克這一難題,來自微軟和清華的研究團隊提出了 DIFF Transformer,一種基於差分注意力機制的創新基礎模型架構
論文標題
Differential Transformer
論文連結
https://openreview.net/pdf?id=OvoCm1gGhN
程式碼連結
https://aka.ms/Diff-Transformer
其核心思想是透過計算兩組 Softmax 注意力圖的差值來放大對關鍵上下文的關注,同時消除注意力噪聲干擾。DIFF Transformer 具備以下顯著優勢:
在語言建模任務中,DIFF Transformer 在模型大小、訓練 token 數量等方面展現出了卓越的可擴充套件性,僅需約 65% 的模型規模或訓練 token 數量即可達到與傳統 Transformer 相當的效能,大幅提升了語言模型通用表現。
在長文字建模、關鍵資訊檢索、數學推理、對抗幻覺、上下文學習、模型啟用值量化等一系列任務中,DIFF Transformer 展現了獨特優勢,相比傳統 Transformer 有顯著提升。
DIFF Transformer 的特性使其在自然語言處理領域具有廣闊的應用前景,有望成為推動語言模型發展的新動力。此外,已有跟進研究初步驗證方法在視覺、多模態等領域中的有效性,顯示出其跨模態通用的潛力。該研究已被 ICLR 2025 接收,並獲選為 Oral 論文(入選比例 1.8%)。

方法

本文提出了一種名為 Differential Transformer(DIFF Transformer) 的基礎模型架構,旨在解決傳統 Transformer 在長文字建模中對無關上下文過度分配注意力的問題。該方法透過差分注意力機制(Differential Attention)放大對關鍵上下文的關注,同時消除注意力噪聲,從而顯著提升模型在多種任務中的效能。
差分注意力機制
傳統 Transformer 的注意力機制透過 Softmax 函式對輸入序列中的不同 token 進行加權,但 Softmax 的性質導致模型難以完全消除無關上下文的影響。為了克服這一問題,DIFF Transformer 引入了差分注意力機制。
具體而言,該機制將查詢向量(Query)和鍵向量(Key)在注意力頭(Head)維度分為兩組,分別計算兩組的 Softmax 注意力圖,然後計算兩者的差值作為最終的注意力分數。這一設計類似於電子工程中的差分放大器,以及降噪耳機,透過兩組訊號相減以消除共有噪聲。
差分注意力的數學表達如下:

其中,

分別是兩組查詢和鍵向量,

是值向量,

是一個可學習的標量引數,用於調節兩組注意力圖的權重。計算過程如圖 1 所示。

▲ 圖1. 差分注意力機制圖示與虛擬碼
為了同步學習速率,將

重引數化為:

其中,

是可學習的向量,而

是用於初始化的常數。

多頭差分注意力
為了進一步提升模型的表達能力,DIFF Transformer 採用了多頭機制。每個注意力頭獨立計算差分注意力,並將多頭輸出拼接為最終結果。具體實現如下:

其中

是注意力頭的數量,

是輸出投影矩陣。為了保持與 Transformer 梯度一致,DIFF Transformer 在每個頭的輸出後應用了獨立的歸一化層,採用 RMSNorm 實現。

▲ 圖2. Transformer與DIFF Transformer注意力分數分佈視覺化
圖 2 展示了 DIFF Transformer 和傳統 Transformer 在注意力分數分配上的顯著差異。作者將一段關鍵資訊插入大段不相關文字的中間位置,並對模型抽取關鍵資訊時的注意力分數分配進行視覺化。
傳統 Transformer 的注意力分數被廣泛分配到整個上下文中,只有極少分數分配至關鍵資訊;而 DIFF Transformer 能夠將更高的分數集中在目標答案上,並且幾乎不向無關上下文分配注意力。
注意力分數分配的稀疏性與精準性也使得 DIFF Transformer 在處理長文字關鍵資訊檢索任務時顯著優於 Transformer。
實驗
作者透過一系列實驗驗證了 DIFF Transformer 在多個方面的卓越效能,證明了其在大語言模型中應用的獨特潛力與優勢。
語言建模
作者研究了 DIFF Transformer 在擴充套件模型規模和訓練資料量時的效能,如圖 3 所示。實驗表明,DIFF Transformer 僅需約 65% 的引數規模或訓練資料量即可達到與 Transformer 相當的語言建模效能。例如,6.8B 引數規模的 DIFF Transformer 在語言建模損失上與 11B 引數規模的 Transformer 相當。

▲ 圖3. 語言建模上的模型引數、訓練資料量可擴充套件性實驗
長文字建模
作者將模型擴充套件到 64K 上下文長度,並在長文字書籍資料上進行了評估。結果顯示,考慮累積平均負對數似然(NLL)指標, DIFF Transformer 在不同序列位置上均優於 Transformer,能夠更有效地利用長上下文資訊。

▲ 圖4. 長文字書籍資料模型效能評估
關鍵資訊檢索
作者透過「多針檢索」(Multi-Needle Retrieval)實驗評估了模型從大量上下文中提取關鍵資訊的能力,如圖 5 所示。
實驗表明,DIFF Transformer 在不同上下文長度和答案深度下均表現出更高的準確率,尤其是在文字較長以及答案位於文字更靠前位置時,優勢更為明顯。例如,在 64K 上下文中,DIFF Transformer 在答案位於 25% 深度時的準確率比 Transformer 高出 76%。
此外,統計資訊顯示,DIFF Transformer 在注意力分數分配上也表現出更高的聚焦能力,能夠準確定位關鍵資訊,並展現了更高的信噪比。

▲ 圖5. 多針檢索評估
上下文學習
作者從兩個角度評估了 DIFF Transformer 的上下文學習能力:多樣本上下文學習和樣本順序魯棒性測試。 如圖 6 所示,在多樣本上下文學習任務中,作者使用了 4 個不同的資料集(TREC、TREC-fine、Banking-77 和 Clinic-150),並逐步增加示例數量,直到總長度達到 64K tokens。
結果顯示,DIFF Transformer 在不同資料集上均優於 Transformer,平均準確率提升顯著。

▲ 圖6. 多樣本上下文學習
在魯棒性測試中,作者透過打亂示例順序的方式評估了模型的效能穩定性。如圖 7 所示,DIFF Transformer 在不同示例排列下的效能方差顯著低於 Transformer,表明其對輸入順序的敏感性更低,具有更強的魯棒性。

▲ 圖7. 樣本順序魯棒性測試
幻覺評測
作者利用文字摘要和問答任務作為兩個典型的幻覺評測場景,評估了 DIFF Transformer 在降低大模型幻覺(hallucination)方面的表現。
結果如圖 8 所示,DIFF Transformer 在生成摘要和回答問題時顯著提升了準確率,減少了幻覺現象。這是因為差分注意力機制能夠準確定位重要文段,避免無關上下文對模型預測的干擾。

▲ 圖8. 利用文字摘要、問答任務進行幻覺評測
異常啟用值分析
作者還發現 DIFF Transformer 能夠顯著減少模型啟用中的異常值,這為模型啟用值的量化提供了新的可能性。實驗表明,DIFF Transformer 在注意力啟用值(attention logits)和隱藏狀態(hidden states)中的最大啟用值顯著低於 Transformer。
例如,在注意力啟用值的 Top-1 啟用值上,DIFF Transformer 比 Transformer 低了近 8 倍。利用這一性質,DIFF Transformer 在注意力啟用值的低位元量化下的效能也優於 Transformer,如圖 9 所示。

▲ 圖9. 注意力啟用值的低位元量化
數學推理能力
作者在數學推理任務上進一步驗證了 DIFF Transformer 的效能。作者採用兩階段訓練,在 3B 預訓練模型的基礎上進行有監督微調,並在 MATH 等 8 個數學資料集上評測模型效能。
在第一階段,採用 20B token 合成數學資料對模型進行微調,使模型獲得基礎數學能力,評測結果如圖 10 所示。從 15B token 開始,DIFF Transformer 展現出了顯著優於 Transformer 的數學能力,至 20B token 結束的時候,準確率的差距達到了 11% 左右。

▲ 圖10. 第一階段數學合成數據微調
在第二階段,作者利用 Deepseek-R1 輸出所構造的資料集 OpenThoughts-114K-Math 對模型進行蒸餾,使模型更強大的深度推理能力。
如圖 11 所示,在 8 個數據集上,DIFF Transformer 相較 Transformer 均有不同程度的提升,平均準確率提升了 7.5%,這表明差分注意力機制更強大的上下文建模能力在推理任務中也至關重要。

▲ 圖11. 第二階段深度推理能力評測
討論與未來工作
DIFF Transformer 自發布以來獲得了較大關注與討論。作者在 Hugging Face 論文討論平臺、alphaXiv 平臺上與社群開展了深入的探討。
在 X 平臺(原 Twitter)上,Google DeepMind 高階研究科學家(Senior Staff Research Scientist)Petar Veličković 與作者就文章中的理論分析展開討論,ViT 核心作者 Lucas Beyer 也在閱讀文章後撰寫了一篇深入的論文總結,相關發帖已獲得數十萬瀏覽。
目前 DIFF Transformer 也已整合至 Hugging Face 的 transformers 庫中。
Hugging Face:
https://huggingface.co/papers/2410.05258
alphaXiv:
https://www.alphaxiv.org/abs/2410.05258
Petar Veličković:
https://x.com/PetarV_93/status/1874820028975267866
Lucas Beyer:
https://x.com/giffmana/status/1873869654252544079

transformers庫:

https://github.com/huggingface/transformers/tree/main/src/transformers/models/diffllama
未來工作方面,作者認為可以利用 DIFF Transformer 的性質設計低位元注意力運算元,以及利用差分注意力的稀疏特性進行鍵值快取(key-value cache)的剪枝。
此外,將 DIFF Transformer 應用在除語言以外的其他模態上也值得探索。近期工作 DiffCLIP 將差分注意力擴充套件至視覺、多模態領域,揭示了 DIFF Transformer 在不同模態任務中的更多結構特性與應用潛力。
DiffCLIP:
https://arxiv.org/abs/2503.06626
總結
本文的貢獻主要在兩個方面:
1. DIFF Transformer 透過創新的差分注意力機制,有效解決了傳統 Transformer 在處理文字時受到噪聲干擾、注意力分配不準確的問題; 
2. 憑藉對關鍵資訊的關注和對噪聲的抵禦能力,DIFF Transformer 在語言建模、長文字建模、關鍵資訊檢索、數學推理、對抗幻覺、上下文學習、模型啟用值量化等任務中表現出色,有望在自然語言處理、多模態等領域作為基礎模型架構。
更多閱讀
#投 稿 通 道#
 讓你的文字被更多人看到 
如何才能讓更多的優質內容以更短路徑到達讀者群體,縮短讀者尋找優質內容的成本呢?答案就是:你不認識的人。
總有一些你不認識的人,知道你想知道的東西。PaperWeekly 或許可以成為一座橋樑,促使不同背景、不同方向的學者和學術靈感相互碰撞,迸發出更多的可能性。 
PaperWeekly 鼓勵高校實驗室或個人,在我們的平臺上分享各類優質內容,可以是最新論文解讀,也可以是學術熱點剖析科研心得競賽經驗講解等。我們的目的只有一個,讓知識真正流動起來。
📝 稿件基本要求:
• 文章確係個人原創作品,未曾在公開渠道發表,如為其他平臺已發表或待發表的文章,請明確標註 
• 稿件建議以 markdown 格式撰寫,文中配圖以附件形式傳送,要求圖片清晰,無版權問題
• PaperWeekly 尊重原作者署名權,並將為每篇被採納的原創首發稿件,提供業內具有競爭力稿酬,具體依據文章閱讀量和文章質量階梯制結算
📬 投稿通道:
• 投稿郵箱:[email protected] 
• 來稿請備註即時聯絡方式(微信),以便我們在稿件選用的第一時間聯絡作者
• 您也可以直接新增小編微信(pwbot02)快速投稿,備註:姓名-投稿
△長按新增PaperWeekly小編
🔍
現在,在「知乎」也能找到我們了
進入知乎首頁搜尋「PaperWeekly」
點選「關注」訂閱我們的專欄吧
·


相關文章