新智元報(bào)道
編輯：英智 桃子
【新智元導(dǎo)讀】簡單的任務(wù)，傳統(tǒng)的Transformer卻錯誤率極高。Meta FAIR團(tuán)隊(duì)重磅推出多token注意力機(jī)制（MTA），精準(zhǔn)捕捉復(fù)雜信息，帶來模型性能飛升！
注意力機(jī)制就像大模型的「慧眼」，幫模型在大量文本中找到關(guān)鍵信息。
不過，每個注意力權(quán)重只取決于單個查詢和鍵向量的相似度。
找信息時，一次只能關(guān)注一個「小線索」，好比在龐大的圖書館里找一本書，卻只能檢索一個關(guān)鍵詞。
最近，Meta FAIR團(tuán)隊(duì)提出了多token注意力機(jī)制，為LLM的性能帶來質(zhì)的飛躍。
論文鏈接：http://arxiv.org/abs/2504.00927
作者還特別提到，這不是愚人節(jié)玩笑，而是一篇真實(shí)的論文。
在實(shí)際應(yīng)用中，很多時候相關(guān)的上下文信息沒辦法通過單一token確定。
比如，想查找一個同時提到「Alice」和「rabbit」的句子時，按照傳統(tǒng)注意力機(jī)制，得把這兩個信息壓縮到一個查詢向量里。
但這樣做不僅增加了向量編碼的難度，還可能導(dǎo)致信息丟失。
就算用不同注意力頭分別查找「Alice」和「rabbit」，也沒法有效整合這些注意力權(quán)重，這極大地限制了模型處理復(fù)雜信息的能力。
MTA是為了解決傳統(tǒng)注意力機(jī)制的這個問題而設(shè)計(jì)的。
MTA三大關(guān)鍵創(chuàng)新
MTA讓模型能同時依據(jù)多個查詢和鍵向量來確定注意力權(quán)重，利用更豐富的信息，實(shí)現(xiàn)精準(zhǔn)的注意力分配。
主要有三個關(guān)鍵部分：鍵-查詢卷積、頭混合卷積和帶深度縮放的組歸一化。
鍵-查詢卷積
鍵-查詢卷積是MTA的一大創(chuàng)新亮點(diǎn)。
通過二維卷積操作，模型可以把多個查詢和鍵token的信息放在一起，計(jì)算出更準(zhǔn)確的注意力權(quán)重：
其中是二維卷積操作，有自己的核權(quán)重和大小。
在傳統(tǒng)的注意力計(jì)算中，注意力權(quán)重僅由當(dāng)前的查詢和鍵向量對決定。
而鍵-查詢卷積打破了這種局限，會綜合考慮附近查詢和鍵的信息。
就拿「Where did Alice see the rabbit?」這個問題來說，要是用傳統(tǒng)注意力機(jī)制，很難同時關(guān)注到「Alice」和「rabbit」這兩個關(guān)鍵信息。
鍵-查詢卷積通過合理設(shè)置卷積核大小，能讓模型同時捕捉到這兩個信息。
具體計(jì)算時，它會在鍵和查詢的長度維度上進(jìn)行卷積操作。為了不讓未來的信息搗亂，計(jì)算時只使用過去的查詢。
實(shí)際操作中，為了簡化流程，采用了一種更簡便的雙重掩碼方法。
鍵-查詢卷積有兩種方式，一種是在softmax之前進(jìn)行卷積，另一種是在softmax之后。
在softmax之前卷積，能更好地融合不同查詢和鍵的信息。
在softmax之后卷積，注意力權(quán)重之間的相互作用就變成加法了：
在實(shí)驗(yàn)中，默認(rèn)用的是softmax之前卷積的方式。
信息共享：頭混合卷積
除了鍵-查詢卷積，MTA還有個很特別的頭混合卷積，它的作用是在不同注意力頭之間共享信息。
在傳統(tǒng)的多頭注意力機(jī)制里，各個頭之間的信息相對獨(dú)立。
頭混合卷積就像一座橋梁，打破了這種孤立，讓不同頭的注意力權(quán)重可以分享信息。
具體計(jì)算時，混合注意力權(quán)重有兩種方式，一種是在softmax之后，像
是卷積核權(quán)重。
另一種是在softmax之前，對注意力的對數(shù)幾率進(jìn)行混合：
通過信息共享，模型能從多個角度綜合信息，理解文本的能力就更強(qiáng)了。
整合關(guān)鍵組件
前面介紹了兩種混合注意力權(quán)重的方式，MTA把鍵-查詢卷積和頭混合卷積結(jié)合，形成了強(qiáng)大的信息處理系統(tǒng)。
如果兩種混合方法都是在softmax之前，那么它們可以通過單個三維卷積操作實(shí)現(xiàn)，如圖所示。
模型訓(xùn)練時，隨著層數(shù)越來越多，會出現(xiàn)一個問題，就是殘差流會變得很大，這讓梯度傳遞不太順暢，影響模型的訓(xùn)練效果。
帶深度縮放的組歸一化就是解決這個問題的。它會對每個頭分別進(jìn)行歸一化操作，而且會根據(jù)層的深度進(jìn)行調(diào)整。
這樣能讓模型訓(xùn)練更穩(wěn)定，梯度傳遞得更好。
MTA機(jī)制革新Transformer
理論上MTA機(jī)制很厲害，實(shí)際效果如何呢？
為了驗(yàn)證MTA的有效性，研究人員開展了一系列實(shí)驗(yàn)，涵蓋了從簡單任務(wù)到復(fù)雜的長上下文任務(wù)等。
簡單任務(wù)：小任務(wù)凸顯大優(yōu)勢
有一個簡單的任務(wù)，MTA在這個任務(wù)里的表現(xiàn)比傳統(tǒng)Transformer好太多了。
模型需要在由多個隨機(jī)字母組成的塊序列中，找到包含特定幾個字母的目標(biāo)塊，并輸出相關(guān)信息。
這個任務(wù)看似簡單，卻能精準(zhǔn)暴露傳統(tǒng)注意力機(jī)制的問題。
傳統(tǒng)Transformer不能很好地把多個信息整合起來，要把兩個問題字母的信息壓縮到一個查詢向量中，這對它來說很困難，錯誤率極高。
MTA在這個任務(wù)上的表現(xiàn)堪稱驚艷。
它借助鍵-查詢卷積，先分別找到每個問題字母在序列中的位置，然后通過卷積把這些信息整合起來，精準(zhǔn)定位目標(biāo)塊。
實(shí)驗(yàn)結(jié)果令人驚喜，MTA幾乎以零錯誤率完成了任務(wù)，這充分展現(xiàn)了MTA在處理多信息檢索任務(wù)時的實(shí)力，與傳統(tǒng)注意力機(jī)制相比，有著跨越式的進(jìn)步。
語言建模：全面提升
在語言建模實(shí)驗(yàn)中，研究團(tuán)隊(duì)訓(xùn)練了880M參數(shù)的模型，并與傳統(tǒng)Transformer模型、差分Transformer（DIFF Transformer）做比較。
所有模型都在SlimPajama數(shù)據(jù)集上，用Lingua框架訓(xùn)練。
為提高訓(xùn)練效率，MTA在每第4層應(yīng)用鍵-查詢卷積，頭卷積應(yīng)用于所有層，并固定了卷積核的維度。
實(shí)驗(yàn)結(jié)果令人眼前一亮。在驗(yàn)證困惑度方面，用MTA訓(xùn)練的模型在各個驗(yàn)證數(shù)據(jù)集上都表現(xiàn)更好。
帶層縮放的組歸一化對MTA和DIFF Transformer的性能提升特別重要。
在一些流行的基準(zhǔn)測試中，如BoolQ、PIQA等，MTA模型也超過了基線模型，平均得分更高。
這說明MTA能有效提升模型性能，無論是理解文本含義，還是回答問題，都更出色。
長文本處理的好幫手
研究團(tuán)隊(duì)對模型進(jìn)行了長上下文微調(diào)，把上下文的長度從2048增加到4096，還調(diào)整了一些訓(xùn)練參數(shù)。
用了MTA的模型在困惑度評估里比基線模型好很多。
在LAMBADA任務(wù)里，MTA模型預(yù)測下一個單詞的時候，能更好地利用上下文信息，預(yù)測得更加準(zhǔn)確。
在寫作助手、摘要生成等場景中，MTA能幫助模型更好地理解文章內(nèi)容和邏輯，生成更符合要求的回復(fù)。
精準(zhǔn)找到長距離信息
像「大海撈針」（Needle-In-A-Haystack）和BabiLong這樣的長距離依賴任務(wù)里，MTA的優(yōu)勢就更明顯了。
「大海撈針」任務(wù)要在很長的文本里找到特定的信息。
MTA模型在這個任務(wù)中表現(xiàn)特別好，無論是在2k還是4k的上下文窗口中，準(zhǔn)確率都有顯著提升。
尤其是在查找隱藏較深的目標(biāo)信息時，優(yōu)勢更為突出。
BabiLong任務(wù)評估模型能不能理解長文本里分散的各種事實(shí)，并進(jìn)行推理。
QA1-5任務(wù)中，MTA模型在有很多干擾文本的情況下，也能保持較高的準(zhǔn)確率，準(zhǔn)確找到關(guān)鍵信息，并進(jìn)行推理。
消融實(shí)驗(yàn)
為了搞清楚MTA各個組件的作用，研究人員做了消融實(shí)驗(yàn)。
在鍵-查詢卷積實(shí)驗(yàn)里，就算只有2層用了MTA增強(qiáng)，模型就能超過強(qiáng)大的基線模型，6層MTA在性能和復(fù)雜程度之間達(dá)到了較好的平衡。
在卷積核初始化的實(shí)驗(yàn)里，用單位矩陣初始化的MTA模型訓(xùn)練時收斂得更快，性能也更好。
組歸一化和指數(shù)深度縮放對于提升模型性能很關(guān)鍵，不同大小的卷積核雖然會影響評估結(jié)果，但總體的卷積核模式相似。
改變卷積操作和softmax的先后順序，對模型性能的影響較小。
作者介紹
Meta多token注意力論文中，論文二作是一位華人研究科學(xué)家Tianlu Wang。
她曾獲得了弗吉尼亞大學(xué)計(jì)算機(jī)科學(xué)博士學(xué)位，導(dǎo)師是Vicente Ordó?ez Román教授。在此之前，她還獲得了浙大計(jì)算機(jī)科學(xué)學(xué)士學(xué)位。
Tianlu Wang研究興趣在于，與探索機(jī)器學(xué)習(xí)模型中的公平性、魯棒性和問責(zé)制相關(guān)話題，尤其是在計(jì)算機(jī)視覺和自然語言處理系統(tǒng)方面。
參考資料：
http://arxiv.org/abs/2504.00927
https://x.com/jaseweston/status/1907260086017237207
熱門跟貼