🤖 machine learning

SAKI: Score-Aware Low-Rank Key Indexing for Long-Context KV Retrieval

SAKI는 폐쇄형 비대칭 인수분해를 통해 어텐션 점수 왜곡을 직접 최소화함으로써 KV 캐시 압축을 최적화하는 학습이 필요 없는 점수 인지형 저계수 키 인덱싱 방법으로, 이를 통해 여러 거대 언어 모델의 롱 컨텍스트 검색 재현율 측면에서 PCA와 같은 기존의 키 재구성 기반 방식들을 크게 능가한다.

원저자: Lin Zhang

게시일 2026-08-05

📖 2 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Lin Zhang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 아주 작은 도시 크기만 한 건초더미 속에서 특정한 바늘 하나를 찾아야 한다고 상상해 보십시오. 그런데 질문을 할 때마다 매번 이 작업을 수행해야 합니다. 이것이 현대 인공지능이 긴 대화나 방대한 문서를 기억하려고 할 때 마주하는 일상적인 현실입니다. AI의 "기억"(KV 캐시라고 불림)은 너무 거대해져서 일반적인 노트북이 담을 수 있는 공간보다 더 많은 공간을 차지하게 되며, 이로 인해 검색 속도가 느려지고 비용이 많이 듭니다. 이를 해결하기 위해 엔지니어들은 하나의 기술을 사용합니다. 모든 짚단을 일일이 살펴보는 대신, 중요한 바늘이 어디에 숨어 있을지 추측할 수 있는 빠르고 대략적인 지도(인덱스)를 만드는 것입니다. 만약 지도가 좋다면 AI는 바늘을 빠르게 찾아내지만, 지도가 나쁘다면 AI는 혼란에 빠져 잘못된 답을 내놓게 됩니다.

오랫동안 과학자들은 두 가지 요소를 보고 이 지도를 만들려고 시도했습니다. 바로 AI 두뇌의 "설계도"(가중치)이거나, 혹은 그것이 보유한 데이터의 "형태"(키의 분산)였습니다. 이것은 책의 제목을 전혀 읽지 않은 채, 오직 책등의 색깔이나 페이지의 두께만을 보고 도서관을 정리하려는 것과 같습니다. 문제는 AI가 책등의 색깔이나 페이지의 두께에는 관심이 없다는 점입니다. AI는 특정 질문이 특정 답변과 얼마나 잘 맞물리는지에 관심을 가집니다. SAKI라는 제목의 이 논문은 기존의 지도들이 중요도를 측정하는 데 있어 잘못된 자를 사용했다고 주장합니다. 저자는 완벽한 지도를 만들기 위해서는 일반적인 형태나 설계도를 바탕으로 추측하는 것이 아니라, 특정 질문과 답변이 정확히 얼마나 서로 "클릭(맞물림)"되는지를 측정해야 한다는 사실을 깨달았습니다.

이 논문은 SAKI(Score-Aware Low-Rank Key Indexing)라는 새로운 방법을 소개합니다. SAKI는 일반적인 자를 사용하는 대신, AI의 질문이 저장된 기억과 정확히 얼마나 잘 일치할지를 예측하는 맞춤형 "점수 인식형(score-aware)" 지도를 구축합니다. 저자는 LLaMA-3.1-8B 및 Qwen2.5-7B를 포함한 여러 인기 있는 AI 모델을 대상으로 이 새로운 지도를 테스트했습니다. 그 결과, SAKI는 이전의 최선책들보다 올바른 바늘을 찾는 능력이 훨씬 뛰어나다는 것을 발견했습니다. 예를 들어, 지도를 작은 크기(rank 32)로 압축했을 때, SAKI는 기존 방식들이 여전히 범하던 오류를 13%에서 30%까지 줄임으로써 AI의 정보 회상 능력을 개선했습니다. LLaMA-3.1-8B 모델의 경우 성공률을 0.748에서 0.799로 높였고, Qwen2.5-7B에서는 0.786에서 0.850으로 끌어올렸습니다.

저자는 기존 방식들이 AI의 메모리를 정적인 데이터 더미처럼 취급하여, 질문에 따라 데이터의 중요성이 변한다는 사실을 간과했기 때문에 실패했다고 설명합니다. 저자는 AI의 내부 "스코어링 머신(scoring machine)"이 기이하고 불균형하다는 것(수학적으로 "비정규적(non-normal)"임)을 보여주었습니다. 이는 표준적인 데이터 압축 방식(PCA 등)이 잘못된 부분을 잘라내게 된다는 것을 의미합니다. SAKI는 질문과 답변을 동시에 고려하는 특별한 수학적 지름길을 사용하여 이 문제를 해결합니다. 이 논문은 이 새로운 접근 방식이 단순히 운 좋은 추측이 아님을 증명합니다. 수학적 예측은 거의 완벽한 정확도(상관관계 0.997)로 결과를 예측합니다. 저자는 아직 모든 가능한 유형의 텍스트나 전체 엔드 투 엔드(end-to-end) 대화에 대해 테스트하지 않았다고 언급했지만, 측정 결과 SAKI가 모델을 재학습시키지 않고도 AI의 메모리를 더 빠르고 똑똑하게 만드는 데 있어 중요한 진전임을 보여줍니다.

