Atlas: 在測試時學習最優記憶上下文

Atlas: Learning to Optimally Memorize the Context at Test Time

Atlas: 在測試時學習最優記憶上下文

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

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

1 引言

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

近期的研究致力于通過設計更高效且有效的循環神經網絡來克服 Transformer（即純基于注意力的架構）在長上下文建模中的局限（Behrouz, Zhong, 等人，2024；Peng, Zhang, 等人，2025；Schlag 等人，2021）。這些現代循環架構可以統一為優化一個稱為“注意力偏置”的內部目標的聯想記憶模塊（Behrouz, Razaviyayn, 等人，2025）。與 Transformer 不斷增長的 KV 緩存不同，這些模型使用固定大小的記憶，因此需要改進記憶管理。因此，人們越來越關注通過更有效的方式來增強 RNN 的記憶管理：（i）學習規則，從加法學習（Katharopoulos 等人，2020）到 DeltaNet 的 Delta 規則（Schlag 等人，2021）；（ii）遺忘（保留）門，從 RetNet 的與輸入無關的門控（Sun, Dong, 等人，2023）到 Titans（Behrouz, Zhong, 等人，2024）和 RWKV7（Peng, Zhang, 等人，2025）中的自適應門控；以及（iii）記憶架構，從向量值記憶（Peng, Alcaide, 等人，2023；Sun, Dong, 等人，2023）到神經深度記憶模塊（Behrouz, Zhong, 等人，2024；Sun, Li, 等人，2024）。

盡管這些改進模型在多種下游基準測試中取得了成功，但它們通常在長上下文理解、上下文檢索和外推到更長序列方面存在困難（Arora, Eyuboglu, Zhang, 等人，2024；Behrouz, Zhong, 等人，2024；Wen 等人，2024；Yang, Kautz, 等人，2024）。我們觀察到這些缺點源于三個設計方面：（1）記憶更新的在線性質，即記憶基于當前令牌進行優化，同時保留過去的記憶狀態，導致記憶單個令牌而不考慮更廣泛的上下文；（2）記憶容量有限，架構和鍵值特征映射限制了可完全映射的鍵值對數量；以及（3）記憶管理的表達能力（即內部目標的優化器），因為最新的模型大多使用依賴于令牌動態一階信息的梯度下降，導致記憶收斂到虛假的局部最小值并學習到效果較差的鍵值映射。

記憶視角

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

在整篇論文中，我們使用術語“測試時記憶”，因為該過程涉及在嚴格限定于全局上下文中存儲和檢索信息，而不更新模型的核心學習參數（即外循環）或預訓練得到的初始狀態。通常，一旦記憶被清除，不會將持久學習或技能獲取延續到新的、獨立的全局上下文中。因此，我們更傾向于使用“測試時記憶”而非“測試時訓練”。

貢獻

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

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

新的表達性學習規則。為了克服近期循環模型的在線性質，本文提出了一種滑動窗口更新規則，稱為 Omega 規則。該規則基于給定上下文窗口中所有過去的令牌（而非僅最后一個）來優化和更新記憶。這使得模型能夠更好地管理其固定大小的記憶，并記憶局部上下文而非單個令牌。

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

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

在多樣化下游任務上的改進。廣泛的實驗驗證了我們的模型設計和所提出的技術，包括對現代架構的消融研究。我們在多種基準測試上評估了 DeepTransformers、OmegaNet 和 Atlas，包括語言建模、常識推理、密集召回和“大海撈針”任務，結果表明它們優于現代線性 RNN、局部注意力（SWA）和 Transformer。此外，我們研究了記憶架構、特征映射、記憶管理算法（內部優化器）以及 Omega 規則對記憶模塊容量和在長上下文理解任務中性能的影響。

證明、額外的實驗結果、相關工作討論以及實驗細節詳見附錄。

2 預備知識

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

2.1 背景知識

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

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

使用迭代算法（例如梯度下降）優化這個目標，便產生了記憶更新規則。因此，該序列模型是一個具有兩個優化層級的元上下文學習者：

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

