KEEP: A KV-Cache-Centric Memory Management System for Efficient Embodied Planning
이 논문은 텍스트 기반 메모리 방식의 긴 프롬프트와 높은 지연 시간 문제를 해결하기 위해 KV 캐시 재계산을 줄이고 메모리 상호작용을 동적으로 최적화하는 'KEEP' 시스템을 제안하여 embodied planning 의 속도와 성공률을 크게 향상시킨다고 요약할 수 있습니다.
원저자:Zebin Yang, Tong Xie, Baotong Lu, Shaoshan Liu, Bo Yu, Meng Li
상상해 보세요. 로봇이 "감자를 씻어서 싱크대에 넣으세요"라는 명령을 받았습니다. 로봇은 과거에 배운 경험 (예: 칼이 어디 있는지, 감자가 어떻게 생겼는지) 을 기억에서 찾아서 계획을 세웁니다.
하지만 기존 방식에는 두 가지 큰 문제가 있었습니다.
기억이 너무 길어서 머리가 아픕니다 (긴 문장 문제): 로봇이 모든 기억을 글자 그대로 텍스트로 저장하면, 명령을 내릴 때마다 수만 자에 달하는 긴 글을 다시 한 번 읽어야 합니다. 마치 매번 새로운 책을 처음부터 끝까지 읽어야만 다음 장을 이해할 수 있는 상황이라서, 로봇이 행동을 결정하는 데 시간이 너무 오래 걸립니다.
기억이 자주 바뀌어서 다시 계산해야 합니다 (기억 갱신 문제): 로봇이 "우유를 테이블에서 집어 올랐다"는 행동을 하면, 테이블 위에는 더 이상 우유가 없습니다. 이때 로봇은 우유가 있던 부분부터 끝까지 모든 기억을 다시 계산해야 합니다. 마치 레고 블록을 하나만 떼어냈는데, 그 뒤로 쌓인 모든 탑을 다 부수고 다시 쌓아야 하는 상황이라서 매우 비효율적입니다.
💡 해결책: KEEP 시스템의 3 가지 마법
KEEP 시스템은 이 문제를 해결하기 위해 로봇의 기억을 'KV 캐시 (Key-Value Cache, 기억의 핵심 데이터)' 형태로 저장하고 관리하는 3 가지 똑똑한 전략을 사용합니다.
1. 정적/동적 기억 분리 (Static-Dynamic Memory Construction)
비유: "변하지 않는 책장"과 "수시로 바뀌는 포스트잇"
기존 방식: 모든 기억을 똑같은 크기의 블록으로 쪼개서 관리했습니다. 하나만 바뀌어도 그 블록 전체를 다시 계산해야 했습니다.
KEEP 의 방식: 기억을 변하는 빈도에 따라 두 가지로 나눕니다.
정적 기억 (Static): "테이블은 항상 거실에 있다", "냉장고는 문이 닫혀 있다"처럼 거의 변하지 않는 기억은 한 번 계산해 두면 오랫동안 그대로 둡니다. (책장에 꽂힌 책처럼)
동적 기억 (Dynamic): "우유가 테이블 위에 있다", "감자가 씻겨 있다"처럼 자주 변하는 기억은 따로 관리합니다. (수시로 적어지는 포스트잇처럼)
효과: 우유를 치울 때, 테이블이나 냉장고 같은 '책장'은 건드리지 않고 '포스트잇'만 교체하면 되므로, 계산량을 획기적으로 줄일 수 있습니다.
2. 다단계 기억 재계산 (Multi-hop Memory Re-computation)
비유: "수색 게임의 단서 연결하기"
기존 방식: 중요한 기억을 찾기 위해 무작위하거나 고정된 위치만 다시 계산했습니다.
KEEP 의 방식: 로봇이 "문을 열어야 한다"는 명령을 받으면, 단순히 '문'만 찾는 게 아니라 연관된 기억들을 단계별로 찾아냅니다.
1 단계: "문을 열려면 열쇠가 필요하다."
2 단계: "열쇠는 테이블 위에 있었다."
3 단계: "테이블은 주방에 있었다."
이 **연결고리 (다단계)**를 따라가면서 가장 중요한 기억들만 다시 계산하고, 나머지는 건너뜁니다.
효과: 로봇이 중요한 단서들을 놓치지 않으면서도, 불필요한 기억은 다시 계산하지 않아 속도가 빨라지고 정확도도 유지됩니다.
3. 층별 균형 로딩 (Layer-balanced Memory Loading)
비유: "공장 조립 라인의 병목 현상 해결"
기존 방식: 로봇의 뇌 (AI 모델) 는 여러 층 (Layer) 으로 되어 있는데, 앞쪽 층은 계산이 빠르고 뒤쪽 층은 데이터가 많아 로딩이 느립니다. 이렇게 되면 앞쪽은 기다리고, 뒤쪽은 바쁘게 일하는 불균형이 생겨 전체 공정이 멈추는 '병목 현상'이 발생합니다.
KEEP 의 방식: 앞쪽 층이 계산하는 동안, 이미 앞쪽에서 계산이 끝난 뒤쪽 층의 데이터를 미리 미리 준비해 둡니다. (미리 로딩)
효과: 공장의 모든 작업자가 동시에 일할 수 있게 되어, 기다리는 시간이 사라지고 전체 처리 속도가 극대화됩니다.
🚀 결과: 얼마나 빨라졌나요?
이 시스템을 적용한 결과, 로봇은 다음과 같은 놀라운 성과를 거두었습니다.
속도: 기존 방식보다 약 2.7 배 더 빠릅니다. (첫 번째 행동을 결정하는 시간이 크게 단축됨)
성공률: 다른 최신 기술 (CacheBlend) 보다 성공 확률이 4% 이상 더 높습니다.
정확도: 속도를 높였지만, 로봇이 일을 잘못하는 일은 거의 없습니다.
📝 한 줄 요약
KEEP은 로봇의 기억을 '변하지 않는 것'과 '자꾸 변하는 것'으로 나누고, 중요한 단서끼리 연결해서 다시 계산하며, 데이터를 미리 준비시켜서 로봇이 더 빠르고 똑똑하게 집안일을 할 수 있게 해주는 지능형 기억 관리자입니다.
이 기술은 앞으로 우리가 집에서 로봇과 함께 살 때, 로봇이 더 자연스럽게 우리와 소통하고 일을 도와줄 수 있는 기반이 될 것입니다.
1. 문제 정의 (Problem)
신체화된 계획 (Embodied Planning) 작업에서 에이전트는 장기적인 목표를 달성하기 위해 과거 경험과 환경 상태를 기억 (Memory) 하여 전역적인 관점을 유지해야 합니다. 그러나 기존 접근 방식은 다음과 같은 한계가 있었습니다:
텍스트 기반 메모리의 비효율성: 메모리를 원시 텍스트 (Raw Text) 로 저장하여 프롬프트 길이가 과도하게 길어지고, 이로 인해 프리필 (Prefill) 지연 시간이 계획의 주요 병목 현상이 됩니다.
KV 캐시 재사용의 어려움: LLM 추론에서 KV(Key-Value) 캐시를 재사용하면 계산 비용을 줄일 수 있지만, 신체화된 환경에서는 에이전트의 행동에 따라 메모리 (객체 상태, 위치 등) 가 빈번하게 동적으로 업데이트됩니다.
기존 KV 재사용 기법의 실패:
고정된 크기의 블록 단위로 KV 를 관리하면, 메모리 업데이트 시 해당 블록 이후의 모든 KV 캐시가 무효화되어 재계산이 필요합니다.
반대로 블록 크기를 줄이면 KV 캐시 무효화는 줄어들지만, 서로 다른 메모리 블록 간의 크로스 어텐션 (Cross-attention) 관계가 무시되어 계획 정확도가 떨어집니다.
기존 KV 재계산 (Re-computation) 기법들은 정적인 규칙 (예: 블록의 시작/끝 토큰) 에 의존하여, 현재 문맥에 따라 중요도가 변하는 메모리 간의 상호작용을 효과적으로 포착하지 못합니다.
2. 제안 방법: KEEP (Methodology)
이 논문은 효율적인 신체화된 계획을 위해 KV 캐시 중심의 메모리 관리 시스템인 KEEP을 제안합니다. KEEP 은 세 가지 핵심 혁신을 통해 메모리 업데이트에 따른 KV 캐시 무효화를 최소화하고, 중요한 메모리 간 상호작용을 복원합니다.
가. 정적 - 동적 메모리 구성 (Static-Dynamic Memory Construction)
동기: 메모리 세그먼트마다 업데이트 빈도가 다릅니다 (예: 테이블 위치는 자주 변하지 않으나, 우유 위치는 자주 변함).
방식:
정적 그룹 (Static Groups): 최근 t단계 동안 변경되지 않은 메모리 세그먼트들을 그룹화하여, 그룹 내부에서 전체 크로스 어텐션을 포함한 KV 캐시를 한 번에 계산합니다. 이는 메모리 간 풍부한 상호관계를 유지합니다.
동적 그룹 (Dynamic Groups): 최근 t단계 동안 업데이트된 메모리 세그먼트들은 개별 세그먼트 단위로 KV 캐시를 계산합니다. 이는 하나의 업데이트가 전체 그룹의 캐시를 무효화하는 것을 방지합니다.
전환: 메모리의 상태가 안정화되거나 변경되면 그룹의 정적/동적 상태를 동적으로 전환하고 KV 캐시를 재구성합니다.
나. 멀티홉 메모리 재계산 (Multi-hop Memory Re-computation)
동기: 메모리의 중요성은 쿼리와 전체 문맥에 따라 동적으로 변합니다. 고정된 위치의 토큰만 재계산하는 기존 방식은 중요한 간접적인 메모리 연결 (예: "잠긴 문" -> "열쇠" -> "테이블") 을 놓칠 수 있습니다.
방식:
중요도 전파 (Importance Propagation): 현재 쿼리에 중요한 메모리를 식별한 후, 해당 메모리에 중요한 2 차 메모리를 찾아내는 반복적 (Iterative) 과정을 거칩니다.
선택적 재계산: 각 레이어에서 어텐션 분포를 기반으로 중요도가 높은 메모리 세그먼트들을 동적으로 식별하고, 해당 세그먼트들의 KV 캐시만 재계산합니다.
세그먼트 단위: 토큰 단위가 아닌 메모리 세그먼트 단위로 재계산을 수행하여 I/O 오버헤드를 줄이고 연속적인 로딩을 가능하게 합니다.
다. 레이어 균형 메모리 로딩 (Layer-balanced Memory Loading)
동기: KV 캐시를 CPU RAM 에서 GPU 로 로드할 때, 레이어별 계산 부하와 로딩 부하가 불균형합니다 (초기 레이어는 재계산 비율이 높아 로딩 데이터가 적고, 후기 레이어는 로딩 데이터가 많음). 이는 파이프라인 버블 (Pipeline Bubbles) 을 유발합니다.
방식:
선제적 프리로딩 (Proactive Pre-loading): 초기 레이어의 계산이 완료되기 전에, 향후 레이어에서 재계산되지 않을 것으로 보장된 메모리 KV 캐시를 미리 로드합니다.
부하 분산: 초기 레이어의 유휴 로딩 엔진을 활용하여 후기 레이어의 로딩 부하를 분산시켜, 로딩과 계산이 병목 없이 겹쳐지도록 (Overlap) 최적화합니다.
3. 주요 기여 (Key Contributions)
메모리 관리 분석: 신체화된 계획과 기존 LLM 추론 간의 메모리 구성 및 관리의 근본적인 차이를 상세히 분석했습니다.
새로운 아키텍처 제안:
KV 캐시 무효화를 최소화하는 정적 - 동적 메모리 구성.
메모리 간 상호 연결성을 보존하는 멀티홉 메모리 재계산.
하드웨어 활용률을 극대화하는 레이어 균형 메모리 로딩 스케줄러를 제안했습니다.
성능 입증: ALFRED 및 WAH-NL 벤치마크에서 기존 방법 대비 압도적인 속도 향상과 정확도 개선을 입증했습니다.
4. 실험 결과 (Results)
ALFRED 데이터셋에서 수행된 실험 결과는 다음과 같습니다:
텍스트 기반 메모리 대비: 기존 텍스트 기반 메모리 방법 대비 **2.68 배의 속도 향상 (Speedup)**을 달성했으며, 정확도 손실은 미미했습니다.
KV 재계산 방법 (CacheBlend) 대비:
성공률 (Success Rate, SR):4.13% 향상.
첫 번째 토큰 도달 시간 (TTFT):1.90 배 단축.
기타: Qwen-14B 및 Qwen-32B 모델 모두에서 일관된 성능 개선을 보였으며, 검색된 메모리 세그먼트 수가 증가할수록 지연 시간 증가 폭이 다른 방법들보다 완만했습니다.
5. 의의 및 결론 (Significance)
KEEP 은 신체화된 에이전트가 복잡한 환경에서 장기 계획을 수립할 때 발생하는 **메모리 관리의 역설 (긴 프롬프트 vs. 빈번한 업데이트)**을 해결합니다.
효율성: 불필요한 KV 재계산을 줄이고, 하드웨어 파이프라인을 최적화하여 실시간 계획 속도를 획기적으로 높였습니다.
정확성: 동적인 중요도 전파를 통해 중요한 메모리 간의 맥락적 연결을 유지하여, 단순한 속도 향상뿐만 아니라 계획의 성공률도 높였습니다.
확장성: KV 캐시를 대규모 메모리 (RAM) 에 저장하고 필요한 부분만 로드하는 전략을 통해, 제한된 GPU 메모리 환경에서도 대규모 LLM 을 활용한 복잡한 에이전트 작업을 가능하게 합니다.
이 연구는 신체화된 인공지능 (Embodied AI) 분야에서 LLM 기반 계획의 실용성을 높이는 중요한 이정표로 평가됩니다.