기술 요약: SAKI (Score-Aware Low-Rank Key Indexing)

문제 정의

긴 컨텍스트 디코딩(예: 1M 토큰)에서 8B 파라미터 모델의 KV 캐시는 상당한 메모리(bf16 기준 약 130 GB)를 소비하며, 이는 검색의 핵심적인 병목 구간이 됩니다. 기존의 저차원(low-rank) KV 캐시 인덱싱 전략은 두 가지 범주로 나뉘지만, 둘 다 어텐션 메커니즘에서 실제로 사용되는 메트릭을 최적화하는 데 실패합니다:

가중치 측 절단(Weight-side truncation): 투영 행렬을 절단하는 방식(예: Weight-SVD)은 연산자의 기하학적 구조는 최적화하지만 데이터 분포를 무시하여 낮은 재현율(rank $r=32$ 에서 0.366)을 초래합니다.
데이터 인식 재구성(Data-aware reconstruction): Key-PCA와 같은 방식은 캐시된 키의 분산을 보존합니다. 그러나 어텐션 스코어는 $s = q^\top k$ 로 계산됩니다. 높은 키 분산을 가진 방향이 쿼리에 의해 읽히지 않을 수 있으며, 반대로 쿼리 투영( $W_Q$ )에 의해 증폭된 완만한 분산의 방향이 스코어를 지배할 수도 있습니다. PCA는 재구성 오차( $\|k - Pk\|^2$ )를 최소화할 뿐, 스코어 왜곡(score distortion)을 최소화하지는 않습니다.

핵심 문제는 두 범주 모두 KV 검색의 진정한 목적 함수인 **어텐션 스코어 랭킹(attention score ranking)**을 최적화하지 못한다는 점입니다.

방법론: SAKI

본 논문은 rank- $r$ 키 압축으로 인해 발생하는 기대 어텐션 스코어 왜곡을 도출하고, 이를 직접 최적화하는 훈련 불필요(training-free) 인덱스인 SAKI(Score-Aware Low-Rank Key Indexing)를 제안합니다.

1. 이론적 도출

저자는 rank- $r$ 선형 사상 $P$ 가 키에 적용될 때의 기대 스코어 왜곡 $L(P)$ 를 정의합니다. 2차 모멘트 $\Sigma_q$ 와 $\Sigma_k$ 를 갖는 독립적인 쿼리/키 모델과 어텐션 연산자 $A = W_Q^\top W_K$ 하에서:
$L(P) = \mathbb{E}[(q^\top A (I-P) k)^2] = \|\Sigma_q^{1/2} A (I-P) \Sigma_k^{1/2}\|_F^2$
이 공식은 최적의 인덱스가 **양방향 공분산 가중 저차원 근사(two-sided covariance-weighted low-rank approximation)**임을 보여줍니다. 단일 측면 가중치( $\Sigma_k$ )를 사용하는 PCA와 달리, SAKI는 모델의 스코어링 연산자 $A$ 에 의해 매개되는 쿼리 통계량( $\Sigma_q$ )과 키 통계량( $\Sigma_k$ ) 모두에 의해 가중치를 부여합니다.

2. 정확한 해법 (Exact Solution)

본 논문은 $L(P)$ 에 대한 rank- $r$ 최적해가 투영사(projector, PCA와 같은 방식)가 아니라 **비대칭 저차원 이선형 인수분해(asymmetric low-rank bilinear factorization)**임을 증립합니다.

$C = \Sigma_q^{1/2} \Sigma_k^{1/2}$ 라고 정의합니다 (연산자 $A$ 가 흡수된 헤드 좌표계 기준).
최적의 사상 $M_r$ 은 $C$ 의 절단된 SVD로부터 도출됩니다: $C_r = U_r \Lambda_r V_r^\top$ .
해는 $M_r = \Sigma_q^{-1/2} C_r \Sigma_k^{-1/2}$ 입니다.
핵심 차이점: $M_r$ 은 일반적으로 항등적이지 않습니다 ( $M_r^2 \neq M_r$ ). 이는 $M_r$ 이 멱등한 투영사가 아니라 최적의 선형 스코어 맵임을 의미합니다. 이를 통해 SAKI는 스팬 기반 방식(PCA 포함)이 표현할 수 없는 방향을 포착할 수 있습니다.

3. 알고리즘 (SAKI)

이 방법은 훈련이 필요 없으며(training-free) 단 한 번의 캘리브레이션 패스만 요구합니다:

