본문 바로가기

AI 모델 파헤치기

Transformer Self-Attention 수식 완전 정리 — Query Key Value 직관적 이해

Transformer 파헤치기 (2) — Softmax와 Self-Attention: Transformer의 심장

https://thisisgonnabegreat.tistory.com/entry/%EB%AA%A8%EB%8D%B8-%ED%8C%8C%ED%97%A4%EC%B9%98%EA%B8%B0-Transformer%ED%8E%B8

 

Word2Vec 임베딩 수식 완전 정리 — Skip-Gram부터 네거티브 샘플링까지

Transformer 파헤치기 (1) — 단어는 어떻게 벡터가 되고, 모델은 어떻게 배우는가GPT나 Claude 같은 모델이 텍스트를 처리하려면, 가장 먼저 "단어"라는 추상적인 기호를 컴퓨터가 계산할 수 있는 숫자

thisisgonnabegreat.tistory.com

 

1편에서 우리는 단어가 어떻게 고차원 벡터로 변환되는지, 그리고 그 벡터들이 어떻게 학습되는지를 다뤘다. 그런데 거기서 한 가지 중요한 질문을 남겨두고 왔다. "나는 오늘 학교에서"라는 문장이 존재할 때, 임베딩 직후의 "학교에서" 벡터는 그 단어 자체의 의미만 담고 있을 뿐, "나는 오늘"이라는 앞 문맥을 알 리가 없다. 그렇다면 문장 안의 단어들은 어떻게 서로의 의미를 반영하게 되는 걸까?

이 질문에 대한 답이 바로 이 글에서 다룰 Self-Attention이다. 그 전에, Self-Attention의 핵심 구성 요소인 소프트맥스 함수부터 정확히 짚고 넘어가자.


1. 소프트맥스 함수 (Softmax Function)

소프트맥스는 어떤 숫자들의 나열이든 확률 분포의 형태로 바꾸어주는 매우 효율적인 함수다. 1편에서 Skip-Gram의 출력에서도 등장했지만, Self-Attention에서 핵심 역할을 하기 때문에 좀 더 깊이 이해할 필요가 있다.

정의는 다음과 같다.

 

각 입력값을 자연상수 e의 지수 형태로 변환한 뒤, 전체 합으로 나누는 구조다. 이렇게 하면 두 가지 성질이 보장된다.

  • 모든 출력값이 0∼1 사이에 놓인다.
  • 출력값의 합이 정확히 1이 된다.

즉 어떤 점수 벡터가 들어오든 소프트맥스를 거치면 확률분포가 된다. 입력값이 클수록 대응하는 출력 확률도 커진다.

소프트맥스 함수가 점수 벡터를 확률 분포로 변환하는 모습

1.1 Temperature $T$

소프트맥스는 입력값이 클수록 출력 확률을 크게 만든다. 그런데 이 "크게 만드는" 정도가 선형적이지 않고 지수(exponential)적이다. 소프트맥스는 각 입력값을 $e^{x_n}$ 형태로 변환하기 때문에, 입력값 사이의 차이가 조금만 벌어져도 출력 확률의 차이는 기하급수적으로 커진다.

예를 들어 $x_1 = 10$ 과 $x_2 = 9$ 의 차이는 $x_3 = 2$ 와 $x_4 = 1$ 의 차이와 같이 1이지만, 소프트맥스를 거치면 $e^{10} - e^{9} \approx 12{,}000$인 반면 $e^{2} - e^{1} \approx 4.7$로 두 경우의 출력 차이는 약 2,500배가 넘는다. 결과적으로 점수가 가장 높은 값 하나가 확률을 거의 독점하게 되고, 나머지 값들의 확률은 0에 가까워진다.

 

 

이 특성은 어떤 방면에서는 단점이 된다. 여러 단어에 고르게 주목해야 할 상황에서도 소프트맥스가 가장 높은 점수의 단어에만 집중해버리는 것이다. 이를 조절하기 위해 등장한 파라미터가 temperature $T$다.

 

