3장. 해부: 어텐션

오늘날 언어 모델, 이미지 생성 모델, 단백질 구조 예측 모델의 거의 모두가 트랜스포머이고, 트랜스포머의 중심에는 식 하나가 있습니다. 2017년 논문 “Attention Is All You Need”(Vaswani 외)의 scaled dot-product attention입니다. 이 장에서는 그 식을 2장의 다섯 질문으로 한 줄씩 해부합니다. 필요한 것은 행렬곱(선형대수 5장), 소프트맥스(다변수 미적분 11장), 분산의 성질(확률과 통계 3장)뿐입니다.

3.1 해부할 식

• $Q \in \mathbb R^{n \times d_k}$: 쿼리(query) $n$개. “무엇을 찾는가”.
• $K \in \mathbb R^{m \times d_k}$: 키(key) $m$개. “나는 무엇에 관한 것인가”.
• $V \in \mathbb R^{m \times d_v}$: 값(value) $m$개. “찾으면 돌려줄 내용”.
• softmax는 행마다 적용합니다.

질문 1 (무엇의 함수인가). 입력 행렬 세 개의 함수이고, 학습하는 매개변수는 이 식에 없습니다. 매개변수는 $Q = XW_Q$, $K = XW_K$, $V = XW_V$처럼 입력 $X$에서 세 행렬을 만드는 투영 $W_Q, W_K, W_V$에 있습니다.

질문 2 (모양). $QK^\top$는 $n \times m$, softmax 뒤에도 $n \times m$, $V$를 곱하면 $n \times d_v$. 쿼리마다 $d_v$차원 벡터 하나가 나옵니다(2장 문제 2.2).

3.2 내적은 유사도

$(QK^\top)_{ij} = \mathbf q_i^\top\mathbf k_j$는 $i$번째 쿼리와 $j$번째 키의 내적입니다. 내적은 $\lVert\mathbf q\rVert\lVert\mathbf k\rVert\cos\theta$이므로(선형대수 13장), 방향이 비슷할수록 큽니다. “이 쿼리에 이 키가 얼마나 관련 있는가”의 점수입니다.

3.3 소프트맥스: 점수를 가중치로

정의 3.1. 점수 행 $\mathbf s = (s_1, \dots, s_m)$에 대해 어텐션 가중치(attention weights)는 $a_j = \dfrac{e^{s_j}}{\sum_le^{s_l}}$입니다.

$a_j > 0$이고 $\sum_ja_j = 1$이므로 확률분포입니다. 출력은 값들의 가중 평균 $\sum_ja_j\mathbf v_j$, 곧 기댓값 $\mathbb E_{j \sim \mathbf a}[\mathbf v_j]$입니다.

질문 3 (극단값). 한 점수가 다른 점수보다 훨씬 크면 가중치가 거의 그 하나에 몰려, 출력은 그 값 벡터 하나를 그대로 가져옵니다(사전에서 키로 찾아 값을 꺼내는 것과 같음). 점수가 모두 같으면 값들의 단순 평균입니다.

예제 3.2. 키 두 개 $\mathbf k_1 = (1, 0)$, $\mathbf k_2 = (0, 1)$, 값 $\mathbf v_1 = (10, 0)$, $\mathbf v_2 = (0, 10)$, 쿼리 $\mathbf q = (2, 0)$, $d_k = 2$.

• 점수: $\mathbf q^\top\mathbf k_1/\sqrt2 = \sqrt2 \approx 1.414$, $\mathbf q^\top\mathbf k_2/\sqrt2 = 0$.
• 가중치: $a_1 = \frac{e^{\sqrt2}}{e^{\sqrt2} + 1} \approx 0.804$, $a_2 \approx 0.196$.
• 출력: $0.804(10, 0) + 0.196(0, 10) = (8.04, 1.96)$.

쿼리가 첫째 키 쪽을 향하므로 첫째 값을 주로 가져옵니다.

