专题解读 | 大语言模型中的记忆设计

发布日期：2024-07-05 20:07:00 浏览次数： 5246

作者：北邮 GAMMA Lab

微信搜一搜，关注“北邮 GAMMA Lab”

大语言模型中的记忆设计

前言

Wang 等人^[1]提出基于 LLM 的自主智能体架构设计应该包括四个组成部分:分析模块、记忆模块、规划模块和动作模块（如图1所示）。记忆模块在智能体架构设计中起着非常重要的作用。它不仅存储从环境中感知到的信息，还能通过分析这些信息来帮助智能体更好地理解和适应其所处的环境。这种信息的处理和存储机制，使得智能体能够从过去的行为中学习，并基于历史数据预测和规划未来的行动。因此，记忆模块是链接感知信息与智能决策的关键桥梁。

图1：基于大语言模型的自主智能体架构设计的统一框架

此外，记忆模块通过积累的经验促进智能体的自我进化。例如，通过对成功和失败的行动进行回顾和总结，智能体可以优化其行为策略，逐渐提升处理复杂情况的能力。这种进化过程是动态的，反映了智能体在与环境互动中逐步完善自身的能力。

记忆模块还有助于智能体维持行为的一致性和逻辑性。通过长期存储重要的决策依据和行为准则，智能体能够在面对类似情境时，快速提取相关记忆，确保行为的一致性和合理性。这不仅提高了响应的效率，也保证了行动的适当性和有效性。

基于注意力的记忆设计

在大语言模型中，主要有两种记忆设计方案。一种方案是基于注意力的记忆设计，这种设计深入探究大语言模型的内部机制，通过保存模型过去输出的中间状态作为记忆内容，并在未来的处理中根据需要检索这些记忆来辅助当前响应的生成。

MemTRM

具体而言，如图2所示的 MemTRM^[2]，对 Transformer 模型的顶层进行了创新性的修改。该系统采用了k最近邻和局部注意力机制的结合，形成了一种混合注意力机制。在这种机制下，模型在处理输入数据时，会首先将顶层的键（K）和值（V）向量存储于一个外部的记忆库中。随后，在模型需要生成响应时，通过k近邻搜索算法检索记忆库中与当前输入最相关的内容（token 级别），并使用这些内容来生成混合注意力层的语义嵌入。

这一设计不仅使得混合层能够关注到当前的输入信息，还能够“回忆”起模型在处理之前相似情境时的状态，从而有效扩充了Transformer的输入上下文长度。此外，通过这种方式，模型能够更加精准地捕捉到与当前任务相关的历史信息，从而提高整体的处理效率和响应质量。

图2：通过访问之前见过的子序列的键值对来扩展 transformer

LongMEM

然而，MemTRM系统在训练过程中遭遇了记忆陈旧的问题。具体来说，随着模型训练的继续进行，后期保存的键值对与前期保存的键值对在数据分布上出现了显著差异。这种分布的变化会导致模型在检索记忆时的效率和准确性降低，对于模型的性能是一大挑战。

为了解决这一问题，LongMEM^[3] 提出了一种创新的解决方案，即将记忆的存储和检索过程进行解耦（如图3所示）。该方法中，LongMEM 选择冻结 Transformer 的某些参数，从而保持了键值对的分布稳定。这种设计有效避免了由于模型持续学习导致的记忆陈旧问题。

此外，LongMEM 还引入了一个可训练的边缘网络（SideNet），以及残差连接来融合历史信息。边缘网络专门用来处理和整合来自记忆库的信息，而残巻连接则帮助模型在保持前向传播的同时，加入更多历史上下文。这种结构不仅解决了记忆陈旧的问题，还通过扩展Transformer模型的输入上下文长度，进一步提高了模型对长期依赖信息的处理能力。

图3：LONGMEM 的整体架构，“MemAug”表示记忆增强层

基于RAG的记忆设计

另一种记忆设计方案是使用 RAG（Retrieval-Augmented Generation）的方式来提供记忆支持。这种方法的特点是记忆模块与模型主体相对独立，不需要直接访问大模型的内部内容，因此特别适用于闭源模型的记忆增强。与基于注意力的设计不同，RAG 方案主要保存句子级别的语义信息，而非单个 token 的信息。

MemoryBank

