랭킹 결정 전에 유사성 검토: 효율적인 긴 컨텍스트 어텐션을 위한 유사성 기반 상위-K 재사용

Recall Before You Rank: Similarity-Guided Top-$K$ Reuse for Efficient Long-Context Attention

상위-K 희소 어텐션은 키-값(KV) 항목의 작은 부분에만 집중하여 Softmax 계산 및 값 집계 비용을 줄입니다. 하지만 여전히 현재 쿼리에 대한 전체 KV 캐시를 평가하고 전역 상위-K 선택을 수행해야 하므로, 선택기 비용이 컨텍스트 길이에 따라 선형적으로 증가하며, 긴 컨텍스트 디코딩에 대한 희소 어텐션의 실제 효율성을 제한합니다. 본 논문에서는 훈련 과정이 필요 없는 ReTopK라는 방법을 제안합니다. ReTopK는 과거 검색 결정을 재사용하여 동적 상위-K 어텐션을 가속화합니다. ReTopK는 유사한 쿼리가 종종 중복된 지원(supports)에 집중한다는 점과, 부분적으로 중복되는 지원이 여전히 정확한 상위-K 어텐션의 대부분을 유지할 수 있다는 점을 기반으로 합니다. 각 어텐션 헤드에서 ReTopK는 과거 쿼리-지원 쌍의 제한적인 캐시를 유지하고, 새로운 쿼리에 대해 가장 유사한 캐시된 쿼리를 검색하며, 저장된 지원과 최근 정보를 결합하여, 정확한 현재 쿼리 점수를 사용하여 결과적으로 생성된 작은 후보 집합만을 재평가합니다. 유사성 기반 백업 메커니즘은 재사용이 불확실할 때 전체 이력을 사용한 정확한 상위-K 선택을 수행하며, 주기적인 정확한 업데이트는 캐시의 변화를 제한합니다. ReTopK는 전체 KV 캐시를 유지하고 과거 점수, 어텐션 가중치 또는 출력이 아닌 선택된 인덱스만 재사용합니다. 16K에서 128K 컨텍스트 길이에서 ReTopK는 평가된 근사 방법 중 가장 낮은 PG19 퍼플렉서티(perplexity)를 달성하고, NIAH 및 LongBench 점수가 가장 높습니다. 128K 길이와 K=512의 경우, ReTopK는 정확한 상위-K에 비해 0.50%p의 작은 퍼플렉서티 증가만을 보이며, 어텐션 계산 속도를 3.07배 향상시켰습니다.

Top-$K$ sparse attention reduces the cost of Softmax and value aggregation by attending to only a small subset of key--value (KV) entries. However, identifying this subset still requires scoring the current query against the full KV cache and performing global Top-$K$ selection, leaving selector cost linear in context length and limiting the practical efficiency of sparse attention for long-context decoding. In this paper, we introduce ReTopK, a training-free method that accelerates dynamic Top-$K$ attention by reusing historical retrieval decisions. ReTopK builds on the observation that similar queries often attend to overlapping supports and that partially overlapping supports can still preserve most of the Exact Top-$K$ attention mass. For each attention head, it maintains a bounded cache of historical query--support pairs, retrieves the most similar cached queries for each new query, unions their stored supports with a recent window, and reranks only the resulting compact candidate set using exact current-query scores. A similarity-based fallback invokes full-history Exact Top-$K$ when reuse is unreliable, while periodic exact refreshes limit cache drift. ReTopK retains the complete KV cache and reuses only selected indices, rather than historical scores, attention weights, or outputs. Across 16K--128K contexts, ReTopK achieves the lowest PG19 perplexity and the highest NIAH and LongBench scores among the evaluated approximate methods. At 128K with $K=512$, ReTopK incurs only a 0.50\% perplexity increase over Exact Top-$K$ while accelerating attention computation by $3.07\times$.