본문 바로가기
TYLER SONGBlog
← Blog index
LLM Research9 min read

Scaled Dot-Product Attention 수식: sqrt(d_k)로 나누지 않으면 softmax가 한 곳에 몰린다

softmax(QK^T/sqrt(d_k))V 수식의 항을 하나씩 뜯어서 이해하는 날이다. 왜 sqrt(d_k)로 나누는지, 미래 토큰을 못 보게 하는 causal mask와 padding mask가 어떻게 작동하는지를 NumPy로 직접 구현해 확인한다. 스케일을 빼면 softmax가 한 곳에 몰리는 현상도 숫자로 관찰한다.

By 송민성

오늘의 목표

Attention의 핵심 수식 softmax(QK^T / sqrt(d_k)) V에 나오는 모든 항(Q, K, V, 내적, 스케일, softmax, 마스크)이 각각 무슨 역할인지 설명할 수 있게 된다. 그리고 NumPy로 스케일과 마스크가 들어간 attention을 직접 만들고, 마스크가 미래 토큰을 정말로 막는지 숫자로 확인한다.

개념과 직관

1. 세 가지 재료: Q, K, V

문장 속 토큰(token, 문장을 잘게 쪼갠 단위) 하나하나가 숫자 벡터(vector, 숫자를 일렬로 늘어놓은 것)로 표현되어 있다고 하자. Attention은 각 토큰 벡터에서 세 가지 벡터를 만들어 쓴다.

  • Q (Query, 질의): "나는 지금 어떤 정보를 찾고 있어?"
  • K (Key, 키): "나는 이런 정보를 갖고 있어."
  • V (Value, 값): "나를 선택하면 실제로 건네줄 내용은 이거야."

도서관 비유가 쉽다. Q는 내가 들고 간 검색어, K는 책 표지의 제목, V는 책 속 내용이다. 검색어와 제목이 잘 맞을수록 그 책 내용을 많이 가져온다.

(Q, K, V를 어떻게 만드는지는 이후 Day에서 다룬다. 오늘은 이미 만들어져 있다고 가정하고 수식에만 집중한다.)

2. 수식을 왼쪽에서 오른쪽으로

Attention(Q, K, V) = softmax( Q K^T / sqrt(d_k) ) V

모양(shape)을 먼저 정해 두자. 토큰이 n개, 각 벡터의 길이가 d_k라고 하면:

Q: (n, d_k) K: (n, d_k) V: (n, d_v)

  1. Q K^T: K^T는 K를 뒤집은(전치, transpose) 행렬이다. 결과는 (n, n) 크기의 점수표다. i번째 행, j번째 열의 값은 "i번째 토큰의 Q와 j번째 토큰의 K의 내적(dot product, 같은 위치끼리 곱해서 모두 더한 값)"이다. 내적이 크면 두 벡터가 비슷한 방향이라는 뜻이므로 "관련이 깊다"는 점수로 쓴다.
  2. / sqrt(d_k): 점수를 d_k의 제곱근으로 나눠 크기를 줄인다. 이유는 아래에서 설명한다.
  3. softmax: 각 행의 점수를 "합이 1인 확률 같은 가중치"로 바꾼다. softmax(x_i) = exp(x_i) / (exp(x_1) + ... + exp(x_n)). 큰 점수는 더 크게, 작은 점수는 거의 0으로 만든다.
  4. × V: 그 가중치로 V를 섞는다(가중 평균). 결과는 (n, d_v)로, 각 토큰이 "주변 토큰의 정보를 가중치대로 모아 온" 새 벡터가 된다.

3. 왜 sqrt(d_k)로 나눌까

핵심은 "차원이 커질수록 내적 값이 커진다"는 점이다.

Q와 K의 각 원소가 평균 0, 분산(variance, 값이 평균 주위에 퍼진 정도) 1인 값이라고 가정하자. 내적은 d_k개의 곱을 더한 것이다. 곱 하나의 분산이 1이고, 서로 독립이면 더한 값의 분산은 더한 개수만큼 커진다.

Var(q·k) = d_k → 표준편차 = sqrt(d_k)

d_k = 64라면 내적은 대략 -8 ~ +8 근처, d_k = 512라면 -22 ~ +22 근처까지 퍼진다. 이 큰 값이 softmax에 들어가면 문제가 생긴다. exp는 값이 조금만 차이 나도 폭발적으로 벌어지기 때문에 가장 큰 점수 하나가 거의 1, 나머지가 거의 0이 된다. 이것을 포화(saturation)라고 부른다. 포화된 softmax는 사실상 "한 토큰만 고르는 스위치"가 되어, 나머지 토큰의 정보를 거의 못 쓰고 학습 신호도 약해진다. (학습 신호가 약해지는 이유는 역전파를 배울 때 다시 보게 된다. 오늘은 "한 곳에 몰린다"는 현상까지만 확인한다.)

