Self-Attention: Q·K·V로 문맥을 만드는 계산

지난 글에서 Transformer의 배치도를 봤다. 인코더와 디코더가 층으로 쌓이고, 층마다 어텐션 블록과 피드포워드 블록이 들어간다는 그림이었다. 이 글은 그 그림에서 이름만 적혀 있던 칸 하나를 열어 본다. Self-Attention은 한 시퀀스 안의 모든 토큰 쌍을 비교해 각 토큰이 다른 토큰을 얼마나 참고할지를 숫자로 정하는 계산이다.

계산 전체가 식 한 줄이다.

$$\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V$$

짧지만 이 안에 결정이 셋 들어 있다. 왜 입력 하나에서 행렬을 셋이나 만드는가, 왜 내적을 $\sqrt{d_k}$ 로 나누는가, 그리고 왜 이 계산이 긴 문장에서 비싸지는가. 세 결정을 하나씩 풀면 BERT가 문맥을 어떻게 잡는지, GPT가 다음 토큰을 어떻게 고르는지, 그리고 수많은 변형 모델이 정확히 어느 자리를 손대고 있는지가 같은 그림 위에서 보인다.

순차 처리의 한계

장거리 의존

문장을 이해하는 데 단어를 순서대로 읽는 것만으로는 부족하다. 앞에서 등장한 주어가 한참 뒤의 동사와 연결되고, 멀리 떨어진 수식어가 특정 명사 하나를 설명한다. 「어제 서점에서 산 그 두꺼운 책은 아직 반도 못 읽었다」에서 「책」과 「못 읽었다」 사이에는 어절이 둘 끼어 있고, 「어제」와 「샀다」의 관계는 문장 전체를 건너뛴다. 읽는 사람은 이 연결을 순서와 무관하게 잡아내지만, 왼쪽에서 오른쪽으로 한 칸씩 밀고 가는 계산에는 그런 통로가 없다.

순환 신경망은 이 관계를 은닉 상태 하나에 눌러 담아 전달했다. 그러면 두 토큰이 멀수록 그 사이에 낀 스텝이 많아지고, 스텝마다 상태가 다시 쓰이면서 앞쪽 정보가 희석된다. 거리가 곧 손실이 되는 구조다. Self-Attention은 이 문제를 다르게 푼다. 정보를 계속 들고 가는 대신, 필요한 자리에서 필요한 토큰을 직접 꺼내 온다. 「못 읽었다」를 계산할 때 「책」을 한 번에 참조하므로 사이에 몇 개가 끼어 있든 상관이 없다.

Self와 크로스 어텐션

어텐션이라는 아이디어 자체는 Transformer보다 먼저 나왔다. Bahdanau 어텐션은 기계 번역에서 디코더가 출력 단어를 하나 만들 때마다 인코더의 어느 입력 단어를 볼지 정하는 장치였다. 여기서 비교되는 것은 서로 다른 두 시퀀스다 — 한쪽은 원문, 다른 쪽은 번역문이다. 이런 어텐션을 크로스 어텐션(cross-attention)이라 부른다.

Self-Attention은 이름이 말하는 그대로 비교 대상이 자기 자신이다. 하나의 시퀀스 안에서 각 토큰이 같은 시퀀스의 다른 모든 토큰과 얼마나 관련 있는지를 스스로 계산한다. 「나는 사과를 먹었다」를 처리한다면, 「먹었다」라는 토큰이 「나는」·「사과를」·「먹었다」 각각과 얼마나 관련 있는지 세 개의 숫자가 나온다. 자기 자신도 후보에 들어간다는 점이 중요하다 — 어떤 토큰은 문맥보다 자기 정보를 그대로 들고 가는 편이 나을 때가 있고, 그 선택도 같은 계산 안에서 이뤄진다.

두 어텐션은 계산식이 완전히 같다. 다른 것은 $Q$ 를 어디서 가져오고 $K$ 와 $V$ 를 어디서 가져오느냐뿐이다. 셋을 모두 같은 시퀀스에서 뽑으면 Self-Attention이고, $Q$ 만 다른 시퀀스에서 뽑으면 크로스 어텐션이다. Transformer의 디코더가 두 종류의 어텐션 블록을 나란히 갖고 있는 것도 이 차이 하나 때문이다.

