토큰 축 압축: 여러 위치의 KV를 요약해서 읽기

앞선 글에서는 인덱서가 읽을 KV 위치를 골랐습니다. 이번에는 선택하기 전에 여러 토큰 위치의 정보를 하나의 요약 KV로 합치는 방법을 살펴보겠습니다. 위치마다 항목을 하나씩 읽는 대신, 여러 위치를 대표하는 항목을 읽으면 긴 문맥을 더 적은 항목으로 다룰 수 있습니다.

먼저 이 압축이 MLA에서 배운 압축과 어느 축에서 다른지 구별하겠습니다. 이어 두 위치의 입력을 하나의 요약 벡터로 만드는 계산을 따라가고, 완성된 요약과 최근 위치별 KV를 함께 읽는 이유를 살펴보겠습니다. 마지막으로 압축된 요약 중 일부를 인덱서로 선택하는 방식까지 연결하겠습니다.

벡터 폭과 토큰 위치 수

KV 저장량을 줄인다는 말만으로는 무엇이 줄어드는지 알 수 없습니다. 위치별 벡터를 행으로 쌓았을 때, 한 행의 성분 수를 줄이는 것과 여러 행을 하나로 합치는 것은 서로 다른 변화입니다.

그림 1은 두 방식 모두 8행 × 4성분에서 출발합니다. 행은 $p_0$부터 $p_7$까지의 토큰 위치이고, 열은 벡터 성분입니다. 헤드 수를 나타내는 격자가 아닙니다.

왼쪽은 각 위치의 표현을 네 성분에서 두 성분으로 줄입니다. 위치는 여전히 여덟 개이므로 결과는 8행 × 2성분입니다. MLA의 저장 구조에서 배운 잠재 표현도 위치마다 저장할 벡터의 폭을 줄인다는 관점으로 이해할 수 있습니다. 다만 실제 MLA에는 별도의 위치 성분도 있으므로, 이 격자 자체가 MLA의 전체 캐시 크기를 나타내는 것은 아닙니다.

오른쪽은 인접한 두 위치를 묶어 하나의 요약으로 만듭니다. $p_0$·$p_1$은 $b_0$, $p_2$·$p_3$은 $b_1$이 됩니다. 각 요약의 성분 수는 네 개로 유지하고, 항목 수를 여덟 개에서 네 개로 줄여 4행 × 4성분을 얻습니다. 이것이 이번 편에서 다룰 토큰 축 압축입니다. 여기서 토큰 수를 줄인다는 것은 모델이 생성할 출력 토큰 수가 아니라, Attention이 읽을 KV 항목 수를 줄인다는 뜻입니다.

요약 $b_0$는 원래 $p_0$의 KV를 그대로 남겨 둔 것이 아닙니다. 두 위치의 정보를 합쳐 만든 새로운 표현입니다. 따라서 읽을 위치를 고르는 인덱서와도 구별해야 합니다. 인덱서는 후보 중 일부를 고르고, 압축은 여러 후보의 정보를 합쳐 새 항목을 만듭니다. 그림 양쪽의 최종 칸 수가 같더라도 담는 정보나 계산이 같다는 뜻은 아닙니다.

두 투영으로 요약 벡터 만들기

두 위치를 합칠 때 모든 성분을 단순 평균할 수도 있지만, 여기서는 각 성분에서 어느 위치를 얼마나 반영할지 학습하는 방식을 살펴보겠습니다. 같은 입력에서 두 경로가 나옵니다. 내용 투영은 요약에 담을 벡터를 만들고, 점수 투영은 그 내용을 섞을 비율을 정하는 데 사용할 점수를 만듭니다.

그림 2의 입력은 두 위치의 벡터 $h_0$ = [1, 0], $h_1$ = [0, 1]입니다. 두 벡터를 행으로 쌓은 행렬을 H라고 하겠습니다. 이 예시는 두 성분을 유지하면서 위치 두 개를 하나로 합칩니다.

첫 단계에서는 H에 서로 다른 투영 행렬 $W_u$와 $W_s$를 곱합니다. $W_u$는 내용 벡터를, $W_s$는 성분별 점수를 만드는 학습 파라미터입니다. 모델이 학습하는 것은 투영 행렬이며, 투영 결과는 입력에 따라 달라집니다. 그림의 행렬 값은 계산을 쉽게 따라가도록 정한 교육용 예시입니다.

내용 투영 $W_u$의 두 행을 [2, 0], [0, 4]로 두었습니다. $h_0$에 곱하면 [2, 0], $h_1$에 곱하면 [0, 4]가 나오며, 이 결과를 행으로 쌓은 것이 내용 행렬 U입니다. 이 단계에서는 아직 위치를 합치지 않습니다. 두 위치 각각의 내용을 변환했을 뿐입니다.

점수 투영 $W_s$의 두 행은 [ln 3, 0], [0, ln 3]입니다. 같은 입력 H를 곱하면 점수 행렬 S도 이 값이 됩니다. 입력을 단위행렬로 잡았기 때문에 이번 예시에서는 투영 행렬과 출력이 같아 보이지만, 일반적인 입력에서는 그렇지 않습니다. ln 3과 0은 내용 벡터 [2, 0]에서 계산한 값이 아니라, 같은 입력을 별도의 점수 행렬로 투영한 결과입니다.

