谷歌又有新的注意力了！

他們提出的新架構(gòu)參數(shù)減少40%，訓(xùn)練速度較RNN提升5-8倍，在某些任務(wù)上性能甚至Transformer好7.2%！

在大語言模型（LLMs）中，他們引入了新的注意力偏向策略，并重新構(gòu)想了「遺忘」這一過程，用「保留」來取而代之。

所謂的「注意力偏向」現(xiàn)象，是指人類天然傾向于優(yōu)先處理特定事件或刺激

受人類認知中的「關(guān)聯(lián)記憶」（associative memory）與「注意力偏向」（attentional bias）概念啟發(fā)，谷歌的團隊提出了統(tǒng)一視角：

Transformer與RNN，都可以被看作是優(yōu)化某種「內(nèi)在記憶目標」（即注意力偏向），從而學習鍵值映射的關(guān)聯(lián)記憶系統(tǒng)。

他們發(fā)現(xiàn)：

1.幾乎所有現(xiàn)代序列模型的底層學習過程，都可以歸結(jié)為關(guān)聯(lián)記憶機制；

2.所謂的遺忘機制，本質(zhì)上是一種對注意力偏向的正則化操作；

3.不同模型之間的差異，可以用「注意力偏向+保留機制」這一組合來解釋。

為此，他們把這一切都被整合進了名為Miras的新框架中，提供四個關(guān)鍵設(shè)計維度，指導(dǎo)下一代序列模型的構(gòu)建。

1.記憶架構(gòu) — 如何構(gòu)建記憶，決定了模型的記憶能力，比如向量、矩陣、MLP等；

2.注意力偏向 — 模型如何集中注意力，負責建模潛在的映射模式；

3.保留門控 — 如何平衡學習新概念和保留已學概念；

4.記憶學習算法 — 模型如何訓(xùn)練，負責記憶管理，比如梯度下降、牛頓法等。

圖1：Miras框架概述

這次他們，一口氣提出了三種新型序列模型，在某些任務(wù)上甚至超越了超越Transformer。

這三種新模型——Moneta、Yaad和Memora，超越了現(xiàn)有線性遞歸神經(jīng)網(wǎng)絡(luò)的能力，同時保持快速可并行訓(xùn)練的過程。

新模型各有所長，在特定任務(wù)中表現(xiàn)卓越：

? Moneta：在語言建模任務(wù)中PPL指標提升23%  

? Yaad：常識推理準確率達89.4%（超越Transformer7.2%）  

? Memora：記憶密集型任務(wù)召回率提升至91.8%

在多個任務(wù)上，新模型提升明顯：

? 在PG19長文本建模任務(wù)中，參數(shù)量減少40%情況下保持相當性能

? 線性計算復(fù)雜度使訓(xùn)練速度較傳統(tǒng)RNN提升5-8倍

? 在CLUTRR關(guān)系推理基準上創(chuàng)造92.3%的新SOTA紀錄

論文鏈接：https://arxiv.org/abs/2504.13173

模型沒有失憶，但也有問題 

研究者定義并形式化了注意力偏向的概念，作為序列模型的內(nèi)部記憶目標，旨在學習輸入（即鍵和值）之間的潛在映射。

廣義上講，關(guān)聯(lián)記憶是將一組鍵K映射到一組值V的操作符（Operator）。

為了學習數(shù)據(jù)中的潛在映射模式，它需要一個目標，該目標針對某種類型的記憶并衡量學習到的映射質(zhì)量：

研究人員不再用「遺忘」（forget）這個詞，而是提出了「保留」（retention）的概念。

因此，「遺忘門」（forget gate）也就變成了「保留門」（retention gate）。

模型并不會真的清除過去的記憶——  

它只是選擇對某些信息不那么「上心」而已。

此外，研究人員提供了一套全新的替代保留門控（忘記門）用于序列模型，帶來了新的洞察，幫助平衡學習新概念和保留先前學到的概念。

現(xiàn)有的深度學習架構(gòu)中的遺忘機制，可以重新解釋為一種針對注意力偏向的??正則化。

比如，softmax注意力是Miras的一個實例，利用Nadaraya-Watson估計器找到MSE損失的非參數(shù)解時，無需保留項。

