KVEraser: 효율적인 국소적 컨텍스트 삭제를 위한 KV 캐시 제어 학습

KV 캐시(KV cache)에 대한 사후 컨텍스트 삭제(Post-hoc context erasing)는 국소적인 편집이 전역적인 결과를 초래하기 때문에 매우 어렵습니다. 즉, 특정 구간(span)이 한 번 처리되면, 그 영향력이 이후 모든 토큰의 캐시된 상태로 전파됩니다. 이러한 문제는 긴 컨텍스트(long-context) LLM 애플리케이션에서 자연스럽게 발생하는데, 오래된 검색 사실, 잘못된 도구 관찰 결과, 철회된 사용자 선호도 또는 유해한 프롬프트 주입(prompt injections) 등이 프리필(prefill) 단계 이후에야 식별될 수 있기 때문입니다. 정확한 삭제를 수행하려면 삭제된 구간 이후의 모든 토큰을 다시 계산해야 하며, 이로 인해 계산 비용이 삭제된 구간의 길이가 아닌 접미사(suffix) 길이에 따라 결정됩니다.

우리는 효율적인 국소적 컨텍스트 삭제를 위한 학습된 KV 캐시 편집 방법인 KVEraser를 소개합니다. 처리된 컨텍스트와 제거할 구간이 주어지면, KVEraser는 나머지 캐시는 변경 없이 재사용하면서 삭제된 구간의 KV 상태만을 학습된 제어 상태(steering states)로 교체합니다. 전이 가능한 삭제 메커니즘을 학습하기 위해, 우리는 2단계 훈련 파이프라인을 구축했습니다. 일반적인 구간-이웃(span-neighbor) 사전 학습(pre-training)은 삭제된 구간의 영향력을 억제하도록 삭제기(eraser)를 가르치며, 작업 특화 미세 조정(task-specific fine-tuning)은 이러한 능력을 다운스트림 시나리오에 적응시킵니다.

실험 결과, KVEraser는 1K32K 컨텍스트 길이에 걸쳐 인도메인(in-domain) 작업에서 삭제 후 성능이 전체 재계산(full recomputation)과 거의 일치하는 것을 보여주었습니다. 반면 지연 시간(latency)은 전체 재계산이 17.6배 증가하는 것에 비해 단 24%만 증가했습니다. 또한 KVEraser는 유해한 사실적 방해 요소(harmful factual distractors)가 포함된 미학습 긴 문서 QA 작업에도 일반화되어, 근사치 베이스라인(approximate baselines) 중 최고의 성능을 달성하는 동시에 전체 재계산 대비 34배의 속도 향상을 이루었습니다.