SIFT: 어텐션 불변성(Attention Invariance)을 활용한 RAG 프리필(Prefill)의 빠른 연산을 위한 선택적 인덱스

검색 증강 생성 (Retrieval-Augmented Generation, RAG)은 응답 품질을 향상시키기 위해 LLM 쿼리에 관련 문서를 주입합니다. 이러한 주입은 프롬프트 길이를 증가시키고 첫 번째 토큰 생성 시간 (Time To First Token, TTFT)을 늦춥니다. 표준 쿼리와 달리, RAG 쿼리는 동일한 문서가 사용자 쿼리 전반에 걸쳐 반복되는 컨텍스트 재사용 (Context Reuse)이라는 독특한 특성을 가집니다. 따라서 모든 RAG 쿼리에 대해 문서를 완전히 다시 계산하는 것은 불필요한 연산을 수행하며 TTFT를 증가시킵니다. 기존 연구들은 RAG 문서의 KV 텐서 (KV tensors)를 오프라인에서 미리 계산하고, 온라인 프리필 (Online Prefill) 중에 일부 토큰을 거칠게 (Coarsely) 다시 계산합니다. 그러나 이러한 KV 재사용은 높은 지연 시간을 갖는 디스크 전송으로 인해 최신 GPU에서는 전체 재계산보다 느린 경우가 많습니다. 또한, 이러한 거친 입도 (Coarse-grained)의 재계산은 정확도를 저하시킵니다. 이러한 한계를 해결하기 위해, 본 논문은 어텐션 불변성 (Attention Invariance)을 활용하여 RAG 프리필의 빠른 연산을 위한 선택적 인덱스인 SIFT (Selective-Index For Fast Compute of RAG Prefill by Exploiting Attention Invariance)를 제안합니다. SIFT는 문서를 오프라인으로 처리하고 각 문서에 대해 높은 어텐션 점수 (Attention scores)를 가진 미세한 위치 (Fine-grained locations)를 추출합니다. 다음으로, 우리는 런타임 중에 추출된 위치를 활용할 수 있게 해주는 다음과 같은 어텐션 불변성 통찰을 식별합니다: (1) 로컬 어텐션 불변성 (Local-Attention Invariance): 문서 내 높은 어텐션 점수의 위치는 주변 문서에 대해 불변합니다. 이는 문서가 자기 자신에게 어텐션을 주는 높은 점수의 위치를 예측하는 데 도움이 됩니다. (2) 교차 어텐션 일관성 (Cross-Attention Consistency): 문서 내 어텐션이 높은 키 (Keys)는 이후의 문서로부터의 교차 어텐션 (Cross-attention)도 유도합니다. 이는 문서가 미래의 문서에 어텐션을 주는 높은 점수의 위치를 예측하는 데 도움이 됩니다. 결정적으로, SIFT는 KV 데이터를 저장하지 않으며 오직 두 개의 압축된 비트 벡터 (Bit vectors) 형태로 높은 점수의 위치만을 저장합니다. SIFT의 저장 용량은 KV 텐서보다 최대 24,000배 작아 비용이 많이 드는 디스크 전송을 방지합니다. 프리필 과정에서 SIFT는 표시된 위치에 대해서만 어텐션을 계산하며, 전체 재계산 대비 정확도를 1% 이내로 유지하면서 TTFT를 1.71배 개선합니다.