新智元報道

編輯：英智 桃子

【新智元導讀】簡單的任務，傳統的Transformer卻錯誤率極高。Meta FAIR團隊重磅推出多token注意力機制（MTA），精準捕捉復雜信息，帶來模型性能飛升！

注意力機制就像大模型的「慧眼」，幫模型在大量文本中找到關鍵信息。

不過，每個注意力權重只取決于單個查詢和鍵向量的相似度。

找信息時，一次只能關注一個「小線索」，好比在龐大的圖書館里找一本書，卻只能檢索一個關鍵詞。

最近，Meta FAIR團隊提出了多token注意力機制，為LLM的性能帶來質的飛躍。

論文鏈接：http://arxiv.org/abs/2504.00927

作者還特別提到，這不是愚人節玩笑，而是一篇真實的論文。

在實際應用中，很多時候相關的上下文信息沒辦法通過單一token確定。

比如，想查找一個同時提到「Alice」和「rabbit」的句子時，按照傳統注意力機制，得把這兩個信息壓縮到一個查詢向量里。

但這樣做不僅增加了向量編碼的難度，還可能導致信息丟失。

就算用不同注意力頭分別查找「Alice」和「rabbit」，也沒法有效整合這些注意力權重，這極大地限制了模型處理復雜信息的能力。

MTA是為了解決傳統注意力機制的這個問題而設計的。

MTA三大關鍵創新

MTA讓模型能同時依據多個查詢和鍵向量來確定注意力權重，利用更豐富的信息，實現精準的注意力分配。

主要有三個關鍵部分：鍵-查詢卷積、頭混合卷積和帶深度縮放的組歸一化。

鍵-查詢卷積
鍵-查詢卷積是MTA的一大創新亮點。

通過二維卷積操作，模型可以把多個查詢和鍵token的信息放在一起，計算出更準確的注意力權重：

其中是二維卷積操作，有自己的核權重和大小。

在傳統的注意力計算中，注意力權重僅由當前的查詢和鍵向量對決定。

而鍵-查詢卷積打破了這種局限，會綜合考慮附近查詢和鍵的信息。

就拿「Where did Alice see the rabbit?」這個問題來說，要是用傳統注意力機制，很難同時關注到「Alice」和「rabbit」這兩個關鍵信息。

鍵-查詢卷積通過合理設置卷積核大小，能讓模型同時捕捉到這兩個信息。

具體計算時，它會在鍵和查詢的長度維度上進行卷積操作。為了不讓未來的信息搗亂，計算時只使用過去的查詢。

實際操作中，為了簡化流程，采用了一種更簡便的雙重掩碼方法。

鍵-查詢卷積有兩種方式，一種是在softmax之前進行卷積，另一種是在softmax之后。

在softmax之前卷積，能更好地融合不同查詢和鍵的信息。

在softmax之后卷積，注意力權重之間的相互作用就變成加法了：

在實驗中，默認用的是softmax之前卷積的方式。

信息共享：頭混合卷積

除了鍵-查詢卷積，MTA還有個很特別的頭混合卷積，它的作用是在不同注意力頭之間共享信息。

在傳統的多頭注意力機制里，各個頭之間的信息相對獨立。

頭混合卷積就像一座橋梁，打破了這種孤立，讓不同頭的注意力權重可以分享信息。

具體計算時，混合注意力權重有兩種方式，一種是在softmax之后，像

是卷積核權重。

另一種是在softmax之前，對注意力的對數幾率進行混合：

通過信息共享，模型能從多個角度綜合信息，理解文本的能力就更強了。

整合關鍵組件

前面介紹了兩種混合注意力權重的方式，MTA把鍵-查詢卷積和頭混合卷積結合，形成了強大的信息處理系統。

如果兩種混合方法都是在softmax之前，那么它們可以通過單個三維卷積操作實現，如圖所示。

模型訓練時，隨著層數越來越多，會出現一個問題，就是殘差流會變得很大，這讓梯度傳遞不太順暢，影響模型的訓練效果。

帶深度縮放的組歸一化就是解決這個問題的。它會對每個頭分別進行歸一化操作，而且會根據層的深度進行調整。

這樣能讓模型訓練更穩定，梯度傳遞得更好。

MTA機制革新Transformer

