Scaled Dot-Product Attention

앞에서 Query와 Key가 “어디를 읽을지”를 정하고 Value가 “읽을 내용”을 제공한다는 역할을 보았다. 이제 이 선택을 실제로 계산한다. Query마다 모든 allowed Keys의 score를 구하고, scale과 mask를 적용한 뒤, 그 Query의 주의를 Keys 사이에 배분하는 weights로 바꾼다. 마지막에는 그 비중대로 Values를 섞어 Query 하나당 output vector 하나를 만든다.

$QK^\top$
÷ $\sqrt{d_{k}}$ + Mask
Softmax over keys
$A$ $V$
query별 output

Computation

각 query에 대해 Score, Scaling과 Mask, Softmax, Value의 가중합 순서로 계산한다.

\operatorname{Attention}(Q,K,V)=\operatorname{Softmax}_{\mathrm{keys}}\!\left(\frac{QK^\top}{\sqrt{d_k}}+M\right)V

Score:

S=QK^\top\quad(n_q\times n_k)

Scaling & Mask:

L=\frac{S}{\sqrt{d_k}}+M

Softmax:

A=\operatorname{Softmax}_{\mathrm{keys}}(L)

Output:

O=AV\quad(n_q\times d_v)

AV:\quad(n_q\times n_k)(n_k\times d_v)=n_q\times d_v

$n_q$ 는 Query positions, $n_k$ 는 memory positions의 수다. Score matrix의 한 row는 Query 하나가 모든 Keys를 비교한 결과이므로 Softmax도 Key 축을 따라 적용한다. 한 row의 weights 합이 1이라는 것은 그 Query의 전체 비중을 접근 가능한 positions에 나눈다는 뜻이다. Self-Attention에서는 $n_q=n_k$ 이지만 Cross-Attention에서는 두 sequence 길이가 달라도 된다.

Mask는 matrix의 shape를 줄이지 않는다. 접근 불가 위치의 logit에 -∞를 더하면 exponential이 0이 되어 해당 Value가 가중합에 기여하지 않는 것이다. 아래 예시에서는 raw score가 가장 큰 세 번째 Key도 mask 때문에 선택되지 않는다.

$d_k=4$ 이고 $M=[0,0,-\infty]$ 일 때 계산은 다음과 같다.

Raw Scores:

[1.5,2,3]

Scaled & Masked:

[0.75,1,-\infty]

Weights:

[0.44,0.56,0]

Weighted Sum:

0.44v_1+0.56v_2

Scaling

Initialization에서 components가 서로 independent이고 평균이 0, 분산이 1이라고 가정한다.

\mathbb{E}[q_i]=\mathbb{E}[k_i]=0,\quad \operatorname{Var}=1

이때 dot product의 variance는 $d_k$ 이다.

\operatorname{Var}(q\cdot k)=d_k

\operatorname{Var}\!\left(\frac{q\cdot k}{\sqrt{d_k}}\right)=1

$\sqrt{d_k}$ 로 나누면 variance가 1이 된다. $d_k=64$ 인 예를 비교하면 scaling의 영향을 볼 수 있다.

Unscaled Logits:

[16,8,0]

Unscaled Weights:

\operatorname{Softmax}([16,8,0])\approx[0.9997,0.0003,\approx0]

Scaled Logits:

[2,1,0]

Scaled Weights:

\operatorname{Softmax}([2,1,0])\approx[0.665,0.245,0.090]

Scaling 전에는 거의 모든 비중이 첫 Key로 몰리지만, scaling 뒤에는 첫 Key가 여전히 가장 크면서 다른 Keys도 비중을 가진다. Scale은 순위를 뒤집는 연산이 아니라 Softmax로 들어가는 차이를 완화하는 연산이다. Softmax의 자기 logit에 대한 미분은 다음과 같다.

\frac{\partial p_i}{\partial \ell_i}=p_i(1-p_i)

$p_i$ 가 0이나 1에 매우 가까우면 이 기울기가 작아진다. 그래서 초기부터 near-one-hot으로 포화되면 다른 관계를 조정하는 학습 신호가 약해질 수 있다. 분산은 상수 $c$ 로 나누면 $c^2$ 로 나뉘므로, 분산 $d_k$ 인 dot product를 $\sqrt{d_k}$ 로 나누어 1로 맞추는 것이다.

단 이 유도는 초기 독립·평균 0·분산 1이라는 가정을 사용한다. 학습 후에는 Q/K의 분포와 norm이 달라질 수 있으므로, 고정된 scaling 하나가 training 내내 score의 크기를 통제한다고 보장하지 않는다. 이 한계가 뒤의 QK Norm과 연결된다.

Score는 Softmax 이전의 값이고 weight는 Softmax 이후의 값이다. Mask의 −∞는 Softmax에서 probability 0이 된다.

Worked Example

입력과 Key dimension은 다음과 같다.

q=[1,0,1,0]

d_k=4

Raw Scores:

[2,0,1]

Scaling:

[1,0,0.5]

Softmax:

[0.51,0.19,0.31]

Values:

V=\begin{bmatrix}1&0\\0&1\\0.5&0.5\end{bmatrix}

여기서는 $d_v=2$ 다. Q/K의 4차원은 compatibility 계산에 쓰이고, 2차원 Value는 전달할 내용을 정하므로 두 폭이 달라도 된다. 세 Value의 기여를 따로 계산하면 다음과 같다.

0.51[1,0]=[0.51,0]

0.19[0,1]=[0,0.19]

0.31[0.5,0.5]\approx[0.15,0.15]

