Atlas: 在測(cè)試時(shí)學(xué)習(xí)最優(yōu)記憶上下文

Atlas: Learning to Optimally Memorize the Context at Test Time

Atlas: 在測(cè)試時(shí)學(xué)習(xí)最優(yōu)記憶上下文

https://arxiv.org/pdf/2505.23735v1

摘要
Transformer 已成為序列建模中最流行的骨干網(wǎng)絡(luò)，這主要得益于其在上下文檢索任務(wù)中的有效性以及大規(guī)模學(xué)習(xí)的能力。然而，其二次方的內(nèi)存和時(shí)間復(fù)雜度限制了其在較長(zhǎng)序列中的應(yīng)用，因此促使研究人員探索有效的替代架構(gòu)，例如現(xiàn)代循環(huán)神經(jīng)網(wǎng)絡(luò)（亦稱長(zhǎng)期循環(huán)記憶模塊）。盡管最近在各種下游任務(wù)中取得了成功，但它們?cè)谛枰L(zhǎng)上下文理解和外推至更長(zhǎng)序列的任務(wù)上仍然存在困難。我們發(fā)現(xiàn)這些缺點(diǎn)源于其設(shè)計(jì)中的三個(gè)相互分離的方面：（1）受限于記憶架構(gòu)和輸入特征映射的有限記憶容量；（2）更新的在線性質(zhì)，即僅針對(duì)最后一個(gè)輸入優(yōu)化記憶；以及（3）對(duì)其固定大小記憶的管理表達(dá)能力不足。為了改善所有這三個(gè)方面，我們提出了 Atlas，一個(gè)具有高容量的長(zhǎng)期記憶模塊，它能夠基于當(dāng)前和過去的標(biāo)記優(yōu)化記憶，從而學(xué)習(xí)記憶上下文，克服了長(zhǎng)期記憶模型的在線性質(zhì)。基于這一見解，我們提出了一類新的類 Transformer 架構(gòu)，稱為 DeepTransformers，它們是原始 Transformer 架構(gòu)的嚴(yán)格泛化。我們?cè)谡Z言建模、常識(shí)推理、密集召回和長(zhǎng)上下文理解任務(wù)上的實(shí)驗(yàn)結(jié)果表明，Atlas 超越了 Transformer 和近期線性循環(huán)模型的性能。Atlas 進(jìn)一步提升了 Titans 的長(zhǎng)上下文性能，在 BABILong 基準(zhǔn)測(cè)試的 1000 萬上下文長(zhǎng)度上實(shí)現(xiàn)了 +80% 的準(zhǔn)確率。

1 引言

注意力模塊（Bahdanau 等人，2014）是現(xiàn)代深度學(xué)習(xí)架構(gòu)（Achiam 等人，2023；Behrouz, Zhong, 等人，2024；Kamath 等人，2025；Vaswani 等人，2017）中的一個(gè)關(guān)鍵構(gòu)建模塊，其卓越性能歸功于其在上下文檢索任務(wù)中的可擴(kuò)展性和表現(xiàn)。原則上，注意力作為一種聯(lián)想記憶，通過計(jì)算直接的成對(duì)令牌依賴關(guān)系來存儲(chǔ)鍵值映射，并通過查詢-鍵相似度檢索它們。然而，計(jì)算這種成對(duì)依賴關(guān)系雖然精確，卻帶來了二次方的空間和時(shí)間復(fù)雜度，限制了其在長(zhǎng)上下文理解、記憶或建模方面的應(yīng)用（Dalal 等人，2025；Li, Huang, 等人，2024；Liu, Lin, 等人，2024）。

近期的研究致力于通過設(shè)計(jì)更高效且有效的循環(huán)神經(jīng)網(wǎng)絡(luò)來克服 Transformer（即純基于注意力的架構(gòu)）在長(zhǎng)上下文建模中的局限（Behrouz, Zhong, 等人，2024；Peng, Zhang, 等人，2025；Schlag 等人，2021）。這些現(xiàn)代循環(huán)架構(gòu)可以統(tǒng)一為優(yōu)化一個(gè)稱為“注意力偏置”的內(nèi)部目標(biāo)的聯(lián)想記憶模塊（Behrouz, Razaviyayn, 等人，2025）。與 Transformer 不斷增長(zhǎng)的 KV 緩存不同，這些模型使用固定大小的記憶，因此需要改進(jìn)記憶管理。因此，人們?cè)絹碓疥P(guān)注通過更有效的方式來增強(qiáng) RNN 的記憶管理：（i）學(xué)習(xí)規(guī)則，從加法學(xué)習(xí)（Katharopoulos 等人，2020）到 DeltaNet 的 Delta 規(guī)則（Schlag 等人，2021）；（ii）遺忘（保留）門，從 RetNet 的與輸入無關(guān)的門控（Sun, Dong, 等人，2023）到 Titans（Behrouz, Zhong, 等人，2024）和 RWKV7（Peng, Zhang, 等人，2025）中的自適應(yīng)門控；以及（iii）記憶架構(gòu)，從向量值記憶（Peng, Alcaide, 等人，2023；Sun, Dong, 等人，2023）到神經(jīng)深度記憶模塊（Behrouz, Zhong, 等人，2024；Sun, Li, 等人，2024）。

盡管這些改進(jìn)模型在多種下游基準(zhǔn)測(cè)試中取得了成功，但它們通常在長(zhǎng)上下文理解、上下文檢索和外推到更長(zhǎng)序列方面存在困難（Arora, Eyuboglu, Zhang, 等人，2024；Behrouz, Zhong, 等人，2024；Wen 等人，2024；Yang, Kautz, 等人，2024）。我們觀察到這些缺點(diǎn)源于三個(gè)設(shè)計(jì)方面：（1）記憶更新的在線性質(zhì)，即記憶基于當(dāng)前令牌進(jìn)行優(yōu)化，同時(shí)保留過去的記憶狀態(tài)，導(dǎo)致記憶單個(gè)令牌而不考慮更廣泛的上下文；（2）記憶容量有限，架構(gòu)和鍵值特征映射限制了可完全映射的鍵值對(duì)數(shù)量；以及（3）記憶管理的表達(dá)能力（即內(nèi)部目標(biāo)的優(yōu)化器），因?yàn)樽钚碌哪Ｐ痛蠖嗍褂靡蕾囉诹钆苿?dòng)態(tài)一階信息的梯度下降，導(dǎo)致記憶收斂到虛假的局部最小值并學(xué)習(xí)到效果較差的鍵值映射。