병렬 계산

순차 처리에 덜 의존한다는 점이 부수 효과가 아니라 핵심이다. 순환 구조에서는 $h_t$ 를 계산하려면 $h_{t-1}$ 이 먼저 나와야 해서 시간 축을 GPU로 나눌 방법이 없었다. 길이 100인 시퀀스는 100번의 스텝을 순서대로 밟아야 했다.

Self-Attention에는 그런 의존이 없다. 모든 토큰 쌍의 점수를 한 번의 행렬 곱으로 동시에 만들고, 가중합도 한 번의 행렬 곱으로 끝낸다. 길이가 100이든 1,000이든 순서대로 밟아야 하는 단계는 그대로 둘이다. 뒤에서 볼 $O(n^2)$ 이라는 비용은 이 병렬성을 사고 치른 값이다 — 연산량 자체는 늘었지만, 그 연산이 전부 한 번에 흩어질 수 있는 종류가 되었다.

Q·K·V의 역할 분리

Query·Key·Value

각 토큰은 세 가지 표현으로 바뀐다. Query는 「지금 내가 찾고 싶은 것」, Key는 「내가 어떤 정보를 갖고 있는지 알리는 표지」, Value는 「실제로 전달할 내용」이다. 데이터베이스에 빗대면 Q는 검색어, K는 인덱스, V는 인덱스가 가리키는 데이터다. 검색어가 인덱스와 얼마나 맞는지로 점수를 매기고, 그 점수에 비례해 데이터를 가져온다.

굳이 셋으로 나누는 이유가 있다. 「무엇을 찾는가」와 「무엇을 갖고 있는가」와 「실제 내용이 무엇인가」는 서로 다른 것이기 때문이다. 「먹었다」라는 동사가 찾는 것은 목적어이고, 「사과를」이라는 명사가 내거는 표지는 「나는 목적어 자리에 있는 먹을 수 있는 것」이며, 정작 전달할 내용은 사과의 의미 표현이다. 셋을 하나로 쓰면, 곧 $Q = K = V = X$ 로 두면 어떻게 되는지도 계산해 보면 나온다. 모든 벡터의 길이가 같다고 하면 코시-슈바르츠 부등식에 따라 자기 자신과의 내적이 최댓값이므로, 점수 행렬의 대각선이 언제나 그 행에서 가장 큰 값이 된다. 문맥을 보라고 만든 계산에 「자기를 보라」는 편향이 박혀 버리는 셈이다.

$Q$ 와 $K$ 를 따로 두는 것에도 이유가 있다. 관계는 방향을 갖는다. 「그것」이 「고양이」를 가리키는 관계와 「고양이」가 「그것」을 가리키는 관계는 같은 세기가 아니다. 두 행렬을 나눠 두면 $q_i \cdot k_j$ 와 $q_j \cdot k_i$ 가 다른 값이 되어 이 비대칭이 표현된다. 하나로 묶으면 관계가 무조건 대칭이 되어, 「대명사가 명사를 가리킨다」와 그 반대를 구별할 수 없다.

선형 투영

입력은 토큰 임베딩을 쌓은 행렬 $X \in \mathbb{R}^{n \times d_{model}}$ 하나다. 여기에 학습되는 가중치 행렬 셋을 곱해 $Q$, $K$, $V$ 를 만든다.

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

$W_Q$ 와 $W_K$ 는 $\mathbb{R}^{d_{model} \times d_k}$, $W_V$ 는 $\mathbb{R}^{d_{model} \times d_v}$ 다. 원 논문의 기본 모델은 $d_{model} = 512$, 헤드마다 $d_k = d_v = 64$ 를 썼다. 곧 512차원 임베딩을 64차원 공간 셋으로 각각 눌러 보내는 셈이다.

여기서 눈여겨볼 것은 학습되는 것이 이 세 행렬뿐이라는 점이다. 토큰이 몇 개인지는 파라미터에 전혀 영향을 주지 않는다. 다섯 토큰짜리 문장이든 5,000 토큰짜리 문서든 같은 $W_Q$ 가 각 행에 똑같이 적용된다. 길이가 달라도 같은 가중치가 돈다는 이 성질이, 문장 길이를 미리 고정하지 않아도 되는 이유다.

