谷歌又有新的注意力了！

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

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

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

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

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

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

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

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

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

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

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

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

3.保留門控 — 如何平衡學(xué)習(xí)新概念和保留已學(xué)概念；

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

圖1：Miras框架概述

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

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

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

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

? Yaad：常識推理準(zhǔn)確率達(dá)89.4%（超越Transformer7.2%）  

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

但它存在幾個問題：

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

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

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

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

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

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

這一公式可以被重新解釋為一種瞬時驚訝度度量，其中模型記憶那些違反目標(biāo)預(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)不同的算法。

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

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

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

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

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

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

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

??更抗異常值，

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

?∞ 聚焦于最大誤差。

Huber損失：「應(yīng)對異?！剐睦頇C(jī)制

Huber損失具備容錯機(jī)制的記憶模塊。

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

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

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

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

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

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

• 不只是優(yōu)化當(dāng)前值，而是對 可能擾動做最壞情況準(zhǔn)備
• 提高模型應(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è)計目標(biāo)：

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

這就是Retention Gate的作用核心。

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

這種形式可以推廣為：

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

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

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

其中：

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

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

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

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

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

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

三個新模型

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

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

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

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

在實驗中，這些新模型在以下任務(wù)中表現(xiàn)優(yōu)于現(xiàn)有最強(qiáng)模型： 語言理解、常識推理、發(fā)現(xiàn)罕見事實（像「大海撈針」那樣找出隱藏信息）、 在長文本中保留細(xì)節(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三個模型變體，以及一些基準(zhǔn)模型（參數(shù)量為340M、760M 和 1.3B）的結(jié)果。

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

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

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

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

擴(kuò)展模式分析（Scaling Pattern）

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

上下文長度

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

Miras的三個變體在上下文長度增加時的擴(kuò)展能力均優(yōu)于當(dāng)前最先進(jìn)的基準(zhǔn)模型。

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

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

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

模型大小

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

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

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

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

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

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

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

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

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

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

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

作者介紹

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

他在哥倫比亞大學(xué)獲得了博士學(xué)位。

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

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

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