긴 추론을 위한 정보 인지형 KV 캐시 압축(InfoKV)

대형 언어 모델의 긴 문맥 처리에서 KV 캐시 메모리는 입력·생성 길이에 따라 선형으로 증가하여 실무적 제약을 초래한다. 이 논문은 단순히 최근 attention에 의존한 토큰 보존 기준이 장거리 생성에 취약하다는 관찰을 바탕으로, 예측 불확실성(entropy)과 레이어별 표현 변화를 결합해 미래 기여도를 정량화함으로써 필요 토큰만 선별적으로 보전한다. 그 결과 장거리 prefilling과 decoding 시나리오 모두에서 기존 attention 기반 압축 방법보다 전반적으로 우수한 성능을 달성했다.

Transformer 기반 생성 모델에서는 KV 캐시가 과거 토큰의 키/값을 저장하여 이후 쿼리에서 빠르게 참조하도록 한다. 이 캐시는 입력 길이나 생성 길이에 비례해 선형적으로 증가하므로 긴 문맥에서는 메모리와 대역폭 제약이 심해진다. 기존의 KV 압축 기법은 최근 관찰 창에서의 attention 가중치를 토큰 중요도의 근사치로 사용했는데, 이 방식은 주로 단기적·역방향적 신호에 의존한다. 토큰이 미래 생성에 미치는 기여는 단순히 최근 쿼리의 attention이 클 것이라는 가정으로는 포착되지 않는 경우가 많다. 본 논문은 예측 분포의 엔트로피가 높은 토큰이 미래의 불확실성을 줄이거나 생성 경로를 크게 바꿀 가능성이 크다는 점에 주목했다. 따라서 예측 엔트로피는 토큰이 보유한 '정보량'을 반영하는 보완적 신호로 작동하며, 특히 먼 미래 컨텍스트에 대한 영향력이 크다. 엔트로피 신호만으로는 표현 수준의 변화나 레이어별 처리 과정을 반영하지 못하므로, 초기 레이어와 최종 레이어의 히든 표현 간 거리를 추가로 측정해 표현이 아직 수렴되지 않은 토큰을 가려냈다. 이러한 레이어간 표현 변화는 해당 토큰이 모델 내부에서 의미적으로 계속 정제되고 있음을 나타내며, 이 값과 엔트로피를 결합할 때 장거리 유틸리티를 보다 정확히 추정할 수 있었다. 최종적으로 attention과 결합해 단기 관련성도 보존하는 가중합 전략이 실용적 균형을 제공했다.

전체 접근 방식은 각 토큰에 대해 세 가지 신호를 계산하여 압축 우선순위를 매기는 것이다. 첫째, 관찰 창(observation window)에서의 마지막 레이어 attention 가중치를 평균해 A_i^(l)을 얻는다. 둘째, 최종 레이어의 예측 확률 분포로부터 Top-k 제한 엔트로피 H_i를 계산해 토큰 불확실성을 평가한다. 레이어별로는 초기(또는 중간) 레이어 l의 히든 표현 h_i^(l)와 최종 레이어 h_i^(L) 사이의 코사인 거리를 D_i^(l)=1-cos(h_i^(l),h_i^(L))로 측정해 표현 변화량을 얻는다. 레이어별 엔트로피 점수 E_i^(l)=D_i^(l) * H_i로 정의하고 Softmax 정규화를 적용해 동일 레이어 내에서 비교 가능하도록 만든다. 마지막으로 각 레이어에서 S_i^(l)=α·A_i^(l)+(1-α)·Softmax(E^(l)) 형태로 attention과 엔트로피 기반 점수를 결합해 상위 토큰을 보존한다. 추가적으로 전체 KV 예산 B를 레이어별 누적 엔트로피 비중 Ē^(l)로 분배하는 adaptive budget 방안을 실험했다. 이때 k_l = (Ē^(l)/Σ_m Ē^(m))·B로 할당량을 계산하되, 균일 할당과 비교해 모델별 안정성과 성능 변동을 평가했다. 하이퍼파라미터로는 α, τ(표현 거리 바이어스), top-k 값이 있으며 논문에서는 α=0.9, τ=1, k=256이 기본 설정으로 채택되었다.

