LookaheadKV: 생성 과정 없이 미래를 예측하여 빠르고 정확한 KV 캐시 제거 수행

긴 문서를 처리할 때 LLM의 메모리 사용량이 급증하는 문제를 해결하기 위해, 무거운 예측 모델 없이도 미래에 중요할 정보를 정확히 골라내는 기술이다. 기존 방식보다 훨씬 가벼운 모듈을 사용하여 추론 속도는 유지하면서도 긴 문맥 이해도를 획기적으로 높였다.

Transformer의 Self-Attention은 모든 토큰의 관계를 계산하기 위해 KV 캐시를 저장하며, 이는 문맥이 길어질수록 메모리 점유율이 선형적으로 증가하는 원인이 된다. 기존에는 중요도가 낮은 토큰을 제거하기 위해 작은 모델로 미래 응답을 미리 생성해보는 방식을 썼지만, 이 과정 자체가 연산 비용을 크게 높인다는 한계가 있었다.

LookaheadKV는 미래 응답을 실제로 생성하는 대신, 'Lookahead 토큰'이라는 특수 임베딩을 통해 미래의 어텐션 패턴을 미리 엿본다. 이 토큰들은 실제 모델이 생성할 미래 응답과 입력 프롬프트 사이의 어텐션 분포를 모방하도록 학습되어, 마치 미래를 미리 본 것과 같은 효과를 낸다. 즉, 무거운 생성 과정 없이도 어떤 토큰이 나중에 중요하게 쓰일지 미리 알 수 있게 된다.

여기에 Lookahead 토큰에만 선택적으로 적용되는 LoRA 어댑터를 결합하여 모델의 원래 지식은 보존하면서도 예측 성능만 정교하게 최적화한다. 결과적으로 매우 적은 추가 연산만으로도 어떤 KV 캐시를 버려야 할지 정확하게 판단하여 메모리 효율과 추론 속도를 동시에 잡았다.

LookaheadKV는 학습 가능한 Lookahead 토큰 P와 전용 LoRA 모듈을 기존 LLM 아키텍처에 추가한다. Prefill 단계에서 입력 시퀀스 X 뒤에 P를 붙여 Forward Pass를 수행하며, 이때 P에서 생성된 Query가 입력 토큰들의 중요도를 결정하는 지표가 된다.

중요도 점수 계산은 Lookahead 토큰의 Query와 입력 토큰의 Key 사이의 Cross-Attention 점수를 기반으로 한다. [Lookahead 토큰의 Query Q_LKV와 입력 Key K_LKV를 입력으로] → [내적 연산 후 Softmax를 적용해 Attention Matrix A_LKV를 생성하고] → [Lookahead 토큰 차원에 대해 평균을 내어] → [각 입력 토큰이 미래 응답에 기여할 확률적 중요도 점수를 얻는다].

학습 시에는 실제 모델이 생성한 응답 Y와의 Ground-truth 어텐션 분포를 타겟으로 삼는다. [실제 응답의 어텐션 분포와 Lookahead 토큰이 예측한 분포를 입력으로] → [KL Divergence Loss를 계산하여] → [두 분포 사이의 차이를 줄이는 방향으로] → [Lookahead 토큰과 LoRA 파라미터만 업데이트하고 나머지 LLM 레이어는 동결한다].

LongBench 벤치마크에서 Llama 3.1-8B 및 Qwen 3-8B 모델을 대상으로 실험한 결과, LookaheadKV는 모든 캐시 예산 범위에서 기존의 SnapKV나 드래프트 기반의 LAQ보다 우수한 성능을 보였다. 특히 매우 제한된 캐시 예산 상황에서도 성능 저하가 가장 적어 효율적인 정보 보존 능력을 입증했다.

효율성 측면에서 32K 문맥 길이 기준, 드래프트 기반 방식인 LAQ 대비 KV 제거 오버헤드를 14.5배 감소시켰다. 실제 TTFT 지연 시간 증가율은 2.16% 미만으로 나타나, 단순 휴리스틱 방식인 SnapKV와 유사한 수준의 빠른 속도를 기록하면서도 정확도는 훨씬 높았다.

RULER 벤치마크를 통한 긴 문맥 이해도 평가에서도 128K 길이까지 성능 우위를 유지했다. 또한 학습 시 사용한 문맥 길이인 16K보다 훨씬 긴 문맥에서도 효과적으로 작동함을 확인하여, 미학습 문맥 길이에 대한 강력한 일반화 성능을 보여주었다.

LookaheadKV는 파라미터 효율적인 튜닝(PEFT) 기법을 KV 캐시 관리 영역에 적용하여 전체 파라미터의 0.5% 미만인 모듈만으로 추론 전략을 최적화한다. 핵심 아키텍처는 'Selectively Activated Low-rank Adapters'를 사용하여 일반 입력 토큰의 연산에는 영향을 주지 않고 Lookahead 토큰의 표현력만 강화하는 구조를 가진다.

학습 과정에서 FlashAttention을 사용하여 메모리 효율을 높였으며, 역전파 시에는 전체 어텐션 행렬을 생성하지 않고 필요한 Cross-Attention 부분만 계산하는 최적화를 통해 16K 이상의 긴 시퀀스 학습 시 발생하는 OOM 문제를 해결했다. 이는 O((|X|+|Y|)^2)의 메모리 요구량을 O(|X|*|Y| + |Y|^2)로 줄이는 효과를 낸다.

기존 드래프트 기반 방식이 별도의 모델 실행이나 반복적인 생성을 요구하여 I/O 및 연산 병목을 일으키는 것과 달리, 단 한 번의 Prefill Forward Pass 내에서 모든 예측이 완료되는 구조를 채택했다. 실험 결과 Lookahead 토큰의 개수를 32개로 설정했을 때 성능과 오버헤드 사이의 최적의 균형점을 찾을 수 있음을 확인했다.

제한된 컴퓨팅 자원으로 인해 70B 이상의 초대형 모델에 대한 실험은 수행되지 못했다. 또한 현재는 Prefill 단계의 KV 제거에 집중하고 있으며, 생성(Decoding) 단계에서 발생하는 동적 KV 캐시 증가를 관리하는 기법은 향후 연구 과제로 남아있다.