記憶視角

聯(lián)想記憶——即映射不同實(shí)體或事件的能力——是人類學(xué)習(xí)中不可分割的組成部分（Terry 2017），因此激發(fā)了最近多項(xiàng)研究通過其視角來理解最先進(jìn)的深度學(xué)習(xí)架構(gòu)（Behrouz, Razaviyayn, 等人，2025；Behrouz, Zhong, 等人，2024；Ramsauer 等人，2021；Wang 等人，2025）。從這一視角看，記憶被定義為由輸入引起的神經(jīng)更新；輸入越令人驚奇，它對(duì)記憶的影響就越大，也就越容易被記住。因此，找到一個(gè)有效的“驚奇度度量”是設(shè)計(jì)此類記憶模塊的關(guān)鍵一步。正如 Behrouz, Razaviyayn, 等人（2025）和 Behrouz, Zhong, 等人（2024）早先所討論的，幾乎所有現(xiàn)有架構(gòu)都使用一種基于當(dāng)前輸入更新記憶的驚奇度度量。然而，一個(gè)事件（作為令牌序列）在很長(zhǎng)一段時(shí)間內(nèi)可能并不總是令人驚奇，盡管它是值得記憶的。為了克服這個(gè)問題，Behrouz, Zhong, 等人（2024）建議將驚奇度度量分解為“瞬時(shí)”和“過去”驚奇度兩部分，在根據(jù)當(dāng)前輸入更新記憶時(shí)納入過去輸入的累積驚奇度。然而，這種設(shè)計(jì)可能因?yàn)橛洃泦蝹€(gè)令牌而錯(cuò)過上下文。為此，在這項(xiàng)工作中，我們提出一個(gè)長(zhǎng)期神經(jīng)記憶模塊，它度量局部（或全局）上下文窗口的驚奇度，這意味著它在測(cè)試時(shí)學(xué)習(xí)如何記憶（令牌）上下文。

在整篇論文中，我們使用術(shù)語“測(cè)試時(shí)記憶”，因?yàn)樵撨^程涉及在嚴(yán)格限定于全局上下文中存儲(chǔ)和檢索信息，而不更新模型的核心學(xué)習(xí)參數(shù)（即外循環(huán)）或預(yù)訓(xùn)練得到的初始狀態(tài)。通常，一旦記憶被清除，不會(huì)將持久學(xué)習(xí)或技能獲取延續(xù)到新的、獨(dú)立的全局上下文中。因此，我們更傾向于使用“測(cè)試時(shí)記憶”而非“測(cè)試時(shí)訓(xùn)練”。

貢獻(xiàn)

本文旨在通過設(shè)計(jì)一個(gè)高容量的長(zhǎng)期神經(jīng)記憶模塊（能夠記憶上下文而非單個(gè)令牌）來克服上述局限——即（1）在線性質(zhì)，（2）有限的記憶容量，以及（3）表達(dá)能力較弱的記憶管理。我們進(jìn)一步基于這些見解，提出了一類嚴(yán)格意義上更強(qiáng)大的 Transformer 架構(gòu)。更具體地說：

更好地理解記憶容量及其瓶頸。為了改善有限的記憶容量，我們建議在輸入令牌上使用高階特征映射（例如多項(xiàng)式特征核函數(shù)）。我們從理論上論證了為什么更深層的記憶模塊和/或更高階的特征映射能夠提升記憶容量——即記憶能夠完美映射的線性無關(guān)鍵值關(guān)聯(lián)的最大數(shù)量。

新的表達(dá)性學(xué)習(xí)規(guī)則。為了克服近期循環(huán)模型的在線性質(zhì)，本文提出了一種滑動(dòng)窗口更新規(guī)則，稱為 Omega 規(guī)則。該規(guī)則基于給定上下文窗口中所有過去的令牌（而非僅最后一個(gè)）來優(yōu)化和更新記憶。這使得模型能夠更好地管理其固定大小的記憶，并記憶局部上下文而非單個(gè)令牌。

對(duì) Transformer 的嚴(yán)格泛化。接下來，我們展示了我們的 Omega 規(guī)則公式如何與全局及局部 softmax 注意力（即滑動(dòng)窗口注意力 - SWA）相聯(lián)系，并提出了一個(gè)新的類 Transformer 架構(gòu)家族，稱為 DeepTransformers 及其滑動(dòng)窗口變體 SWDT。這些架構(gòu)嚴(yán)格泛化了 Transformer（Vaswani 等人，2017）。我們進(jìn)一步提出了一個(gè)新穎的基線模型 Deep Linear Attention，以展示深層記憶的作用。

具有更好記憶管理的新記憶模塊。基于上述改進(jìn)，我們提出了 OmegaNet，這是一個(gè)在其鍵和查詢上使用多項(xiàng)式特征，并基于 Omega 規(guī)則和梯度下降更新記憶的新架構(gòu)。為了進(jìn)一步增強(qiáng)記憶管理，我們引入了 Atlas，它利用流行的 Muon 優(yōu)化器來更新內(nèi)部記憶。我們展示了 OmegaNet 和 Atlas 都能利用可并行化的訓(xùn)練算法，從而在不產(chǎn)生顯著開銷的情況下實(shí)現(xiàn)快速訓(xùn)練（與在線版本，即上下文窗口=1 相比）。據(jù)我們所知，Atlas 是第一個(gè)利用二階信息（近似）優(yōu)化記憶（即具有局部最優(yōu)記憶模塊）的可并行化循環(huán)架構(gòu)。

