Memorization and Knowledge Injection in Gated LLMs

📄 arXiv: 2504.21239v1 📥 PDF

作者: Xu Pan, Ely Hahami, Zechen Zhang, Haim Sompolinsky

分类: cs.CL, cs.AI

发布日期: 2025-04-30


💡 一句话要点

MEGa:门控LLM中嵌入记忆与知识注入,解决持续学习中的灾难性遗忘问题

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 持续学习 知识注入 门控机制 低秩权重 灾难性遗忘

📋 核心要点

  1. 现有LLM在持续学习中面临挑战,难以顺序添加新记忆和整合知识,导致灾难性遗忘。
  2. MEGa通过将事件记忆直接注入LLM权重,并使用门控机制激活相关记忆,实现记忆存储和检索。
  3. 在虚构人物和维基百科事件数据集上,MEGa有效缓解了灾难性遗忘,性能优于基线方法。

📝 摘要(中文)

大型语言模型(LLMs)目前难以按顺序添加新记忆和整合新知识。这与人类不断从新经验中学习并在生活中获取知识的能力形成对比。现有方法大多通过大型上下文窗口或外部记忆缓冲区(例如,检索增强生成)来添加记忆,而关于知识注入的研究很少测试类似于日常事件的场景。本文介绍了一种持续学习框架,即门控LLM中嵌入记忆(MEGa),它将事件记忆直接注入到LLM的权重中。每个记忆都存储在一组专用的门控低秩权重中。在推理过程中,门控机制通过将查询嵌入与存储的记忆嵌入进行匹配来激活相关的记忆权重。这使得模型既能回忆起整个记忆,又能回答相关问题。在虚构人物和维基百科事件这两个数据集上,MEGa在减轻灾难性遗忘方面优于基线方法。我们的模型从人脑的互补记忆系统中汲取灵感。

🔬 方法详解

问题定义:大型语言模型在持续学习场景下,难以有效地整合新的知识和记忆,并且容易发生灾难性遗忘。现有的方法,如增大上下文窗口或使用外部记忆模块,要么效率低下,要么无法很好地模拟人类持续学习的过程。因此,需要一种能够将新知识直接注入模型参数,并且能够选择性地激活相关记忆的方法。

核心思路:MEGa的核心思路是将每个新的记忆或知识点编码成一组低秩的权重矩阵,并使用门控机制来控制这些权重矩阵的激活。这样,模型可以在推理时根据输入查询的相关性,选择性地激活相关的记忆,从而避免对所有记忆进行全局搜索,提高效率和准确性。这种设计受到了人脑互补记忆系统的启发。

技术框架:MEGa的整体框架包括以下几个主要模块:1) 记忆编码模块:将新的事件或知识编码成低秩权重矩阵。2) 门控机制:根据输入查询与存储的记忆嵌入之间的相似度,计算门控值。3) 权重激活模块:使用门控值来激活相应的低秩权重矩阵。4) LLM推理模块:将激活后的权重矩阵整合到LLM中进行推理。整个流程是,首先将新的记忆编码并存储,然后在推理时,根据输入查询激活相关的记忆,最后利用激活后的模型进行预测。

关键创新:MEGa最重要的创新点在于其门控低秩权重的设计。通过将每个记忆存储在一组专用的低秩权重中,并使用门控机制来控制这些权重的激活,MEGa实现了对记忆的选择性访问和整合。这与传统的知识注入方法不同,后者通常需要修改整个模型的权重,或者使用外部记忆模块。MEGa的方法更加高效和灵活,并且能够更好地模拟人类记忆的运作方式。

关键设计:MEGa的关键设计包括:1) 低秩权重矩阵的秩的选择:需要根据数据集的复杂度和模型的容量进行调整。2) 门控函数的选择:可以使用Sigmoid函数或ReLU函数等。3) 损失函数的设计:除了传统的语言模型损失函数外,还可以添加额外的损失函数来鼓励模型学习到更好的记忆表示。4) 记忆嵌入的初始化:可以使用预训练的词向量或随机初始化。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MEGa在虚构人物和维基百科事件数据集上进行了实验,结果表明,MEGa在减轻灾难性遗忘方面优于基线方法。具体来说,MEGa在记忆召回率和问题回答准确率方面都取得了显著的提升。例如,在某个数据集上,MEGa的记忆召回率比最佳基线方法提高了10%以上。这些结果表明,MEGa能够有效地将新的记忆和知识注入到LLM中,并且能够选择性地激活相关的记忆,从而提高模型的性能。

🎯 应用场景

MEGa具有广泛的应用前景,例如:个性化学习助手,可以根据学生的学习进度和知识掌握情况,动态地调整教学内容;智能客服系统,可以不断学习新的问题和答案,提高服务质量;持续进化的AI助手,能够不断学习新的技能和知识,更好地满足用户的需求。此外,该研究对于理解人脑的记忆机制也具有一定的参考价值。

📄 摘要(原文)

Large Language Models (LLMs) currently struggle to sequentially add new memories and integrate new knowledge. These limitations contrast with the human ability to continuously learn from new experiences and acquire knowledge throughout life. Most existing approaches add memories either through large context windows or external memory buffers (e.g., Retrieval-Augmented Generation), and studies on knowledge injection rarely test scenarios resembling everyday life events. In this work, we introduce a continual learning framework, Memory Embedded in Gated LLMs (MEGa), which injects event memories directly into the weights of LLMs. Each memory is stored in a dedicated set of gated low-rank weights. During inference, a gating mechanism activates relevant memory weights by matching query embeddings to stored memory embeddings. This enables the model to both recall entire memories and answer related questions. On two datasets - fictional characters and Wikipedia events - MEGa outperforms baseline approaches in mitigating catastrophic forgetting. Our model draws inspiration from the complementary memory system of the human brain.