키밸류 캐시 자동 관리를 위한 강화학습 기반 정책

초록

본 논문은 대형 언어 모델의 자동 회귀 생성 과정에서 메모리 사용량을 크게 차지하는 키‑밸류(KV) 캐시를 효율적으로 관리하기 위해, 토큰의 미래 활용도를 예측하는 강화학습(RL) 기반 순위 정책을 제안한다. 각 어텐션 헤드마다 경량 RL 에이전트를 학습시켜 키와 밸류만을 입력으로 토큰을 중요도 순으로 정렬하고, 원하는 메모리 예산에 따라 낮은 순위의 토큰을 삭제한다. 제안 방법(KVP)은 사전 계산된 생성 트레이스를 이용해 오프라인으로 학습되며, LLM 구조나 추론 과정에 어떠한 변경도 가하지 않는다. 실험 결과 RULER와 OASST2‑4k 같은 장문·다중턴 벤치마크에서 기존 휴리스틱 기반 방법들을 크게 앞섰으며, LongBench, BOOLQ, ARC 등 다운스트림 태스크에서도 제로샷 일반화 능력을 보였다.

상세 분석

이 연구는 KV 캐시 관리 문제를 “토큰을 미래 디코딩에 얼마나 유용할지”라는 직접적인 목표로 재정의하고, 이를 순위 학습 문제로 전환한다는 점에서 혁신적이다. 기존 방법들은 최근성, 어텐션 스코어, 혹은 압축 기법 등 간접적인 신호에 의존했으며, 이러한 신호는 토큰의 실제 미래 중요도를 완전히 반영하지 못한다. 논문은 먼저 캐시 내 모든 토큰에 대해 미래 어텐션 합계(즉, 향후 토큰들이 해당 토큰에 할당하는 어텐션 양)를 “진정한 중요도”로 정의하고, 이 값을 기반으로 전체 예산에 대해 최소 손실을 주는 순위를 학습한다.

핵심 설계는 두 가지이다. 첫째, 각 어텐션 헤드마다 별도의 경량 MLP 기반 스코어링 함수를 두어 헤드별 특성을 반영한다. 입력은 키, 밸류, 그리고 토큰 위치 정보만을 사용하므로, 쿼리나 실제 어텐션 값에 접근할 필요가 없으며, 이는 기존 방법이 요구하던 추가 연산을 크게 절감한다. 둘째, 순위 정책을 Plackett‑Luce 모델로 확률화하고, Gumbel‑Sort 기법을 이용해 한 번의 포워드 패스로 전체 순열을 샘플링한다. 이는 순차적 샘플링의 비효율성을 없애고 GPU에서 완전 병렬화가 가능하도록 만든다.

보상 설계도 주목할 만하다. 전체 보상은 모든 가능한 캐시 예산 b(1≤b≤n‑1)에 대해, 해당 예산으로 상위 b개의 토큰만 유지했을 때 손실되는 미래 어텐션 합계의 총합으로 정의된다. 즉, 하나의 순열에 대해 “예산 전반에 걸친 평균 손실”을 최소화하도록 학습한다. 이 보상은 오프라인 RL 설정에서 LLM을 재추론할 필요 없이, 미리 저장된 어텐션 매트릭스만으로 계산할 수 있다.

실험에서는 두 종류의 모델(예: LLaMA‑2와 Falcon)과 두 벤치마크(RULER, OASST2‑4k)에서 KVP가 기존의 FIFO, LRU, StreamingLLM, H2O 등과 비교해 퍼플렉시티 및 정확도 측면에서 일관되게 우수함을 보였다. 특히 메모리 예산이 극히 제한된 상황에서도 중요한 토큰을 보존함으로써 성능 저하를 최소화한다. 제로샷 일반화 실험에서는 LongBench, BOOLQ, ARC 등 다양한 다운스트림 작업에서 KVP가 사전 학습된 정책을 그대로 적용했음에도 불구하고, 기존 휴리스틱 대비 평균 3~5%의 성능 향상을 기록했다. 이는 미래 유틸리티 예측이 모델‑특정 혹은 데이터‑특정 편향에 크게 좌우되지 않음을 시사한다.

한계점으로는 현재 각 헤드마다 독립적인 에이전트를 학습하기 때문에 전체 파라미터 수가 헤드 수에 비례한다는 점과, 트레이닝 데이터가 충분히 다양하지 않을 경우 특정 도메인에 과적합될 가능성이 있다는 점을 들 수 있다. 향후 연구에서는 헤드 간 파라미터 공유, 메모리 계층 관리와의 연계, 그리고 요약 토큰 생성과 같은 고차원 압축 전략과의 결합을 통해 더욱 효율적인 캐시 관리 프레임워크를 구축할 수 있을 것이다.