헤드별 공분산 $\Sigma_q$ (pre-RoPE 쿼리의 비중심 2차 모멘트) 및 $\Sigma_k$ (중심화된 키)를 계산합니다.
행렬 $C = \Sigma_q^{1/2} \Sigma_k^{1/2}$ 를 계산합니다.
$C$ 에 대해 SVD를 수행하여 $U_r, \Lambda_r, V_r$ 을 얻습니다.
키를 코드 $c_j = B_k^\top (k_j - \mu)$ 로 저장하고 $(B_q^\top q)^\top c_j$ 를 통해 스코어를 재구성하는 압축 행렬 $B_k$ 와 $B_q$ 를 구성합니다.
이 맵들은 오프라인에서 모델의 $W_Q$ 및 $W_K$ 가중치로 폴딩(fold)될 수 있습니다.

주요 기여

문제 정식화: 기대 스코어 왜곡을 KV 검색의 올바른 목적 함수로 식별하여, 양방향 공분산 가중 근사 문제를 도출했습니다.
정확한 해법: 스팬 제한적(projector-based) 방식보다 엄격하게 우수한 비대칭 저차원 이선형 인수분해의 폐쇄형(closed-form) 해를 도출했습니다.
경험적 검증: 이론적 독립성 가정(폐쇄형 도출에 사용됨)이 매우 높은 정확도로 성립함을 입증했습니다 ( $r=0.997$ 의 예측 및 실제 스코어-MSE 감소 간 피어슨 상관계수).
연산자 기하학 진단: 어텐션 스코어 연산자가 매우 비정규적(non-normal)이며(중앙값 Henrici departure 0.95–0.98), 두꺼운 꼬리를 가진 스펙트럼을 가짐을 분석했습니다. 이는 왜 가중치 전용 절단이 실패하는지(살아있는 방향을 버림), 그리고 왜 불변 부분 공간(invariant-subspace) 방식이 실패하는지(허브 형태의 교차 부분 공간 결합을 끊음)를 설명합니다.

실험 결과

저자는 네 가지 모델(LLaMA-3.1-8B, Qwen2.5-7B, Mistral-7B-v0.1, Llama-3.2-3B)에 대해 4,096 토큰의 자연어 텍스트와 정확한 쿼리를 사용하는 프로토콜로 top-64 재현율을 측정하여 SAKI를 평가했습니다.

성능: SAKI는 모든 랭크( $r=16, 32, 64$ $r = 16, 32, 64$ )와 모든 모델에서 Key-PCA를 능가합니다.
- $r=32$ 에서, SAKI는 PCA가 남긴 **잔여 재현율 오차의 13~30%**를 제거합니다.
- 예시 (LLaMA-3.1-8B): 재현율이 0.748 (PCA)에서 **0.799 (SAKI)**로 향상되었습니다.
- 예시 (Qwen2.5-7B): 재현율이 0.786 (PCA)에서 **0.850 (SAKI)**로 향상되었습니다.
헤드별 개선: 모델당 어텐션 헤드의 68~89%에서 개선이 나타났습니다.
계층별 분포: 이득은 쿼리와 키의 기하학적 구조가 가장 크게 벌어지는 깊은 층(예: 레이어 16–31)에 집중되어 있습니다.
어블레이션(Ablation): 논문은 이득의 원천을 격리했습니다. "가중 스팬(weighted span)"을 사용하는 변형(SAP-map)이 일반 스팬(PCA)보다 성능이 좋았으나, 정확한 비대칭 최적해(SAKI-opt)가 엄격하게 더 우수했습니다. 이는 이득이 단순히 공분산 행렬의 삽입 때문이 아니라, 목적 함수와 비대칭 인수분해에서 온다는 것을 확인시켜 줍니다.

의의 및 주장

본 논문은 기존의 접근 방식(가중치 측 또는 데이터 재구성)이 잘못된 양을 최적화하기 때문에 실패한다고 주장합니다. 어텐션은 키의 분산뿐만 아니라 쿼리와 키 통계 사이의 상호작용에 의존합니다.

이론적 타당성: 도출된 폐쇄형 솔루션은 단순한 휴리스틱이 아니라, 스코어 왜곡 목적 함수에 대한 정확한 최적해입니다.
실질적 영향: SAKI는 훈련이 필요 없는 드롭인(drop-in) 교체 방식으로, 특히 검색이 가장 어려운 깊은 층에서 KV 인덱스의 검색 충실도를 크게 향상시킵니다.
한계점: 저자는 SAKI가 재현율을 높이지만, $r \le 32$ 에서 아직 0.95 재현율에 도달하지는 못했다고 언급했습니다. 엔드 투 엔드 생성 품질과 더 긴 컨텍스트(4K 이상) 또는 다른 도메인에서의 성능은 추가 검증이 필요합니다. 독립성 가정은 RoPE 회전에 대한 근사치이지만 경험적으로 검증되었습니다.

결론적으로, 본 논문은 미래의 인덱싱이 어텐션 메커니즘의 특정 "연산자 기하학", 즉 SAKI가 포착하도록 설계된 비정규적이고 교차 부분 공간적인 결합을 보존해야 한다고 강조합니다.