표준편차가 sqrt(d_k)이므로 sqrt(d_k)로 나누면 분산이 다시 1 정도로 돌아온다. d_k가 얼마든 softmax 입력의 크기가 일정하게 유지된다는 것이 핵심이다.

4. 마스크 1: causal mask (미래 가리기)

문장을 앞에서부터 한 단어씩 생성하는 모델은 3번째 단어를 예측할 때 4번째 단어를 보면 안 된다. 시험 답안을 쓰는데 정답지를 보는 것과 같다. 그래서 점수표의 "미래 위치"를 가린다.

점수표가 (n, n)이므로 i번째 행의 j번째 열이 j > i이면 미래다. 4개 토큰이라면 다음처럼 오른쪽 위 삼각형을 가린다.

허용(O) 가림(X) O X X X O O X X O O O X O O O O

가리는 방법은 해당 점수를 -inf(음의 무한대)로 바꾸는 것이다. exp(-inf) = 0이므로 softmax를 거치면 가중치가 정확히 0이 된다. 0을 넣으면 안 된다. 점수 0도 exp(0) = 1이라 가중치를 받기 때문이다. 반드시 softmax 전에 가려야 한다.

5. 마스크 2: padding mask (빈칸 가리기)

문장마다 길이가 다르므로 여러 문장을 묶어 처리하려면 짧은 문장 뒤에 의미 없는 채움 토큰(padding, <pad>)을 붙여 길이를 맞춘다. 이 채움 토큰은 진짜 단어가 아니므로 attention이 보면 안 된다. 방법은 같다. pad 위치의 "열"(K 쪽)에 해당하는 점수를 -inf로 만든다.

코드로 직접 해보기

NumPy만 쓴다. 터미널에서 pip install numpy로 설치한 뒤 attention.py로 저장해 실행하자.

1단계: softmax (안정적으로)

exp는 큰 값에서 오버플로가 나기 쉬워서, 각 행의 최댓값을 먼저 빼 준다. 결과는 수학적으로 같다.

python
import numpy as np def softmax(x): x = x - np.max(x, axis=-1, keepdims=True) # 행마다 최댓값을 빼서 안정화 e = np.exp(x) return e / np.sum(e, axis=-1, keepdims=True)

-inf가 섞여 있어도 최댓값(유한한 값)을 빼면 -inf는 그대로 -inf이고 exp는 0이 되므로 잘 작동한다. 단, 한 행이 전부 -inf이면 nan이 나온다(아래 "막히는 지점" 참고).

2단계: attention 함수

python
def attention(Q, K, V, causal=False, pad_mask=None, scale=True): d_k = Q.shape[-1] scores = Q @ K.T # (n, n) 점수표 if scale: scores = scores / np.sqrt(d_k) # 스케일링 n = scores.shape[0] if causal: future = np.triu(np.ones((n, n), dtype=bool), k=1) # 대각선 위쪽 = 미래 scores = np.where(future, -np.inf, scores) if pad_mask is not None: # pad_mask: 길이 n, pad 위치가 True scores = np.where(pad_mask[None, :], -np.inf, scores) weights = softmax(scores) # (n, n), 각 행의 합은 1 return weights @ V, weights

np.triu(..., k=1)은 대각선 바로 위부터의 삼각형만 True로 남긴다. 위 그림의 X 자리와 같다.

3단계: 마스크가 미래를 정말 막는지 확인

python
rng = np.random.default_rng(0) n, d_k = 4, 8 Q = rng.normal(size=(n, d_k)) K = rng.normal(size=(n, d_k)) V = rng.normal(size=(n, d_k)) out, w = attention(Q, K, V, causal=True) print(np.round(w, 3)) # 오른쪽 위 삼각형이 전부 0이어야 한다 print(w.sum(axis=-1)) # 각 행의 합 = 1 # 진짜 검증: 마지막 토큰(미래)의 값을 완전히 바꿔도 앞 토큰의 출력이 같아야 한다 K2, V2 = K.copy(), V.copy() K2[-1] = rng.normal(size=d_k) * 100 V2[-1] = rng.normal(size=d_k) * 100 out2, _ = attention(Q, K2, V2, causal=True) print(np.allclose(out[:-1], out2[:-1])) # True면 미래를 못 본 것

가중치 표에서 0을 확인하는 것에 더해, "미래를 바꿔도 과거 출력이 안 변한다"를 확인하는 것이 더 강한 검증이다. 마지막 줄이 False이면 마스크가 잘못된 것이다.

padding mask도 같은 방식으로 확인한다.

python
pad = np.array([False, False, True, True]) # 뒤의 2개는 pad out, w = attention(Q, K, V, pad_mask=pad) print(np.round(w, 3)) # 3, 4번째 열이 전부 0

오늘의 실험

