사전 희소성 기반의 근사 오라클 KV 선택을 통한 장문 컨텍스트 추론

Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference

대규모 언어 모델(LLM) 추론의 핵심적인 병목 현상은 지속적으로 증가하는 키-값(KV) 캐시에 대한 어텐션 연산 비용입니다. 근사 오라클 방식의 상위 k개 KV 선택은 어텐션의 품질을 유지하면서 계산 및 대역폭을 크게 줄일 수 있지만, 기존의 희소성 방법은 일반적으로 어텐션 또는 프록시 점수를 기반으로 하는 사후 휴리스틱에 의존합니다. 이러한 조건 설정은 사후 편향을 유발하며, 이는 실제 토큰의 중요성을 왜곡하고 중요한 토큰을 놓쳐 장거리 추론 능력을 저하시킵니다. 이러한 문제를 해결하기 위해, 우리는 어텐션 점수 계산 전에 KV 항목을 선택하고 명시적인 정확도 제어를 제공하는 사전 희소성(PrHS)을 제안합니다. 버려지는 항목의 어텐션 가중치를 델타(삭제된 가중치)라고 할 때, 우리는 정보 손실에 대한 상한을 유도하는데, 이는 삭제된 가중치에만 의존합니다. 이 관계는 사후 휴리스틱의 실패 원인을 설명하고, 삭제된 가중치를 미리 제어함으로써 검증 가능한 보장을 제공합니다. PrHS 내에서, 우리는 시간, 깊이 및 레이어를 기준으로 세 가지 직교하는 사전 선택기를 구현했습니다. LLaMA 및 Mistral 패밀리에 대한 광범위한 실험 결과, PrHS의 효과가 입증되었습니다. GSM8K 및 CoQA 데이터셋에서 PrHS는 검색 오버헤드를 90% 이상 줄이고, 동일하거나 더 나은 정확도를 유지하면서 HShare보다 3배 높은 검색 희소성을 달성했습니다. LongBench 데이터셋에서 평균 1% 미만의 성능 저하를 보였으며, 기존의 희소성 기준선 대비 어텐션 FLOPs를 약 15% 줄였습니다. 또한 NVIDIA A100-80GB GPU에서 어텐션 연산 지연 시간을 9.9배 단축하고 처리량을 2.8배 향상시켰습니다.

A core bottleneck in large language model (LLM) inference is the cost of attending over the ever-growing key-value (KV) cache. Although near-oracle top-k KV selection can preserve the quality of dense attention while sharply reducing computation and bandwidth, existing sparse methods generally rely on posterior heuristics, i.e., selectors conditioned on observed attention or proxy scores. Such conditioning introduces posterior bias: it tends to distort true token importance and miss salient tokens, thereby impairing long-range reasoning. To tackle this problem, we propose Pre-hoc Sparsity (PrHS), which selects KV entries before attention scoring and provides explicit accuracy control. Let the attention mass of discarded entries be delta (the dropped mass). Through a marginal-to-mutual-information analysis, we derive an upper bound on the mutual-information loss that depends only on the dropped mass. This relation explains failure modes of posterior heuristics and enables verifiable guarantees by controlling the dropped mass in advance. Within PrHS, we instantiate three orthogonal pre-hoc selectors along the axes of time, depth, and layer. Extensive experiments on LLaMA and Mistral families validate PrHS. Across GSM8K and CoQA, PrHS reduces retrieval overhead by over 90%, achieving 3x higher retrieval sparsity than HShare at matched or better accuracy. It incurs under 1% average degradation on LongBench, lowers attention FLOPs by about 15% versus prior sparse baselines, and yields a 9.9x speedup in attention-operator latency and 2.8x higher throughput on NVIDIA A100-80GB GPUs than the dense baseline.