It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization
作者: Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, Vahab Mirrokni
分类: cs.LG, cs.AI
发布日期: 2025-04-17
💡 一句话要点
提出Miras框架，通过可定制的记忆、注意力偏置和遗忘机制，设计高性能序列模型。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 序列模型 注意力机制 遗忘机制 关联记忆 注意力偏置 深度学习架构 保留正则化
📋 核心要点
- 现有序列模型架构设计缺乏统一视角，且在注意力机制和遗忘机制上存在局限性。
- 论文提出Miras框架，将序列模型视为关联记忆模块，通过定制注意力偏置、遗忘门和记忆学习算法来设计模型。
- 实验表明，基于Miras设计的模型在语言建模、常识推理和回忆任务上超越了现有线性RNN和Transformer。
📝 摘要（中文）
为了提升基础模型的能力，研究的重点一直放在设计高效的架构骨干上。受到人类认知中注意力偏置现象的启发，我们将包括Transformer、Titan和现代线性循环神经网络在内的神经架构重新概念化为关联记忆模块，这些模块使用内部目标（称为注意力偏置）来学习键和值的映射。令人惊讶的是，我们观察到大多数现有的序列模型要么利用(1)点积相似性，要么利用(2)L2回归目标作为它们的注意力偏置。为了超越这些目标，我们提出了一组替代的注意力偏置配置以及它们的有效近似，以稳定它们的训练过程。然后，我们将现代深度学习架构中的遗忘机制重新解释为一种保留正则化形式，为序列模型提供了一组新的遗忘门。基于这些见解，我们提出了Miras，这是一个通用框架，用于基于以下四个选择来设计深度学习架构：(i)关联记忆架构，(ii)注意力偏置目标，(iii)保留门，以及(iv)记忆学习算法。我们提出了三种新的序列模型——Moneta、Yaad和Memora——它们超越了现有线性RNN的能力，同时保持了快速并行化的训练过程。我们的实验表明，Miras中不同的设计选择会产生具有不同优势的模型。例如，Miras的某些实例在诸如语言建模、常识推理和回忆密集型任务等特殊任务中表现出色，甚至优于Transformer和其他现代线性循环模型。
🔬 方法详解
问题定义：现有序列模型，如Transformer和线性RNN，在架构设计上缺乏统一的理论框架，并且在注意力机制和遗忘机制的选择上较为单一，限制了模型在不同任务上的性能表现。现有方法通常采用点积相似度或L2回归作为注意力偏置，遗忘机制也缺乏灵活性。
核心思路：论文的核心思路是将序列模型视为关联记忆模块，通过学习键值对的映射来实现序列建模。借鉴人类认知中的注意力偏置现象，将模型的内部目标定义为注意力偏置，并探索不同的注意力偏置函数。同时，将遗忘机制视为一种保留正则化，通过引入可定制的遗忘门来控制信息的保留和遗忘。
技术框架：Miras框架包含四个主要组成部分：(1)关联记忆架构：选择合适的记忆模块，如Transformer或线性RNN；(2)注意力偏置目标：定义键值对之间的相似度度量，如点积、L2距离或其他可学习的函数；(3)保留门：引入遗忘门机制，控制信息的保留和遗忘，实现保留正则化；(4)记忆学习算法：选择合适的优化算法来训练模型，学习键值对的映射。基于Miras框架，论文提出了三种新的序列模型：Moneta、Yaad和Memora。
关键创新：论文的关键创新在于提出了Miras框架，将序列模型的设计问题转化为四个可配置的选择：关联记忆架构、注意力偏置目标、保留门和记忆学习算法。通过探索不同的配置组合，可以设计出适用于不同任务的序列模型。此外，论文还提出了新的注意力偏置函数和遗忘门机制，进一步提升了模型的性能。
关键设计：注意力偏置目标的设计是关键。论文探索了多种注意力偏置函数，包括点积、L2距离、可学习的线性变换等。为了稳定训练过程，论文还提出了这些注意力偏置函数的有效近似。保留门的设计也至关重要，论文提出了一组新的遗忘门，可以根据输入动态地调整信息的保留和遗忘。
🖼️ 关键图片
📊 实验亮点
实验结果表明，基于Miras框架设计的Moneta、Yaad和Memora模型在语言建模、常识推理和回忆密集型任务上表现出色，甚至优于Transformer和其他现代线性循环模型。例如，在某些任务上，Miras模型能够达到与Transformer相当甚至更高的性能，同时保持了更快的训练速度和更低的计算成本。
🎯 应用场景
该研究成果可应用于自然语言处理领域的各种任务，如语言建模、机器翻译、文本摘要、对话生成等。通过选择合适的Miras配置，可以针对特定任务设计出高性能的序列模型。此外，该框架还可以推广到其他序列数据处理领域，如语音识别、时间序列预测等。
📄 摘要（原文）
Designing efficient and effective architectural backbones has been in the core of research efforts to enhance the capability of foundation models. Inspired by the human cognitive phenomenon of attentional bias-the natural tendency to prioritize certain events or stimuli-we reconceptualize neural architectures, including Transformers, Titans, and modern linear recurrent neural networks as associative memory modules that learn a mapping of keys and values using an internal objective, referred to as attentional bias. Surprisingly, we observed that most existing sequence models leverage either (1) dot-product similarity, or (2) L2 regression objectives as their attentional bias. Going beyond these objectives, we present a set of alternative attentional bias configurations along with their effective approximations to stabilize their training procedure. We then reinterpret forgetting mechanisms in modern deep learning architectures as a form of retention regularization, providing a novel set of forget gates for sequence models. Building upon these insights, we present Miras, a general framework to design deep learning architectures based on four choices of: (i) associative memory architecture, (ii) attentional bias objective, (iii) retention gate, and (iv) memory learning algorithm. We present three novel sequence models-Moneta, Yaad, and Memora-that go beyond the power of existing linear RNNs while maintaining a fast parallelizable training process. Our experiments show different design choices in Miras yield models with varying strengths. For example, certain instances of Miras achieve exceptional performance in special tasks such as language modeling, commonsense reasoning, and recall intensive tasks, even outperforming Transformers and other modern linear recurrent models.