기술 요약: KVEraser

문제 정의

장기 문맥(Long-context) 거대 언어 모델(LLM) 애플리케이션에서, 컨텍스트는 검색된 문서, 도구 관찰 결과, 또는 대화 메모리로부터 온라인으로 조립되는 동적인 특성을 가집니다. 특정 처리된 컨텍스트의 구간이 프리필(prefill) 단계 이후에 유효하지 않거나, 틀렸거나, 혹은 해로운 것(예: 철회된 사용자 선호도, 잘못된 도구 결과, 또는 프롬프트 인젝션)으로 식별될 때 중대한 문제가 발생합니다.

"사후 컨텍스트 삭제(post-hoc context erasing)"의 핵심적인 어려움은 KV 캐시(Key-Value Cache)의 본질에 있습니다. 인과적 셀프 어텐션(Causal self-attention) 하에서, 임의의 토큰 표현은 모든 이전 토큰에 의존합니다. 결과적으로, 특정 구간이 편집되거나 제거되면, 해당 구간을 포함하여 계산되었던 모든 후속 토큰들의 캐시 상태는 무효화됩니다.

• 정확한 삭제 (Exact Erasing): 표준적인 해결책은 유효한 프리픽스 캐시를 재사용하고 남은 서픽스(suffix) 전체에 대해 프리필 패스를 다시 실행하는 것입니다. 그러나 이 방식의 계산 비용은 삭제된 구간의 길이가 아니라 서픽스의 길이에 비례하여 증가합니다. 짧은 오류 구간 뒤에 긴 컨텍스트가 이어지는 시나리오에서, 이 접근법은 과도한 지연 시간을 초래합니다.
• 지시문 기반 망각 (Instruction-Only Forgetting): 단순히 해당 구간을 "잊으라"는 지시문을 추가하는 것은 효과가 없습니다. 왜냐로 서픽스의 KV 상태가 이미 오염되었기 때문이며, 모델은 디코딩 중에 캐시된 증거를 스스로 다시 쓸 수 없습니다.
• 기존 방법의 한계: KV 캐시 재사용(예: 공유 프리픽스용)이나 제거(예: 어텐션이 낮은 토큰 제거)에 관한 기존 연구들은, 특정 구간이 존재하지 않았던 것처럼 행동하도록 유도하는 반사실적(counterfactual) 동작 요구 사항을 해결하지 못합니다.

방법론: KVEraser

저자들은 전체 서픽스를 재계산하지 않고도 효율적이고 국소적인 컨텍스트 삭제를 수행하도록 설계된 학습 가능한 KV 캐시 편집 방법인 KVEraser를 소개합니다.

핵심 메커니즘

KVEraser는 편집된 캐시의 정확한 재구축 대신 **함수적 근사(functional approximation)**를 도입합니다. 서픽스를 다시 계산하는 대신, 다음과 같이 대리 KV 캐시($\widehat{KV}$)를 구축합니다:

1. 유효한 프리픽스 캐시($KV_{1:m-1}$)를 유지합니다.
2. 삭제된 구간($e$)의 KV 상태를 학습된 스티어링 상태($KV_{steer}$)로 교체합니다.
3. 원래의 오염된 서픽스 캐시($KV_{n+1:T}$)를 변경 없이 재사용합니다.

목표는 스티어링 상태가 누락된 구간과 서픽스의 오염을 보상하여, 고정된 생성기(generator)가 해당 구간이 처음부터 없었던 것처럼 동일한 출력을 생성하도록 유도하는 것입니다.

모델 아키텍처

• 생성기 ($p_\theta$): 고정된 LLM 백본.
• 이레이저 (Eraser, $E_\phi$): 삭제된 구간에 대한 교체 KV 상태를 예측하는 학습 가능한 모듈 (LM 헤드를 제외한 생성기의 백본을 복사하여 초기화됨).
• 조건화 (Conditioning): 이레이저는 보존된 프리픽스 KV 캐시와 삭제된 구간 자체를 조건으로 합니다. 이는 다운스트림 쿼리나 서픽스에는 조건화되지 않으며, 이를 통해 스티어링 블록이 다양한 쿼리에 대해 재사용 가능하도록 보장합니다.

학습 파이프라인

컨텍스트 삭제를 위한 대규모 주석 데이터의 부재로 인해, 저자들은 2단계 학습 전략을 제안합니다.

1. 일반적 구간-이웃 사전 학습 (Generic Span-Neighbor Pre-training): 이레이저는 위키피디아 데이터를 사용한 대규모 합성 태스크에서 사전 학습됩니다. 모델은 삽입된 "방해꾼(distractor)" 구간의 영향을 억제하고 올바른 이웃 텍스트를 복구하는 법을 배웁니다. 이는 캐시 편집을 위한 일반적인 메커니즘을 가르칩니다.
2. 태스크 특화 미세 조정 (Task-Specific Fine-tuning): 사전 학습된 이레이저는 다음의 다운스트림 태스크에서 미세 조정됩니다:
   • Needle-in-a-Haystack (NIAH): 모델이 삭제된 "바늘(needle)"을 무시하고 유지된 바늘을 기억해야 하는 통제된 합성 벤치마크.
   • 사실적 방해꾼 제거 (Factual Distractor Removal): 긴 문서 내에 삽입된 오도하는 사실적 덩어리들을 모델이 무시해야 하는 QA 설정.