論文鏈接：https://arxiv.org/abs/2407.04620

實際上，這次谷歌團隊發(fā)現(xiàn)大多數(shù)現(xiàn)有模型（如Transformer、RetNet、Mamba等）都采用了類似的注意力偏向目標，即嘗試最小化鍵值對之間的?? 范數(shù)（均方誤差）。

但它存在幾個問題：

• 對異常值敏感：極端或錯誤輸入可能嚴重干擾記憶更新
• 不支持可調(diào)節(jié)的保留策略：不同任務(wù)/token 重要性不同，不能一視同仁
• 無法應(yīng)對復(fù)雜上下文需求：長文檔、多語義層、跨段落推理等任務(wù)對注意力機制要求更高

表1：基于Miras框架視角的近期序列模型概覽

目標函數(shù)：注意力偏向策略

基于關(guān)聯(lián)記憶概念的神經(jīng)架構(gòu)設(shè)計，被轉(zhuǎn)化為學習鍵值之間的基本映射，可以利用最小化目標函數(shù)L來實現(xiàn):

為了求解上述優(yōu)化問題，最簡單的方法就是利用梯度下降。

具體來說，給定一對新的鍵值對，可以通過以下方式更新記憶（一下叫做更新方程）：

這一公式可以被重新解釋為一種瞬時驚訝度度量，其中模型記憶那些違反目標預(yù)期的token。

更新方程可以看作是在線梯度下降的一步，涉及損失函數(shù)序列的優(yōu)化：

眾所周知，在線梯度下降可以被視為跟蹤正則化領(lǐng)導(dǎo)者（Follow-The-Regularized-Leader, FTRL） 算法的一個特例。

這其實對應(yīng)于某些特定選擇的損失函數(shù)。

具體來說，假設(shè)W? = 0，則更新方程中的更新規(guī)則等價于下列方程（以后稱為二次更新方程）：

以上方程使用了損失函數(shù)的線性近似和二次正則化。

然而，從原則上講，也可以使用其他損失函數(shù)的近似以及其他正則化函數(shù)。

更具體地說，可以將二次更新方程推廣到如下形式：

其中：

• 第一項是注意力偏向（Attentional Bias）的和；
• 最后一項是記憶穩(wěn)定性（Memory Stability） 正則化項。

不同的損失函數(shù)和正則化項，對應(yīng)不同的算法。

在這種情況下，記憶的更新不僅依賴于當前輸入數(shù)據(jù)的特征，還受到記憶結(jié)構(gòu)的影響，正則化項在其中起到了平衡學習和記憶穩(wěn)定性的作用。

Miras提出的三類新型注意力偏向策略。

??范數(shù)：記憶精度可調(diào)

如正文所述?2回歸損失通常是自然選擇，但其對數(shù)據(jù)噪聲較為敏感。

自然的擴展是采用???范數(shù)目標函數(shù)類。

具體而言，設(shè)M為記憶模塊，k為鍵集合，v為值集合，???注意力偏向定義為：

不同的范數(shù)對應(yīng)對噪聲的敏感度：

??更抗異常值，

??是常規(guī)選擇，

?∞ 聚焦于最大誤差。

Huber損失：「應(yīng)對異常」心理機制

Huber損失具備容錯機制的記憶模塊。

盡管?2范數(shù)目標是許多統(tǒng)計與機器學習任務(wù)的常見選擇，但其對異常值和極端樣本的敏感性眾所周知。

這種敏感性同樣存在于將?2損失用于注意力偏向的場景。

為解決該問題，并借鑒穩(wěn)健回歸的思路，研究者建議采用Huber損失類型作為注意力偏向，從而降低異常數(shù)據(jù)對記憶學習過程的負面影響。

Huber損失結(jié)合了??（正常情況下）和??（出現(xiàn)大誤差時），在面對異常值時也能保持學習的穩(wěn)定性。

魯棒優(yōu)化：考慮最壞情況

魯棒優(yōu)化（Robust Optimization）的核心思想：最小化最壞情況下的損失；在一個不確定性集合（uncertainty set）內(nèi)優(yōu)化性能。