在外部记忆库的实现上，这种方案具有多种形式。例如，MemoryBank^[4]，它采用向量库的形式存储了过去的对话记录、事件摘要和用户特征，利用向量相似度计算技术来加速记忆的检索过程。如图4所示，MemoryBank 将对话数据作为基本的记忆单元，不仅存储简单的对话文本，还从这些对话中动态提取并总结出关键事件和用户画像，这些高阶的记忆信息为智能体提供了丰富的背景知识。

此外，MemoryBank 还引入了类似艾宾浩斯遗忘曲线的机制来管理记忆的存储。这意味着每条记忆都有一个与之相关的“记忆强度”，该强度会随时间逐渐衰减。当记忆强度降至某一阈值以下时，该记忆将被视为低价值信息并从记忆库中清除，从而为新的记忆释放空间。

在实际应用中，当智能体如 SiliconFriend 接收到一个查询（Query）时，它会从 MemoryBank 中检索与该查询相关的事件总结、用户画像及相关对话。这些信息被用来填充“Meta Prompt”，进而帮助模型生成更加精准和具有建议性的回答。这种基于检索的记忆增强方法不仅提高了响应的相关性和质量，也使得模型能够更好地适应用户的具体需求和历史背景。

图4：MemoryBank 的整体架构以及 MemoryBank 增强的人工智能伙伴—SiliconFriend

AI-town

AI-town⁵ 成功地构建了一个名为 Smallville 的虚拟小镇，这里居住着25个 AI 智能体，他们不仅有各自的工作，还能进行社交活动，如八卦交流、结交新朋友，甚至举办情人节派对。每个智能体都拥有独特的个性和背景故事，这些特征丰富了小镇的社会动态和互动复杂性。

记忆检索是AI-town运行的关键环节。如图5所示，该系统在检索记忆时考虑三个主要因素：相关度、新近度和重要性。这三个因素的得分相乘，形成了记忆片段的最终检索得分。得分越高的记忆片段越容易被检索出来，从而被用于指导智能体当前的行动和决策。

相关度衡量的是查询（Query）和记忆片段之间的语义匹配程度。高相关度意味着记忆内容与当前查询紧密相关，因此更有可能提供有用的信息。
新近度偏好最近的记忆片段，基于这样一个假设：最近的事件可能与当前情况更相关，因此赋予这些记忆更高的权重。
重要性则是根据智能体对记忆内容的重要性评价来分配分数。这允许智能体区分哪些记忹是关键的核心记忆，哪些则是可有可无的普通记忆。

这种综合评分机制确保了智能体能够高效地从庞大的记忆库中检索出最相关、最有用的信息，从而提高响应的准确性和适时性。通过这种方式，AI-town中的智能体能够更加自然和有效地与环境及其他居民互动，展现出更贴近真实世界的社会行为。

图5：AI-town 的检索方式

AI-town 在其智能体设计中采取了一种创新性的反思机制，进一步提升了智能体处理信息和做出决策的深度。这个机制不仅仅停留在简单的观察层面，而是将这些观察转化为在特定条件下的更抽象、更高阶的反思。如图6所展示的，智能体 Klaus Mueller 的记忆被构建成一棵反思树，其中叶节点代表对世界的直接观察，而这些观察经过层层递归的总结和综合，最终形成了表达 Klaus 高阶自我观念的节点，即他对研究工作的高度投入与专注。

在记忆检索阶段，AI-town 的系统会平等地处理这两种记忆：基础的观察和高阶的反思。这种方法使得智能体在处理查询或决策时，不仅考虑具体事件的直接信息，还融入了从过去经验中提炼出的深层洞察和策略。通过这种方式，智能体能够展现出更加复杂和成熟的行为模式，更好地适应和响应其社会环境。

图6：Klaus Mueller 的反思树

ChatDB

在 ChatDB^[6] 中，记忆支持采用了数据库形式的组织结构，这种结构为历史信息的精确记录和管理提供了强有力的支持。数据库作为外部记忆库，不仅能够存储大量数据，还能够通过结构化查询语言（SQL）高效地进行数据检索和操作，从而大大提高了信息处理的准确性和效率。

此外，ChatDB 特别引入了有监督微调技术（Supervised fine-tuning, SFT），这一技术使得 ChatDB 能够理解和转换用户的自然语言查询。具体来说，当用户通过自然语言提出查询请求时，ChatDB 可以自动将这些查询转换为适当的 SQL 语句。这一过程不仅简化了用户与系统的交互，也使得非技术用户无需了解复杂的 SQL 语法即可直接利用数据库资源。