MemDecay: Region-Aware KV Cache Eviction for Efficient LLM Agent Inference

기술 요약: MemDecay

문제 정의

대규모 언어 모델(LLM)은 다회차 상호작용, 도구 사용, 반복적인 계획을 포함하는 장기적 워크플로우를 실행하는 에이전트로 점점 더 많이 배치되고 있습니다. 이러한 에이전트 트레이스(trace)는 시스템 지침, 안전 제약 조건, 진화하는 계획, 사용자 턴, 검색된 문서, 도구 출력, 중간 추론 상태와 같은 이질적인 구성 요소들을 축적합니다.

자기 회귀적 디코딩에 필요한 키-값(KV) 캐시는 이 트레이스와 함께 선형적으로 증가하며, 이는 상당한 메모리 오버헤드, 대역폭 요구량 및 지연 시간을 초래합니다. 기존의 KV 캐시 제거 정책은 일반적으로 모델 중심의 신호(예: 누적 어텐션, 최신성 윈도우 또는 어텐션 싱크)에 의존하며 전체 시퀀스에 대해 균일한 유지 규칙을 적용합니다. 이러한 접근 방식은 오케스트레이션된 프롬프트의 내재적인 의미 구조를 활용하지 못합니다. 에이전트 워크로드에서 서로 다른 구성 요소들은 각기 매우 다른 기능적 역할과 기대 수명을 가집니다. 예를 들어, 시스템 지침은 에피소드 전반에 걸쳐 지속되어야 하는 반면, 스크래치패드(scratchpad) 추론이나 오래된 도구 출력은 짧은 구간 동안만 유효할 수 있습니다. 모든 토큰을 동일한 기대 수명을 가진 것으로 취급하는 것은 이러한 구조를 무시하는 것이며, 이는 최적화되지 못한 메모리 관리를 초래합니다.

방법론: MemDecay

저자들은 장기적 LLM 에이전트를 위해 설계된 MemDecay라는, 학습이 필요 없는(training-free) 영역 인식형(region-aware) KV 캐시 제거 프레임을 제안합니다. 이 방법은 고정된 캐시 예산 하에서 작동하며, 오케스트레이션 레이어의 구조적 사전 지식(structural priors)을 어텐션 기반 중요도와 통합합니다.

핵심 구성 요소

1. 영역 레이블(Region Labels): 오케스트레이션 레이어는 연속적인 프롬프트 구간에 대한 영역 레이블(예: system, plan, user, tool_in, tool_out, retrieval, scratchpad)을 제공합니다.
2. 유지 점수(Retention Score): 각 토큰 $i$는 시간 $t$에서의 유지 점수 $s_i(t)$를 할당받으며, 다음과 같이 계산됩니다:
   $$s_i(t) = b_{r(i)} \exp(-\lambda_{r(i)} \Delta t_i(t)) + \alpha a_i(t)$$
   • 구조적 사전 지식: $b_{r(i)}$는 해당 토큰의 영역에 특화된 기본 우선순위입니다.
   • 시간적 감쇠(Temporal Decay): $\lambda_{r(i)}$는 영역별 감쇠율입니다. $\Delta t_i(t)$ 항은 토큰이 마지막으로 "새로고침(refreshed)"된 이후 경과된 시간을 나타냅니다.
   • 어텐션 중요도: $a_i(t)$는 토큰이 받은 어텐션 질량(attention mass)의 지수 가중 이동 평균(EWMA)입니다.
   • 새로고침 메커니즘: 토큰에 대한 어텐션이 임계값 $\tau$를 초과하면, 감쇠 시계가 리셋되어($\Delta t_i$가 0이 됨) 지속적인 사용에 기반하여 수명을 효과적으로 연장합니다.
3. 고정(Pinning): 시스템 영역(및 선택적으로 계획 세그먼트)은 "고정"될 수 있으며, 이는 점수에 관계없이 절대 제거되지 않음을 의미합니다.
4. 제거 정책(Eviction Policy): 캐시는 페이지 단위(예: 16개 토큰당 1페이지)로 관리됩니다. 예산을 초과할 경우, 페이지를 구성하는 토큰들의 평균 유지 점수에 따라 점수를 매깁니다. 점수가 가장 낮은 페이지가 먼저 제거됩니다.

캘리브레이션(Calibration)

MemDecay는 모델 가중치에 대해서는 학습이 필요 없지만, 각 영역의 감쇠율($\lambda_r$)을 맞추기 위해 가벼운 통계적 캘리브레이션을 사용합니다. 이 과정은 에이전트 트레이스로부터 영역별 어텐션 수명을 추정하여, 시스템이 LLM을 재학습시키지 않고도 특정 워크로드 특성에 적응할 수 있도록 합니다.

주요 기여