이 점수를 바로 가중치로 쓰지는 않습니다. 각 성분에서 두 위치의 점수에 softmax를 적용합니다. 그림의 열 하나를 위아래로 정규화하는 것입니다. 첫 성분의 점수는 [ln 3, 0]이므로 지수값은 [3, 1]이고, 합 4로 나누면 [0.75, 0.25]가 됩니다. 둘째 성분의 점수는 [0, ln 3]이므로 가중치는 [0.25, 0.75]입니다. 이처럼 계산이 간단한 비율을 얻기 위해 예시 점수로 ln 3과 0을 골랐습니다.

결과인 가중치 행렬 A는 각 열의 합이 1입니다. 첫 성분은 위치 0을 75%, 위치 1을 25% 반영하고, 둘째 성분은 반대로 반영합니다. 두 위치를 섞는 비율을 벡터 전체에 하나만 적용하는 것이 아니라, 성분마다 따로 정합니다.

두 번째 단계에서는 U와 A의 대응하는 성분끼리 곱한 뒤 위치 방향으로 더합니다. 여기서 두 행렬의 곱은 앞선 투영의 행렬 곱과 달리 성분별 곱입니다. 첫 성분은 2 × 0.75 + 0 × 0.25 = 1.5, 둘째 성분은 0 × 0.25 + 4 × 0.75 = 3입니다. 최종 요약 $c_0$는 [1.5, 3]이 됩니다. 성분 수는 두 개로 유지되면서 위치별 항목 두 개가 요약 하나로 줄었습니다.

이 가중치는 본 Attention에서 현재 Query와 Key를 비교해 만드는 가중치와 역할이 다릅니다. 압축 가중치는 요약을 만드는 데 쓰이고, Attention 가중치는 만들어진 요약을 읽을 때 쓰입니다. 그림 2의 요약 생성에는 현재 Query가 참여하지 않습니다. 완성된 요약은 캐시에 저장해 이후 Query들이 다시 사용할 수 있습니다.

DeepSeek-V4 공식 구현의 Compressor에서도 같은 입력에 내용용 투영과 점수용 투영을 적용하고, 위치 방향 softmax와 가중합으로 압축합니다. 그림은 그 계산을 작은 예시로 분리한 것입니다. 실제 구현의 위치별 편향, 겹치는 압축 범위, 후속 정규화와 위치 정보 처리는 생략했습니다. 이미 만들어진 임의의 KV를 이 계산으로 합치면 원래 Attention 결과가 그대로 보존된다는 뜻은 아닙니다.

완성된 요약과 최근 위치별 KV

여러 위치를 묶으려면 그 위치들의 입력이 모두 도착해야 합니다. 토큰을 하나씩 처리하는 중에는 아직 완성되지 않은 묶음이 생깁니다. 동시에 최근 문맥은 요약에만 의존하지 않고 각 위치를 따로 읽을 수 있도록 남길 수 있습니다.

그림 3에서는 두 위치마다 요약을 하나 만들고, 현재 위치를 포함한 최근 세 위치의 KV를 별도로 유지합니다. 이를 최근 위치별 KV라고 부르겠습니다. 여러 위치를 합치지 않은 KV라는 뜻이며, 모델에 들어온 입력 벡터 자체를 말하는 것은 아닙니다. 그림의 ‘원본’도 이처럼 토큰 축으로 합치기 전의 항목을 가리킵니다.

현재 $p_6$을 처리한다면 $p_0$·$p_1$을 합친 $b_0$, $p_2$·$p_3$을 합친 $b_1$, $p_4$·$p_5$를 합친 $b_2$까지 완성되어 있습니다. 하지만 $p_6$·$p_7$의 묶음 $b_3$는 아직 만들 수 없습니다. 미래 위치 $p_7$을 포함한 요약을 $q_6$이 읽으면 미래 정보를 미리 사용하는 셈이 됩니다.

그래서 $q_6$은 완성된 요약 $b_0$·$b_1$·$b_2$와 최근 위치별 KV $p_4$·$p_5$·$p_6$을 함께 읽습니다. 그림의 Summary bank는 완성된 요약을 모아 둔 캐시입니다. 여기서는 선택 과정을 잠시 제외하고, 완성된 요약은 모두 읽는다고 가정합니다. 아직 요약에 들어가지 못한 현재 $p_6$의 정보는 최근 위치별 KV를 통해 읽을 수 있습니다.

여기서 $p_4$·$p_5$는 이미 $b_2$로 합쳐졌는데 왜 개별 KV도 남길까요? $b_2$는 두 위치가 섞인 하나의 항목입니다. Query는 $b_2$에 가중치를 줄 수 있지만, 그 안에서 $p_4$와 $p_5$를 다시 분리해 각각 다른 Attention 가중치를 줄 수는 없습니다. 최근 위치별 KV를 유지하면 가까운 문맥을 각 위치별로 구별해서 읽을 수 있습니다.