• 不只是優(yōu)化當前值，而是對 可能擾動做最壞情況準備
• 提高模型應(yīng)對微小輸入變化的魯棒性，適用于噪聲或?qū)剐暂斎氕h(huán)境

類似「備份記憶」策略——即使現(xiàn)實偏離，也不崩盤。

魯棒優(yōu)化使模型在輸入有小幅變動時也能保持穩(wěn)定。

正則化：保留門策略

在多數(shù)傳統(tǒng)模型中（如 LSTM、Mamba、Transformer），信息的遺忘或記憶更新是隱式的，模型只是不斷地「覆蓋」舊狀態(tài)。

但現(xiàn)實中，大家知道：

并不是所有信息都值得被長期記住，有些應(yīng)該快速遺忘，有些則必須深深保留。

因此，Miras 框架提出了一個明確的設(shè)計目標：

引入可控的、可設(shè)計的保留機制 Retention Gate，使模型顯式判斷是否保留舊記憶。

這就是Retention Gate的作用核心。

另一種解讀的方法是，將更新方程視為從最新的鍵值對（k?, v?）中學習（通過使用其梯度或驚訝度度量），同時保持接近先前狀態(tài) W???，以保留先前記憶的token。

這種形式可以推廣為：

其中，右側(cè)第一項是?(W; k_t, v_t) 的近似，最小化它對應(yīng)于從新概念（k?, v?）中學習。

第二項則對W的變化進行正則化，以使學習動態(tài)穩(wěn)定，并保留先前學到的知識。

Retention函數(shù)可能包括局部和全局組件：

其中：

• 第一項是一種預(yù)度量，用于控制W_{t-1}的偏差，旨在保留先前學到的知識。
• 系數(shù)η?可以被視為一種元上下文學習率，其中較大的η?值意味著從新概念中學習更多，同時允許對先前學到的概念有更高的遺忘率。
• 第二項是全局保留，它控制與記憶大小相關(guān)的變化。

從目標函數(shù)角度，保留門對應(yīng)正則項。

基于概率的機制：將記憶處理為概率分布（比如用KL散度）來保持其穩(wěn)定性。

彈性網(wǎng)（Elastic net）：結(jié)合了軟遺忘（??）和硬遺忘（??）的方法。

Lq穩(wěn)定性：可調(diào)節(jié)記憶對變化的抵抗程度。

Bregman散度：引入非線性、能感知數(shù)據(jù)結(jié)構(gòu)形狀的記憶更新方式。

三個新模型

研究人員利用 Miras 框架構(gòu)建了三個新模型：

? Moneta —— 靈活且表達力強。它采用可定制的 ?p/?q范數(shù)來靈活控制記憶更新的精度。

? Yaad —— 抗噪和抗極端值能力強。它使用Huber損失和自適應(yīng)更新機制來保持模型的穩(wěn)定性。 

? Memora —— 穩(wěn)定且規(guī)范的記憶控制。它通過KL散度和Softmax更新方法，確保記憶在合理范圍內(nèi)波動。

在實驗中，這些新模型在以下任務(wù)中表現(xiàn)優(yōu)于現(xiàn)有最強模型： 語言理解、常識推理、發(fā)現(xiàn)罕見事實（像「大海撈針」那樣找出隱藏信息）、 在長文本中保留細節(jié)信息。

實驗表明，Miras中的不同設(shè)計選擇產(chǎn)生了具有不同優(yōu)勢的模型。

Moneta專注于記憶更新中的可定制精度，使用靈活的??/?q 范數(shù)。

Yaad使用Huber損失和自適應(yīng)更新來保持穩(wěn)定性。

Memora利用KL散度和Softmax更新來保持記憶的邊界。

實驗結(jié)果

首先關(guān)注語言建模中的困惑度（perplexity）以及常識推理任務(wù)的表現(xiàn)。

研究者在表2中報告了Memora、Yaad、Moneta三個模型變體，以及一些基準模型（參數(shù)量為340M、760M 和 1.3B）的結(jié)果。

