프리헥 KV 선택으로 장기 컨텍스트 추론을 가속화하다

초록

본 논문은 LLM 추론 시 KV 캐시를 전부 스캔하는 비용을 줄이기 위해, 어텐션 점수를 계산하기 전에 KV를 선택하는 ‘Pre‑hoc Sparsity(PrHS)’ 프레임워크를 제안한다. 선택된 KV의 누락된 어텐션 질량 δ를 기준으로 상호정보 손실의 상한을 이론적으로 도출하고, 시간·깊이·레이어 축을 따라 세 가지 사전 선택기(CIS, PSAW, ETF)를 설계한다. 실험 결과, GSM8K·CoQA 등에서 90 % 이상의 KV 검색 오버헤드를 감소시키고, LongBench에서 1 % 미만의 정확도 저하와 15 % 정도의 FLOP 절감을 달성했으며, A100‑80GB GPU에서 어텐션 연산 지연을 9.9배, 처리량을 2.8배 가속한다.

상세 분석

논문은 먼저 KV 캐시가 토큰 수 L에 비례해 선형적으로 성장하면서, 매 디코딩 단계마다 전체 KV에 대해 어텐션 스코어를 계산해야 하는 근본적인 병목을 지적한다. 기존의 후행(Post‑hoc) 방식은 이미 관측된 어텐션 분포나 토큰 연령, 스케치와 같은 후행 정보를 이용해 후보 KV를 제한하지만, 이러한 ‘posterior bias’는 실제 토큰 중요도를 왜곡하고 장기 의존성을 손상시킨다. 저자들은 이를 ‘Posterior‑conditioned Sparsity(PoHS)’라 명명하고, 어텐션 질량 손실 δ와 상호정보 손실 I_full – I_S 사이의 함수 g(δ)=2