$\text{softmax}\left(\frac{x}{T}\right)_i=\frac{e^{x_i/T}}{\sum_{n}e^{x_n/T}}$

 

- $T$가 높을 때: 모든 점수를 $T$로 나누면 값들 사이의 차이가 줄어들어, 지수 함수의 증폭 효과가 약해진다. 결과적으로 출력 확률 분포가 균일해지고, 어느 단어에도 골고루 주목하게 된다.

 

- $T$가 낮을 때: 반대로 값들 사이의 차이가 상대적으로 커져 지수 함수의 증폭 효과가 강해진다. 가장 높은 점수를 가진 단어에 집중하는 첨예한 분포가 된다.

T가 높게 설정된 경우

 

T가 낮게 설정된 경우

 

이 temperature 개념은 언어 모델의 생성 단계에서도 그대로 쓰인다. gemini와 같은 대부분의 LLM은 API를 통해 temperature 파라미터를 직접 설정할 수 있도록 지원하는데, 값을 높이면 출력이 다양하고 창의적으로, 낮추면 일관적이고 결정론적으로 변하는 이유가 바로 여기 있다.

 

2. Self-Attention — 단어들이 서로를 참조하는 메커니즘

2.1 문제 설정: 임베딩만으로는 부족하다

1편에서 만들어진 임베딩 벡터는 단어의 평균적인 의미만 담고 있다. 예를 들어 한국어 "배"라는 단어는 선박, 복부, 과일 등 전혀 다른 의미로 쓰인다. 그런데 임베딩은 학습 데이터에서 "배"가 등장하는 모든 맥락을 평균낸 하나의 벡터이므로, "배를 타고 떠났다"의 "배"와 "배가 고프다"의 "배"가 임베딩 공간에서 같은 점으로 표현된다. 문맥에 따라 의미가 확정되지 않는 것이다.

 

Self-Attention은 이 문제를 해결한다. 각 단어의 임베딩 벡터가 문장 안의 다른 단어들을 참조해서 자신의 표현을 업데이트하도록 만드는 것이다. "배를 타고 떠났다"에서 "타고"라는 단어가 주변에 있다면, "배"의 임베딩은 선박 쪽 의미로 업데이트되고, "배가 고프다"에서 "고프다"가 주변에 있다면 복부 쪽 의미로 업데이트된다.

 

각 단어의 초기 임베딩은 단어의 의미와 위치 정보를 담고 있는 고차원 벡터다(위치 정보를 포함한다는 뜻에서 positional embedding이라고도 한다).

positional embedding

 

2.2 Query: "나는 어떤 단어에 주목해야 하는가?"

 

Self-Attention에는 세 가지 핵심 행렬이 등장한다 — Query, Key, Value. 하나씩 이해해보자.

 

Query 행렬은 각 단어가 던지는 질문이다.

 

구체적으로, 가중치 행렬 $W_Q$에 각 단어의 임베딩 벡터를 곱해서 Query 벡터를 구한다. $W_Q$는 모든 토큰에 동일하게 적용되는 하나의 학습 파라미터로, 한 토큰에 대해서는 다음과 같이 표현된다.

 

$q_i = W_Q \cdot e_i, \quad W_Q \in \mathbb{R}^{d_k \times d_{\text{model}}}$

 

전체 문장의 토큰을 한꺼번에 처리할 때는 임베딩 행렬 $E \in \mathbb{R}^{d_{\text{model}} \times n}$에 곱해 Query 행렬 $Q \in \mathbb{R}^{d_k \times n}$을 구한다. $Q = W_Q \cdot E$ 여기서 $d_{\text{model}}$은 임베딩 공간의 차원(GPT-3 기준 12288), $d_k$는 Query/Key 공간의 차원(GPT-3 기준 128), $n$은 토큰 수다. $W_Q$가 임베딩 공간($d_{\text{model}}$차원)에서 Query/Key 공간($d_k$차원)으로 차원을 줄이는 선형 변환이기 때문에, 각 토큰의 Query 벡터는 임베딩보다 훨씬 작은 128차원 공간에 표현된다.