3 在測試時學習記憶上下文

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

本節將進一步討論這些挑戰，并介紹 Omega 規則：一種具有表達力的記憶更新規則，可以直接訪問局部上下文窗口中的令牌，從而記憶上下文而非單個令牌。

3.1 具有超線性容量的聯想記憶

3.2 具有上下文記憶的長期記憶

如前所述，大多數現有循環模型的一個關鍵缺陷在于其在線性質，即它們僅基于當前輸入優化內部目標（注意力偏置），同時保留記憶的先前狀態（Behrouz, Razaviyayn 等人，2025；Liu, Wang 等人，2024），即：

其中 Ret(·,·) 是保留門。這種在線性質雖然使記憶的優化更簡單、更快速，但可能導致對上下文的記憶次優，因為記憶會貪婪地記憶單個令牌。然而，在更一般的情況下，我們可以在每個時間步基于整個上下文（輸入序列）來優化記憶，即：

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

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

3.3 Omega 規則的并行化

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

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

4.1 記憶的在線與局部上下文優化

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

這等價于滑動窗口注意力（SWA）。這種關聯為注意力機制與循環模型之間的差異提供了重要洞見：注意力不僅是一種非參數解（與循環模型的參數化特性相反），它還全局優化其內部目標（注意力偏置），而大多數最近的現代循環模型是在線學習器 (Behrouz, Razaviyayn, et al. 2025; Peng, Zhang, et al. 2025; Sun, Li, et al. 2024; Yang, Kautz, et al. 2024)。我們對滑動窗口RNN和Omega規則的形式化填補了這一空白，通過基于參數化方法優化相對于過去token上下文窗口的記憶，有效地記憶上下文而非單個token。

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

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

4.2 記憶容量與指數核

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

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

該公式從另一個重要角度揭示了注意力機制與（核）循環模型的差異：作為聯想記憶的 softmax 注意力擁有無界記憶，因而能夠將更大的上下文信息更充分地編碼進其參數中。基于這一洞見，我們在深度線性注意力公式（公式 19）中用 §*() 核替換多項式核，提出 DEEPTRANSFORMER，得到未歸一化的表達式：

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

4.3 Deep Omega Transformer（Dot）：采用 Omega 學習規則的 Transformer

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

上述（未歸一化）公式可視為采用 Delta 規則的 Transformer 的泛化。因此，憑借無界記憶，Dot 不僅像原始 Transformer 那樣追加新的鍵和值，還會用前一狀態對該新值的預測結果來替換它。

5 Atlas：一種具有高容量的局部最優記憶

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

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

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

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

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

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

6 實驗

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

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

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

3. Omega 規則是否帶來更具表達力的記憶更新？（見表 2 與表 6——OMEGANET 與 ATLAS 的表現）

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

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

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

7. 內部優化器對記憶有何影響？（見圖 6）

   
   

實驗設置
我們在 FineWeb 數據集（Penedo et al. 2024）上以 4K 上下文窗口訓練模型，參數規模分別為 340 M、400 M、790 M 與 1.3 B，對應訓練 token 量為 15 B、15 B、30 B 與 100 B。基線結果取自 Behrouz, Razaviyayn 等（2025）、Behrouz, Zhong 等（2024）以及 Yang, Kautz 等（2024）。困惑度在預留驗證集上測量。下游任務包括 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）。實驗設置及其他數據集的詳細信息見附錄 E。

6.1 語言建模與常識推理

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

6.2 長上下文：大海撈針

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

6.3 長上下文：BABILong 基準

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

此前章節已展示類 Transformer 架構（DeepTransformers 與 Dot）在語言建模與長上下文大海撈針任務中的有效性。自此以下，我們聚焦循環架構（Atlas 與 OmegaNet），以進一步說明所提出改進的關鍵作用。

6.4 可學習性實驗

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

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

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

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

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

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