理論上MTA機制很厲害，實際效果如何呢？

為了驗證MTA的有效性，研究人員開展了一系列實驗，涵蓋了從簡單任務到復雜的長上下文任務等。

簡單任務：小任務凸顯大優勢

有一個簡單的任務，MTA在這個任務里的表現比傳統Transformer好太多了。

模型需要在由多個隨機字母組成的塊序列中，找到包含特定幾個字母的目標塊，并輸出相關信息。

這個任務看似簡單，卻能精準暴露傳統注意力機制的問題。

傳統Transformer不能很好地把多個信息整合起來，要把兩個問題字母的信息壓縮到一個查詢向量中，這對它來說很困難，錯誤率極高。

MTA在這個任務上的表現堪稱驚艷。

它借助鍵-查詢卷積，先分別找到每個問題字母在序列中的位置，然后通過卷積把這些信息整合起來，精準定位目標塊。

實驗結果令人驚喜，MTA幾乎以零錯誤率完成了任務，這充分展現了MTA在處理多信息檢索任務時的實力，與傳統注意力機制相比，有著跨越式的進步。

語言建模：全面提升

在語言建模實驗中，研究團隊訓練了880M參數的模型，并與傳統Transformer模型、差分Transformer（DIFF Transformer）做比較。

所有模型都在SlimPajama數據集上，用Lingua框架訓練。

為提高訓練效率，MTA在每第4層應用鍵-查詢卷積，頭卷積應用于所有層，并固定了卷積核的維度。

實驗結果令人眼前一亮。在驗證困惑度方面，用MTA訓練的模型在各個驗證數據集上都表現更好。

帶層縮放的組歸一化對MTA和DIFF Transformer的性能提升特別重要。

在一些流行的基準測試中，如BoolQ、PIQA等，MTA模型也超過了基線模型，平均得分更高。

這說明MTA能有效提升模型性能，無論是理解文本含義，還是回答問題，都更出色。

長文本處理的好幫手

研究團隊對模型進行了長上下文微調，把上下文的長度從2048增加到4096，還調整了一些訓練參數。

用了MTA的模型在困惑度評估里比基線模型好很多。

在LAMBADA任務里，MTA模型預測下一個單詞的時候，能更好地利用上下文信息，預測得更加準確。

在寫作助手、摘要生成等場景中，MTA能幫助模型更好地理解文章內容和邏輯，生成更符合要求的回復。

精準找到長距離信息

像「大海撈針」（Needle-In-A-Haystack）和BabiLong這樣的長距離依賴任務里，MTA的優勢就更明顯了。

「大海撈針」任務要在很長的文本里找到特定的信息。

MTA模型在這個任務中表現特別好，無論是在2k還是4k的上下文窗口中，準確率都有顯著提升。

尤其是在查找隱藏較深的目標信息時，優勢更為突出。

BabiLong任務評估模型能不能理解長文本里分散的各種事實，并進行推理。

QA1-5任務中，MTA模型在有很多干擾文本的情況下，也能保持較高的準確率，準確找到關鍵信息，并進行推理。

消融實驗

為了搞清楚MTA各個組件的作用，研究人員做了消融實驗。

在鍵-查詢卷積實驗里，就算只有2層用了MTA增強，模型就能超過強大的基線模型，6層MTA在性能和復雜程度之間達到了較好的平衡。

在卷積核初始化的實驗里，用單位矩陣初始化的MTA模型訓練時收斂得更快，性能也更好。

組歸一化和指數深度縮放對于提升模型性能很關鍵，不同大小的卷積核雖然會影響評估結果，但總體的卷積核模式相似。

改變卷積操作和softmax的先后順序，對模型性能的影響較小。

作者介紹

Meta多token注意力論文中，論文二作是一位華人研究科學家Tianlu Wang。

她曾獲得了弗吉尼亞大學計算機科學博士學位，導師是Vicente Ordó?ez Román教授。在此之前，她還獲得了浙大計算機科學學士學位。

Tianlu Wang研究興趣在于，與探索機器學習模型中的公平性、魯棒性和問責制相關話題，尤其是在計算機視覺和自然語言處理系統方面。

參考資料：

http://arxiv.org/abs/2504.00927

https://x.com/jaseweston/status/1907260086017237207