01. 어텐션 한 줄부터 · Annotated Transformer 한 단계씩

원본: 어텐션

Transformer를 한 문장으로 줄이면 "모든 토큰이 다른 모든 토큰을 얼마나 참고할지 스스로 정해서 섞는 모델"이다. 그 "섞는" 연산이 어텐션이고, 나머지 부품은 전부 이 연산을 잘 돌리기 위한 장치다. 그래서 첫 줄은 어텐션부터 쓴다.

Scaled Dot-Product Attention

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

query와 key의 내적 QK^T: "i번째 토큰이 j번째 토큰과 얼마나 관련 있는가"를 점수로 매긴다.
softmax: 점수를 행마다 합이 1인 가중치로 바꾼다.
value에 곱하기: 그 가중치로 value 벡터들을 섞는다. 출력 토큰 하나는 입력 토큰들의 가중 평균이다.
/ math.sqrt(d_k): 벡터 차원이 커질수록 내적값이 커져서 softmax가 한 칸에 쏠린다. 차원 수의 제곱근으로 나눠 이를 막는다.

함수가 결과와 함께 가중치 p_attn도 돌려주는 점을 기억해 두자. 마지막 단계에서 이 가중치를 그림으로 그린다.

실행 결과에서 볼 것

가중치 행렬의 각 행을 더하면 1이다.
출력 첫 번째 토큰을 가중치로 직접 섞어 계산한 값과 비교하면 똑같다.
나누지 않으면 d_k=512에서 softmax가 한두 칸에 거의 전부를 몰아준다. 나누면 고르게 퍼진다.

입력 x : (1, 3, 4) (batch, 토큰 수, d_k) 출력 : (1, 3, 4) 가중치 p_attn: (1, 3, 3) (batch, query 토큰, key 토큰) 어텐션 가중치 — i행은 'i번째 토큰이 각 토큰을 얼마나 참고하는가' tensor([[0.9630, 0.0100, 0.0270], [0.0450, 0.7620, 0.1930], [0.2560, 0.3900, 0.3550]]) 행마다 합: ['1.000', '1.000', '1.000'] 출력 0번 토큰 = 가중치로 섞은 value attention() : ['1.4535', '-0.3072', '-2.1104', '0.5605'] 직접 계산 : ['1.4535', '-0.3072', '-2.1104', '0.5605'] 왜 √d_k로 나누는가 — 성분이 평균 0·분산 1인 q, k의 내적은 분산이 d_k가 된다 d_k= 4 q·k 표준편차 2.00 → √d_k로 나누면 1.00 d_k= 64 q·k 표준편차 8.11 → √d_k로 나누면 1.01 d_k=512 q·k 표준편차 22.85 → √d_k로 나누면 1.01 d_k=512, key 10개일 때 softmax 결과 나누지 않으면: 0.000 0.000 0.000 0.000 0.000 0.652 0.348 0.000 0.000 0.000 √d_k로 나누면: 0.085 0.033 0.093 0.154 0.076 0.213 0.208 0.054 0.062 0.022 → 나누지 않으면 한 칸에 거의 다 몰려서(포화) 기울기가 사라진다

1	+	import math
2	+
3	+	import torch
4	+
5	+
6	+	def attention(query, key, value):
7	+	"Compute 'Scaled Dot Product Attention'"
8	+	d_k = query.size(-1)
9	+	scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
10	+	p_attn = scores.softmax(dim=-1)
11	+	return torch.matmul(p_attn, value), p_attn

1	import math
2
3	import torch
4
5
6	def attention(query, key, value):
7	"Compute 'Scaled Dot Product Attention'"
8	d_k = query.size(-1)
9	scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
10	p_attn = scores.softmax(dim=-1)
11	return torch.matmul(p_attn, value), p_attn

1	import math
2
3	import torch
4
5	from transformer import attention
6
7	torch.manual_seed(0)
8
9	# 토큰 3개짜리 문장 하나. 토큰마다 4차원 벡터.
10	x = torch.randn(1, 3, 4)
11
12	# self-attention: query, key, value를 모두 같은 x로
13	out, p_attn = attention(x, x, x)
14
15	print("입력 x :", tuple(x.shape), " (batch, 토큰 수, d_k)")
16	print("출력 :", tuple(out.shape))
17	print("가중치 p_attn:", tuple(p_attn.shape), " (batch, query 토큰, key 토큰)")
18	print()
19	print("어텐션 가중치 — i행은 'i번째 토큰이 각 토큰을 얼마나 참고하는가'")
20	print(p_attn[0].round(decimals=3))
21	print("행마다 합:", [f"{v:.3f}" for v in p_attn[0].sum(-1).tolist()])
22	print()
23	print("출력 0번 토큰 = 가중치로 섞은 value")
24	manual = sum(p_attn[0, 0, j] * x[0, j] for j in range(3))
25	print(" attention() :", [f"{v:.4f}" for v in out[0, 0].tolist()])
26	print(" 직접 계산 :", [f"{v:.4f}" for v in manual.tolist()])
27	print()
28
29	print("왜 √d_k로 나누는가 — 성분이 평균 0·분산 1인 q, k의 내적은 분산이 d_k가 된다")
30	for d_k in (4, 64, 512):
31	q, k = torch.randn(10000, d_k), torch.randn(10000, d_k)
32	dots = (q * k).sum(-1)
33	print(
34	f" d_k={d_k:3d} q·k 표준편차 {dots.std():6.2f}"
35	f" → √d_k로 나누면 {(dots / math.sqrt(d_k)).std():.2f}"
36	)
37	print()
38	q, k = torch.randn(1, 512), torch.randn(10, 512)
39	raw = (q @ k.T).softmax(-1)
40	scaled = (q @ k.T / math.sqrt(512)).softmax(-1)
41	print("d_k=512, key 10개일 때 softmax 결과")
42	print(" 나누지 않으면:", " ".join(f"{v:.3f}" for v in raw[0].tolist()))
43	print(" √d_k로 나누면:", " ".join(f"{v:.3f}" for v in scaled[0].tolist()))
44	print(" → 나누지 않으면 한 칸에 거의 다 몰려서(포화) 기울기가 사라진다")

어텐션 한 줄부터

실행 결과에서 볼 것

코드 변화

실행 결과

참고

`key.transpose` 인자 `(-2, -1)`의 의미

실행 결과에서 볼 것

코드 변화

실행 결과

참고

key.transpose 인자 (-2, -1)의 의미

`key.transpose` 인자 `(-2, -1)`의 의미