Query 행렬의 역할
W_Q 행렬은 특정 토큰의 임베딩 벡터를 Query/Key space로 매핑시켜주는 역할을 한다

 

2.3 Key: "나는 어떤 정보를 제공할 수 있는가?"

Key 행렬은 Query에 대한 답변이라고 생각하면 편하다.

 

$K = W_K \cdot E, \quad W_K \in \mathbb{R}^{d_k \times d_{\text{model}}}$ 

 

$W_K$ 역시 $W_Q$와 마찬가지로 모든 토큰에 동일하게 적용되는 학습 파라미터다. Query 벡터와 Key 벡터의 내적값이 클수록 두 단어 사이의 연관성이 높다고 판단한다.

W_K 행렬은 W_Q와 마찬가지로 특정 토큰의 임베딩 벡터를 Query/Key space로 매핑시켜주는 역할을 한다.

 

이쯤에서 한 가지 자연스러운 의문이 생긴다. $W_Q$와 $W_K$는 그냥 행렬인데, 왜 의미적으로 연관된 단어 쌍에서 내적값이 크게 나오는 걸까?

결론부터 말하면 처음부터 그런 게 아니라, 학습을 통해 그렇게 되도록 만들어지는 것이다.

$W_Q$와 $W_K$는 학습 전에는 랜덤한 값이다. 내적을 수식으로 풀어쓰면 다음과 같다.

$$
q_i \cdot k_j = (W_Q e_i) \cdot (W_K e_j) = e_i^\top W_Q^\top W_K e_j
$$
이를 보면 $W_Q^\top W_K$가 사실상 임베딩 $e_i$와 $e_j$ 사이의 학습된 유사도 행렬 역할을 한다는 것을 알 수 있다. 즉 "어떤 두 단어가 서로 연관이 높다고 판단할 것인가"의 기준 자체가 학습되는 것이다.

예를 들어 "배를 타고"에서 "배"와 "타고"는 수백억 건의 텍스트에서 함께 등장하는 패턴이 반복적으로 학습된다. 모델은 역전파를 반복적으로 진행하는 과정에서 이런 쌍에 대해 내적값이 크게 나오도록 $W_Q$와 $W_K$를 조정한다. 반면 "배가 고프다"에서 "배"와 "타고"는 함께 등장하지 않으므로, 내적값이 낮게 유지되는 방향으로 학습된다.

정리하면, $W_Q$는 "이 단어와 연관이 있는 단어가 존재하는가"를, $W_K$는 "$W_Q$에서 말한 단어와 연관된 단어의 유무"를 각각 표현하는 공간으로 매핑하는데, 이 두 공간이 서로 잘 맞게 설계되도록 학습되는 것이 핵심이다. "질문-답변"이라는 비유는 이 학습된 결과를 직관적으로 표현한 것이다.

 

2.4 Query·Key 내적: 연관성 측정

이제 각 단어의 Query 벡터와 모든 다른 단어의 Key 벡터를 내적해서 얼마나 잘 맞는지를 측정한다.

 

$\vec{Q}_i \cdot \vec{K}_j = \langle Q_i, K_j \rangle$

 

예를 들어 "The fluffy blue creature sat on the mat"라는 문장에서, 학습을 통해 $W_Q$와 $W_K$는 형용사-명사 수식 관계에 있는 단어 쌍의 내적값이 크게 나오도록 조정된다. 그 결과 "creature"의 Query 벡터와 이를 수식하고 있는 "fluffy", "blue"의 Key 벡터를 내적하면 큰 양수값이 나오고, 구조적 단어인 "the"의 Key 벡터와는 음수값이 나온다. 이렇게 "fluffy"와 "blue"가 "creature"라는 명사에 주목(attend to)하고 있다고 말한다. 이것이 "Attention"이라는 이름의 유래다.

 