계산 흐름

$Q$, $K$, $V$ 가 준비되면 나머지는 네 걸음이다.

1. $Q$ 와 $K$ 의 내적으로 관련도를 계산한다. $QK^\top$ 은 $n \times n$ 짜리 점수 행렬이 된다.
2. $\sqrt{d_k}$ 로 나눠 점수의 크기를 안정화한다. 다음 절의 주제다.
3. softmax로 가중치를 정규화한다. 행마다 합이 1인 확률 분포가 나온다.
4. 그 가중치로 $V$ 를 섞는다. 문맥을 반영한 새 표현이 여기서 나온다.

디코더에서는 2번과 3번 사이에 마스킹이 한 걸음 더 낀다. 아직 생성하지 않은 뒤쪽 토큰의 점수를 $-\infty$ 로 덮어 softmax가 그 자리에 0을 주도록 만드는 조작인데, 계산의 뼈대는 그대로다. 네 걸음 중 무게가 실린 곳은 1번과 2번이므로 그 둘을 차례로 본다.

Scaled Dot-Product Attention

내적

내적(dot product)은 두 벡터의 대응하는 원소끼리 곱해 모두 더한 값이다. 기하로 옮기면 $q \cdot k = \lVert q \rVert \lVert k \rVert \cos\theta$ 이므로, 두 벡터가 같은 방향을 가리킬수록 값이 크고 직교하면 0이다. 어텐션이 내적을 쓰는 이유가 여기 있다 — 「Query가 찾는 방향」과 「Key가 내건 방향」이 일치하는 정도를 숫자 하나로 재 준다.

이 값을 모든 쌍에 대해 구한 것이 $QK^\top$ 이다. $i$ 번째 행 $j$ 번째 열에는 토큰 $i$ 의 Query와 토큰 $j$ 의 Key를 내적한 값이 들어간다. 토큰이 여섯 개면 36개의 점수가 한 번의 행렬 곱으로 동시에 나온다. 반복문도 순서도 없다.

스케일링과 포화

문제는 이 점수의 크기다. $Q$ 와 $K$ 의 각 원소가 평균 0, 분산 1을 따른다고 하면, 독립인 $d_k$ 개의 곱을 더한 내적의 분산은 $d_k$ 가 된다. 표준편차는 $\sqrt{d_k}$ 다. $d_k = 64$ 라면 점수가 대략 $\pm 8$ 범위에서 흩어진다는 뜻이다.

그 값을 그대로 softmax에 넣으면 무슨 일이 생기는지 숫자로 보자. 점수가 $8, 0, -8$ 이면 softmax 결과는 대략 $0.9997$, $0.0003$, $0.0000001$ 이다. 사실상 첫 번째에 전부 쏠린 one-hot이 된다. 이 상태를 포화(saturation)라 부른다. 그리고 softmax 출력이 자기 입력에 대해 갖는 미분은 $p(1-p)$ 라, $p = 0.9997$ 인 자리의 기울기는 $0.0003$ 이다. 거의 0이고, 학습 신호가 이 블록을 통과하지 못한다.

같은 점수를 $\sqrt{64} = 8$ 로 나누면 $1, 0, -1$ 이 되고 softmax는 $0.665$, $0.245$, $0.090$ 이다. 최댓값 자리의 기울기는 $0.665 \times 0.335 \approx 0.223$ 으로, 스케일링 전보다 700배 넘게 크다. 순위는 하나도 바뀌지 않았다 — 여전히 첫 번째가 가장 크다. 바뀐 것은 차이의 폭뿐이고, 그 폭이 학습 가능 여부를 갈랐다.

나누는 값이 하필 $\sqrt{d_k}$ 인 것도 여기서 나온다. 분산이 $d_k$ 에 비례하니 표준편차는 그 제곱근에 비례하고, 표준편차로 나누면 분산이 1로 돌아간다. $d_k$ 로 나누면 지나치게 눌려 점수 차이가 사라지고, 나누지 않으면 위처럼 포화된다. 차원이 커질수록 이 조작의 무게도 커진다 — $d_k = 512$ 라면 표준편차가 22.6이라 포화가 훨씬 심하다.

