用键值缓存优化大模型长时程智能体规划,提速近3倍且精度不降。
KEEP: A KV-Cache-Centric Memory Management System for Efficient Embodied Planning
- 以键值缓存为核心构建混合粒度记忆,减少重复计算
- 多跳重计算机制动态优化跨记忆组关联,提升推理效率
- 分层均衡加载避免计算负载不均,适合长期规划场景
带有记忆的大型语言模型在复杂、长时程的具身规划任务中表现出色。通过记录过往经验与环境状态,记忆使模型能保持全局视野,避免重复探索。然而,现有方法通常以原始文本形式存储记忆,导致提示过长、预填充延迟高。尽管可存储和复用键值(KV)缓存,但频繁更新严重削弱了其效率。本文提出KEEP,一种面向高效具身规划的键值缓存中心型内存管理机制。主要创新包括:(1) 静态-动态记忆构建算法,通过混合粒度记忆组减少KV缓存重计算;(2) 多跳记忆重计算算法,动态识别不同记忆组间的重点交叉注意力并迭代重构;(3) 分层均衡记忆加载,消除各层间KV缓存加载与交叉注意力计算的不平衡。大量实验表明,在ALFRED数据集上,KEEP相较文本记忆方法实现2.68倍加速,精度损失可忽略;相比同为重计算方法的CacheBlend(EuroSys'25),成功率达4.13%提升,首次生成时间(TTFT)降低1.90倍。代码已开源:https://github.com/PKU-SEC-Lab/KEEP_Embodied_Memory。
原文摘要 · Abstract (English)
Memory-augmented Large Language Models (LLMs) have demonstrated remarkable capability for complex and long-horizon embodied planning. By keeping track of past experiences and environmental states, memory enables LLMs to maintain a global view, thereby avoiding repetitive exploration. However, existing approaches often store the memory as raw text, leading to excessively long prompts and high prefill latency. While it is possible to store and reuse the KV caches, the efficiency benefits are greatly undermined due to frequent KV cache updates. In this paper, we propose KEEP, a KV-cache-centric memory management system for efficient embodied planning. KEEP features 3 key innovations: (1) a Static-Dynamic Memory Construction algorithm that reduces KV cache recomputation by mixed-granularity memory group; (2) a Multi-hop Memory Re-computation algorithm that dynamically identifies important cross-attention among different memory groups and reconstructs memory interactions iteratively; (3) a Layer-balanced Memory Loading that eliminates unbalanced KV cache loading and cross-attention computation across different layers. Extensive experimental results have demonstrated that KEEP achieves 2.68x speedup with negligible accuracy loss compared with text-based memory methods on ALFRED dataset. Compared with the KV re-computation method CacheBlend (EuroSys'25), KEEP shows 4.13% success rate improvement and 1.90x time-to-first-token (TTFT) reduction. Our code is available on https://github.com/PKU-SEC-Lab/KEEP_Embodied_Memory.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。