Query와 Key 벡터 간의 내적을 통한 단어 관계 파악
세부 사진

2.5 Value와 가중합: 임베딩을 업데이트하기

단어 쌍마다 내적값을 구했으면, 이제 이 값들로 임베딩을 실제로 업데이트해야 한다. 여기서 Value 벡터 $\vec{V}_j$가 등장한다. 각 단어 $j$의 Value 벡터에 해당 내적값(가중치)을 곱하고 합산하면, 단어 $i$의 임베딩에 더해줄 업데이트 벡터 $\Delta\vec{E}_i$를 구할 수 있다.

 

$\Delta \vec{E}_i = \sum_j w_{ij} \cdot \vec{V}_j$

 

여기서 가중치 $w_{ij}$는 $\vec{Q}_i \cdot \vec{K}_j$의 값이다. 그런데 이 가중치들의 합이 1이 되도록 정규화하면 계산이 훨씬 수월해진다. 이 정규화를 담당하는 함수가 바로 앞서 배운 소프트맥스다. $w_{ij} = \text{softmax}\left(\vec{Q}_i \cdot \vec{K}_j\right)$ 소프트맥스를 거쳐 나온 이 가중치 행렬 — 어떤 단어가 어떤 단어에 얼마나 주목(attention)하는지를 나타내는 표 — 을 Attention Pattern(Map)이라고 부른다.

Attention Pattern(Map)

 

마지막으로 $\Delta\vec{E}_i$를 원래 단어의 임베딩 벡터 $\vec{E}_i$에 더해주면, "fluffy blue creature"를 표현하는 새로운 임베딩 벡터 $\vec{E}'_i$를 얻을 수 있다.

 

$ \vec{E}'_i = \vec{E}_i + \Delta\vec{E}_i $

 

이처럼 원래 값에 변화량을 덧셈으로 더하는 구조를 Residual Connection(잔차 연결)이라고 부른다. ResNet(2015)에서 제안된 기법으로, 그래디언트가 역전파될 때 덧셈 경로를 통해 깊은 레이어까지 온전히 전달되기 때문에 학습이 훨씬 안정적으로 이루어진다. (후에 ResNet과 관련된 글도 작성하도록 하겠다) Transformer는 이 Residual Connection을 매 Attention 블록과 MLP 블록마다 반복적으로 사용한다. 

잔차 연결의 예시

 

3. 최종 수식 : 

지금까지 직관적으로 설명한 과정을 행렬 연산으로 압축하면 다음 한 줄이 된다.

 

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

 

각 항이 무엇을 의미하는지 정리하면:

  • $Q \in \mathbb{R}^{d_k \times n}$ : 문장 내 모든 토큰의 Query 벡터를 열로 쌓은 행렬 ($n$은 토큰 수, $d_k$는 Query/Key 공간의 차원)
  • $K \in \mathbb{R}^{d_k \times n}$: 모든 토큰의 Key 벡터를 열로 쌓은 행렬
  • $V \in \mathbb{R}^{d_v \times n}$: 모든 토큰의 Value 벡터를 열로 쌓은 행렬
  • $K^\top Q \in \mathbb{R}^{n \times n}$: 모든 토큰 쌍의 Query-Key 내적을 담은 행렬. $(i, j)$ 성분이 $\vec{Q}_j \cdot \vec{K}_i$에 해당한다.
  • $\sqrt{d_k}$로 나누는 이유: $d_k$가 커질수록 내적값의 분산이 커져서 소프트맥스 출력이 극단적으로 치우치는 문제가 생긴다. $\sqrt{d_k}$로 스케일링하면 이를 안정화할 수 있다.
  • 마지막 $V$를 곱하면: 소프트맥스로 정규화된 가중치를 Value 벡터들과 결합해서 최종 $\Delta E$ 행렬을 얻는다.

