QK Norm

Pre-LN으로 입력 $X$를 정규화해도 그 뒤의 learned projections가 만든 Q/K의 크기까지 고정되지는 않는다. 또 $\sqrt{d_k}$ scaling은 초기 분포에 대한 조정이었다. Training 중 Q/K norm이 커지면 dot-product scores가 커지고 Softmax가 한 위치로 포화될 수 있다. Block 입력의 scale, 내적에 들어가는 vector의 scale, 내적 후 score의 범위는 서로 다른 위치의 문제다.

1. $Q$/$K$ magnitude ↑
2. QK logits ↑
3. Softmax near one-hot
4. Entropy ↓
5. loss 불안정
6. vector Norm / score cap

Entropy Collapse

ViT-22B recipe의 8B ViT ablation에서 Attention logits 폭주 → one-hot → $\mathrm{Entropy}\approx0$ → gradient/loss 불안정 사례가 소개됐다. $|\mathrm{logit}|>50{,}000$은 해당 case의 관측이며 모든 8B 모델의 보편적 threshold가 아니다. $\sqrt{d_{h}}$는 $q$,$k$ 가정 아래의 initial scale을 조정할 뿐 learned norms를 고정하지 않는다.

Norm & Soft-Capping

QK Norm은 token/head별 $Q$/$K$ projections에 RMSNorm·LayerNorm·L2 Norm 등을 적용한 뒤 dot product한다. 강의 예시는 learned gain $\gamma$가 있는 per-head RMSNorm이다. $\gamma$=1, $d_{h}=128$ 가정의 scaled score bound는 $\sqrt{128}\approx11.31$이며 learned gains에 따라 달라진다.

Soft-capping은 score 이후 $a\leftarrow50\tanh\!\left(\frac{a}{50}\right)$로 $(-50,50)$에 넣는다. Vector 정규화와 score 범위 제한의 작용 위치를 구분한다.

Magnitude & Direction

내적은 방향의 관계뿐 아니라 두 vector의 크기에도 영향을 받는다.

$$q^\top k=\lVert q\rVert\lVert k\rVert\cos\theta$$

방향 관계가 그대로여도 norms만 커지면 score가 커질 수 있다. Raw $\lVert q\rVert=\lVert k\rVert=40$인 예시에서는 내적의 절댓값이 최대 1600까지 가능하다. QK Norm은 projections 뒤 token/head별 Q와 K를 먼저 정규화하여 이 raw magnitude 의존을 줄인다.

Per-head RMSNorm에서 learned gain을 1로 놓고 $d_h=128$이라고 가정하면 scaled score의 bound는 약 $\sqrt{128}=11.31$이다. Learned gains가 달라지면 같은 고정 bound를 그대로 주장할 수 없다. Henry et al.의 L2 Norm과 learned scalar, ViT-22B의 LayerNorm, Qwen3의 RMSNorm은 같은 목적을 다른 구체적 방법으로 다루는 예시다.

Entropy가 거의 0이라는 것은 Attention 분포가 거의 하나의 position만 선택한다는 뜻이다. ViT ablation은 text가 아닌 image patches에서도 같은 Attention 연산의 포화가 training instability와 연결될 수 있음을 보여 준다. 관측된 50,000은 그 Adam 1e-3 실험의 수치이지 모델 크기가 8B면 발생하는 공통 threshold가 아니다.

Bounding the Score

Soft-capping은 vector가 아니라 계산된 score를 $50\tanh(a/50)$으로 바꾼다. Tanh의 출력이 (-1,1)이므로 결과는 (-50,50)에 놓인다. 이후 learned scale 등으로 score가 커질 수 있는 설계에서도 score 자체의 범위를 제한하려는 접근이다.

QK Norm은 내적 전 vector를, soft-cap은 내적 후 score를 다룬다. 서로 같은 switch가 아니므로 함께 사용할 수도 있다. 어느 Norm을 썼다는 이름보다 적용 위치와 제어 대상을 확인해야 한다.

$$q_n=\operatorname{RMSNorm}(q),\quad k_n=\operatorname{RMSNorm}(k)$$ $$a_{ij}=\frac{q_n^\top k_n}{\sqrt{d_h}}$$ $$a\leftarrow50\tanh\!\left(\frac{a}{50}\right)\in(-50,50)$$ $$\gamma=1,\quad d_h=128\quad\text{가정}$$ $$|a|\le\sqrt{128}\approx11.31$$

QK Norm과 block 입력 Norm은 위치가 다르다. Learned gains가 있는 Norm에 같은 고정 score bound를 무조건 적용하지 않는다.

QK Norm은 내적 전 vector, soft-cap은 내적 후 score다.

√d_k Scaling · QK Norm · Soft-cap

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

Check

$\sqrt{d_{k}}$ scaling이 있는데 QK Norm이 필요할 수 있는 이유를 cause→effect로 설명하라.

Answer

학습 중 $Q$/$K$ magnitude가 증가하면 dot product가 커질 수 있다. 고정 scaling만으로 막지 못해 Softmax가 sharp해지고 Entropy/gradient가 불안정해질 수 있어 vector를 Norm한다.

QK Norm과 Soft-cap의 작용 위치와 보장 범위를 구분하라.

Answer

QK Norm은 내적 전 vector scale을 제어하지만 learned gains의 영향이 남는다. Soft-cap은 내적 후 score 자체를 유한한 구간으로 매핑한다.

Q/K의 방향이 그대로인데도 Attention이 near-one-hot으로 변할 수 있는 이유는?

Answer

$q^\top k=\lVert q\rVert\lVert k\rVert\cos\theta$이므로 norms가 커지면 score 차이도 커질 수 있다. 고정 초기 scaling만으로 learned magnitudes를 통제하지 못한다. QK Norm은 내적 전 vector를, soft-capping은 내적 후 score를 제어한다.

Attention score를 만들기 전 projection된 $Q$/$K$의 vector scale을 제어하는 기법은?

Answer

QK Norm. Score에 사후 적용하는 soft-cap과 작용 위치가 다르다.

학습 중 $Q$/$K$ magnitude 증가 이후 학습 불안정까지의 경로와 대응 위치를 말하라.

Answer

$Q$/$K$ magnitude↑ → QK logit magnitude↑ → Softmax near one-hot → Entropy↓ → gradient/loss 불안정. QK Norm은 dot product 전, soft-cap은 score 후.

Related

• Scaled Dot-Product Attention
• Post-LN vs Pre-LN
• Contrastive Training