요약 생성과 최근 KV 보관은 별도의 규칙입니다. 묶음이 완성되면 요약을 추가하고, window 안에 있는 동안은 개별 KV도 유지합니다. 따라서 압축이 끝났다고 즉시 그 위치의 개별 KV를 지우지는 않습니다. $b_2$와 $p_4$·$p_5$는 같은 시기의 정보를 포함하지만 서로 다른 읽기 항목입니다. 앞선 글에서 동일한 원본 위치를 window와 인덱서가 함께 골라 중복을 제거한 경우와는 다릅니다.

오른쪽처럼 $p_8$까지 진행하면 $p_6$·$p_7$의 요약 $b_3$도 완성됩니다. 최근 위치별 KV는 $p_6$·$p_7$·$p_8$로 이동합니다. 이제 $p_4$·$p_5$는 최근 window 밖으로 나갔으므로 이 경로의 개별 KV에서는 빠지고, 그 정보는 $b_2$에 요약된 형태로 남습니다. 아직 미래인 $p_9$을 포함할 $b_4$는 읽지 않습니다.

이 예시는 두 위치씩 요약하는 방식이므로 캐시가 일정한 크기로 고정되는 것은 아닙니다. 최근 위치별 KV는 세 개로 유지되지만, 과거 요약은 입력 두 위치마다 하나씩 늘어납니다. 또한 요약을 만드는 동안 필요한 중간 상태도 있습니다. 토큰 축 압축은 과거 정보를 저장하는 항목 수의 증가를 완만하게 만들지만, 긴 문맥 전체를 고정된 상태 하나로 바꾸는 구조는 아닙니다.

압축과 인덱서 결합하기

요약이 많아지면 그중 일부만 골라 읽을 수도 있습니다. 앞선 글의 인덱서가 토큰 위치를 골랐다면, 이제는 압축된 요약 항목을 선택 대상으로 삼는 것입니다. 반대로 충분히 적은 수의 요약을 만든 뒤 모두 읽는 설계도 가능합니다.

그림 4는 현재 $p_{16}$에서 완료된 과거 위치 $p_0$부터 $p_{15}$까지를 다루는 두 경로입니다. 양쪽 모두 최근 위치별 KV $p_{14}$·$p_{15}$·$p_{16}$을 함께 읽습니다. 바뀌는 것은 과거 위치를 묶는 크기와 요약을 선택하는 방식입니다.

왼쪽은 두 위치씩 묶어 요약 여덟 개를 만듭니다. 인덱서는 이 여덟 항목 중 Top-2인 $b_1$과 $b_5$를 선택합니다. $b_1$은 위치 2·3의 요약이고, $b_5$는 위치 10·11의 요약입니다. 선택 후에 원래 네 위치의 KV를 다시 펼쳐 읽는 것이 아니라, 선택된 요약 두 개 자체를 읽습니다. 최근 위치별 KV 세 개까지 합치면 이 예시의 읽기 항목은 다섯 개입니다.

오른쪽은 네 위치씩 묶어 요약 네 개를 만든 뒤 모두 읽습니다. 여기에 같은 최근 위치별 KV 세 개를 더하므로 읽기 항목은 일곱 개입니다. 더 크게 압축했는데도 이번 예시에서는 왼쪽보다 많은 항목을 읽습니다. 왼쪽에는 별도의 선택 단계가 있기 때문입니다. 압축으로 몇 개를 저장하는지와, 그중 몇 개를 읽는지는 구별해야 합니다.

왼쪽도 지금 읽는 요약 두 개만 보관하는 것은 아닙니다. 다음 Query가 다른 요약을 선택할 수 있으므로 여덟 요약을 후보로 유지합니다. 오른쪽은 네 요약을 저장하고 모두 읽습니다. 이 비교는 저장량과 읽는 양의 차이를 보여 주는 교육용 예시이며, 다섯 개를 읽는 경로가 일곱 개를 읽는 경로보다 반드시 빠르다는 성능 비교는 아닙니다. 왼쪽에는 인덱서의 점수 계산과 선택 비용도 있습니다.

DeepSeek-V4의 공식 Attention 구현은 압축된 항목의 인덱스와 최근 window의 위치를 결합합니다. 그림의 2개·4개 묶음은 실제 모델의 압축률을 재현한 수치가 아닙니다. 여기서 익힐 관계는 압축된 항목을 일부 선택하는 경로와 더 적은 요약을 모두 읽는 경로를 구별하는 것입니다.

묶음을 크게 만들수록 요약 수는 줄지만, 한 벡터에 더 많은 위치의 정보를 담아야 합니다. 이미 요약 과정에서 구별하기 어려워진 세부 정보를 인덱서가 복원해 주지는 않습니다. 따라서 압축 정도, 선택 개수, 최근 위치별 KV의 window 크기를 함께 정해야 합니다. 더 적은 항목을 저장하고 읽는 이점과, 필요한 정보를 얼마나 유지할 수 있는지 사이의 절충입니다.