Attention(Q,K,V) 최종 공식

 

4. 마스킹 (masking)

Self-Attention에는 중요한 제약이 하나 있다. 언어 모델은 기본적으로 앞의 단어들을 보고 다음 단어를 예측하는 구조다. 그런데 Attention을 그냥 계산하면, 예측하려는 단어 이후에 나오는 단어들도 참조할 수 있게 된다. 정답을 미리 보는 것과 다름없다. 이를 막기 위해 마스킹(masking)을 적용한다. 미래 위치에 해당하는 Query-Key 내적값을 소프트맥스 이전에 $-\infty$로 설정하는 것이다.

 

$w_{ij} = \text{softmax}\!\left(\frac{\vec{Q}_i \cdot \vec{K}_j}{\sqrt{d_k}} + M_{ij}\right)$

 

여기서 $M_{ij}$는 마스킹 행렬로, $j > i$인 경우(미래 토큰) $-\infty$, 아닌 경우 $0$이다. $-\infty$를 소프트맥스에 통과시키면 $e^{-\infty} = 0$이 되어 해당 위치의 가중치가 정확히 0이 된다 — 즉 미래 단어가 현재 단어에 영향을 미칠 수 없도록 해버린다.

 

Attention에서는 앞의 문장을 토대로 다음에 올 단어를 예측하는 것이 메커니즘 → 이렇게 하기 위해서는 미래의 단어들이 이전의 단어에 영향을 미치지 않도록 해야 함. 이전의 단어에 영향을 줄 수 있게 되면 다음에 무슨 단어가 올 지 정답을 알고 있는 것과 다름 없기 때문
마스킹 (masking)

4.1 마스킹의 두 종류: Causal vs Padding

마스킹에는 두 가지 유형이 있다. 어떤 마스킹을 쓰느냐에 따라 모델의 성격이 달라진다.

 

Causal Masking(인과 마스킹) : 위에서 설명한 방식으로, 현재 위치 이후의 토큰들을 모두 $-\infty$로 막는다. "앞만 보고 다음을 예측"하는 구조이므로 텍스트 생성에 적합하다. GPT 계열 모델이 이 방식을 쓴다.

 

Padding Masking : 배치 학습 시에는 길이가 다른 문장들을 같은 크기로 맞추기 위해 짧은 문장 끝에 패딩(PAD) 토큰을 붙이는데, 이 패딩 위치가 Attention 계산에 영향을 주지 않도록 막는 것이 Padding Masking이다.

 

BERT 같은 인코더 모델은 Causal Masking을 쓰지 않는다. 문장 전체를 양방향으로 다 볼 수 있도록 설계되어 있기 때문이다. 대신 학습 시 일부 토큰을 [MASK]로 가려두고 맞히는 방식(Masked Language Modeling)으로 문맥을 학습한다. 이 설계 차이가 GPT(생성 특화)와 BERT(이해 특화)의 근본적인 차이를 만들어낸다. 참고로 Attention 연산의 계산 복잡도는 시퀀스 길이의 제곱 $O(n^2)$에 비례한다. 긴 문서를 처리할 때 병목이 되는 이유가 여기 있고, 이를 개선하기 위한 다양한 효율적 Attention 연구들(Sparse Attention, Flash Attention 등)이 진행되어 있다.

 

5. Multi-Head Attention

 

지금까지 설명한 것은 Attention 한 번의 계산, 즉 One head of Attention이다. 그런데 실제 Transformer에서는 이 과정을 여러 번 병렬로 수행한다. 이것이 Multi-Head Attention이다.

 

One head of attention

 

왜 여러 개의 head를 쓸까? 문장 안에서 단어들 사이의 관계는 하나가 아니기 때문이다. 예를 들어 "그 선수는 마이클이다. 그는 농구를 한다"라는 문장에는 다음과 같은 여러 유형의 관계가 동시에 존재한다.

 

