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 模型的同时，回答以下问题：
深度记忆对 softmax 注意力是否有效？（见表 2——Transformer++ 与 DEEPTRANSFORMERS 的比较）
引入 Omega 能否提升 softmax 注意力性能？（见表 2——Transformer++、DEEPTRANSFORMERS 与 DoT 的比较）
Omega 规则是否带来更具表达力的记忆更新？（见表 2 与表 6——OMEGANET 与 ATLAS 的表现）
局部最优记忆更新是否有效？（见表 2 与表 6——OMEGANET 与 ATLAS 的比较）
非线性特征映射是否有效？（见表 6）
所提改进能否在上下文回忆任务上缩小与 Transformer 的差距？（见表 5）
内部优化器对记忆有何影响？（见图 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ⱼ‖² / ‖oⱼ‖²。
该实验帮助我们理解用于表示记忆的模型的表示能力，以及优化算法快速学习底层序列映射的能力。
我们研究了五种不同的序列到序列函数：
低秩映射：我们随机采样一个低秩矩阵 W = XY，其中 X ∈ ℝᵈˣᵏ，Y ∈ ℝᵏˣᵈ。然后从高斯分布中随机采样 i₁, …, iₜ，并对所有 j ∈ [t] 设 oⱼ = W iⱼ。
MLP 映射：我们采样一个具有 1 个输入层、1 个隐藏层和 1 个输出层的 MLP M，使用 GELU 非线性，并将隐藏维度设为 d 以避免扩展。然后从高斯分布中随机采样 i₁, …, iₜ，并对所有 j ∈ [t] 设 oⱼ = M(iⱼ)。
注意力 + 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′ⱼ)。
注意力输出作为输入：与上述设置相同，只是我们令 o′₁, …, o′ₜ 作为输入序列，o₁, …, oₜ 作为输出序列。
滑动窗口注意力 + MLP 映射：与“注意力 + MLP 映射”设置相同，只是我们使用滑动窗口注意力而非全注意力。实验中滑动窗口大小设为 512。
请注意，设置 3 和 5 的学习难度远高于其余设置，因为它们需要（部分）记住先前的输入与输出，才能学会将 iⱼ 映射到 oⱼ 的函数；而设置 1、2 和 4 无需记忆任何历史输入-输出对，只需学会把输入映射到输出的低秩矩阵或 MLP 即可。
设置 4 与设置 2 略有不同：其每一步的输入并非独立，而是通过我们用以计算输入的注意力机制彼此关联。因此，若学习算法足够强，就能利用这种潜在相关性，在设置 4 中比设置 2 更快地学会映射。
我们设 d = 256，并在图 6 中展示了全部五种设置下，随序列位置变化的损失曲线；函数学习 MLP M 在不同设置中定义并训练。可以看到，在所有设置中，模型最终都能学会非平凡的输入-输出映射，损失 lⱼ = ‖iⱼ − oⱼ‖² / ‖oⱼ‖² 最终均小于 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。实验表明：
更强大的记忆架构（如带门控的 MLP）可进一步提升 Atlas 性能；
混合变体进一步带来增益，其中 MAG 架构的改进幅度大于 MAL；
当采用上下文记忆（即 Omega 规则）时，多项式映射与深度记忆尤为关键。
图 5 还展示了局部上下文长度 c 对模型性能的影响：随着 c 增大，性能提升，主要得益于可按需剪枝上下文的门控参数 γ。
打开网易新闻 查看精彩图片
模型规模 图 8 给出了 Atlas 与 OmegaNet 随参数规模变化的扩展曲线，并与基线对比。两模型随规模增大均呈现良好的扩展性，在各规模下困惑度均低于基线。
上下文长度 图 8 同时展示了 Atlas 与 OmegaNet 随训练上下文长度的扩展曲线。凭借高记忆容量，两模型在上下文长度增加时均能良好扩展。
7 结论
我们提出 Atlas——一种新型长期记忆模块，旨在解决现代循环模型在长上下文理解中的核心局限：记忆容量有限、仅支持在线更新，以及记忆管理薄弱。所提出的滑动窗口学习规则、高阶特征映射与先进记忆优化器，为克服这些挑战提供了原则化且可扩展的方案。实证表明，我们的模型——OmegaNet、Atlas、DeepTransformers 与 Dot——在多种基准上均相对 Transformer 及最新 RNN 变体实现一致提升。理论上，我们深入分析了记忆容量与优化动态，为先前工作中观察到的上下文长度限制提供了阐释。
原文链接：https://arxiv.org/pdf/2505.23735v1
热门跟贴