万物互联：测时记忆、注意偏好、信息留存与在线优化之旅
It’s All Connected: A Journey Through Test-Time Memorization,Attentional Bias, Retention, and Online Optimization
摘要：
设计高效且有效的架构主干，一直是增强基础模型能力的核心研究方向。受人类注意力偏向认知现象的启发——即自然倾向于优先处理某些事件或刺激——我们重新概念化了神经架构，包括 Transformers、Titans 及现代线性循环神经网络，将它们视为联想记忆模块，这些模块使用内部目标（称为注意力偏向）来学习键与值的映射。令人惊讶的是，我们观察到大多数现有序列模型要么利用 (1) 点积相似性，要么利用 (2) ℓ2 回归目标作为其注意力偏向。超越这些目标，我们提出了一组替代的注意力偏向配置及其有效近似，以稳定其训练过程。接着，我们将现代深度学习架构中的遗忘机制重新解释为一种留存正则化的形式，从而为序列模型提供了一套新颖的遗忘门。基于这些见解，我们提出了 Miras，这是一个基于以下四种选择来设计深度学习架构的通用框架：(i) 联想记忆架构，(ii) 注意力偏向目标，(iii) 留存门，以及 (iv) 记忆学习算法。我们提出了三个新颖的序列模型——Moneta、Yaad 和 Memora——它们超越了现有线性 RNN 的能力，同时保持了快速可并行的训练过程。我们的实验表明，Miras 中不同的设计选择会产生具有不同优势的模型。例如，Miras 的某些实例在特定任务（如语言建模、常识推理和召回密集型任务）中取得了卓越的性能，甚至超越了 Transformers 和其他现代线性循环模型。
1 引言
为序列建模设计高效的架构主干是增强基础模型在语言、计算机视觉、计算生物学和神经科学等领域能力的关键。虽然 Transformers 主要得益于其上下文学习能力和规模化学习能力，已在序列建模中稳固确立为最先进的模型，但其二次时间与空间复杂度限制了其在需要长上下文建模任务中的应用。
近期的努力旨在通过设计高效的循环替代方案来克服 Transformers 在长上下文建模中的局限性。与 Transformers 线性增长的内存不同，这些模型将上下文压缩到固定大小的内存中，需要改进内存管理以获得可比的性能。为了设计更有效的架构，研究重点在于通过使用/设计更具表现力的以下方面来提升内存容量及其管理：(1) 学习规则：从赫布规则到 Delta 规则；(2) 遗忘门：从 LSTM 到 Mamba2 再到 Titan 的遗忘门；(3) 更具表现力的内存架构：从 RetNet 和 LRU 的向量值内存，到 Titans 和 TTT 的神经深度内存。
这些进步的核心存在一个关键问题：“这些序列模型背后的基本设计框架是什么，以及如何增强这些模型？”。受神经心理学文献中关于联想记忆和学习的广义定义的启发，多项研究讨论了 Transformers 与（线性）循环神经网络同联想记忆之间的联系。然而，这些研究要么 (1) 缺乏能完整阐明底层学习算法的普适性解释，要么 (2) 局限于特定的联想记忆定义而缺乏通用性，和/或 (3) 无法描述标准且广泛使用的组件，如遗忘门。
贡献：受人类注意力偏向认知现象的启发——即自然倾向于优先处理某些事件或刺激——我们基于具有注意力偏向的广义联想记忆定义，重新概念化了神经架构，包括 Transformers、Titans 和其他现代线性循环神经网络。我们将注意力偏向定义并形式化为序列模型的内部记忆目标，旨在学习输入之间的底层映射。我们的表述揭示，几乎所有现有的序列模型都是利用同一类型注意力偏向的联想记忆。我们将现代深度学习架构中现有的遗忘机制重新解释为注意力偏向的一种留存 ℓ2 正则化形式，然后为序列模型提供了一套新颖的替代留存门，为如何平衡学习新概念与保留已学概念提供了新的见解。
基于我们对记忆和遗忘门的表述，我们提出了 Miras¹，这是一个通过以下四种选择来设计新颖序列建模架构的基础框架：(1) 注意力偏向（即记忆目标），(2) 留存门，(3) 记忆架构，以及 (4) 记忆学习算法（即优化器）。我们提出并讨论了几种新颖的设计选择，从而产生了超越现有序列建模架构的新架构。
最后，我们专注于 Miras 的三个新颖变体——Moneta、Yaad 和 Memora——它们基于超越简单 ℓ2 回归目标的注意力偏向，以及比现有机制更稳健的新型留存门控机制。我们进一步在语言建模、常识推理、大海捞针和召回密集型任务上对这三个变体进行了实验评估。结果表明这些变体性能优越，超越了最先进的序列模型。
路线图：在第 2 节中，我们回顾了文献并讨论了本文使用的相关概念。在第 3 节中，我们提出并讨论了联想记忆的广义定义，并正式定义了注意力偏向的概念。接着，我们讨论了两种视角——学习-保留和跟随正则化领导者——以通过优化透镜来解读序列建模，并证明学习-保留相对于 FTRL 的普适性。在第 4 节中，我们介绍了 Miras 框架并讨论了它如何统一现代序列模型。在第 5 节中，为了展示 Miras 框架的潜力，我们讨论了针对 (1) 注意力偏向和 (2) 留存门（遗忘门）的多种新颖设计选择。随后在第 5.3 节，我们介绍了作为 Miras 变体的三个新颖序列模型，并讨论了如何以可并行化的方式训练它们。最后，我们的实验评估在第 6 节中报告。
2 预备知识与背景
在本节中，我们将回顾全文所涉及的相关研究和背景概念。
为了解决上述限制，近年来，使用矩阵值内存并采用 Delta 学习规则的循环模型变得越来越受欢迎。尽管有显著优势，即使是这些基于 Delta 规则的循环模型也面临理论限制，并且在实践中性能一般。最近，多项研究旨在通过添加标量或通道级遗忘门机制、使用负特征值以及多步学习来提高此类模型的性能。然而，由于内存架构的表达能力不足，它们在长上下文任务中仍然存在性能下降的问题。
深度记忆模块：Titans 与测试时训练 为了克服有限的内存并扩展深度序列模型的有效上下文长度，近期的研究聚焦于具有深度记忆模块的新一代架构。这些架构建立在元学习视角之上，其中记忆是一个使用梯度下降（可能带动量）更新的 MLP 架构。Sun 等人 (2024) 进一步提供了一个统一的视角，阐释了线性和 softmax 注意力如何分别是（核）回归损失的参数化和非参数化解，但将其他现代线性 RNNs 视为该类模型之外的方法。最近，在我们的一项并行工作中，Wang 等人 (2025) 表明，通过对现代 RNNs（如 RetNet、Mamba）进行额外简化，它们大致可归为内部优化回归损失的同类模型。然而，“这些序列模型背后能够准确统一现有架构的基本设计框架是什么？”这个问题仍未得到解答。此外，遗忘门的作用及其在现代序列模型中的替代选择，目前的研究仍显不足。
3 联想记忆、注意力偏向与留存
联想记忆是人类学习中不可或缺的组成部分，一直是文献中许多人造神经架构的灵感来源。然而，这些研究只定义了联想记忆概念的特定实例，将架构限制在实体间某种特定类型的相似性度量范围内。广义而言，联想记忆是一种将一组键 映射到一组值 的算子。因此，为了学习数据中潜在的映射模式，它需要一个针对某种记忆类型并衡量所学映射质量的目标：
需注意以下几点：
备注1.当我们用参数 W W参数化记忆时，我们使用 M ( W , k )
。在这种参数化设置下，公式 (4) 中的优化问题应在参数 W W上进行。此外，在参数化设置中，我们可能会使用额外的正则化项 R ( W )
来控制对过去数据的保留。
备注2.学习键与值之间的映射（公式 4）是一个元学习问题，其中注意力偏向在内循环中被优化，而神经网络的所有其他参数（例如线性投影、卷积等）则在外循环中被优化。因此，模型学习如何在测试时将数据存储到其参数中。
3.1 通过优化的视角学习记忆与留存
定义 3.1 将基于联想记忆概念的神经架构设计，转化为通过最小化目标函数 L L来学习键与值之间的底层映射。为了优化公式 4，一个简单的方法是采用梯度下降的思想。具体来说，给定一个新的键值对，我们按如下方式更新记忆：
3.2 视角一：在线回归与跟随正则化领导者
公式 (5) 可视为对损失函数序列执行一步在线梯度下降：
公式 (7) 使用了损失函数的线性近似和二次正则化。然而，原则上我们可以使用损失函数的其他近似以及其他正则化函数，正如过去在线性优化或一般优化中所使用的那样。这种改变是开发其他优化算法（如镜像下降）背后的思想。更具体地说，我们可以将公式 (7) 中的更新规则推广为以下形式：
3.3 视角二：学习最新令牌同时保留先前信息
我们在附录 B 中提供了证明。上述命题表明，在一些温和假设下，（学习-保留视角）也可以解释由（跟随正则化领导者视角）得到的方法。因此，（学习-保留视角）可被视为一个更通用的版本。这就是为什么我们在接下来的大部分推导中主要关注这一视角。
备注3：根据上述视角，我们可以看到，即使使用额外的全局正则化，也不存在记忆擦除或遗忘过程（现代架构中的常见术语），但模型可能会决定不保留记忆的过去状态。有趣的是，这一观察也与人类的记忆过程相符，即大脑不会擦除记忆，但记忆可能因提取失败而变得无法访问。因此，我们随后将使用 “留存门”来指代这一项，而非遗忘门。
备注4：如我们在第4节讨论并总结在表1中的那样，大多数现有的现代序列模型都在使用梯度下降来优化联想记忆目标（公式4中的注意力偏向）。因此，为了进一步理解现有序列模型及其在线学习解释之间的联系，我们讨论了上述两种限于基于梯度下降更新规则的视角。然而，我们在公式4中对注意力偏向和联想记忆的初始定义更为宽泛，并且可以通过任何优化算法（例如，甚至牛顿法或非参数解）进行优化。
4 MIRAS：学习用鲁棒且富有表达力的记忆进行记忆
基于我们关于联想记忆、注意力偏差的定义以及之前的观点，我们提出了 MIRAS 框架，它不仅精确地统一了现有的骨干架构，还为如何设计下一代序列模型提供了见解。如第3节早先讨论的，学习联想记忆可以解释为一个元学习任务，在此任务中，联想记忆学习如何在测试时压缩数据并将其存储到其参数中。在这类任务中，记忆的架构尤为重要，因为在较长上下文中，记忆结构的表达能力可能会限制其学习潜在模式的能力。因此，设计序列模型的第一个选择是记忆的结构。给定由一组参数 W W参数化的记忆结构，如前所述，我们旨在通过一个学习算法（例如梯度下降）来最小化带有记忆保留正则项 Ret ( ⋅ ) 的损失函数 ℓ ( W ; ⋅ )
。相应地，MIRAS 需要四个设计选择：
5 超越现有的注意力偏差与保留门机制
5.2 替代的保留门机制
变体 1：通过 -散度在缩放概率单纯形上进行记忆化。在学习过程中，一种防止数值不稳定和值爆炸的常见技术是将搜索空间限制在一个有界域内。遵循这一原则，为了避免数值不稳定性，我们可以将变量 约束在（缩放的）概率单纯形内。换句话说，我们可以将状态限制在如下约束集内：
5.3 MIRAS 的变体：MONETA、YAAD 和 MEMORA
在上一节中，我们讨论了注意力偏差和保留门的不同潜在选择，以展示 MIRAS 的通用性和潜力。在本节中，基于我们的框架，我们提出了三种新颖的序列模型，每种模型都基于不同的动机设计，并讨论了它们如何利用快速并行训练。
YAAD。基于我们关于保护记忆免受极端事件（令牌）影响的鲁棒性记忆重要性的讨论，我们基于 Huber 目标设计了 YAAD。即在 MIRAS 框架中，对于记忆结构的选择，我们遵循 MONETA，使用与上述相同架构的 MLP；对于注意力偏差的选择，我们使用 Huber 损失（由公式 16 定义）；对于保留门的选择，为简洁起见，我们结合使用局部和全局保留项，表示为
这等价于 Behrouz 等人（2024c）引入的“遗忘门”机制；最后，我们直接使用梯度下降作为记忆学习算法。根据以上选择，我们可以将得到的记忆学习过程描述如下：
Memora。 最后，在 Memora 中，我们采用了弹性网络正则化的思想（即硬保留和软保留）。为此，在 MIRAS 框架中：(1) 在记忆架构的选择上，与上述变体类似，我们使用 MLP（架构与之前的变体相同）。(2) 在注意力偏差的选择上，我们使用简单的 ℓ₂ 回归损失。(3) 在保留门的选择上，我们使用如公式 21 所示的 KL 散度。(4) 最后，我们使用梯度下降来优化记忆，从而得到以下更新规则：
5.4 架构骨干与快速训练
架构骨干。对于架构骨干，我们完全遵循近期研究（Behrouz 等人 2024c; Yang 等人 2024a）的做法：在 Llama 的宏观架构中，使用带有 SwiGLU() 激活函数的 MLP、旋转位置编码（RoPE）（Su 等人 2024）和 RMSNorm（Zhang 等人 2019），将注意力模块替换为我们 MIRAS 的变体。对于 MIRAS 层模块，我们遵循近现代线性循环模型（Behrouz 等人 2024c; Yang 等人 2024a），在查询、键和值投影之后分别加入一维深度可分离卷积层（卷积核大小为 4）。为了训练稳定性，我们同样对 q 和 k 使用 ℓ 2
归一化。MIRAS 层模块的输出经过归一化，并通过一个线性层进行门控（Mehta 等人 2023）。
6 实验
在我们的实验评估中，我们旨在回答三个主要问题：(1) 不同的注意力偏差在实践中是否会导致不同的架构？(2) 不同类型的保留门（即保留门机制）如何影响模型在长上下文中的性能？(3) 与基线模型相比，Memora、Moneta 和 Yaad 在下游任务中的表现如何？
实验设置。我们使用训练上下文窗口大小为 4096 来训练我们的模型，使用的数据集为 FineWeb-Edu (Penedo 等人 2024)（用于语言建模和常识推理任务）或 C4 数据集 (Raffel 等人 2020)（用于扩展规律分析）。我们使用的模型参数量级为 120M、340M、760M 和 1.3B。小模型（120M 和 340M）在从数据集中采样的 15B tokens 上训练，中等模型（760M）在 30B tokens 上训练，大模型在 100B tokens 上训练。基线结果由 Behrouz 等人 (2024c) 报告。
6.1 语言建模与常识推理
我们遵循近期研究 (Behrouz 等人 2024c; Yang 等人 2024a,c)，首先关注语言建模的困惑度以及常识推理任务。Memora、Yaad、Moneta 以及参数量为 340M、760M 和 1.3B 的基线模型的结果在表 2 中报告。我们所有的变体都优于所有基线，包括 Transformer++、现代线性循环模型和混合方法。相较于混合模型的优越性能尤为重要，因为我们所有的变体都是纯循环的（无注意力机制）。在 Miras 的三个变体中，虽然 Moneta 的表现略弱于 Memora 和 Yaad，但其他两个变体表现接近，并且根据任务和模型规模，最佳模型会有所不同。
6.2 扩展规律
为了评估模型的扩展规律并与基线进行比较，本节中，我们绘制了它们在不同模型规模和上下文窗口下的性能表现。
上下文长度。我们首先将训练上下文长度从 2K 变化到 32K，对规模为 340M 和 760M 的两个模型版本进行评估。结果报告在图 3（中和右）中。当增加上下文长度时，Miras 的所有三个变体都比最先进的基线模型扩展得更好。我们将这种优越性能归因于：(1) 富有表达力的记忆架构。与使用向量值和矩阵值记忆的 Mamba2 和 GSA 等基线不同，我们的变体使用具有更强表达能力的 2 层 MLP 来从更长的序列中学习。(2) 保留门和注意力偏差的选择：我们所有的三个变体都超越了标准的注意力偏差和保留门。这些选择可以帮助记忆更好地管理其固定大小的容量。
模型规模。我们还在图 3（左）中报告了我们的模型和基线在 FLOPs 与困惑度方面的表现。在几乎相同的 FLOPs 预算下，所有三个变体都优于所有基线。这些结果再次支持了强大记忆设计的重要性。
6.3 大海捞针
为了评估我们的模型和基线的有效上下文窗口，我们使用“大海捞针”任务。在该任务中，我们评估模型从长干扰文本（即“干草堆”）中检索特定信息（即“针”）的能力。我们专注于 RULER 基准测试 (Hsieh 等人 2024) 中的单针任务 (S-NIAH)，并在长度为 1K、2K、4K 和 8K 的序列上评估我们的模型和基线。结果报告在表 3 中。我们所有的变体都以相当大的优势优于所有基线。有趣的是，当数据是合成噪声（S-NIAH-PK）时，Moneta 表现出比其他模型更好的性能。这一观察结果验证了 -范数目标和保留门的有效性，因为它们对噪声更加鲁棒。
6.4 消融实验
在本节中，我们进行消融实验，以验证本文中讨论的不同设计选择是否对实现更好的结果有积极贡献。
p值对性能的影响。我们首先评估p值对Moneta性能的影响。我们将p值在{1, 1.5, 2, 2.8, 3, 3.2, 4}范围内变化，并将上下文窗口从2K调整至16K。结果如图4所示。有趣的是，随着p值增加，性能并未呈现单调变化趋势，当p=3时达到最佳性能，而p=4时性能最差。此外，尽管不同p值会导致具有不同性能的记忆模块，但随着上下文长度增加，扩展规律几乎保持一致。
q值对性能的影响。类似地，我们通过在{2, 3, 4, 5}范围内变化q值来评估其影响。有趣的是，与p值不同，q值能够改变上下文长度增加时的扩展规律。这一现象的主要原因是q值决定了保留门的特性，一个强大的保留门可以改善记忆管理，从而带来更好的性能表现。
设计选择的影响。为了评估架构设计选择，我们对Yaad进行了消融实验。结果见表4。第一行报告了Yaad的原始性能，而（1）第二行移除了保留机制（即设β=1），（2）第三行使δ参数独立于输入，（3）第三行从Huber损失中移除ℓ2损失项，（4）第四行移除ℓ1条件约束，（5）最后一行用线性层替换MLP。这些结果表明所有设计选择都对模型性能具有贡献。
7 结论
本文提出了一种通用框架 Miras，用于解释在线优化与测试时记忆化之间的关联。Miras 框架能够阐释文献中多种标准架构选择（例如遗忘门）的作用，并有助于设计能够更有效管理记忆的新一代架构。基于我们的框架，我们提出了三种新颖的序列模型，每种模型都具有其各自的优势与局限性。实验评估表明，在各种下游任务中，所有这些变体均优于 Transformer 和线性 RNN。本研究通过 Miras 展示了一系列多样化的变体。未来，探索这些替代架构在不同下游任务中的应用是一个值得关注的研究方向。
原文链接：https://arxiv.org/pdf/2504.13173
热门跟贴