TokMem: Tokenized Procedural Memory for Large Language Models

📄 arXiv: 2510.00444v1 📥 PDF

作者: Zijun Wu, Yongchang Hao, Lili Mou

分类: cs.CL

发布日期: 2025-10-01


💡 一句话要点

提出TokMem:一种用于大型语言模型的Token化过程记忆,提升任务泛化与效率。

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

关键词: 过程记忆 大型语言模型 Token化 提示工程 函数调用 知识表示 持续学习

📋 核心要点

  1. 现有大语言模型依赖提示工程,效率低、泛化差,且缺乏模块化复用机制。
  2. TokMem将重复过程编码为可训练的Token嵌入,实现恒定开销下的目标行为控制。
  3. 实验表明,TokMem在原子回忆和函数调用任务上优于检索增强生成,且参数量更少。

📝 摘要(中文)

大型语言模型严重依赖提示来指定任务、回忆知识和指导推理。然而,这种依赖效率低下,因为提示必须在每个步骤中重新读取,跨任务的可扩展性差,并且缺乏模块化重用机制。我们引入了TokMem,一种Token化的过程记忆,它将重复出现的过程存储为紧凑的可训练嵌入。每个记忆Token编码一个过程的地址和一个控制信号,以恒定大小的开销实现目标行为。为了支持持续适应,TokMem保持骨干模型冻结,允许添加新过程而不干扰现有过程。我们在1000个原子回忆任务和函数调用组合回忆任务上评估了TokMem,它始终优于检索增强生成,同时避免了重复的上下文开销,并且使用更少的参数进行微调。这些结果表明,TokMem是提示工程和微调的可扩展和模块化替代方案,为LLM提供了一种显式的过程记忆。

🔬 方法详解

问题定义:大型语言模型依赖提示(prompts)来完成任务,包括知识回忆、推理引导等。然而,这种方式存在效率问题,因为每次都需要重新读取提示,并且在不同任务之间的泛化能力较差,缺乏模块化和可重用性。现有方法如检索增强生成(RAG)虽然可以缓解部分问题,但仍然存在上下文冗余和计算开销。

核心思路:TokMem的核心思想是将重复使用的过程(procedural knowledge)存储为可训练的Token嵌入,这些Token嵌入可以被视为一种“过程记忆”。每个Token不仅包含指向特定过程的地址,还包含控制信号,用于指导模型的生成过程。通过这种方式,模型可以在需要时快速调用相应的过程,而无需重复读取冗长的提示。

技术框架:TokMem主要包含两个部分:Token化的过程记忆模块和语言模型。过程记忆模块负责存储和管理过程Token,语言模型则负责利用这些Token进行生成。当模型需要执行某个任务时,首先从过程记忆中选择相关的Token,然后将这些Token添加到输入序列中。语言模型根据这些Token的引导,生成相应的输出。关键在于,TokMem允许在不微调整个LLM的情况下添加新的过程,从而实现持续学习和适应。

关键创新:TokMem的关键创新在于它将过程知识显式地存储为可训练的Token,并利用这些Token来控制模型的行为。与传统的提示工程相比,TokMem可以避免重复的上下文开销,并提高模型的效率和泛化能力。与微调相比,TokMem可以在不干扰现有知识的情况下添加新的过程,从而实现持续学习。

关键设计:TokMem的设计包括以下几个关键方面:1) 过程Token的嵌入表示:如何将过程知识编码为紧凑的Token嵌入?2) Token选择机制:如何根据当前任务选择相关的过程Token?3) 控制信号的设计:如何利用Token中的控制信号来指导模型的生成过程?论文中可能使用了特定的损失函数来训练过程Token的嵌入表示,并设计了特定的网络结构来实现Token选择和控制信号的传递。具体的技术细节需要参考论文原文。

📊 实验亮点

TokMem在1000个原子回忆任务和函数调用组合回忆任务上进行了评估,实验结果表明,TokMem始终优于检索增强生成(RAG),同时避免了重复的上下文开销,并且使用更少的参数进行微调。这些结果证明了TokMem作为一种可扩展和模块化替代方案的有效性。

🎯 应用场景

TokMem具有广泛的应用前景,例如可以应用于智能助手、代码生成、对话系统等领域。通过将常用的操作流程存储为过程记忆,可以显著提高这些系统的效率和智能化水平。此外,TokMem还可以用于构建可扩展的知识库,支持持续学习和知识更新。未来,TokMem有望成为构建更智能、更高效的大型语言模型的重要技术。

📄 摘要(原文)

Large language models rely heavily on prompts to specify tasks, recall knowledge and guide reasoning. However, this reliance is inefficient as prompts must be re-read at each step, scale poorly across tasks, and lack mechanisms for modular reuse. We introduce TokMem, a tokenized procedural memory that stores recurring procedures as compact, trainable embeddings. Each memory token encodes both an address to a procedure and a control signal that steers generation, enabling targeted behavior with constant-size overhead. To support continual adaptation, TokMem keeps the backbone model frozen, allowing new procedures to be added without interfering with existing ones. We evaluate TokMem on 1,000 tasks for atomic recall, and on function-calling tasks for compositional recall, where it consistently outperforms retrieval-augmented generation while avoiding repeated context overhead, and fine-tuning with far fewer parameters. These results establish TokMem as a scalable and modular alternative to prompt engineering and fine-tuning, offering an explicit procedural memory for LLMs.