在多樣化下游任務(wù)上的改進(jìn)。廣泛的實(shí)驗(yàn)驗(yàn)證了我們的模型設(shè)計(jì)和所提出的技術(shù)，包括對(duì)現(xiàn)代架構(gòu)的消融研究。我們?cè)诙喾N基準(zhǔn)測(cè)試上評(píng)估了 DeepTransformers、OmegaNet 和 Atlas，包括語言建模、常識(shí)推理、密集召回和“大海撈針”任務(wù)，結(jié)果表明它們優(yōu)于現(xiàn)代線性 RNN、局部注意力（SWA）和 Transformer。此外，我們研究了記憶架構(gòu)、特征映射、記憶管理算法（內(nèi)部?jī)?yōu)化器）以及 Omega 規(guī)則對(duì)記憶模塊容量和在長(zhǎng)上下文理解任務(wù)中性能的影響。

證明、額外的實(shí)驗(yàn)結(jié)果、相關(guān)工作討論以及實(shí)驗(yàn)細(xì)節(jié)詳見附錄。

2 預(yù)備知識(shí)

在本節(jié)中，我們首先討論整篇論文使用的符號(hào)，然后回顧背景概念和相關(guān)工作。關(guān)于相關(guān)研究的更多討論見附錄 A。

2.1 背景知識(shí)

其中 ? 是任意的結(jié)合性運(yùn)算符， A t
是一個(gè)與數(shù)據(jù)（不）相關(guān)的對(duì)角矩陣或低秩加單位矩陣（Yang, Wang, Zhang 等人，2024）。盡管這些模型具有高效的線性循環(huán)特性，但其記憶可能溢出，尤其是在上下文長(zhǎng)度增加的情況下。盡管遺忘門最近顯著改善了這些架構(gòu)中的記憶管理（Peng, Zhang 等人，2025；Sun, Dong 等人，2023），但其記憶的表達(dá)能力仍然受限于其線性結(jié)構(gòu)。

深度記憶模塊。為了克服記憶表達(dá)能力有限的問題，并增強(qiáng)循環(huán)模型的有效上下文長(zhǎng)度，近期的研究聚焦于一類具有深度記憶模塊的新架構(gòu)（Behrouz, Razaviyayn 等人，2025；Behrouz, Zhong 等人，2024；Irie 等人，2021；Sun, Li 等人，2024）。這些架構(gòu)建立在元學(xué)習(xí)的視角上，其中記憶是一個(gè)通過梯度下降（可能帶有動(dòng)量）更新的深度 MLP 架構(gòu)。最近，Behrouz, Razaviyayn 等人（2025）提出了一個(gè)框架，準(zhǔn)確地統(tǒng)一了流行的序列模型，將其視為測(cè)試時(shí)記憶的實(shí)例。也就是說，序列模型是聯(lián)想記憶模塊，旨在通過優(yōu)化一個(gè)稱為注意力偏置的內(nèi)部記憶目標(biāo)來學(xué)習(xí)給定鍵與值之間的底層映射。這種優(yōu)化基于迭代優(yōu)化算法，如梯度下降。更正式地說，聯(lián)想記憶定義如下：

使用迭代算法（例如梯度下降）優(yōu)化這個(gè)目標(biāo)，便產(chǎn)生了記憶更新規(guī)則。因此，該序列模型是一個(gè)具有兩個(gè)優(yōu)化層級(jí)的元上下文學(xué)習(xí)者：

我們的術(shù)語體系建立在此框架之上。因此，我們不再使用完整的循環(huán)公式來描述模型，而是通過以下方面來描述：（1）記憶架構(gòu)，（2）內(nèi)部目標(biāo)（即注意力偏置），以及（3）記憶學(xué)習(xí)算法（優(yōu)化器）。在大多數(shù)情況下，模型使用矩陣值記憶并進(jìn)行在線梯度下降；為簡(jiǎn)潔起見，在此類情況下，我們僅通過其內(nèi)部記憶目標(biāo)來指代一個(gè)架構(gòu)。更多討論和示例見附錄 B。

3 在測(cè)試時(shí)學(xué)習(xí)記憶上下文

長(zhǎng)期聯(lián)想記憶對(duì)人類學(xué)習(xí)至關(guān)重要（Terry 2017），它啟發(fā)了許多人工神經(jīng)架構(gòu)（Behrouz, Razaviyayn 等人，2025；Behrouz, Zhong 等人，2024；He 等人，2024；Hopfield 1982；Krotov 和 Hopfield 2016；Ramsauer 等人，2021；Schmidhuber 和 Hochreiter 1997）。雖然許多此類模型使用矩陣值或向量值記憶來壓縮過去的數(shù)據(jù)（Schlag 等人，2021；Von Oswald 等人，2023；Yang, Kautz 等人，2024），但近期研究提倡使用深度非線性神經(jīng)記憶，將過去的抽象概念編碼到其參數(shù)中（Behrouz, Razaviyayn 等人，2025；Behrouz, Zhong 等人，2024；Dalal 等人，2025；Sun, Li 等人，2024）。然而，對(duì)于長(zhǎng)上下文推理/理解，這些長(zhǎng)期神經(jīng)記憶模塊仍然需要：（1）高容量——參數(shù)中可存儲(chǔ)的（鍵，值）對(duì)的最大數(shù)量（見 §3.1）；（2）強(qiáng)大的內(nèi)部記憶目標(biāo)（即注意力偏置）以學(xué)習(xí)鍵與值之間的復(fù)雜映射（見 §3.2）；（3）強(qiáng)大的記憶管理能力以實(shí)現(xiàn)更好的固定大小記憶管理（見 §3.2）；以及（4）高效的并行訓(xùn)練過程，以便在現(xiàn)代加速器上進(jìn)行大規(guī)模訓(xùn)練（見 §3.3）。