질문: sqrt(d_k)로 나누지 않으면 정말 softmax가 한 곳에 몰리는가?

방법: d_k를 바꿔 가며 scale=False와 scale=True를 비교한다. 몰림 정도는 "각 행의 최대 가중치의 평균"으로 잰다. 1에 가까울수록 한 곳에 몰린 것이다.

python
rng = np.random.default_rng(0) n = 16 print("d_k 스케일없음 스케일있음") for d_k in [4, 16, 64, 256, 1024]: Q = rng.normal(size=(n, d_k)) K = rng.normal(size=(n, d_k)) V = rng.normal(size=(n, d_k)) _, w_raw = attention(Q, K, V, scale=False) _, w_scl = attention(Q, K, V, scale=True) print(d_k, round(w_raw.max(axis=-1).mean(), 3), round(w_scl.max(axis=-1).mean(), 3))

추가로 내적의 표준편차도 찍어 보자.

python
for d_k in [4, 64, 1024]: q = rng.normal(size=(10000, d_k)); k = rng.normal(size=(10000, d_k)) print(d_k, np.std((q * k).sum(axis=1)), np.sqrt(d_k))

관찰할 것:

  • 스케일 없음 열의 값이 d_k가 커질수록 1에 가까워지는가?
  • 스케일 있음 열은 d_k와 상관없이 비슷한 수준을 유지하는가?
  • 내적의 표준편차가 sqrt(d_k)와 거의 같은가?

기록: 아래 같은 표를 노트에 남기고, 한 줄로 결론을 쓴다. 정확한 숫자는 실행 환경과 난수 시드에 따라 달라지므로 직접 나온 값을 적는다.

d_k스케일 없음 (최대 가중치 평균)스케일 있음내적 표준편차sqrt(d_k)
4
64
1024

결론 예시 형식: "d_k가 커질수록 ___. 이유는 내적의 분산이 ___에 비례하기 때문이다."

왜 이걸 배우나 (LLM / Post-training 연결)

LLM의 핵심 블록인 Transformer는 이 수식을 층마다 여러 번 반복하고, GPT 같은 생성 모델은 causal mask를 기본으로 쓴다. 파인튜닝(fine-tuning, 이미 학습된 모델을 추가 데이터로 더 학습하는 것)이나 Post-training 업무에서 배치 학습 시 padding mask 누락, 마스크 실수로 인한 미래 정보 누출(학습은 잘 되는데 실제 생성은 망가지는 버그)은 흔한 디버깅 대상이다. 오늘 직접 구현해 본 검증 방식(미래를 바꿔도 과거 출력이 같은가)은 이후 실제 모델 코드를 점검할 때 그대로 쓸 수 있다.

초보자가 막히는 지점

  • 마스크를 softmax 뒤에 적용하는 실수: 가중치를 0으로 만든 뒤에는 합이 1이 아니게 된다. 반드시 softmax "전"의 점수에 -inf를 넣는다.
  • -inf 대신 0을 넣는 실수: exp(0) = 1이라 가려지지 않는다. 위 3단계의 검증 코드로 바로 잡아낼 수 있다.
  • 전부 -inf인 행에서 nan: 한 행의 모든 칸이 가려지면 0/0이 되어 nan이 나온다. 전부 pad인 문장에서 생길 수 있다. 지금은 "이런 경우가 있다"만 알아 두자.
  • sqrt(d_k)와 d_k 혼동: 분산은 d_k에 비례하지만 나누는 값은 표준편차인 sqrt(d_k)다.
  • shape 오류: Q @ K.T에서 에러가 나면 print(Q.shape, K.shape)를 먼저 찍어 보자. 이 코드는 토큰 1개 문장(배치 없음) 기준이라 3차원 입력은 다루지 못한다.
  • "내적의 분산이 d_k" 가정: 원소가 평균 0, 분산 1로 독립이라는 가정에서 나온 결론이다. 실제 학습된 모델의 값이 정확히 그렇다는 뜻은 아니고, "초기에 크기를 안정시키는 설계 이유"로 이해하면 충분하다.

알고리즘 (병행 1시간)

LeetCode Medium 그래프 문제 2문제를 푼다. (예: BFS/DFS로 탐색하는 섬 개수, 코스 스케줄처럼 그래프 기본기를 쓰는 문제 중 골라도 좋다. 이미 푼 문제는 피한다.)

문제마다 노트에 기록할 것:

  • 접근: 그래프를 어떻게 표현했는가(인접 리스트 등), BFS와 DFS 중 무엇을 왜 골랐는가
  • 시간복잡도 / 공간복잡도: 노드 수 V, 간선 수 E 기준으로 적는다
  • 틀린 이유: 방문 체크 누락, 경계 조건, 사이클 처리 등 틀렸거나 막힌 지점을 한 줄로. 틀리지 않았다면 "처음 떠올린 접근에서 헷갈린 부분"을 적는다