구현 코드

앞의 네 걸음을 코드로 옮기면 다섯 줄이다.

import math
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)   # (B, n, n)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    alpha = F.softmax(scores, dim=-1)                   # 행마다 합이 1
    return alpha @ V, alpha

dim=-1이 마지막 걸림돌이다. softmax를 마지막 축에 걸어야 「토큰 $i$ 가 모든 $j$ 에 나눠 준 가중치의 합이 1」이 된다. 여기를 dim=-2로 잘못 두면 오류는 나지 않고 모양도 그대로인데, 정규화 방향이 뒤집혀 학습만 조용히 망가진다.

마지막 줄의 가중합을 숫자로 한 번 따라가 보자. 「나는 사과를 먹었다」 세 토큰의 Value가 각각 $[1, 0]$, $[0, 1]$, $[1, 1]$ 이고 「먹었다」의 어텐션 가중치가 $[0.2, 0.5, 0.3]$ 이면, 출력은 $0.2[1,0] + 0.5[0,1] + 0.3[1,1] = [0.5, 0.8]$ 이다. 「먹었다」 자리에 원래 있던 $[1,1]$ 이 문맥을 섞은 $[0.5, 0.8]$ 로 바뀌었다.

가중치의 합이 1이라는 점이 여기서 한 가지를 보장한다. 출력은 언제나 Value 벡터들의 볼록 조합이다 — 음이 아닌 가중치를 합이 1이 되게 걸어 섞은 값이라는 뜻이고, 그래서 결과가 Value 벡터들이 이루는 범위를 벗어나지 않는다. 어텐션은 새 값을 지어내는 것이 아니라 이미 있는 값들 사이에서 위치를 고르는 계산이다.

어텐션 행렬의 해석

행과 열의 뜻

softmax를 통과한 $\alpha \in \mathbb{R}^{n \times n}$ 이 어텐션 행렬이고, 이 글에서 유일하게 사람이 그대로 들여다볼 수 있는 중간 산출물이다. $\alpha[i][j]$ 는 토큰 $i$ 가 토큰 $j$ 를 얼마나 참고했는지를 뜻한다.

읽을 때 헷갈리기 쉬운 자리가 둘이다. 첫째, 합이 1인 것은 행뿐이다. 한 토큰이 여러 곳에 나눠 준 관심의 총량은 1로 고정되지만, 한 토큰이 남들에게서 받는 관심의 총량에는 제한이 없다. 문장의 첫 토큰처럼 여러 토큰이 함께 쳐다보는 자리는 열 합이 1을 훌쩍 넘고, 아무도 참고하지 않는 토큰은 열 합이 0에 가깝다. 둘째, 행렬이 대칭이 아니다. $\alpha[i][j]$ 와 $\alpha[j][i]$ 는 다른 값이고, 그래야 앞에서 말한 방향 있는 관계가 담긴다.

이 두 성질 때문에 어텐션 행렬은 관계를 「누가 누구를 본다」로 읽어야 한다. 대각선이 밝으면 각 토큰이 자기 정보를 지키고 있다는 뜻이고, 특정 열이 통째로 밝으면 그 토큰이 문장 전체의 정보를 모으는 자리라는 뜻이다.

헤드별 패턴

실제 모델에서 이 행렬을 꺼내 보는 것도 어렵지 않다. BERT-base는 12개 층에 층마다 12개의 헤드(어텐션을 독립적으로 한 번 수행하는 단위)를 두고 있으므로, 문장 하나를 넣으면 144개의 어텐션 행렬이 나온다.

from transformers import BertModel, BertTokenizer

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased", output_attentions=True)

inputs = tokenizer("The cat sat on the mat", return_tensors="pt")
attn = model(**inputs).attentions[0][0]   # 첫 층: (12헤드, n, n)

이 행렬들을 분석한 연구들은 헤드마다 성격이 다르다는 것을 보고했다. 어떤 헤드는 바로 앞이나 뒤 토큰에 몰리는 위치 패턴을 보이고, 어떤 헤드는 동사와 그 목적어를 잇고, 어떤 헤드는 대명사와 그것이 가리키는 명사를 잇는다 — 「it」이 「cat」을 향하는 식이다. 아무도 그렇게 하라고 지시하지 않았는데 학습만으로 갈라졌다는 점이 이 관찰의 요지다.