本節(jié)將進(jìn)一步討論這些挑戰(zhàn)，并介紹 Omega 規(guī)則：一種具有表達(dá)力的記憶更新規(guī)則，可以直接訪問局部上下文窗口中的令牌，從而記憶上下文而非單個(gè)令牌。

3.1 具有超線性容量的聯(lián)想記憶

3.2 具有上下文記憶的長(zhǎng)期記憶

如前所述，大多數(shù)現(xiàn)有循環(huán)模型的一個(gè)關(guān)鍵缺陷在于其在線性質(zhì)，即它們僅基于當(dāng)前輸入優(yōu)化內(nèi)部目標(biāo)（注意力偏置），同時(shí)保留記憶的先前狀態(tài)（Behrouz, Razaviyayn 等人，2025；Liu, Wang 等人，2024），即：

其中 Ret(·,·) 是保留門。這種在線性質(zhì)雖然使記憶的優(yōu)化更簡(jiǎn)單、更快速，但可能導(dǎo)致對(duì)上下文的記憶次優(yōu)，因?yàn)橛洃洉?huì)貪婪地記憶單個(gè)令牌。然而，在更一般的情況下，我們可以在每個(gè)時(shí)間步基于整個(gè)上下文（輸入序列）來優(yōu)化記憶，即：

這種嚴(yán)格的全局優(yōu)化公式通常存在兩個(gè)關(guān)鍵限制：（1）效率：循環(huán)架構(gòu)的一個(gè)重要優(yōu)勢(shì)在于其在長(zhǎng)上下文訓(xùn)練和推理中的高效性。然而，基于所有過去令牌（整個(gè)上下文）優(yōu)化記憶會(huì)（i）在每次記憶更新步驟引入額外的優(yōu)化約束，導(dǎo)致在極長(zhǎng)序列上效率低下，以及（ii）需要在測(cè)試時(shí)緩存過去的鍵和值，增加內(nèi)存消耗；（2）上下文剪枝：在長(zhǎng)上下文任務(wù)中，使用所有過去令牌進(jìn)行優(yōu)化可能導(dǎo)致性能次優(yōu)，這主要是由于輸入序列中間的上下文變化（或無關(guān)上下文）。這一觀察結(jié)果推動(dòng)了具有保留（遺忘）門的架構(gòu)設(shè)計(jì)，使模型能夠在不需要過去上下文時(shí)清除記憶（Behrouz, Razaviyayn 等人，2025；Behrouz, Zhong 等人，2024；Peng, Zhang 等人，2025；Sun, Dong 等人，2023；Yang, Wang, Shen 等人，2024）。

從記憶的角度，Omega 規(guī)則（OmegaNet）并不衡量一個(gè) token 的驚奇度，而是基于上下文內(nèi)單個(gè) token 的上下文感知組合來衡量局部上下文的驚奇度。

3.3 Omega 規(guī)則的并行化

4 深度Transformer：具有深度記憶的Transformer

近期研究已通過聯(lián)想記憶（associative memory）的視角廣泛討論了Transformer架構(gòu) (Behrouz, Razaviyayn, et al. 2025; Sun, Li, et al. 2024; Wang et al. 2025)。因此，很自然地會(huì)問到我們對(duì)記憶容量以及Omega規(guī)則的討論會(huì)如何影響Transformer。在本節(jié)中，我們討論我們的Omega規(guī)則形式化與Transformer及其滑動(dòng)窗口變體（即SWA）之間的關(guān)聯(lián)。我們進(jìn)一步為Transformer提供兩種擴(kuò)展，每種擴(kuò)展都是Transformer的嚴(yán)格泛化。

4.1 記憶的在線與局部上下文優(yōu)化

與滑動(dòng)窗口注意力（Sliding Window Attention）的關(guān)聯(lián)。 Softmax注意力塊也可以被重新表述為使用Nadaraya-Watson估計(jì)器 (Fan 2018; Zhang et al. 2022) 求解回歸的非參數(shù)解：

這等價(jià)于滑動(dòng)窗口注意力（SWA）。這種關(guān)聯(lián)為注意力機(jī)制與循環(huán)模型之間的差異提供了重要洞見：注意力不僅是一種非參數(shù)解（與循環(huán)模型的參數(shù)化特性相反），它還全局優(yōu)化其內(nèi)部目標(biāo)（注意力偏置），而大多數(shù)最近的現(xiàn)代循環(huán)模型是在線學(xué)習(xí)器 (Behrouz, Razaviyayn, et al. 2025; Peng, Zhang, et al. 2025; Sun, Li, et al. 2024; Yang, Kautz, et al. 2024)。我們對(duì)滑動(dòng)窗口RNN和Omega規(guī)則的形式化填補(bǔ)了這一空白，通過基于參數(shù)化方法優(yōu)化相對(duì)于過去token上下文窗口的記憶，有效地記憶上下文而非單個(gè)token。

深度線性注意力（Deep Linear Attention）。 作為一種新穎的基線方法，我們提出了深度（門控）線性注意力（DLA），它將（門控）線性注意力 (Katharopoulos et al. 2020; Yang, Wang, Shen, et al. 2024) 中的矩陣值記憶替換為一個(gè)深度神經(jīng)網(wǎng)絡(luò)（例如，k層MLP）。如前文（Hebbian Rule）所討論的，使用點(diǎn)積相似度作為內(nèi)部注意力偏置會(huì)產(chǎn)生線性注意力。因此，利用最近的深度記憶模塊 (Behrouz, Razaviyayn, et al. 2025; Behrouz, Zhong, et al. 2024; Sun, Li, et al. 2024)，我們使用梯度下降優(yōu)化記憶，并采用點(diǎn)積注意力偏置：