장거리 prefilling 평가에서 LongReason 데이터셋을 사용해 Llama-3.1-8B-Instruct와 Llama-3.2-3B-Instruct를 대상으로 다양한 컨텍스트 길이(16k, 32k, 64k)와 캐시 보유율(100%, 40%, 20%) 조건으로 비교했다. 표와 본문 결과에서 InfoKV는 다수의 설정에서 SnapKV, PyramidKV, Expected Attention 등 attention 기반 기법을 상회하거나 동등한 성능을 보였고, 시퀀스 길이가 증가할수록 우위가 뚜렷해지는 경향이 관찰되었다. 체인오브생각(CoT)과 직접응답 설정 모두에서 평균 성능 개선이 보고되었다. 장거리 decoding 평가에서는 DeepSeek-R1 계열 모델을 IFEval, AIME 2024, LiveCodeBench에 적용해 최대 출력 길이 32768 토큰 조건으로 평가했다. RPC와 같은 온라인 압축 기법과 비교했을 때 InfoKV는 전체적으로 더 높은 pass@1 성능을 보였고, 일부 경우에는 압축된 캐시가 전체 캐시보다 더 나은 성능을 보이는 사례가 나타났다. 이는 장거리 생성에서 과거의 모든 토큰을 유지하는 것이 오히려 산만함을 유발할 수 있음을 시사한다. 추가 분석에서는 τ와 top-k, α의 민감도 실험을 통해 τ=1과 k=256, α≈0.9 조합이 안정적 성능을 제공함이 확인되었다. adaptive budget 전략은 R1-Distill-Llama-8B에서 이득을 주었으나 R1-Distill-Qwen-7B에서는 성능 저하를 초래해 아키텍처별 민감성을 시사했다.

전체 시스템은 KV 캐시 압축을 레이어별로 독립적으로 수행하는 구조를 따른다. 각 레이어 l에서 토큰 x_i의 attention 점수 A_i^(l)을 관찰 창 내 쿼리-키 상호작용으로부터 추출하며, 엔트로피 H_i는 최종 레이어의 예측 분포에 대해 Top-k 제한을 적용해 계산한다. 표현 변화 D_i^(l)는 코사인 거리를 사용해 계산하며 작은 바이어스 τ를 더해 최종 레이어의 엔트로피 스코어가 0이 되는 것을 방지한다. 중요도 결합 공식은 S_i^(l)=α·A_i^(l)+(1-α)·Softmax(E^(l))이며, Softmax는 해당 레이어 내 토큰들에 대한 정규화 기능을 수행한다. 보존 대상 토큰은 각 레이어에서 상위 k_l 개를 유지하며, 기본으로 균일한 레이어 예산을 사용하되 필요 시 Ē^(l)에 기반한 적응형 예산 k_l = (Ē^(l)/Σ Ē^(m))·B를 적용할 수 있다. KL 발산 기반의 Forward Influence는 문서의 앞부분을 압축한 뒤 특정 미래 청크(l_c:r_c)에 대해 원래 예측 분포와 제거 후 분포의 평균 KL을 계산해 얻었다. 실험적 구현에서는 Llama-3.1-8B-Instruct, Llama-3.2-3B-Instruct, DeepSeek-R1 계열 모델을 사용했으며 압축은 prefilling 단계와 decoding 단계(예: 1024 토큰 마다)에서 적용되었다. 하이퍼파라미터 튜닝 결과 τ=1, top-k=256, α=0.9이 기본값으로 채택되었고 이들 값은 성능과 안정성 측면에서 타협점을 제공했다. 성능 평가는 LongReason, IFEval, AIME 2024, LiveCodeBench 등의 표준화된 벤치마크를 통해 수행되었다.

논문은 엔트로피가 장거리 영향력을 더 잘 포착한다는 사실을 보였으나 엔트로피 자체는 미래 유틸리티를 직접 최적화하는 목적함수가 아니므로 근사치에 불과하다. 레이어별 적응형 예산 할당은 일부 모델에서 성능을 개선했지만 다른 모델에서는 안정성을 해치는 사례가 보고되어 아키텍처 의존성이 존재한다. 또한 논문에서 제시한 설정과 하이퍼파라미터가 모든 모델·업무에 보편적으로 최적이라는 근거는 없으며 추가적인 아키텍처별 조정이 필요하다.