가중치의 해석 한계

다만 이 관찰을 너무 멀리 밀고 가지 않는 편이 좋다. 모든 헤드의 역할이 깔끔한 규칙으로 떨어지지는 않는다. 이름 붙일 만한 패턴을 보이는 것은 일부이고, 나머지 상당수는 사람이 읽을 만한 구조 없이 흩어져 있다. 구문 관계처럼 보이는 패턴도 문장을 바꾸면 흐트러지는 경우가 흔하다.

한 걸음 더 조심할 것은 어텐션 가중치를 곧 모델의 근거로 읽는 일이다. $\alpha[i][j]$ 가 크다는 것은 토큰 $j$ 의 Value가 많이 섞였다는 뜻이지, 모델이 그 토큰 때문에 그 답을 냈다는 뜻이 아니다. 층이 12개 쌓이면 아래층의 출력이 위층의 입력이 되므로 한 층의 가중치만으로는 최종 출력까지의 경로를 복원할 수 없고, 피드포워드 블록과 잔차 연결이 그 사이에서 값을 계속 바꾼다.

그래도 헤드를 여러 개 두는 효과 자체는 분명하다. 해석이 늘 명확하지 않을 뿐, 서로 다른 관계를 나눠 맡게 하면 표현력이 넓어진다는 것은 성능으로 확인되는 사실이다. 이 이야기는 마지막 절에서 다시 이어진다.

n² 비용

행렬 곱의 비용

Self-Attention의 비용은 행렬 곱 두 번에서 나온다.

연산	시간 복잡도	결과 크기
$QK^\top$	$O(n^2 d_k)$	$n \times n$
$\alpha V$	$O(n^2 d_v)$	$n \times d_v$
저장해야 하는 어텐션 행렬	—	$O(n^2)$
최대 경로 길이	$O(1)$	—

$n$ 이 제곱으로 들어가는 것이 전부를 정한다. 참고로 순환 층의 복잡도는 $O(n d^2)$ 이라 $n$ 에 선형이다. 두 값을 비교하면 $n$ 이 $d$ 보다 작을 때는 오히려 Self-Attention이 층당 연산이 적다. $d_{model} = 512$ 인 모델에서 512 토큰짜리 문장은 두 방식의 연산량이 얼추 같고, 그보다 짧으면 어텐션 쪽이 싸다. 흔히 「어텐션은 비싸다」고 말하지만 정확히는 긴 시퀀스에서만 비싸다.

어텐션 행렬의 메모리

문제가 되는 자리를 숫자로 보자. $n = 8192$ 이면 어텐션 행렬 하나가 $8192^2 = 67{,}108{,}864$ 개의 원소다. 16비트 부동소수점으로 저장하면 한 행렬에 약 134MB이고, 헤드가 12개인 층 하나만 해도 1.6GB가 된다. 층이 12개면 이 값이 다시 12배로 늘어난다. 학습할 때는 역전파를 위해 이 행렬을 붙들고 있어야 하므로 부담이 그대로 메모리 한계가 된다.

증가 속도가 더 나쁜 소식이다. 시퀀스 길이를 두 배로 늘리면 시간과 메모리가 네 배로 뛴다. 4,096에서 8,192로 가는 데 드는 추가 비용이 0에서 4,096까지 오는 데 든 비용의 세 배다. 컨텍스트 길이를 늘리는 일이 그토록 어려운 이유가 이 제곱에 있다.

여기를 공략하는 방법이 크게 둘이다. 하나는 어텐션 행렬을 통째로 만들지 않는 것이다. FlashAttention은 행렬을 작은 타일로 쪼개 GPU의 빠른 메모리 안에서 계산하고 중간 결과를 버려, 메모리 사용을 $O(n)$ 으로 낮춘다. 계산하는 값 자체는 똑같아서 결과가 바뀌지 않는다. 다른 하나는 비교하는 쌍의 수를 줄이는 것으로, 각 토큰이 가까운 이웃이나 정해진 몇 자리만 보게 만드는 효율적 어텐션 계열이 여기 속한다. 이쪽은 계산 결과가 달라지므로 정확도와 맞바꾸는 선택이 된다.

