Sparse Attention의 Indexer: 내용에 따라 읽을 위치 고르기

앞선 글에서는 가까운 위치나 일정한 간격의 위치처럼, 정해진 규칙으로 읽을 범위를 제한했습니다. 하지만 현재 토큰에 필요한 정보가 항상 그 범위 안에 있지는 않습니다. 멀리 떨어진 위치라도 지금의 계산에 중요하다면 직접 읽을 수 있어야 합니다.

이번에는 현재 입력과 각 위치의 내용을 비교해 읽을 위치를 고르는 Indexer(인덱서)를 살펴보겠습니다. 먼저 선택용 Query와 Key로 위치를 고른 뒤 본 Attention에 연결하는 과정을 보겠습니다. 이어 토큰을 하나씩 처리하면서 선택이 어떻게 달라지는지, 최근 위치를 읽는 window와 어떻게 결합하는지 설명하겠습니다.

선택용 Query와 Key로 위치 고르기

전체 Key에 대해 본 Attention의 점수를 모두 계산한 다음 일부만 남기면, 점수를 계산하는 비용은 이미 지불한 셈입니다. 인덱서를 두는 목적은 본 Attention보다 가벼운 계산으로 읽을 후보를 먼저 좁히는 것입니다. 이를 위해 위치 선택에 사용할 Query와 Key를 따로 만듭니다.

그림 1은 현재 위치 $p_7$에서 여덟 위치 중 세 개를 고르는 예시입니다. $p_0$부터 $p_7$까지는 토큰 위치이며, 헤드 번호가 아닙니다. $h_7$은 이 Attention 층에 들어온 현재 토큰의 벡터입니다. 같은 입력에서 선택용 Query와 본 Attention의 Query가 만들어지지만, 서로 다른 투영을 사용합니다.

그림 왼쪽의 $q_7^i$와 $k_j^i$는 각각 선택용 Query와 선택용 Key입니다. 위첨자 i는 인덱서용이라는 표시이며, 아래첨자는 토큰 위치입니다. 두 성분짜리 벡터로 계산을 따라갈 수 있도록 단순화했습니다. 선택용 Query는 [1, 0]이고, 위치 0의 선택용 Key는 [0.9, 0.2]입니다. 내적은 1 × 0.9 + 0 × 0.2 = 0.9가 됩니다.

같은 계산을 여덟 위치에 적용하면 위치 순서대로 선택 점수는 [0.9, 0.2, 0.3, 0.1, 0.8, 0.4, 0.5, 0.7]입니다. 가장 큰 세 점수는 위치 0의 0.9, 위치 4의 0.8, 위치 7의 0.7이므로, Top-3 선택 결과는 위치 0·4·7입니다. 여기서 Top-k는 점수가 큰 k개를 고른다는 뜻입니다.

그림 오른쪽에서는 이 세 위치의 본 Attention용 Key와 Value를 읽습니다. 본 Query $q_7$과 선택된 Key로 Attention 점수를 계산하고, 선택된 위치들에 대한 softmax 가중치로 Value를 합칩니다. 이 부분의 연산은 앞서 배운 Attention과 같습니다. 달라지는 것은 참여할 위치가 먼저 정해졌다는 점입니다.

선택 점수 0.9·0.8·0.7을 그대로 Value의 가중치로 쓰는 것은 아닙니다. 인덱서가 넘겨주는 것은 읽을 위치이고, 그 위치의 정보를 얼마나 반영할지는 본 Attention에서 다시 계산합니다. 선택용 Key 역시 본 Attention의 Key를 대체하지 않습니다.

DeepSeek-V3.2의 공식 구현에서도 인덱서가 별도의 Query·Key 표현으로 점수를 구하고 선택한 위치를 반환합니다. 실제 구현은 여러 인덱서 헤드와 위치 정보 등을 사용하므로, 그림의 두 성분 내적은 전체 연산을 그대로 재현한 것이 아닙니다. 그림은 위치를 고르는 계산과 선택한 KV를 읽는 계산의 역할을 분리해 보여 줍니다. 본 KV를 저장하는 표현도 앞서 배운 MLA처럼 달라질 수 있습니다.

생성이 진행되면 선택도 달라집니다

인덱서가 고른 위치는 이후에도 고정되어 있을까요? 같은 층에서 $p_6$을 처리한 뒤 다음 $p_7$을 처리하는 과정을 보겠습니다. 이전 위치의 선택용 Key와 본 KV는 유지하지만, 현재 입력에서 만드는 Query는 달라집니다. 새 위치의 Key와 KV도 캐시에 추가됩니다.