滑動(dòng)窗口線性注意力（Sliding Window Linear Attention）。基于上述直覺以及我們的公式與 SWA 的聯(lián)系，我們提出滑動(dòng)窗口線性注意力（SWLA）模塊。借鑒線性注意力在聯(lián)想記憶視角下的公式（Behrouz、Razaviyayn 等，2025），我們使用點(diǎn)積相似度（即 t(Mt; Ki, Vi) = ）作為注意力偏置，并利用梯度下降優(yōu)化損失函數(shù)。為清晰起見，此處我們采用線性記憶來推導(dǎo)閉式解：

4.2 記憶容量與指數(shù)核

我們首先回顧 Transformer 中 softmax 注意力的公式（即公式 1）：

其 exp() 核不可分離，因此無法寫成遞推形式。按照 Kacham 等人（2024b）的討論，可以將 exp() 核（與多項(xiàng)式核 p() 相比）視為一種將輸入映射到無限維的特征映射。即，我們定義：

該公式從另一個(gè)重要角度揭示了注意力機(jī)制與（核）循環(huán)模型的差異：作為聯(lián)想記憶的 softmax 注意力擁有無界記憶，因而能夠?qū)⒏蟮纳舷挛男畔⒏浞值鼐幋a進(jìn)其參數(shù)中。基于這一洞見，我們?cè)谏疃染€性注意力公式（公式 19）中用 §*() 核替換多項(xiàng)式核，提出 DEEPTRANSFORMER，得到未歸一化的表達(dá)式：

其輸出與未歸一化的 Transformer 一致。因此，DeepTransformer 是帶有 softmax 注意力的 Transformer（Vaswani 等，2017）的嚴(yán)格泛化。

4.3 Deep Omega Transformer（Dot）：采用 Omega 學(xué)習(xí)規(guī)則的 Transformer

我們上述的 DeepTransformer 公式基于（Hebb 規(guī)則），這也是原始 Transformer 所采用的。然而，如前所述，在聯(lián)想記憶模塊中引入更強(qiáng)大的記憶管理與學(xué)習(xí)規(guī)則可進(jìn)一步提升性能。為此，我們用 Omega 學(xué)習(xí)規(guī)則替換 Hebb 規(guī)則，得到未歸一化的 Deep Omega Transformer（Dot）公式：

上述（未歸一化）公式可視為采用 Delta 規(guī)則的 Transformer 的泛化。因此，憑借無界記憶，Dot 不僅像原始 Transformer 那樣追加新的鍵和值，還會(huì)用前一狀態(tài)對(duì)該新值的預(yù)測(cè)結(jié)果來替換它。

5 Atlas：一種具有高容量的局部最優(yōu)記憶

盡管 Omega 規(guī)則的設(shè)計(jì)使模型能夠記憶整個(gè)上下文而非單個(gè) token，且采用多項(xiàng)式（或指數(shù)）特征映射提升了記憶容量，但記憶管理（即鍵-值映射的優(yōu)化）仍局限于簡(jiǎn)單的梯度下降。該優(yōu)化器的選擇可能使模型陷入局部最優(yōu)的低質(zhì)量解，在長(zhǎng)上下文場(chǎng)景下?lián)p害性能。為解決此問題，我們建議采用 Muon 優(yōu)化器（Jordan 等，2024）（帶權(quán)重衰減），它不僅能近似二階信息，還主要依賴矩陣乘法且可沿序列并行。據(jù)此，用 Muon 優(yōu)化公式 9 的內(nèi)部目標(biāo)，得到如下更新規(guī)則：

其中 c 為局部上下文長(zhǎng)度，k 為 Newton–Schulz 迭代步數(shù)。關(guān)于該算法及此運(yùn)算的進(jìn)一步討論，請(qǐng)參見 Jordan 等（2024）。根據(jù) Muon 優(yōu)化器的相關(guān)文獻(xiàn)，當(dāng) k → ∞ 時(shí)，NewtonSchulz-k(St) 收斂到與動(dòng)量項(xiàng) St 最近的半正交矩陣，從而以更低誤差近似二階信息。因此，有趣的是，參數(shù) k 可視為 ATLAS 內(nèi)部的一種“測(cè)試時(shí)計(jì)算”超參：步數(shù)越多，潛在的記憶效果越好。

5.1 并行訓(xùn)練
本節(jié)討論如何對(duì) ATLAS 的訓(xùn)練過程做并行化。為清晰起見，設(shè) c = 1；將過程推廣到任意 c 值的方法與 3.3 節(jié)相同。我們沿用 3.3 節(jié)的方案，把序列分塊，并針對(duì)前一塊的最終狀態(tài)計(jì)算所有梯度。于是，利用帶動(dòng)量但不含 的 ATLAS 遞推式，可得：

其中 為梯度矩陣， 與 Θ 是元素分別為 與 的對(duì)角矩陣，⊙ 表示廣播乘法。
上述（按塊遞推）公式的主要優(yōu)勢(shì)在于：動(dòng)量的遞推與記憶狀態(tài)無關(guān)。也就是說，我們可以在塊的起始處一次性算出所有動(dòng)量項(xiàng)。對(duì)于 Muon 優(yōu)化器，我們需對(duì)這些動(dòng)量項(xiàng)應(yīng)用 Newton–Schulz 算法，于是得到：

由于所有 S 的計(jì)算可以并行完成，Newton–Schulz?(·) 的計(jì)算同樣能夠并行執(zhí)行。

架構(gòu)主干 至于架構(gòu)主干，我們遵循近期現(xiàn)代循環(huán)模型（Allen-Zhu 2025；Arora、Eyuboglu、Zhang 等 2024；Behrouz、Zhong 等 2024；Yang、Wang、Zhang 等 2024）的做法，先用線性層將鍵、值和查詢投影，再接尺寸為 4 的短卷積層。我們對(duì)鍵和查詢做歸一化以穩(wěn)定訓(xùn)練。同時(shí)，參照 Behrouz、Zhong 等（2024），我們?yōu)?Atlas 模型采用兩種混合變體 MAL 與 MAG，其結(jié)構(gòu)見圖 3。對(duì)于采用深層記憶架構(gòu)的模型，我們使用帶殘差連接的雙層 MLP：