1. 정식화: 본 논문은 영역 인식형 선택적 망각을 LLM 에이전트를 위한 메모리 관리 문제로 정식화하며, 솔루션을 오케스트레이션된 프롬프트의 의미 구조에 근거하여 구축했습니다.
2. 알고리즘: 영역별 기본 우선순위, 시간적 감쇠, 그리고 어텐션 트리거 새로고침을 하나의 해석 가능한 유지 점수로 결합한 정책인 MemDecay를 소개합니다. 이는 기존 서빙 스택(예: PagedAttention)과 호환되는 페이지 단위 레이아웃에서 작동합니다.
3. 측정 프로토콜: 저자들은 영역별 어텐션 수명을 추정하고 감쇠율을 캘리브레이션하기 위한 프로토콜을 명시합니다. 이를 두 가지 컨텍스트 길이(~450 및 ~1,700 토큰)와 두 가지 모델 크기(Qwen2.5-1.5B 및 3B)에 대해 실행하여 총 4,320개의 점수 기반 제거 프로브(eviction probes)를 수행했습니다.

실험 결과

본 연구는 랜덤 제거, 최신성 기반(스트리밍) 유지, 누적 어텐션 유지(H2O 스타일)를 포함한 베이스라인들과 MemDecay를 비교 평가합니다.

주요 발견 사항

• 영역 수명 계층 구조: 측정된 어텐션 수명은 영역에 따라 수 자릿수 차이가 났습니다. 모든 설정에서 system 토큰의 반감기는 148189 디코드 단계였던 반면, scratchpad 토کن은 급격히 감소했습니다(1416 단계). Retrieval 토큰은 이동 가능한 영역 중 가장 긴 수명을 가진 것으로 나타났으며, 이는 이들이 빠르게 감쇠할 것이라는 기본 가정과 상반되는 결과입니다.
• 고정(Pinning)의 효과: 시스템 영역을 고정함으로써 지침의 생존을 보장합니다. MemDecay는 풀 캐시 용량 제한 상황에서 24개의 시스템 프로브 중 2124개를 보존한 반면, 최신성 기반 베이스라인은 02개만을 보존했습니다.
• 최신성 vs. 구조: 최신성 기반 유지(스트리밍)는 컨텍스트 길이가 길어짐에 따라 성능이 붕괴되며, 긴 티어(long tiers)에서 거의 0에 가까운 성능을 보였습니다. 반면, MemDec much의 구조적 사전 지식은 여전히 효과적입니다.
• 누적 어텐션: 고정되지 않은 콘텐츠의 경우, 누적 어텐션(H2O 스타일)은 모델 규모가 커질수록 강화되어 3B 모델 설정에서 MemDecay보다 높은 unpinned recall 성능을 보였습니다.
• MemDecay의 한계: MemDecay는 오래된 사실(예: 첫 번째 턴의 사용자 영역 사실)에 대한 unpinned recall 측면에서 어텐션 기반 베이스라인보다 특정 손실을 보였습니다. 절제 연구(ablation study) 결과, 어텐션 항($\alpha a_i$)의 크기가 이러한 토큰들에 대한 구조적 감쇠를 상쇄하기에는 너무 작았습니다. 저자들은 어텐션 항이 단순히 가중치 튜닝만으로는 효과적이지 않으며, (균등 어텐션 수준과 대비하여) 크기 정규화(magnitude normalization)가 필요하다고 지적합니다.

의의 및 주장

본 논문은 MemDecay를 에이전트 워크플로우의 의미 구조와 캐시 유지를 일치시키는 실용적이고 학습이 필요 없는 솔루션으로 포지셔닝합니다. 그 의의는 다음과 같습니다:

• 해석 가능성: 정책이 명시적인 영역 레이블과 감쇠율에 의해 구동되므로, 유지 결정이 투명합니다.
• 호환성: 모델 재학습이나 보조 모듈 없이 페이지 단위 캐시에서 작동하여 기존 서빙 스택에 적합합니다.
• 보장된 안정성: 핵심 영역을 고정함으로써 시스템 지침과 안전 제약 조건이 메모리 압박으로 인해 손실되지 않도록 보장하며, 이는 어텐션이나 최신성 기반 방식이 제공하지 못하는 보장입니다.
• 적응성: 캘리브레이션 프로토콜은 영역 수명이 워크로드에 의존적이며, 감쇠율을 최적화하기 위해 측정될 수 있음을 입증하며, 단일한 방식의 정적 제거 정책 사용에 도전합니다.

저자들은 MemDecay가 구조적 사전 지식과 고정된 지침을 성공적으로 보존하지만, 향후 연구에서는 unpinned recall을 개선하기 위한 어텐션 항의 크기 정규화 문제를 해결하고, 돌이킬 수 없는 제거를 처리하기 위한 계층적 확장(예: 호스트 메모리로의 오프로딩)을 탐색해야 한다고 결론짓습니다.