$p_6$을 처리할 때는 위치 0부터 6까지만 후보입니다. 아직 미래인 위치 7은 후보에 들어가지 않습니다. 그림의 선택용 Query $q_6^i$는 [0, 1]이므로, 각 선택용 Key의 두 번째 성분이 점수가 됩니다. 점수는 [0.2, 0.9, 0.1, 0.7, 0.3, 0.4, 0.8]이고, 위치 1·3·6이 선택됩니다.

다음 $p_7$을 처리할 때는 새 위치의 선택용 Key와 본 KV를 추가합니다. 이때 선택용 Query $q_7^i$는 [1, 0]이므로 첫 번째 성분으로 점수가 정해집니다. 결과는 그림 1과 같은 위치 0·4·7입니다. 과거 Key를 바꾸지 않아도, 현재 Query가 바뀌면 같은 과거 위치에 매기는 점수가 달라집니다. 실제 생성에서도 선택은 달라질 수 있지만, 매번 반드시 다른 위치를 고르는 것은 아닙니다.

이 예시는 앞선 글에서 구별한 읽는 범위와 캐시 보관량의 차이도 보여 줍니다. $p_6$에서 읽지 않았던 위치 0과 4를 $p_7$에서 다시 읽습니다. 따라서 $p_6$의 Top-3에 들지 않았다는 이유만으로 나머지 KV를 지우면 다음 선택을 수행할 수 없습니다.

그림에서는 $p_7$에서 본 KV 세 위치만 읽지만 캐시에는 여덟 위치가 남아 있습니다. 인덱서가 이후에도 전체 과거 위치를 후보로 삼는 동안은, 해당 위치의 선택용 Key와 다시 읽을 본 KV를 사용할 수 있어야 합니다. 읽을 위치 선택만으로 캐시 보관량까지 같은 비율로 줄어드는 것은 아닙니다.

최근 window와 선택한 위치 결합하기

내용만으로 상위 위치를 고르면 최근 위치도 선택에서 빠질 수 있습니다. 그림 1에서는 직전 위치 5와 6이 Top-3에 들어가지 않았습니다. 가까운 문맥은 항상 읽고 싶다면, 인덱서의 선택 결과에 sliding window를 더할 수 있습니다.

앞선 글과 같이 window 3을 현재 위치와 직전 두 위치로 정의하겠습니다. $p_7$에서는 위치 5·6·7입니다. 그림 3은 이 window와 인덱서가 고른 위치 0·4·7을 합칩니다.

두 집합의 합집합은 위치 0·4·5·6·7입니다. 위치 7은 양쪽에 있지만 같은 원본 KV이므로 한 번만 읽습니다. 이렇게 합친 다섯 위치에 대해 본 Attention을 계산합니다. Window는 최근 세 위치가 빠지지 않게 하고, 인덱서는 내용에 따라 다른 위치에 접근할 수 있게 합니다.

이 그림에서는 window 안의 위치도 인덱서의 후보에 포함했습니다. 따라서 window 크기 3과 Top-3을 더해도 실제로 읽는 위치는 여섯 개가 아니라 다섯 개입니다. 다른 설계로는 window 안을 후보에서 제외하고, window 밖에서 추가로 k개를 고를 수도 있습니다. 후보가 충분하다면 이 방식은 선택 예산 k개를 모두 window 밖에 사용합니다. Window와 인덱서를 결합한다는 말만으로 후보 범위나 중복 처리까지 정해지는 것은 아닙니다.

선택 비용과 읽기 비용

인덱서는 본 Attention이 읽을 위치를 줄이는 대신, 선택을 위한 계산을 추가합니다. 그림 1에서도 본 Attention은 세 위치만 읽지만, 인덱서는 여덟 위치의 선택용 Key를 비교합니다. 문맥이 길어지면 후보 Key를 읽고 점수를 계산하는 비용과 Top-k를 구하는 비용도 커집니다.

따라서 유리한 조건은 선택 비용을 추가해도 본 Attention에서 줄어드는 비용이 더 큰 경우입니다. 선택용 표현의 크기, 인덱서 헤드 수, 고를 위치 수, 선택된 KV를 가져오는 방식과 커널 구현에 따라 결과가 달라집니다. 여덟 위치 중 세 개를 읽는다는 이유만으로 전체 실행 시간이 3/8이 되는 것은 아닙니다.

선택의 정확도도 중요합니다. 본 Attention이 중요하게 사용할 위치를 인덱서가 빠뜨리면, 그 위치의 Value는 이번 출력에 직접 반영되지 않습니다. 선택 개수를 늘리면 더 많은 정보를 읽을 수 있지만 비용도 늘어납니다. 인덱서는 모든 위치를 읽는 Attention과 같은 결과를 보장하는 계산 순서 변경이 아니라, 읽을 정보를 제한하는 모델 설계입니다.