6 實(shí)驗(yàn)

接下來，我們?cè)谡Z言建模、常識(shí)推理、大海撈針（needle-in-haystack）以及上下文回憶任務(wù)上評(píng)估 ATLAS、OMEGANET、DEEPTRANSFORMERS 和 DoT 的性能。盡管我們還討論過若干其他變體（如 SwLA），實(shí)驗(yàn)部分僅聚焦上述模型，從而在對(duì)比 SOTA 模型的同時(shí)，回答以下問題：

1. 深度記憶對(duì) softmax 注意力是否有效？（見表 2——Transformer++ 與 DEEPTRANSFORMERS 的比較）

2. 引入 Omega 能否提升 softmax 注意力性能？（見表 2——Transformer++、DEEPTRANSFORMERS 與 DoT 的比較）

3. Omega 規(guī)則是否帶來更具表達(dá)力的記憶更新？（見表 2 與表 6——OMEGANET 與 ATLAS 的表現(xiàn)）

4. 局部最優(yōu)記憶更新是否有效？（見表 2 與表 6——OMEGANET 與 ATLAS 的比較）

5. 非線性特征映射是否有效？（見表 6）

6. 所提改進(jìn)能否在上下文回憶任務(wù)上縮小與 Transformer 的差距？（見表 5）

7. 內(nèi)部?jī)?yōu)化器對(duì)記憶有何影響？（見圖 6）

   
   

實(shí)驗(yàn)設(shè)置
我們?cè)?FineWeb 數(shù)據(jù)集（Penedo et al. 2024）上以 4K 上下文窗口訓(xùn)練模型，參數(shù)規(guī)模分別為 340 M、400 M、790 M 與 1.3 B，對(duì)應(yīng)訓(xùn)練 token 量為 15 B、15 B、30 B 與 100 B。基線結(jié)果取自 Behrouz, Razaviyayn 等（2025）、Behrouz, Zhong 等（2024）以及 Yang, Kautz 等（2024）。困惑度在預(yù)留驗(yàn)證集上測(cè)量。下游任務(wù)包括 Wikitext（Merity et al. 2017）、LMB（Paperno et al. 2016）、PIQA（Bisk et al. 2020）、HellaSwag（Zellers et al. 2019）、WinoGrande（Sakaguchi et al. 2021）、ARC-easy（ARC-e）與 ARC-challenge（ARC-c）（Clark, Cowhey 等 2018）、SIQA（Sap et al. 2019）以及 BoolQ（Clark, Lee 等 2019）。實(shí)驗(yàn)設(shè)置及其他數(shù)據(jù)集的詳細(xì)信息見附錄 E。

6.1 語言建模與常識(shí)推理

表 2 給出了 760 M 與 1.3 B 規(guī)模的 Atlas、OmegaNet 及其對(duì)應(yīng)基線 SWDT、DLA、DeepTransformers 和 DoT 的結(jié)果（小模型結(jié)果見附錄 F）。在非混合模型中，包括 Transformer++ 在內(nèi)，我們的 Atlas 與 OmegaNet 在困惑度和準(zhǔn)確率兩項(xiàng)指標(biāo)上均取得最佳表現(xiàn)。我們將此歸因于它們“記憶上下文而非單個(gè) token”的能力。
對(duì)比同樣采用瞬時(shí)目標(biāo)（即 ?2 損失）但上下文窗口僅為 1 的 Titans，OmegaNet 的優(yōu)勢(shì)印證了“非在線學(xué)習(xí)規(guī)則”的有效性。另一方面，即便完全不引入注意力，我們的模型也能超越混合模型；而其混合變體 MAG 進(jìn)一步提升了性能。這一增益也與多項(xiàng)式核的引入有關(guān)——它顯著增強(qiáng)了模型的記憶容量。關(guān)于各組件影響的受控研究見表 6。
將 Transformer++ 與我們更泛化的 Transformer（即 DeepTransformers 與 DoT）相比，可觀察到一致的性能提升。我們認(rèn)為這得益于其“深度記憶”，使其對(duì) token 間依賴的建模能力更強(qiáng)。再比較 DoT 與 DeepTransformers，Omega 規(guī)則的優(yōu)勢(shì)顯而易見：它幫助模型更有效地管理記憶。

6.2 長(zhǎng)上下文：大海撈針

我們?cè)O(shè)計(jì) Atlas 的核心動(dòng)機(jī)之一，是在長(zhǎng)上下文任務(wù)中提升神經(jīng)長(zhǎng)期記憶模塊的表現(xiàn)。為驗(yàn)證各項(xiàng)設(shè)計(jì)對(duì)有效上下文長(zhǎng)度與記憶容量的改進(jìn)效果，我們?cè)?RULER 基準(zhǔn)（Hsieh 等，2024）的“大海撈針”任務(wù)上進(jìn)行實(shí)驗(yàn)。表 3 報(bào)告了 Atlas 及其混合變體、類 Transformer 架構(gòu)與基線方法的性能。Atlas 相比循環(huán)基線表現(xiàn)優(yōu)異，超越了 Titans、DeltaNet 等現(xiàn)代循環(huán)網(wǎng)絡(luò)；其混合變體進(jìn)一步延長(zhǎng)了有效上下文長(zhǎng)度，可在訓(xùn)練上下文 4× 的更長(zhǎng)子序列上外推。我們將此歸功于所提出的記憶容量增強(qiáng)方案，并通過消融實(shí)驗(yàn)加以驗(yàn)證。此外，我們的類 Transformer 架構(gòu)在更長(zhǎng)上下文上甚至優(yōu)于 Atlas 的混合變體，表明指數(shù)特征映射在長(zhǎng)序列中的重要性。

