清華微軟最新力作：用物理學(xué)革新Transformer注意力，「大海撈針」精度暴漲30%！

作者：新智元 2024-10-09 13:22:10

隨著諾貝爾物理學(xué)獎(jiǎng)?lì)C給了「機(jī)器學(xué)習(xí)之父」Geoffrey Hinton，另一個(gè)借鑒物理學(xué)概念的模型架構(gòu)也橫空出世——微軟清華團(tuán)隊(duì)的最新架構(gòu)Differential Transformer，從注意力模塊入手，實(shí)現(xiàn)了Transformer的核心能力提升。

隨著近些年來NLP領(lǐng)域研究的不斷深入，我們逐漸發(fā)現(xiàn)，Transformer架構(gòu)中出現(xiàn)的幻覺問題，以及各種下游任務(wù)中的性能不足，都或多或少與注意力缺陷有關(guān)。

雖然上下文窗口可以擴(kuò)展，但是Transformer還是無法真正關(guān)注到有價(jià)值的信息。

最近，微軟研究院和清華大學(xué)的研究人員共同提出了一種新的模型架構(gòu)——Differential Transformer，不僅保留了原始Transformer中的可擴(kuò)展性，也能讓模型更加關(guān)注上下文中與任務(wù)相關(guān)的關(guān)鍵信息。

實(shí)驗(yàn)表明，注意力機(jī)制的改進(jìn)，不僅顯著提升了檢索精度，還能緩解LLM的幻覺。

論文地址：https://arxiv.org/abs/2410.05258

Transformer的困境

眾所周知，Transformer的核心是注意力機(jī)制，采用softmax函數(shù)來衡量序列中各種標(biāo)記的重要性。然而，最近的研究表明，LLM難以從上下文中準(zhǔn)確到檢索關(guān)鍵信息。

比如去年斯坦福Percy Liang團(tuán)隊(duì)的一篇論文就指出，雖然語(yǔ)言模型能夠接受較長(zhǎng)的上下文作為輸入，但并不能穩(wěn)健地利用長(zhǎng)輸入上下文中的信息。

論文地址：https://arxiv.org/abs/2307.03172

比如，實(shí)驗(yàn)中發(fā)現(xiàn)，僅僅改變關(guān)鍵信息在文檔中的出現(xiàn)位置，就可以造成GPT-3.5 Turbo檢索性能的大范圍波動(dòng)。

此外，本篇論文的實(shí)驗(yàn)結(jié)果證明，Transformer經(jīng)常過度關(guān)注不相關(guān)的上下文，本文將其稱之為「注意力噪聲」。

如圖1（左）所示，模型分配給正確答案的注意力分?jǐn)?shù)很低，同時(shí)不成比例地關(guān)注不相關(guān)的上下文，這意味著信噪比很低，最終淹沒了正確答案。

由此看來，我們對(duì)于LLM檢索、利用長(zhǎng)上下文的過程，知之甚少，其注意力過程也需要更多的改進(jìn)。

本文所提出的Differential Transformer（DIFF Transformer）正是希望用「差分注意力」（differential attention）機(jī)制消除注意力噪聲，促使模型關(guān)注上下文中的關(guān)鍵信息。

圖1的對(duì)比結(jié)果可以看出，DIFF Transformer給出的注意力分?jǐn)?shù)的分布明顯不同于傳統(tǒng)Transformer架構(gòu)，給予關(guān)鍵信息更高的注意力分?jǐn)?shù)，進(jìn)而顯著提升了檢索能力。

這種能力的提升，對(duì)于有效利用LLM的長(zhǎng)上下文窗口、緩解幻覺、關(guān)鍵信息檢索等方面都有重要的意義。

模型架構(gòu)

DIFF Transformer也可以用于純Encoder或Encoder-Decoder模型，但本篇論文以純Decoder模型為例進(jìn)行描述。

整個(gè)模型由L個(gè)DIFF Transformer層堆疊而成，每層由一個(gè)差分注意力模塊和前饋網(wǎng)絡(luò)模塊連接形成。

宏觀布局類似于傳統(tǒng)Transformer架構(gòu)，但主要區(qū)別在于修改了注意力的softmax過程，并且采用了pre-RMSNorm、SwiGLU等改進(jìn)。

差分注意力

該模塊的結(jié)構(gòu)示意圖和偽代碼如圖2所示，具體的代碼實(shí)現(xiàn)可參考項(xiàng)目GitHub。

代碼地址：https://github.com/microsoft/unilm/tree/master/Diff-Transformer

除了傳統(tǒng)注意力中的權(quán)重矩陣W^Q、W^K、W^V ∈ ?^{d_model×2?d}，模塊中還加入了可學(xué)習(xí)標(biāo)量λ。

具體來說，給定輸入序列X ∈ ?^{N×d_model}，首先將其投影為Q、K、V矩陣Q_1,Q_2,K_1,K_2 ∈ ?^{N×d} , V ∈ ?^{N×2?d}，然后是差分注意力算子DiffAttn(·)通過公式（1）計(jì)算輸出：

λ被初始化為常量λ_{init} ∈ (0,1)，并依照公式（2）與其他權(quán)重參數(shù)同步更新：

其中，λ_???? , λ_???? , λ_???? , λ_???? ∈ ?^d也都是是可學(xué)習(xí)向量。

之所以命名為「差分注意力」，是指兩個(gè)softmax函數(shù)間的差異可以消除注意力噪音。

這個(gè)想法類似于電氣工程中提出的差分放大器（differential amplifiler），將兩個(gè)信號(hào)之間的差異作為輸出，從而消除輸入中的共模噪聲；降噪耳機(jī)的設(shè)計(jì)也是基于類似的思路。