3.4 왜 √d로 나누는가

이것이 이 식에서 가장 자주 묻는 질문이고, 답은 확률과 통계 3장의 분산 계산입니다.

정리 3.3. $\mathbf q, \mathbf k \in \mathbb R^d$의 성분이 서로 독립이고 평균 0, 분산 1이면 $\mathrm{Var}(\mathbf q^\top\mathbf k) = d$입니다.

증명. $\mathbf q^\top\mathbf k = \sum_{i=1}^dq_ik_i$이고 각 항은 평균 0, 분산 $E[q_i^2]E[k_i^2] = 1$, 서로 독립이므로 분산이 더해져 $d$입니다. □

그래서 $d = 512$이면 점수의 표준편차가 $\sqrt{512} \approx 22.6$나 됩니다. 점수 차이가 이렇게 크면 소프트맥스가 거의 원-핫이 되어(질문 3), 가중치의 그래디언트 $a_j(\delta_{jl} - a_l)$가 거의 0이 되고 학습이 멈춥니다. $\sqrt{d}$로 나누면 분산이 1로 돌아옵니다.

질문 4 (특수한 경우). $d_k = 1$이면 나누지 않은 것과 같습니다. 스케일링은 고차원에서만 의미가 있습니다.

3.5 인과 마스크

언어 모델은 다음 단어를 예측하므로, $i$번째 위치가 $i$보다 뒤의 단어를 보면 안 됩니다. softmax 전에 $j > i$인 점수를 $-\infty$로 바꿉니다($e^{-\infty} = 0$). 이를 인과 마스크(causal mask)라 하며, 가중치 행렬이 아래삼각 행렬이 됩니다.

3.6 다중 헤드

정의 3.4. 다중 헤드 어텐션(multi-head attention)은 $h$개의 서로 다른 투영으로 어텐션을 따로 계산해 이어 붙이고 선형변환합니다.

보통 $d_k = d_{\text{model}}/h$로 잡아 전체 계산량을 헤드 하나일 때와 비슷하게 맞춥니다. 헤드마다 다른 종류의 관계(문법, 지시 대상, 위치)를 본다고 해석합니다.

질문 5 (어디서 왔는가). 최적화나 확률 모델에서 유도한 식이 아니라 설계된 식입니다. 다만 “점수의 소프트맥스로 가중 평균”은 커널 회귀(나다라야–왓슨 추정)와 같은 꼴이고, 소프트맥스 자체는 최대 엔트로피 분포(최적화 6장)에서 나옵니다.

3.7 계산량

길이 $n$인 문장의 자기 어텐션($n = m$)에서 $QK^\top$는 $n \times n$ 행렬이고, 계산량은 $\mathcal O(n^2d)$, 메모리는 $\mathcal O(n^2)$입니다. 문맥을 2배로 늘리면 4배가 듭니다. 긴 문맥 연구(희소 어텐션, 선형 어텐션, 상태공간 모델)가 모두 이 $n^2$을 줄이려는 시도입니다.

3.8 이 장의 요약

• $\mathrm{softmax}(QK^\top/\sqrt{d_k})V$: 내적 점수 → 확률 가중치 → 값의 가중 평균.
• $\sqrt{d_k}$로 나누는 이유: 내적의 분산이 $d_k$라서, 소프트맥스 포화를 막으려고.
• 인과 마스크는 $-\infty$로 미래를 가림. 다중 헤드는 여러 관계를 병렬로.
• 계산량 $\mathcal O(n^2d)$.

쓰이는 곳. 트랜스포머 기반의 모든 논문(GPT, BERT, ViT, 확산 트랜스포머), 효율적 어텐션 논문(FlashAttention은 같은 식을 메모리 효율적으로 계산), 해석 가능성 연구(어텐션 가중치 시각화, 귀납 헤드), 그래프 신경망의 어텐션, 단백질 구조 모델이 이 식을 씁니다.

문제

5

0.8808

8

16