이를 합하면 아래 output을 얻는다. Raw score가 0이던 둘째 Key도 masked된 것은 아니므로 Softmax 뒤에는 약 0.19의 비중으로 참여한다.

\mathrm{output}\approx[0.66,0.34]

같은 logits를 10배하면 weights는 다음처럼 sharp해진다.

\mathrm{weights}\approx[0.993,0.00005,0.007]

반올림 때문에 화면에 표시한 weights의 합이 정확히 1이 아닐 수 있다.

가중합은 세 Values를 하나로 섞으므로 각 slot의 정보를 따로 보존한 목록은 아니다. Single head가 한 관계에 집중하면 다른 관계가 같은 mixture 안에서 약해질 수 있다. 다음의 Multi-Head Attention은 이 계산을 여러 learned subspaces에서 수행하여 서로 다른 관점을 표현할 여지를 만든다. 모든 head가 반드시 다른 유용한 역할을 얻는다는 보장은 없다.

슬라이드의 개별 vector 내적 $q^\top k_i$ 와 matrix 표기 $QK^\top$ 는 다른 algorithm이 아니다. 개별 vector를 column으로 적느냐, token states를 matrix의 rows로 쌓느냐에 따라 transpose의 위치가 달라지는 것이다.

Attention과 FFN의 특징과 차이를 Transformer block 안의 역할과 연결해 설명할 수 있어야 한다.

Check

$Q$ ( $n_{q}\times d_{k}$ ), $K$ ( $n_{k}\times d_{k}$ ), $V$ ( $n_{k}\times d_{v}$ )에서 $QK^\top$ 와 최종 output shape는?

Answer

$QK^\top$ 는 $n_{q}\times n_{k}$ , Softmax는 각 query의 key축에 적용한다. AV는 $n_{q}\times d_{v}$ 다.

왜 $d_{k}$ 가 아니라 $\sqrt{d_{k}}$ 로 나누며 학습 중에는 무엇을 보장하지 못하는가?

Answer

초기 independent unit-variance components의 dot product variance가 $d_{k}$ 여서 표준편차 $\sqrt{d_{k}}$ 로 나누면 variance1이다. Learned $Q$ / $K$ magnitude 증가와 초기 가정의 지속은 보장하지 못한다.

Scaled logits $\begin{bmatrix}1&0&0.5\end{bmatrix}$ 와 $V$ 세 개로 output을 계산하고 logits×10의 결과를 설명하라.

Answer

Softmax≈ $\begin{bmatrix}0.51&0.19&0.31\end{bmatrix}$ , weighted $V\approx \begin{bmatrix}0.66&0.34\end{bmatrix}$ 다. ×10이면 거의 첫 value에 집중하여 다른 slot 정보가 덜 전달된다.

Raw score가 큰 Key라도 output에 전혀 기여하지 않을 수 있는가?

Answer

가능하다. Mask가 그 위치의 logit을 −∞로 만들면 Softmax weight가 0이 되어 해당 Value는 output에 기여하지 않는다. Mask는 score matrix의 shape를 줄이는 연산이 아니며, Softmax는 Query마다 Key 축을 따라 계산한다.

$Q$ / $K$ / $V$ 가 주어졌을 때 output을 얻는 네 연산을 shape와 함께 복원하라.

Answer

$QK^\top$ ( $n_{q}\times n_{k}$ ) → ÷ $\sqrt{d_{k}}$ +M → Softmax_keys weights( $n_{q}\times n_{k}$ ) → AV( $n_{q}\times d_{v}$ ).

Attention · FFN를 축별로 비교하고 그 차이가 필요한 이유를 설명하세요.

Answer

교수님이 안내한 첫 서술형 준비 방향: 두 sublayer의 차이를 Residual Stream과 연결한다.

정보 접근: Attention=Allowed positions 간 혼합 / FFN=각 position 독립 계산

주요 연산: Attention= $QK^\top$ → scale/mask → Softmax → AV / FFN=Up → activation/gate → down

Output: Attention=Value의 weighted mixture / FFN=Token features의 비선형 변환

Vanilla matrix weights: Attention= $Q$ / $K$ / $V$ /O: $4d^{2}$ / FFN= $d$ →4d→ $d$ : $8d^{2}$

Vanilla forward: Attention= $8nd^{2}$ projections + $4n^{2}d$ / FFN= $16nd^{2}$

Residual 역할: Attention=위치 간 통신 update를 더함 / FFN=위치별 계산 update를 더함

확장 개념: Attention=MHA / GQA / MLA / FFN=SwiGLU / MoE

조건: FFN 입력은 이미 Attention의 문맥을 포함할 수 있다.

위 상수는 dense vanilla / FFN4d / multiply-add2 FLOPs 가정이다.

$\sqrt{d_{k}}$ Scaling · QK Norm · Soft-cap를 축별로 비교하고 그 차이가 필요한 이유를 설명하세요.

Answer

서로 다른 위치에서 서로 다른 scale 문제를 제어한다.

위치: Scaling=Dot product 후 / QK Norm=Dot product 전 $Q$ / $K$ / Soft-capping=Dot product 후 score

제어 대상: Scaling=초기 $d_{k}$ 의존 variance / QK Norm=Vector magnitude / Soft-capping=Score 범위

표현: Scaling= $QK^\top$ / $\sqrt{d_{k}}$ / QK Norm=Norm( $Q$ ),Norm( $K$ ) / Soft-capping=50 tanh(a/50)

한계: Scaling=Learned norms 증가 보장 못함 / QK Norm=Learned gains 영향 남음 / Soft-capping=표시한 score를 (-50,50)에 매핑

Computation

Scaling

Worked Example

Check

Related