Core Attention: 토큰 사이의 정보 조합하기

지난 글에서는 같은 입력에서 Q·K·V를 만들고, 각 벡터의 성분을 여러 head로 구분하는 방식을 살펴봤습니다. Core Attention의 결과는 출력 projection을 거쳐 Residual stream에 더해졌습니다. 이번 글에서는 내부 계산을 생략했던 Core Attention을 살펴보겠습니다. Q와 K로 토큰 사이의 점수를 계산하고, 그 점수를 가중치로 바꾼 뒤, V를 가중합하는 과정입니다.

Core Attention은 head마다 독립적으로 계산합니다

지난 글과 같은 MHA 구성을 사용하겠습니다. 토큰 수 T = 3, 모델 차원 d = 8, head 수 h = 2이므로 head 하나의 차원 dh는 4입니다. Q·K·V는 각각 3 × 8 형태이며, 각 토큰의 앞 네 성분은 Head 1, 뒤 네 성분은 Head 2에 해당합니다. 그림 1에서는 이 성분들을 head별로 나누어 Core Attention에 전달하는 흐름을 표시했습니다.

같은 번호의 Q·K·V가 한 묶음으로 계산에 참여합니다. Head 1에서는 Head 1의 Q·K·V로 결과 O1을 만들고, Head 2에서는 Head 2의 Q·K·V로 결과 O2를 만듭니다. 각 head는 세 토큰 모두에 대해 계산하지만, 서로 다른 head의 Q와 K를 곱하지는 않습니다. 계산 순서는 같아도 사용하는 Q·K·V가 다르므로, 토큰별로 정보를 반영하는 비중과 결과도 달라질 수 있습니다.

O1과 O2는 각각 T × dh, 즉 3 × 4 형태입니다. 두 결과를 같은 토큰 위치끼리 이어 붙이면(Concat) T × d, 즉 3 × 8 형태의 H가 됩니다. 이어 붙이기는 각 head의 값을 그대로 나란히 놓는 단계이며, 이후 출력 projection에서 여러 head의 성분을 조합합니다. 그림 2~5에서는 Head 1의 내부 계산만 확대하겠습니다. Head 2도 자신의 Q·K·V를 사용해 같은 순서로 독립적으로 계산합니다. 이후의 수치는 실제 모델에서 가져온 값이 아닌 설명용 예시입니다.

QKᵀ: 토큰 사이의 점수 계산하기

첫 단계는 Q의 각 토큰과 K의 각 토큰을 짝지어 점수를 계산하는 것입니다. Head 하나에서 Q와 K는 모두 T × dh 형태입니다. K의 행과 열을 바꾼 Kᵀ는 dh × T가 되므로, S = Q × Kᵀ의 결과는 T × T입니다. 이 예시에서는 3 × 4와 4 × 3을 곱해 3 × 3 점수 행렬 S를 만듭니다.

점수 하나는 Query 벡터와 Key 벡터의 같은 위치 성분을 곱한 뒤 더하는 내적입니다. 그림 2에서 토큰 3의 Query는 [−1, 1, 1, 2], 토큰 1의 Key는 [1, −1, 2, 1]입니다. 따라서 두 토큰 사이의 점수는 −1 × 1 + 1 × (−1) + 1 × 2 + 2 × 1 = 2입니다. 이 값은 S의 세 번째 행, 첫 번째 열에 들어갑니다. 행은 정보를 모을 Query 토큰, 열은 참조할 Key 토큰에 해당합니다. S의 세 번째 행 [2, 0, 6]은 토큰 3의 Query를 토큰 1·2·3의 Key와 각각 비교한 결과입니다.

이어서 S의 모든 값을 √dh로 나눕니다. 내적은 dh개 성분의 곱을 더하므로, 차원이 커질수록 점수의 크기가 커지기 쉽습니다. 점수 차이가 지나치게 커지면 뒤의 Softmax에서 가중치가 한쪽으로 몰릴 수 있어, head 차원에 맞춰 크기를 조절합니다. 여기서는 dh = 4이므로 모든 점수를 2로 나눕니다. 강조한 점수 2는 1이 되고, 세 번째 행 [2, 0, 6]은 [1, 0, 3]이 됩니다.

마스크: 참조할 수 있는 토큰 구분하기

지금까지는 각 Query와 모든 Key의 점수를 계산했습니다. 하지만 다음 토큰을 예측하는 디코더에서는 현재 위치가 자기 자신과 이전 위치만 참조하도록 제한합니다. 각 위치의 출력으로 다음 토큰을 예측할 때, 뒤에 있는 토큰의 정보가 미리 들어가면 안 되기 때문입니다. 이 참조 범위를 정하는 것이 Causal mask입니다.

