An Evolved Universal Transformer Memory
作者: Edoardo Cetin, Qi Sun, Tianyu Zhao, Yujin Tang
分类: cs.LG, cs.AI, cs.CL
发布日期: 2024-10-17 (更新: 2025-02-13)
备注: Published at ICLR 2025. Source code available at https://github.com/SakanaAI/evo-memory
💡 一句话要点
提出神经注意力记忆模型，提升Transformer长文本处理效率与性能
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer 长文本处理 注意力机制 记忆管理 神经网络 零样本迁移 模型压缩
📋 核心要点
- 现有方法通过人工规则丢弃上下文信息以降低计算成本，但牺牲了模型性能，存在效率与性能的权衡。
- 本文提出神经注意力记忆模型（NAMMs），利用可学习的记忆管理网络，为Transformer提供更高效和相关的上下文信息。
- 实验表明，NAMMs在多个长文本基准测试中显著提升性能，同时大幅减少输入上下文长度，并具备跨模态迁移能力。
📝 摘要（中文）
为了克服现代大型模型不断增长的计算成本，现有方法通常采用手工设计的规则来丢弃上下文中的特定部分，以期在保持原始性能的同时提高效率。本文提出神经注意力记忆模型（NAMMs），通过引入一个可学习的记忆管理网络来改进Transformer的性能和效率，从而克服了这种权衡。NAMMs构建于预训练Transformer之上，为每个层和注意力头提供不同的潜在上下文，专注于最相关的信息。NAMMs具有通用性，可应用于任何使用自注意力机制的模型，因为它们仅以生成的注意力矩阵中的值为条件。通过在少量问题上学习NAMMs，可以在多个长上下文基准测试中实现显著的性能提升，同时将模型的输入上下文减少到原始大小的一小部分。本文还证明了这种条件作用的通用性，使得仅在语言上训练的NAMMs能够零样本迁移到全新的Transformer架构，甚至跨越输入模态，其优势可以扩展到视觉和强化学习领域。
🔬 方法详解
问题定义：现有大型Transformer模型在处理长文本时面临计算成本过高的问题。为了降低成本，一些方法采用人工设计的规则来丢弃部分上下文信息，但这往往会导致模型性能下降，无法兼顾效率和性能。因此，如何有效地管理和利用长文本上下文信息，在降低计算成本的同时保持甚至提升模型性能，是本文要解决的核心问题。
核心思路：本文的核心思路是引入一个可学习的记忆管理网络（NAMMs），该网络能够根据Transformer的每一层和注意力头的需求，动态地选择和保留最相关的上下文信息。通过学习的方式，NAMMs能够更精准地捕捉到对当前任务最重要的信息，从而避免了人工规则的局限性，实现了效率和性能的双重提升。
技术框架：整体框架是在预训练的Transformer模型之上构建NAMMs。NAMMs以Transformer的每一层和注意力头产生的注意力矩阵的值作为输入，学习如何选择和保留上下文信息。具体流程如下：1) Transformer模型处理输入序列，生成注意力矩阵；2) NAMMs以注意力矩阵为条件，学习生成一个记忆选择掩码；3) 该掩码用于过滤掉不重要的上下文信息，保留关键信息；4) 过滤后的上下文信息被用于后续的Transformer层计算。
关键创新：本文最重要的技术创新在于提出了神经注意力记忆模型（NAMMs），它是一种可学习的记忆管理机制，能够动态地选择和保留对当前任务最相关的上下文信息。与传统的手工规则相比，NAMMs能够更精准地捕捉到关键信息，从而在降低计算成本的同时提升模型性能。此外，NAMMs的设计具有通用性，可以应用于任何使用自注意力机制的模型，并且可以跨模态迁移。
关键设计：NAMMs的关键设计包括：1) 以注意力矩阵的值作为条件输入，使得NAMMs能够根据Transformer的内部状态动态地选择上下文信息；2) 使用神经网络学习记忆选择掩码，使得NAMMs能够自动地学习到哪些信息是重要的；3) 设计损失函数，鼓励NAMMs选择对模型性能提升最有效的上下文信息。具体的网络结构和参数设置在论文中有详细描述，但核心思想是通过学习的方式，让NAMMs能够自适应地管理和利用上下文信息。
🖼️ 关键图片
📊 实验亮点
实验结果表明，NAMMs在多个长上下文基准测试中实现了显著的性能提升，同时大幅减少了输入上下文的长度。例如，在某些任务上，NAMMs可以将输入上下文减少到原始大小的1/10，同时保持甚至提升模型性能。此外，NAMMs还展现出了强大的零样本迁移能力，可以将其在语言任务上训练的模型直接应用于视觉和强化学习任务，并取得良好的效果。
🎯 应用场景
该研究成果可广泛应用于需要处理长文本的自然语言处理任务，如机器翻译、文本摘要、问答系统等。通过降低计算成本和提升模型性能，该方法有望推动大型语言模型在资源受限环境下的应用，并促进跨模态学习的发展，例如在视觉和强化学习领域。
📄 摘要（原文）
Prior methods propose to offset the escalating costs of modern foundation models by dropping specific parts of their contexts with hand-designed rules, while attempting to preserve their original performance. We overcome this trade-off with Neural Attention Memory Models (NAMMs), introducing a learned network for memory management that improves both the performance and efficiency of transformers. We evolve NAMMs atop pre-trained transformers to provide different latent contexts focusing on the most relevant information for individual layers and attention heads. NAMMs are universally applicable to any model using self-attention as they condition exclusively on the values in the produced attention matrices. Learning NAMMs on a small set of problems, we achieve substantial performance improvements across multiple long-context benchmarks while cutting the model's input contexts up to a fraction of the original sizes. We show the generality of our conditioning enables zero-shot transfer of NAMMs trained only on language to entirely new transformer architectures even across input modalities, with their benefits carrying over to vision and reinforcement learning.