- 지시 관계: "그"가 "마이클"을 가리킨다.

- 사실 관계: "마이클 조던"이 "농구"와 연관된다.

- 의미론적 관계: "선수"와 "농구"가 같은 도메인에 속한다.

- 통사적 관계: "농구를"과 "한다"가 동사-목적어 관계를 이룬다.

 

Multi - headed attention

 

단일 head는 이 중 하나의 관점에서만 문장을 바라볼 수 있다. 실제로 학습된 모델을 분석해보면, 특정 head는 대명사와 선행사 관계를, 다른 head는 동사와 목적어 관계를, 또 다른 head는 형용사와 명사 수식 관계를 각각 포착하는 방향으로 특화되는 경향이 관찰된다. 각 head가 독립적인 가중치 행렬 $W_Q^{(k)}, W_K^{(k)}, W_V^{(k)}$를 가지기 때문에, 같은 문장을 완전히 다른 관점으로 바라볼 수 있는 것이다.

 

GPT-3는 총 96개의 Attention head를 사용한다. Query·Key 벡터에는 각각 $128 \times 12288$ 크기의 가중치 행렬이 사용된다.

 

Multi-head Attention에서 각 head들이 $\Delta\vec{E_i^{(k)}}$ 를 구하는 모습

5.1 Value의 Low-Rank Transformation

Value 벡터는 원칙적으로 $12288 \times 12288$ 크기의 정사각 행렬로 표현될 수 있지만, 이를 그대로 쓰면 파라미터 수가 너무 커진다. 실제로는 저차원 선형변환(Low-Rank Transformation)을 사용한다.

 

$\text{Value params} \approx \text{Query params} + \text{Key params}$

 

구체적으로 두 행렬의 곱으로 나눈다.

 

- $\text{Value}{\downarrow}$ : 고차원($12288$차원) → 저차원으로 압축하는 행렬 (차원 축소)

- $\text{Value}{\uparrow}$ : 저차원 → 고차원($12288$차원)으로 다시 복원하는 행렬 (차원 확장)

 

Low Rank transformation

 

실제 구현에서의 순서는 다음과 같다. 각 head 내에서 먼저 $\text{Value}{\downarrow}$ 로 차원을 낮춰 Attention 계산을 진행한다. 이후 모든 head가 각자 구한 $\Delta\vec{E}_i^{(k)}$를 합산하고, 최종 Output matrix에서 $\text{Value}{\uparrow}$ 를 적용해 원래 차원($12288$차원)으로 복원한다.

 

최종적으로는 각 head들이 구한 $\Delta\vec{E_i^{(k)}}$ 값을 원래의 $\vec{E_i}$ 에 더해줌으로써 새로운 임베딩으로 넘어감

 

최종적으로는 96개의 head가 각자 구한 $\Delta\vec{E}_i^{(k)}$ 값을 원래의 $\vec{E}_i$에 모두 더해줌으로써, 문맥이 충분히 반영된 새로운 임베딩으로 넘어간다.

 

6. Self-Attention이 전체 파라미터에서 차지하는 비중

Attention은 단순히 문맥을 파악하는 역할을 넘어, Transformer 전체 파라미터의 약 $\frac{1}{3}$을 차지한다. 나머지 $\frac{2}{3}$은 다음 편에서 다룰 MLP(Multi-Layer Perceptron)와 기타 구성 요소가 차지한다. Self-Attention이 Transformer의 "심장"이라고 불리는 이유가 여기 있다. 


 

다음 편에서는 Self-Attention과 다른 데이터 유형 간의 관계를 다루는 Cross-Attention, 모델이 "사실"을 저장하는 방식인 MLP, 그리고 고차원 공간에서 의미가 중첩되는 Superposition 현상을 다룬다.

 

오늘 작성한 내용은 3blue1brown 유튜브 채널의 transformer 관련 영상을 기반으로 작성하였습니다.

수정이 필요한 내용이 있다면 댓글 부탁드립니다.