表2：Miras各個變體與基準模型在語言建模和常識推理任務(wù)中的表現(xiàn)。帶有*標記的為混合模型，高亮的內(nèi)容是表現(xiàn)最好的純模型和混合模型

所有模型變體都優(yōu)于包括Transformer++、現(xiàn)代線性遞歸模型和混合方法在內(nèi)的全部基準方法。

尤其是在與混合模型的比較中取得更好表現(xiàn)更為關(guān)鍵，因為所有模型變體都是純遞歸結(jié)構(gòu)（完全不依賴注意力機制）。

在Miras的三個變體中，雖然Moneta的表現(xiàn)略遜于Memora和Yaad，但這三者的差距并不大，且具體哪個模型效果最好會因任務(wù)類型和模型大小而異。

擴展模式分析（Scaling Pattern）

為了評估新模型的擴展能力，并與基準模型做對比，研究者繪制了模型在不同大小和上下文窗口下的性能變化圖。

上下文長度

研究者將訓(xùn)練時使用的上下文長度從2K擴展到32K，分別在模型大小為340M和760M的兩個版本上進行實驗。結(jié)果如圖3中間和右側(cè)所示。

Miras的三個變體在上下文長度增加時的擴展能力均優(yōu)于當前最先進的基準模型。

這種性能優(yōu)勢主要來自兩個方面：

(1) 更強表達能力的記憶結(jié)構(gòu)。與Mamba2和GSA這些使用向量或矩陣形式記憶的基準模型不同，新模型變體使用了兩層的多層感知機（MLP），能更有效地學習長序列信息；

(2) 保留門（retention gate）和注意力偏向的設(shè)計：新的模型突破了傳統(tǒng)做法，這有助于更高效地管理固定容量的記憶。

模型大小

研究者還在圖3左側(cè)展示了模型的計算量（FLOPs）與困惑度的關(guān)系。

在相同的 FLOPs（計算預(yù)算）下，三個模型變體的表現(xiàn)都超過了所有基準模型。再次證明了強大的記憶機制設(shè)計對模型性能的重要性。

圖3：在C4數(shù)據(jù)集上擴展模型規(guī)模和序列長度時的表現(xiàn)趨勢。（左）隨著模型規(guī)模增加的表現(xiàn)；（中）在模型規(guī)模為340M時，序列長度增加帶來的影響；（右）在模型規(guī)模為760M時，序列長度增加帶來的影響

大海撈針任務(wù)（Needle In Haystack）

為了評估模型在處理長文本時的有效上下文能力，研究者采用了「大海撈針」（Needle In Haystack）任務(wù)。

在「大海撈針」任務(wù)中，模型需要從一段很長的干擾文本中找出一條特定的信息（即「針」）。

在RULER基準中的S-NIAH（單一大海撈針）任務(wù)，在文本長度分別為1K、2K、4K和8K的情境下對新模型和基準模型進行測試，結(jié)果見表3。

所有模型變體都以顯著優(yōu)勢超過了所有基準模型。

值得注意的是，在處理合成噪聲數(shù)據(jù)（S-NIAH-PK）時，Moneta 的表現(xiàn)優(yōu)于其他模型。這一發(fā)現(xiàn)說明 ??-范數(shù)目標函數(shù)和保留門機制在噪聲環(huán)境下更具魯棒性，能更好地保持模型性能。

表3：Moneta、Yaad、Memora以及基準模型在RULER中的NIAH任務(wù)上的表現(xiàn)。最佳結(jié)果用高亮表示。

更多細節(jié)和理論推導(dǎo)，請參閱原文。

作者介紹

Peilin Zhong目前是谷歌紐約的算法與優(yōu)化團隊的研究科學家。

他在哥倫比亞大學獲得了博士學位。

在此之前，他曾是清華大學跨學科信息科學研究院（姚班）的本科生。

他的研究興趣廣泛，主要集中在理論計算機科學領(lǐng)域，特別是算法的設(shè)計與分析。

具體包括并行算法和大規(guī)模并行算法、隱私算法、壓縮算法、流式算法、圖算法、機器學習、高維幾何、度量嵌入、數(shù)值線性代數(shù)、聚類以及與大規(guī)模數(shù)據(jù)計算相關(guān)的其他算法。