DIFF Transformer中也可以使用多頭注意力機(jī)制，在同一層的多個(gè)head間共享參數(shù)λ，將每個(gè)head的輸出進(jìn)行歸一化處理后再拼接、投影，就得到了最終輸出，如公式（3）所示。

公式（3）中的LN（·）是指對(duì)每個(gè)頭使用RMSNorm，但如圖2（左）所示，也可以使用GroupNorm。

加上前饋網(wǎng)絡(luò)模塊，每個(gè)DIFF Transformer層就可以描述為：

實(shí)驗(yàn)

下游任務(wù)

首先，研究人員在1T token上訓(xùn)練3B大小的DIFF Transformer模型，并在各種下游任務(wù)上與之前有競(jìng)爭(zhēng)力的Transformer架構(gòu)模型進(jìn)行比較，結(jié)果如表1所示。

基線模型大小都為3B，其中，StableLM-3B-4E1T的1T結(jié)果取自技術(shù)報(bào)告，而OpenLLaMA-v2-3B和StableLM-base-alpha-3B-v2同樣使用1T數(shù)據(jù)訓(xùn)練，表中分?jǐn)?shù)為Eval Harness基準(zhǔn)測(cè)試上的零樣本準(zhǔn)確率。

結(jié)果顯示，，與之前經(jīng)過精心調(diào)優(yōu)的Transformer語(yǔ)言模型相比，DIFF Transformer取得了良好的性能。

尤其是對(duì)于長(zhǎng)上下文任務(wù)，如圖4所示，隨著上下文長(zhǎng)度不斷增加，累計(jì)平均的負(fù)對(duì)數(shù)似然值（NLL）持續(xù)降低，說明Diff Transformer可以更有效地利用不斷增加的上下文。

關(guān)鍵信息檢索

「大海撈針」（Needle-In-A-Haystack）測(cè)試被廣泛用于評(píng)估LLM提取長(zhǎng)上下文中的關(guān)鍵信息的能力。

本文的實(shí)驗(yàn)遵循LWM和Gemini 1.5的「多針」評(píng)估方案，在不同長(zhǎng)度的上下文中，N根針被插入不同的深度。每根「針」都由一個(gè)簡(jiǎn)潔的句子組成，為特定城市分配一個(gè)獨(dú)特的魔法數(shù)字。

答案針被放置在上下文中的5個(gè)不同深度：0%、25%、50%、75%和100%，同時(shí)隨機(jī)放置其他分散注意力的針。待測(cè)LLM的目標(biāo)，就是是檢索與查詢城市相對(duì)應(yīng)的數(shù)字。

4k上下文檢索的可結(jié)果如表2所示。雖然兩種模型在N=1或N=2時(shí)都取得了良好的準(zhǔn)確率，但隨著N的增加，DIFF Transformer的性能保持相對(duì)一致，Transformer則顯著下降。

4K長(zhǎng)度的平均檢索精度，N代表針數(shù)，R表示查詢城市的數(shù)量

將上下文長(zhǎng)度擴(kuò)展至64k時(shí)，差距就更加明顯，尤其是關(guān)鍵信息位于前半部分時(shí)（即0%、25% 和 50%深度）。

特別是，將針放置在25%深度時(shí)，DIFF Transformer相對(duì)于傳統(tǒng)Transformer實(shí)現(xiàn)了76%的精度提升。

除了檢索精度，表3進(jìn)一步分析了兩種模型為上下文分配的注意力分?jǐn)?shù)。可以看出， DIFF Transformer的確將更多的注意力分配給了有用的信息，并有效地消除注意力噪聲。

值得注意的是，DIFF Transformer在提升檢索精度的同時(shí)也緩解了幻覺現(xiàn)象。

實(shí)驗(yàn)包含模型在總結(jié)（圖4a）和問答（圖4b）兩種任務(wù)上的幻覺評(píng)估。可以發(fā)現(xiàn)，與Transformer相比，DIFF Transformer的上下文幻覺明顯減輕。

這種性能的提高可能源于，改進(jìn)后的注意力模塊能更好第關(guān)注任務(wù)所需信息，而非不相關(guān)的上下文。

這與之前研究中的觀察結(jié)果一致，即Transformer出現(xiàn)上下文幻覺的一個(gè)主要原因是注意力分?jǐn)?shù)的錯(cuò)誤分配。

對(duì)文本摘要和問題回答的幻覺評(píng)估。準(zhǔn)確度越高表示幻覺越少；評(píng)估時(shí)采用GPT-4o進(jìn)行自動(dòng)化的二元判斷

縮放特性

除了下游任務(wù)性能，論文還進(jìn)行了縮放特性的對(duì)比。

擴(kuò)展模型規(guī)模

如圖3a所示，分別使用830M、1.4B、2.8B、6.8B和13.1B參數(shù)訓(xùn)練語(yǔ)言模型，發(fā)現(xiàn)DIFF Transformer依舊遵循Scaling Law。

根據(jù)擬合曲線，68億參數(shù)規(guī)模的DIFF Transformer達(dá)到了與110億參數(shù)規(guī)模Transformer相當(dāng)?shù)尿?yàn)證損失，但僅需62.2%的參數(shù)。

同樣，78億參數(shù)的DIFF Transformer匹配了131億參數(shù)的Transformer的性能，參數(shù)量是后者的59.5%。

擴(kuò)展訓(xùn)練Token

如圖3b所示，訓(xùn)練數(shù)據(jù)的縮放也遵循類似規(guī)律，且擬合曲線表明，使用160B token訓(xùn)練的DIFF Transformer達(dá)到了與使用251B token訓(xùn)練的Transformer相當(dāng)?shù)男阅埽珒H消耗了63.7%的訓(xùn)練數(shù)據(jù)。