학습 목표는 교사 강제(teacher-forced) 부정 로그 가능도 손실(negative log-likelihood loss)이며, 이를 통해 서리게이트 캐시로부터의 생성기 출력이 깨끗한 편집 프롬프트로부터의 출력과 일치하도록 이레이저를 최적화합니다.

주요 기여

1. 국소적 KV 스티어링의 정식화: 본 논문은 컨텍스트 삭제를 글로벌 캐시 오염을 보상하는 로컬 스티어링 블록을 학습하는 문제로 정식화하여, 계산 의존성을 서픽스 길이가 아닌 삭제된 구간의 길이로 전환했습니다.
2. KVEraser 프레임워크: 소량의 교체 KV 블록만 필요로 하며, 기존 LLM 서빙 시스템과 원활하게 통합되는 실용적이고 학습 가능한 아키텍처를 제공합니다.
3. 2단계 학습 전략: 일반적인 사전 학습을 통해 전이 가능한 삭제 능력을 배우고, 특정 다운스트림 시나리오에 맞춰 미세 조정하는 새로운 파이프라인을 활용합니다.
4. 포괄적인 평가: 1K에서 32K 토큰에 이르는 컨텍스트 길이와 합성 및 자연어 장문 QA 태스크 모두에 대해 평가를 수행했습니다.

실험 결과

통제된 벤치마크에서의 성능 (NIAH)

• 정확도: KVEraser는 모든 컨텍스트 크기(1K–32K)에서 완벽에 가까운 Exact Match 성능을 달성하며, 전체 재계산(full recomputation)의 성능과 일치하는 모습을 보였습니다.
• 베이스라인: 근사 베이스라인들(Delete-and-shift, Instruction-only forgetting, Local suffix repair)은 낮은 정확도로 시작하거나 컨텍스트 길이가 길어짐에 따라 급격히 성능이 저하되었습니다.
• 효율성: 전체 재계산의 지연 시간은 1K에서 32K 컨텍스트로 이동할 때 17.6배 증가하는 반면, KVEraser의 지연 시간은 단 **24%**만 증가했습니다. 이는 비용이 서픽스 길이($|s|$)가 아닌 삭제된 구간의 길이($|e|$)에 따라 스케일링되기 때문입니다.

자연어 장문 QA에서의 성능

• 품질-효율성 트레이드오프: 해로운 사실적 방해꾼이 포함된 미학습 QA 데이터셋(2WikiMultiHopQA, MuSiQue, IIRC)에서, KVEraser는 모든 근사 베이스라인 중 최고의 성능을 기록했습니다.
• 전체 재계산과의 비교: 전체 재계산이 가장 높은 절대적 정확도를 보여주기는 하지만, KVEraser보다 3~4배 더 느립니다. KVEraser는 품질과 효율성 사이의 최적의 균형을 제공하는 파레토 프런티어(Pareto frontier)에 위치합니다.

절제 연구 (Ablation Studies)

• 조건화 소스: 저자들은 이레이저를 보존된 프리픽스 KV와 삭제된 구간에 조건화하는 것이 최적임을 발견했습니다. 쿼리나 서픽스 정보를 추가하는 것은 일관된 개선을 가져오지 못했으며, 이는 쿼리 불가지론적(query-agnostic) 설계의 타당성을 입증합니다.
• 실패 사례: 근사 베이스라인들은 삭제된 값과 유지된 값을 동시에 출력하거나(오염의 징후), 무관한 값을 생성하며 실패하는 경우가 많았습니다. KVEraser의 실패는 드물고 고립되어 나타납니다.

의의 및 주장

본 논문은 학습된 로컬 KV 스티어링이 효율적인 컨텍스트 삭제를 위한 유망하고 실행 가능한 메커니즘임을 주장합니다.

• 확장성: KVEraser는 정확한 삭제의 병목 현상인, 비용이 남은 컨텍스트의 길이에 의해 결정되는 문제를 해결합니다.
• 신뢰성: 지시문 기반의 망각이나 단순한 캐시 삭제와 달리, KVEraser는 삭제된 구간의 영향을 효과적으로 억제하여 모델이 오래되거나 해로운 정보에 의존하는 것을 방지합니다.
• 실용성: 이 방법은 전체 재계산보다 3~4배 빠른 속도를 제공하면서도 정확한 방식에 필적하는 성능을 유지하므로, 컨텍스트 업데이트가 빈번한 RAG 시스템, 도구 증강 에이전트, 대화형 AI와 같은 실제 애플리케이션에 적합합니다.

저자들은 자연어 QA 태스크에서 KVEraser와 전체 재계산 사이의 잔여 격차가 학습 데이터의 규모와 다양성에 기인할 수 있음을 인정하며, 향의 스케일링을 통해 이 격차를 좁힐 수 있을 것이라고 시사했습니다.