6.3 長(zhǎng)上下文：BABILong 基準(zhǔn)

為在超大規(guī)模序列上比較 Atlas 與 Titans（Behrouz、Zhong 等，2024）的效果，我們進(jìn)一步在 BABILong 基準(zhǔn)（Kuratov 等，2024）上評(píng)估 Atlas。實(shí)驗(yàn)沿用 Behrouz、Zhong 等（2024）的設(shè)置，采用 MAC 架構(gòu)但不使用持久記憶 token，并按原基準(zhǔn)設(shè)定進(jìn)行微調(diào)。結(jié)果見圖 4：在 1 M 上下文長(zhǎng)度內(nèi)，Atlas 與 Titans 表現(xiàn)相當(dāng)；當(dāng)長(zhǎng)度增至 10 M 時(shí)，Titans 性能下降，而 Atlas 仍保持 80 % 以上的準(zhǔn)確率。我們歸因于 Atlas 擁有更強(qiáng)大的記憶能力：(1) 借助 Muon 實(shí)現(xiàn)更優(yōu)的記憶管理；(2) 多項(xiàng)式核帶來更大的記憶容量；(3) 其本質(zhì)是對(duì)“上下文”而非“單個(gè) token”進(jìn)行記憶。

此前章節(jié)已展示類 Transformer 架構(gòu)（DeepTransformers 與 Dot）在語言建模與長(zhǎng)上下文大海撈針任務(wù)中的有效性。自此以下，我們聚焦循環(huán)架構(gòu)（Atlas 與 OmegaNet），以進(jìn)一步說明所提出改進(jìn)的關(guān)鍵作用。

6.4 可學(xué)習(xí)性實(shí)驗(yàn)

我們還進(jìn)行了一些小型實(shí)驗(yàn)，以在線方式分析小型 MLP 的函數(shù)學(xué)習(xí)能力。在此設(shè)置中，我們有序列元組 (i?, o?), …, (i?, o?)，其中對(duì)所有 j 都有 i?, o? ∈ ??。我們以在線方式訓(xùn)練一個(gè) MLP M，使損失最小化；具體地，我們?cè)跁r(shí)間步 j 計(jì)算梯度為
l? = ‖i? ? o?‖2 / ‖o?‖2。
該實(shí)驗(yàn)幫助我們理解用于表示記憶的模型的表示能力，以及優(yōu)化算法快速學(xué)習(xí)底層序列映射的能力。

我們研究了五種不同的序列到序列函數(shù)：

1. 低秩映射：我們隨機(jī)采樣一個(gè)低秩矩陣 W = XY，其中 X ∈ ????，Y ∈ ????。然后從高斯分布中隨機(jī)采樣 i?, …, i?，并對(duì)所有 j ∈ [t] 設(shè) o? = W i?。

2. MLP 映射：我們采樣一個(gè)具有 1 個(gè)輸入層、1 個(gè)隱藏層和 1 個(gè)輸出層的 MLP M，使用 GELU 非線性，并將隱藏維度設(shè)為 d 以避免擴(kuò)展。然后從高斯分布中隨機(jī)采樣 i?, …, i?，并對(duì)所有 j ∈ [t] 設(shè) o? = M(i?)。

3. 注意力 + MLP 映射：我們從高斯分布中采樣 (i?, …, i?)，并如上采樣 MLP M。此外，我們隨機(jī)采樣三個(gè) d×d 矩陣 W_q、W_k 和 W_v，并對(duì)所有 j ∈ [t] 計(jì)算 q? = W_q i?、k? = Wk i? 和 v? = Wv i?。然后，我們將因果掩碼注意力機(jī)制應(yīng)用于 (q?)?∈[t]、(k?)?∈[t]、(v?)?∈[t]，得到輸出 o′?, …, o′?，最后計(jì)算 o? = M(o′?)。

4. 注意力輸出作為輸入：與上述設(shè)置相同，只是我們令 o′?, …, o′? 作為輸入序列，o?, …, o? 作為輸出序列。

5. 滑動(dòng)窗口注意力 + MLP 映射：與“注意力 + MLP 映射”設(shè)置相同，只是我們使用滑動(dòng)窗口注意力而非全注意力。實(shí)驗(yàn)中滑動(dòng)窗口大小設(shè)為 512。

請(qǐng)注意，設(shè)置 3 和 5 的學(xué)習(xí)難度遠(yuǎn)高于其余設(shè)置，因?yàn)樗鼈冃枰ú糠郑┯涀∠惹暗妮斎肱c輸出，才能學(xué)會(huì)將 i? 映射到 o? 的函數(shù)；而設(shè)置 1、2 和 4 無需記憶任何歷史輸入-輸出對(duì)，只需學(xué)會(huì)把輸入映射到輸出的低秩矩陣或 MLP 即可。

設(shè)置 4 與設(shè)置 2 略有不同：其每一步的輸入并非獨(dú)立，而是通過我們用以計(jì)算輸入的注意力機(jī)制彼此關(guān)聯(lián)。因此，若學(xué)習(xí)算法足夠強(qiáng)，就能利用這種潛在相關(guān)性，在設(shè)置 4 中比設(shè)置 2 更快地學(xué)會(huì)映射。

我們?cè)O(shè) d = 256，并在圖 6 中展示了全部五種設(shè)置下，隨序列位置變化的損失曲線；函數(shù)學(xué)習(xí) MLP M 在不同設(shè)置中定義并訓(xùn)練。可以看到，在所有設(shè)置中，模型最終都能學(xué)會(huì)非平凡的輸入-輸出映射，損失 l? = ‖i? ? o?‖2 / ‖o?‖2 最終均小于 1。最值得注意的是，設(shè)置 4 中由注意力機(jī)制引入的輸入相關(guān)性，使模型比設(shè)置 2 更快學(xué)會(huì)映射；而模型在設(shè)置 1（復(fù)雜度最低的函數(shù)）中通常學(xué)得最好。