그림 3에서 토큰 1은 토큰 1만, 토큰 2는 토큰 1과 2를, 토큰 3은 토큰 1·2·3을 참조합니다. 점수 행렬로 보면 대각선은 자기 자신이고, 그 왼쪽은 이전 토큰, 오른쪽은 미래 토큰입니다. 따라서 대각선 위쪽의 점수를 −∞로 바꾸고 나머지는 유지합니다. 두 번째 행 [3, 1.5, −1]은 [3, 1.5, −∞]가 되고, 세 번째 행 [1, 0, 3]은 그대로입니다.

이를 수식으로는 스케일링한 점수에 마스크 M을 더하는 것으로 나타낼 수 있습니다. M은 참조 가능한 위치에서 0, 제외할 위치에서 −∞인 행렬입니다. 0을 더한 점수는 그대로이고, −∞를 더한 점수는 −∞가 됩니다. 그림의 “마스크 M의 값”은 이때 더하는 값을 뜻합니다.

Softmax: 점수를 가중치로 바꾸기

이제 토큰 사이의 점수를 각 Value에 곱할 가중치로 바꾸겠습니다. Softmax는 한 행의 점수를 0 이상의 가중치로 바꾸고, 그 합이 1이 되도록 합니다. 계산은 P = softmax(S / √dh + M)이며, Query 토큰 하나에 해당하는 한 행 안에서 이루어집니다. P의 각 성분은 해당 열에 있는 토큰의 Value를 얼마나 반영할지 나타냅니다.

Softmax는 각 점수에 지수 함수 exp를 적용한 뒤, 그 행에서 얻은 값들의 합으로 나눕니다. 그림 4의 토큰 3을 보면 점수는 [1, 0, 3]입니다. 각각 [e¹, e⁰, e³]으로 바꾸고 e¹ + e⁰ + e³으로 나누면, 가중치는 약 [0.114, 0.042, 0.844]가 됩니다. 같은 행 안에서는 점수가 클수록 더 큰 가중치를 받습니다. 이 예시에서 토큰 3은 자신의 Value에 가장 큰 비중을 두고, 토큰 1과 2의 Value도 함께 반영합니다.

마스크로 설정한 −∞는 지수 함수를 적용하면 0이 되므로, 해당 위치의 가중치도 0입니다. 토큰 1은 참조 가능한 위치가 하나뿐이라 [1, 0, 0]이 되고, 토큰 2는 약 [0.818, 0.182, 0]이 됩니다. 점수가 0인 토큰도 가중합에 참여합니다. 토큰 3의 두 번째 점수는 0이지만 e⁰ = 1이므로 약 0.042의 가중치를 받습니다. 그림의 P는 소수점 셋째 자리까지 반올림한 값입니다.

PV: 여러 토큰의 Value를 가중합하기

마지막으로 가중치 P를 사용해 Value 벡터들을 합칩니다. O = P × V이며, P는 T × T, V는 T × dh이므로 결과 O는 T × dh입니다. 이 예시에서는 3 × 3과 3 × 4를 곱해 3 × 4 출력을 만듭니다. P의 한 행은 한 Query 토큰이 사용할 가중치이고, V의 각 행은 해당 토큰의 Value 벡터입니다.

그림 5에서 토큰 3의 출력은 토큰 1의 Value × 0.114 + 토큰 2의 Value × 0.042 + 토큰 3의 Value × 0.844로 볼 수 있습니다. 가중치 하나는 해당 Value 벡터의 네 성분 모두에 곱해집니다. 세 결과를 성분별로 더하면 약 [0.114, 0.042, 0.844, 1]이 됩니다. 마지막 성분이 1인 이유는 세 Value의 마지막 성분이 모두 1이고, 가중치의 합도 1이기 때문입니다.

다른 토큰도 같은 방식입니다. 토큰 1은 가중치가 [1, 0, 0]이므로 자신의 Value가 그대로 출력됩니다. 토큰 2는 토큰 1과 2의 Value를 가중합하고, 토큰 3은 세 토큰의 Value를 가중합합니다. 여러 토큰의 정보를 합치더라도 출력 벡터의 차원은 dh로 유지됩니다. 이 예시에서는 각 토큰마다 네 성분이 나오므로, Head 1의 전체 출력 O는 입력 V와 같은 3 × 4 형태입니다.

Core Attention에서는 Q와 K로 토큰 사이의 점수를 계산하고, 마스크로 참조 범위를 정한 뒤, Softmax로 얻은 가중치에 따라 V를 모읍니다. 어떤 토큰을 얼마나 반영할지는 Q와 K로 계산하고, 실제로 가중합하는 정보는 V에서 가져옵니다. 그림 1의 전체 흐름으로 돌아가면, 각 head에서 이렇게 계산한 결과를 같은 토큰 위치끼리 이어 붙인 뒤 출력 projection으로 전달합니다.