최대 경로 길이

$O(n^2)$ 을 치르고 무엇을 얻는지도 같은 표에 적혀 있다. 최대 경로 길이는 시퀀스의 두 토큰 사이에 정보가 흐르려면 계산 단계를 몇 번 거쳐야 하는지를 재는 값이다. 순환 신경망에서는 이 값이 $O(n)$ 이다 — 1번 토큰의 정보가 100번 토큰에 닿으려면 99번의 갱신을 통과해야 한다. Self-Attention에서는 $O(1)$ 이다. 어떤 두 토큰이든 한 층 안에서 직접 이어진다.

이 차이가 학습에서 그대로 드러난다. 신호가 통과해야 하는 곱셈의 사슬이 짧으면 기울기가 사라질 자리도 그만큼 줄어든다. 순환 구조가 기울기 소실과 싸우며 게이트를 발명해야 했던 것과 대비된다. 앞 절에서 「거리가 곧 손실」이라 했던 성질이, 여기서는 아예 사라진다.

정리하면 거래는 이렇다. 순환은 시간을 선형으로 쓰는 대신 먼 관계를 잃고, 어텐션은 제곱의 비용을 내는 대신 거리를 없앤다. 하드웨어가 병렬 연산을 싸게 만들어 준 시점에 뒤쪽 거래가 이겼다.

단일 헤드의 한계

관계의 혼재

여기까지가 Self-Attention 한 벌의 전부다. 그런데 이 한 벌로 문장 안의 모든 관계를 담으려 하면 걸리는 지점이 있다.

softmax를 통과한 각 행은 확률 분포 하나다. 합이 1이라는 것은 관심을 나눠 써야 한다는 뜻이고, 한 곳에 쏠리면 다른 곳에는 남는 것이 없다. 「먹었다」라는 토큰이 동시에 잡아야 할 것이 여럿이라고 해 보자 — 주어가 누구인지, 목적어가 무엇인지, 시제를 정하는 어미가 어디 붙었는지. 분포 하나로 이 셋을 표현하려면 셋에 골고루 나눠 주는 수밖에 없고, 그러면 어느 관계도 뚜렷하지 않은 평평한 가중치가 나온다.

투영 공간이 하나라는 것도 같은 제약의 다른 얼굴이다. $W_Q$ 와 $W_K$ 가 한 벌뿐이면 「무엇이 비슷한가」를 재는 기준도 하나다. 문법적 인접성으로 재는 기준과 의미적 유사성으로 재는 기준이 다른데, 둘을 한 공간에 욱여넣으면 서로를 방해한다.

멀티헤드 어텐션

답은 나눠서 여러 번 하는 것이다. 같은 문장을 서로 다른 투영 공간에서 여러 벌 바라보고 결과를 합치면, 한 벌이 하나의 관계에 집중할 수 있다. 어떤 벌은 가까운 문법 관계를 잡고, 다른 벌은 멀리 떨어진 의미 관계를 잡는다. 앞에서 이미 「헤드」라는 이름으로 여러 번 등장한 그 단위다.

앞에서 그냥 지나간 숫자 하나가 여기서 설명된다. $d_{model} = 512$ 인데 $d_k = 64$ 였던 이유는 $512 \div 8 = 64$ 이기 때문이다. 512차원을 한 벌이 통째로 쓰는 대신 여덟 벌이 64차원씩 나눠 쓰는 배분이고, 그래서 벌을 여덟으로 늘려도 전체 연산량은 거의 그대로다. 관점을 늘리는 값을 차원을 쪼개어 치른 셈이다.

다음 글에서는 이 나눔을 정확히 어떻게 하는지 본다. 몇 벌로 쪼개야 하는지, 나눠 계산한 결과를 어떤 순서로 다시 이어 붙이는지, 그리고 차원을 쪼갠 만큼 각 벌이 얇아지는데도 성능이 오르는 이유가 무엇인지가 그 글의 내용이다.

---

읽어주셔서 감사합니다. 😊