ICLR 2025 Oral|差分注意力機(jī)制引領(lǐng)變革,DIFF Transformer攻克長(zhǎng)序列建模難題
近年來(lái),Transformer 架構(gòu)在自然語(yǔ)言處理領(lǐng)域取得了巨大成功,從機(jī)器翻譯到文本生成,其強(qiáng)大的建模能力為語(yǔ)言理解與生成帶來(lái)了前所未有的突破。
然而,隨著模型規(guī)模的不斷擴(kuò)大和應(yīng)用場(chǎng)景的日益復(fù)雜,傳統(tǒng) Transformer 架構(gòu)逐漸暴露出缺陷,尤其是在處理長(zhǎng)文本、關(guān)鍵信息檢索以及對(duì)抗幻覺等任務(wù)時(shí),Transformer 常常因過(guò)度關(guān)注無(wú)關(guān)上下文而陷入困境,導(dǎo)致模型表現(xiàn)受限。
為攻克這一難題,來(lái)自微軟和清華的研究團(tuán)隊(duì)提出了 DIFF Transformer,一種基于差分注意力機(jī)制的創(chuàng)新基礎(chǔ)模型架構(gòu)。
- 論文標(biāo)題:Differential Transformer
- 論文鏈接:https://openreview.net/pdf?id=OvoCm1gGhN
- 代碼鏈接:https://aka.ms/Diff-Transformer
其核心思想是通過(guò)計(jì)算兩組 Softmax 注意力圖的差值來(lái)放大對(duì)關(guān)鍵上下文的關(guān)注,同時(shí)消除注意力噪聲干擾。DIFF Transformer 具備以下顯著優(yōu)勢(shì):
在語(yǔ)言建模任務(wù)中,DIFF Transformer 在模型大小、訓(xùn)練 token 數(shù)量等方面展現(xiàn)出了卓越的可擴(kuò)展性,僅需約 65% 的模型規(guī)模或訓(xùn)練 token 數(shù)量即可達(dá)到與傳統(tǒng) Transformer 相當(dāng)?shù)男阅埽蠓嵘苏Z(yǔ)言模型通用表現(xiàn)。
在長(zhǎng)文本建模、關(guān)鍵信息檢索、數(shù)學(xué)推理、對(duì)抗幻覺、上下文學(xué)習(xí)、模型激活值量化等一系列任務(wù)中,DIFF Transformer 展現(xiàn)了獨(dú)特優(yōu)勢(shì),相比傳統(tǒng) Transformer 有顯著提升。
DIFF Transformer 的特性使其在自然語(yǔ)言處理領(lǐng)域具有廣闊的應(yīng)用前景,有望成為推動(dòng)語(yǔ)言模型發(fā)展的新動(dòng)力。此外,已有跟進(jìn)研究初步驗(yàn)證方法在視覺、多模態(tài)等領(lǐng)域中的有效性,顯示出其跨模態(tài)通用的潛力。該研究已被 ICLR 2025 接收,并獲選為 Oral 論文(入選比例 1.8%)。
方法
本文提出了一種名為 Differential Transformer(DIFF Transformer) 的基礎(chǔ)模型架構(gòu),旨在解決傳統(tǒng) Transformer 在長(zhǎng)文本建模中對(duì)無(wú)關(guān)上下文過(guò)度分配注意力的問(wèn)題。該方法通過(guò)差分注意力機(jī)制(Differential Attention)放大對(duì)關(guān)鍵上下文的關(guān)注,同時(shí)消除注意力噪聲,從而顯著提升模型在多種任務(wù)中的性能。
差分注意力機(jī)制
傳統(tǒng) Transformer 的注意力機(jī)制通過(guò) Softmax 函數(shù)對(duì)輸入序列中的不同 token 進(jìn)行加權(quán),但 Softmax 的性質(zhì)導(dǎo)致模型難以完全消除無(wú)關(guān)上下文的影響。為了克服這一問(wèn)題,DIFF Transformer 引入了差分注意力機(jī)制。
具體而言,該機(jī)制將查詢向量(Query)和鍵向量(Key)在注意力頭(Head)維度分為兩組,分別計(jì)算兩組的 Softmax 注意力圖,然后計(jì)算兩者的差值作為最終的注意力分?jǐn)?shù)。這一設(shè)計(jì)類似于電子工程中的差分放大器,以及降噪耳機(jī),通過(guò)兩組信號(hào)相減以消除共有噪聲。
差分注意力的數(shù)學(xué)表達(dá)如下:
其中,和
分別是兩組查詢和鍵向量,
是值向量,
是一個(gè)可學(xué)習(xí)的標(biāo)量參數(shù),用于調(diào)節(jié)兩組注意力圖的權(quán)重。計(jì)算過(guò)程如圖 1 所示。
圖 1. 差分注意力機(jī)制圖示與偽代碼
為了同步學(xué)習(xí)速率,將重參數(shù)化為:
其中,是可學(xué)習(xí)的向量,而
是用于初始化的常數(shù)。
多頭差分注意力
為了進(jìn)一步提升模型的表達(dá)能力,DIFF Transformer 采用了多頭機(jī)制。每個(gè)注意力頭獨(dú)立計(jì)算差分注意力,并將多頭輸出拼接為最終結(jié)果。具體實(shí)現(xiàn)如下:
其中是注意力頭的數(shù)量,
是輸出投影矩陣。為了保持與 Transformer 梯度一致,DIFF Transformer 在每個(gè)頭的輸出后應(yīng)用了獨(dú)立的歸一化層,采用 RMSNorm 實(shí)現(xiàn)。
圖 2. Transformer 與 DIFF Transformer 注意力分?jǐn)?shù)分布可視化
圖 2 展示了 DIFF Transformer 和傳統(tǒng) Transformer 在注意力分?jǐn)?shù)分配上的顯著差異。作者將一段關(guān)鍵信息插入大段不相關(guān)文本的中間位置,并對(duì)模型抽取關(guān)鍵信息時(shí)的注意力分?jǐn)?shù)分配進(jìn)行可視化。
傳統(tǒng) Transformer 的注意力分?jǐn)?shù)被廣泛分配到整個(gè)上下文中,只有極少分?jǐn)?shù)分配至關(guān)鍵信息;而 DIFF Transformer 能夠?qū)⒏叩姆謹(jǐn)?shù)集中在目標(biāo)答案上,并且?guī)缀醪幌驘o(wú)關(guān)上下文分配注意力。
注意力分?jǐn)?shù)分配的稀疏性與精準(zhǔn)性也使得 DIFF Transformer 在處理長(zhǎng)文本關(guān)鍵信息檢索任務(wù)時(shí)顯著優(yōu)于 Transformer。
實(shí)驗(yàn)
作者通過(guò)一系列實(shí)驗(yàn)驗(yàn)證了 DIFF Transformer 在多個(gè)方面的卓越性能,證明了其在大語(yǔ)言模型中應(yīng)用的獨(dú)特潛力與優(yōu)勢(shì)。
語(yǔ)言建模
作者研究了 DIFF Transformer 在擴(kuò)展模型規(guī)模和訓(xùn)練數(shù)據(jù)量時(shí)的性能,如圖 3 所示。實(shí)驗(yàn)表明,DIFF Transformer 僅需約 65% 的參數(shù)規(guī)?;蛴?xùn)練數(shù)據(jù)量即可達(dá)到與 Transformer 相當(dāng)?shù)恼Z(yǔ)言建模性能。例如,6.8B 參數(shù)規(guī)模的 DIFF Transformer 在語(yǔ)言建模損失上與 11B 參數(shù)規(guī)模的 Transformer 相當(dāng)。
圖 3. 語(yǔ)言建模上的模型參數(shù)、訓(xùn)練數(shù)據(jù)量可擴(kuò)展性實(shí)驗(yàn)
長(zhǎng)文本建模
作者將模型擴(kuò)展到 64K 上下文長(zhǎng)度,并在長(zhǎng)文本書籍?dāng)?shù)據(jù)上進(jìn)行了評(píng)估。結(jié)果顯示,考慮累積平均負(fù)對(duì)數(shù)似然(NLL)指標(biāo), DIFF Transformer 在不同序列位置上均優(yōu)于 Transformer,能夠更有效地利用長(zhǎng)上下文信息。
圖 4. 長(zhǎng)文本書籍?dāng)?shù)據(jù)模型性能評(píng)估
關(guān)鍵信息檢索
作者通過(guò)「多針檢索」(Multi-Needle Retrieval)實(shí)驗(yàn)評(píng)估了模型從大量上下文中提取關(guān)鍵信息的能力,如圖 5 所示。實(shí)驗(yàn)表明,DIFF Transformer 在不同上下文長(zhǎng)度和答案深度下均表現(xiàn)出更高的準(zhǔn)確率,尤其是在文本較長(zhǎng)以及答案位于文本更靠前位置時(shí),優(yōu)勢(shì)更為明顯。例如,在 64K 上下文中,DIFF Transformer 在答案位于 25% 深度時(shí)的準(zhǔn)確率比 Transformer 高出 76%。此外,統(tǒng)計(jì)信息顯示,DIFF Transformer 在注意力分?jǐn)?shù)分配上也表現(xiàn)出更高的聚焦能力,能夠準(zhǔn)確定位關(guān)鍵信息,并展現(xiàn)了更高的信噪比。
圖 5. 多針檢索評(píng)估
上下文學(xué)習(xí)
作者從兩個(gè)角度評(píng)估了 DIFF Transformer 的上下文學(xué)習(xí)能力:多樣本上下文學(xué)習(xí)和樣本順序魯棒性測(cè)試。 如圖 6 所示,在多樣本上下文學(xué)習(xí)任務(wù)中,作者使用了 4 個(gè)不同的數(shù)據(jù)集(TREC、TREC-fine、Banking-77 和 Clinic-150),并逐步增加示例數(shù)量,直到總長(zhǎng)度達(dá)到 64K tokens。結(jié)果顯示,DIFF Transformer 在不同數(shù)據(jù)集上均優(yōu)于 Transformer,平均準(zhǔn)確率提升顯著。
圖 6. 多樣本上下文學(xué)習(xí)
在魯棒性測(cè)試中,作者通過(guò)打亂示例順序的方式評(píng)估了模型的性能穩(wěn)定性。如圖 7 所示,DIFF Transformer 在不同示例排列下的性能方差顯著低于 Transformer,表明其對(duì)輸入順序的敏感性更低,具有更強(qiáng)的魯棒性。
圖 7. 樣本順序魯棒性測(cè)試
幻覺評(píng)測(cè)
作者利用文本摘要和問(wèn)答任務(wù)作為兩個(gè)典型的幻覺評(píng)測(cè)場(chǎng)景,評(píng)估了 DIFF Transformer 在降低大模型幻覺(hallucination)方面的表現(xiàn)。結(jié)果如圖 8 所示,DIFF Transformer 在生成摘要和回答問(wèn)題時(shí)顯著提升了準(zhǔn)確率,減少了幻覺現(xiàn)象。這是因?yàn)椴罘肿⒁饬C(jī)制能夠準(zhǔn)確定位重要文段,避免無(wú)關(guān)上下文對(duì)模型預(yù)測(cè)的干擾。
圖 8. 利用文本摘要、問(wèn)答任務(wù)進(jìn)行幻覺評(píng)測(cè)
異常激活值分析
作者還發(fā)現(xiàn) DIFF Transformer 能夠顯著減少模型激活中的異常值,這為模型激活值的量化提供了新的可能性。實(shí)驗(yàn)表明,DIFF Transformer 在注意力激活值(attention logits)和隱藏狀態(tài)(hidden states)中的最大激活值顯著低于 Transformer。例如,在注意力激活值的 Top-1 激活值上,DIFF Transformer 比 Transformer 低了近 8 倍。利用這一性質(zhì),DIFF Transformer 在注意力激活值的低比特量化下的性能也優(yōu)于 Transformer,如圖 9 所示。
圖 9. 注意力激活值的低比特量化
數(shù)學(xué)推理能力
作者在數(shù)學(xué)推理任務(wù)上進(jìn)一步驗(yàn)證了 DIFF Transformer 的性能。作者采用兩階段訓(xùn)練,在 3B 預(yù)訓(xùn)練模型的基礎(chǔ)上進(jìn)行有監(jiān)督微調(diào),并在 MATH 等 8 個(gè)數(shù)學(xué)數(shù)據(jù)集上評(píng)測(cè)模型性能。在第一階段,采用 20B token 合成數(shù)學(xué)數(shù)據(jù)對(duì)模型進(jìn)行微調(diào),使模型獲得基礎(chǔ)數(shù)學(xué)能力,評(píng)測(cè)結(jié)果如圖 10 所示。從 15B token 開始,DIFF Transformer 展現(xiàn)出了顯著優(yōu)于 Transformer 的數(shù)學(xué)能力,至 20B token 結(jié)束的時(shí)候,準(zhǔn)確率的差距達(dá)到了 11% 左右。
圖 10. 第一階段數(shù)學(xué)合成數(shù)據(jù)微調(diào)
在第二階段,作者利用 Deepseek-R1 輸出所構(gòu)造的數(shù)據(jù)集 OpenThoughts-114K-Math 對(duì)模型進(jìn)行蒸餾,使模型更強(qiáng)大的深度推理能力。如圖 11 所示,在 8 個(gè)數(shù)據(jù)集上,DIFF Transformer 相較 Transformer 均有不同程度的提升,平均準(zhǔn)確率提升了 7.5%,這表明差分注意力機(jī)制更強(qiáng)大的上下文建模能力在推理任務(wù)中也至關(guān)重要。
圖 11. 第二階段深度推理能力評(píng)測(cè)
討論與未來(lái)工作
DIFF Transformer 自發(fā)布以來(lái)獲得了較大關(guān)注與討論。作者在 Hugging Face 論文討論平臺(tái)、alphaXiv 平臺(tái)上與社區(qū)開展了深入的探討。在 X 平臺(tái)(原 Twitter)上,Google DeepMind 高級(jí)研究科學(xué)家(Senior Staff Research Scientist)Petar Veli?kovi? 與作者就文章中的理論分析展開討論,ViT 核心作者 Lucas Beyer 也在閱讀文章后撰寫了一篇深入的論文總結(jié),相關(guān)發(fā)帖已獲得數(shù)十萬(wàn)瀏覽。目前 DIFF Transformer 也已集成至 Hugging Face 的 transformers 庫(kù)中。
- Hugging Face:https://huggingface.co/papers/2410.05258
- alphaXiv:https://www.alphaxiv.org/abs/2410.05258v1
- Petar Veli?kovi?:https://x.com/PetarV_93/status/1874820028975267866
- Lucas Beyer:https://x.com/giffmana/status/1873869654252544079
- transformers庫(kù):https://github.com/huggingface/transformers/tree/main/src/transformers/models/diffllama
未來(lái)工作方面,作者認(rèn)為可以利用 DIFF Transformer 的性質(zhì)設(shè)計(jì)低比特注意力算子,以及利用差分注意力的稀疏特性進(jìn)行鍵值緩存(key-value cache)的剪枝。此外,將 DIFF Transformer 應(yīng)用在除語(yǔ)言以外的其他模態(tài)上也值得探索。近期工作 DiffCLIP 將差分注意力擴(kuò)展至視覺、多模態(tài)領(lǐng)域,揭示了 DIFF Transformer 在不同模態(tài)任務(wù)中的更多結(jié)構(gòu)特性與應(yīng)用潛力。
- DiffCLIP:https://arxiv.org/abs/2503.06626
總結(jié)
本文的貢獻(xiàn)主要在兩個(gè)方面:
(1)DIFF Transformer 通過(guò)創(chuàng)新的差分注意力機(jī)制,有效解決了傳統(tǒng) Transformer 在處理文本時(shí)受到噪聲干擾、注意力分配不準(zhǔn)確的問(wèn)題;
(2)憑借對(duì)關(guān)鍵信息的關(guān)注和對(duì)噪聲的抵御能力,DIFF Transformer 在語(yǔ)言建模、長(zhǎng)文本建模、關(guān)鍵信息檢索、數(shù)學(xué)推理、對(duì)抗幻覺、上下文學(xué)習(xí)、模型激活值量化等任務(wù)中表現(xiàn)出色,有望在自然語(yǔ)言處理、多模態(tài)等領(lǐng)域作為基礎(chǔ)模型架構(gòu)。