5. 滑動窗口注意力 + MLP 映射：與“注意力 + MLP 映射”設置相同，只是我們使用滑動窗口注意力而非全注意力。實驗中滑動窗口大小設為 512。

請注意，設置 3 和 5 的學習難度遠高于其余設置，因為它們需要（部分）記住先前的輸入與輸出，才能學會將 i? 映射到 o? 的函數；而設置 1、2 和 4 無需記憶任何歷史輸入-輸出對，只需學會把輸入映射到輸出的低秩矩陣或 MLP 即可。

設置 4 與設置 2 略有不同：其每一步的輸入并非獨立，而是通過我們用以計算輸入的注意力機制彼此關聯。因此，若學習算法足夠強，就能利用這種潛在相關性，在設置 4 中比設置 2 更快地學會映射。

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

模型在設置 3 和 5 中表現最差，因為這兩者要求模型（部分）記住輸入和輸出，才能學會注意力機制的輸出。令人驚訝的是，模型在設置 3 中的表現反而優于設置 5，而按理設置 3 的容量需求應高于設置 5。我們推測，這是因為學習算法無法讓模型“遺忘”舊輸入，導致在滑動窗口設置下損失比全局注意力設置更高。需指出的是，我們的注意力計算是在隨機初始化的向量上完成的，因此注意力矩陣通常并不尖銳，這與在 LLM 中訓練好的查詢、鍵、值向量所得的注意力矩陣不同，導致注意力輸出接近上下文中值向量的均值。

6.5 附加實驗：上下文回憶、MAD 合成基準與聯想回憶

本節首先在 MAD 合成基準上評估我們的模型；該基準用于衡量模型在回憶、記憶、壓縮與復制任務中的表現（Poli 等，2024）。結果見表 4。ATLAS 在所有方面均取得最佳成績，尤其在記憶方面，凸顯其各組件對提升記憶容量的重要性。

上下文回憶任務對循環神經網絡而言是最具挑戰性的基準之一。本節遵循 Arora、Eyuboglu、Zhang 等（2024）的做法，在 SWDE（Lockard 等，2019）、NQ（Kwiatkowski 等，2019）、DROP（Dua 等，2019）、FDA（Arora、Yang 等，2023）、SQUAD（Rajpurkar 等，2016）和 TQA（Kembhavi 等，2017）上開展實驗，評估并比較 Atlas 與基線及 Transformer 的表現，結果列于表 5。
盡管 Transformer 仍在上下文回憶任務中拔得頭籌，Atlas 與 OmegaNet 展現出具有競爭力的性能，優于當前最先進的循環模型。我們再次將此歸因于更優的記憶管理與容量。

最后，遵循 Yang, Wang, Zhang 等（2024）以及 Arora, Eyuboglu, Timalsina 等（2023）的方法，我們在多查詢聯想回憶（MQAR）任務（Arora, Eyuboglu, Timalsina 等，2023）上評估 Atlas 與 Dot 的性能，結果見圖 7。相比基線，兩模型均表現良好；就單位記憶大小的性能而言，Atlas 優于 DeltaNet（Yang, Wang, Zhang 等，2024）等當前最優模型。

6.6 消融實驗與擴展規律
本節對 Atlas 的不同組件進行消融實驗，并評估其在參數規模與訓練上下文長度上的擴展規律，結果列于表 6。實驗表明：

1. 更強大的記憶架構（如帶門控的 MLP）可進一步提升 Atlas 性能；

2. 混合變體進一步帶來增益，其中 MAG 架構的改進幅度大于 MAL；

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

   
   

模型規模 圖 8 給出了 Atlas 與 OmegaNet 隨參數規模變化的擴展曲線，并與基線對比。兩模型隨規模增大均呈現良好的擴展性，在各規模下困惑度均低于基線。

上下文長度 圖 8 同時展示了 Atlas 與 OmegaNet 隨訓練上下文長度的擴展曲線。憑借高記憶容量，兩模型在上下文長度增加時均能良好擴展。

7 結論

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

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