模型在設(shè)置 3 和 5 中表現(xiàn)最差，因?yàn)檫@兩者要求模型（部分）記住輸入和輸出，才能學(xué)會(huì)注意力機(jī)制的輸出。令人驚訝的是，模型在設(shè)置 3 中的表現(xiàn)反而優(yōu)于設(shè)置 5，而按理設(shè)置 3 的容量需求應(yīng)高于設(shè)置 5。我們推測(cè)，這是因?yàn)閷W(xué)習(xí)算法無法讓模型“遺忘”舊輸入，導(dǎo)致在滑動(dòng)窗口設(shè)置下?lián)p失比全局注意力設(shè)置更高。需指出的是，我們的注意力計(jì)算是在隨機(jī)初始化的向量上完成的，因此注意力矩陣通常并不尖銳，這與在 LLM 中訓(xùn)練好的查詢、鍵、值向量所得的注意力矩陣不同，導(dǎo)致注意力輸出接近上下文中值向量的均值。

6.5 附加實(shí)驗(yàn)：上下文回憶、MAD 合成基準(zhǔn)與聯(lián)想回憶

本節(jié)首先在 MAD 合成基準(zhǔn)上評(píng)估我們的模型；該基準(zhǔn)用于衡量模型在回憶、記憶、壓縮與復(fù)制任務(wù)中的表現(xiàn)（Poli 等，2024）。結(jié)果見表 4。ATLAS 在所有方面均取得最佳成績(jī)，尤其在記憶方面，凸顯其各組件對(duì)提升記憶容量的重要性。

上下文回憶任務(wù)對(duì)循環(huán)神經(jīng)網(wǎng)絡(luò)而言是最具挑戰(zhàn)性的基準(zhǔn)之一。本節(jié)遵循 Arora、Eyuboglu、Zhang 等（2024）的做法，在 SWDE（Lockard 等，2019）、NQ（Kwiatkowski 等，2019）、DROP（Dua 等，2019）、FDA（Arora、Yang 等，2023）、SQUAD（Rajpurkar 等，2016）和 TQA（Kembhavi 等，2017）上開展實(shí)驗(yàn)，評(píng)估并比較 Atlas 與基線及 Transformer 的表現(xiàn)，結(jié)果列于表 5。
盡管 Transformer 仍在上下文回憶任務(wù)中拔得頭籌，Atlas 與 OmegaNet 展現(xiàn)出具有競(jìng)爭(zhēng)力的性能，優(yōu)于當(dāng)前最先進(jìn)的循環(huán)模型。我們?cè)俅螌⒋藲w因于更優(yōu)的記憶管理與容量。

最后，遵循 Yang, Wang, Zhang 等（2024）以及 Arora, Eyuboglu, Timalsina 等（2023）的方法，我們?cè)诙嗖樵兟?lián)想回憶（MQAR）任務(wù)（Arora, Eyuboglu, Timalsina 等，2023）上評(píng)估 Atlas 與 Dot 的性能，結(jié)果見圖 7。相比基線，兩模型均表現(xiàn)良好；就單位記憶大小的性能而言，Atlas 優(yōu)于 DeltaNet（Yang, Wang, Zhang 等，2024）等當(dāng)前最優(yōu)模型。

6.6 消融實(shí)驗(yàn)與擴(kuò)展規(guī)律
本節(jié)對(duì) Atlas 的不同組件進(jìn)行消融實(shí)驗(yàn)，并評(píng)估其在參數(shù)規(guī)模與訓(xùn)練上下文長(zhǎng)度上的擴(kuò)展規(guī)律，結(jié)果列于表 6。實(shí)驗(yàn)表明：

1. 更強(qiáng)大的記憶架構(gòu)（如帶門控的 MLP）可進(jìn)一步提升 Atlas 性能；

2. 混合變體進(jìn)一步帶來增益，其中 MAG 架構(gòu)的改進(jìn)幅度大于 MAL；

3. 當(dāng)采用上下文記憶（即 Omega 規(guī)則）時(shí)，多項(xiàng)式映射與深度記憶尤為關(guān)鍵。
   圖 5 還展示了局部上下文長(zhǎng)度 c 對(duì)模型性能的影響：隨著 c 增大，性能提升，主要得益于可按需剪枝上下文的門控參數(shù) γ。

   
   

模型規(guī)模 圖 8 給出了 Atlas 與 OmegaNet 隨參數(shù)規(guī)模變化的擴(kuò)展曲線，并與基線對(duì)比。兩模型隨規(guī)模增大均呈現(xiàn)良好的擴(kuò)展性，在各規(guī)模下困惑度均低于基線。

上下文長(zhǎng)度 圖 8 同時(shí)展示了 Atlas 與 OmegaNet 隨訓(xùn)練上下文長(zhǎng)度的擴(kuò)展曲線。憑借高記憶容量，兩模型在上下文長(zhǎng)度增加時(shí)均能良好擴(kuò)展。

7 結(jié)論

我們提出 Atlas——一種新型長(zhǎng)期記憶模塊，旨在解決現(xiàn)代循環(huán)模型在長(zhǎng)上下文理解中的核心局限：記憶容量有限、僅支持在線更新，以及記憶管理薄弱。所提出的滑動(dòng)窗口學(xué)習(xí)規(guī)則、高階特征映射與先進(jìn)記憶優(yōu)化器，為克服這些挑戰(zhàn)提供了原則化且可擴(kuò)展的方案。實(shí)證表明，我們的模型——OmegaNet、Atlas、DeepTransformers 與 Dot——在多種基準(zhǔn)上均相對(duì) Transformer 及最新 RNN 變體實(shí)現(xiàn)一致提升。理論上，我們深入分析了記憶容量與優(yōu)化動(dòng)態(tài)，為先前工作中觀察到的上下文長(zhǎng)度限制提供了闡釋。

原文鏈接：https://arxiv.org/pdf/2505.23735v1