STUDYLLM
‹ Lecture 3 08 / 21

Post-LN vs Pre-LN

Norm의 위치를 바꾸는 것은 같은 계산을 보기 좋게 재배치하는 일이 아니다. Post-LN은 sublayer output과 원래 state를 더한 결과 전체를 정규화하고, Pre-LN은 sublayer가 읽는 입력만 먼저 정규화한다. Residual의 원래 xx가 다음 state로 넘어가는 경로에 Norm이 끼어 있느냐가 핵심 차이다.

  1. Post: x+F(x)
  2. LN / Pre: LN(x)
  3. F
  4. +x

Residual Paths

Post-LN은 xnext=LN⁡(x+F(x))x_{\mathrm{next}}=\operatorname{LN}(x+F(x)), Pre-LN은 xnext=x+F(LN⁡(x))x_{\mathrm{next}}=x+F(\operatorname{LN}(x))다. Post의 Jacobian은 JLN(I+JF)J_{\mathrm{LN}}(I+J_F), Pre는 I+JFJLNI+J_FJ_{\mathrm{LN}}으로 직접 identity term을 갖는다.

Warm-Up & Stability

Xiong의 setting에서 Post는 output 근처 initial gradient가 크고 warm-up이 도움됐다. Pre는 warm-up 없이 비교 가능한 결과를 보였지만 모든 setting에서 warm-up이 불필요하다는 법칙은 아니다. Pre-LN stack 뒤에는 Final Norm을 흔히 추가한다.

Following the Gradient

Jacobian은 입력의 작은 변화가 출력에 어떻게 전달되는지를 나타낸다. Backpropagation에서는 layer를 거슬러 갈 때 이 변화율이 연쇄적으로 곱해지므로, 깊은 stack에서 어떤 경로가 남는지가 중요하다. 여기서 F는 Attention 또는 FFN 같은 sublayer다.

Post-LN의 x+F(x)x+F(x)를 미분하면 I+JFI+J_F지만, 그 바깥에 LN이 있으므로 전체에 JLNJ_{\mathrm{LN}}이 곱해진다. Identity term도 Norm의 영향을 피하지 못한다. 반면 Pre-LN의 원래 xx는 바깥에서 그대로 더해지므로, 미분하면 직접 II가 남고 나머지 branch만 JFJLNJ_FJ_{\mathrm{LN}}을 거친다.

이 차이는 초기 layers 사이 gradient scale을 더 안정적으로 유지하는 데 도움이 된다. Xiong et al.의 Post-LN setting에서는 output 근처 gradient가 커서 초기 learning rate가 크면 불안정할 수 있었고, learning rate를 서서히 올리는 warm-up이 도움이 되었다. Pre-LN에서는 그 실험의 warm-up 없이도 비교 가능한 결과를 얻었다. 다만 Pre-LN이면 어떤 학습에서도 warm-up이 필요 없다는 결론은 아니다. Initial gradient 분석과 training 중 발생할 모든 불안정성을 구분해야 한다.

Pre-LN은 residual 전체를 매 sublayer 뒤에 정규화하지 않으므로 stack 마지막에 Final Norm을 흔히 둔다. 다음 QK Norm은 또 다른 위치, 즉 projection 뒤 Attention의 내적 바로 전에 적용되는 정규화다.

Post-LN:LN⁡(x+F(x))\text{Post-LN:}\quad\operatorname{LN}(x+F(x)) J=JLN(I+JF)J=J_{\mathrm{LN}}(I+J_F) Pre-LN:x+F(LN⁡(x))\text{Pre-LN:}\quad x+F(\operatorname{LN}(x)) J=I+JFJLNJ=I+J_FJ_{\mathrm{LN}}

F는 Attention 또는 FFN 같은 sublayer다. 초기 gradient의 안정성과 학습 중 모든 안정성 문제의 해결을 동일시하지 않는다.

Pre-LN Jacobian에는 직접 I가 있다. 모든 안정성 문제의 해결은 아니다.

Post-LN · Pre-LN

Norm의 위치가 직접 identity 경로를 바꾼다.

AspectPost-LNPre-LN
식LN(x+F(x))x+F(LN(x))
JacobianJLN(I+JF)J_{\mathrm{LN}}(I+J_F)I+JFJLNI+J_FJ_{\mathrm{LN}}
Identity 경로J_LN을 통과직접 I term
인용된 settingOutput 근처 init gradient가 큼; warm-up 도움초기 gradient가 더 안정적; warm-up 없이 비교 결과
Stack 끝Sublayer output 정규화Final Norm을 흔히 추가

해당 분석 결과를 모든 architecture의 warm-up 필요 여부로 일반화하지 않는다.

Check

Pre-LN과 Post-LN의 Jacobian에서 I의 위치가 어떻게 다른가?

Answer

Post는 J_LN이 (I+J_F)를 감싼다. Pre는 I가 직접 더해진다. Identity 경로가 Norm을 통과하는지의 차이다.

Pre-LN을 썼으니 QQ/KK logit 폭주와 warm-up 문제가 모두 해결됐다고 할 수 있는가?

Answer

아니다. 해당 분석·실험 setting의 안정성 결과이며 projection 후 QQ/KK 크기는 별도로 커질 수 있다.

Pre-LN에서 gradient의 직접 경로가 남는 이유를 식으로 설명할 수 있는가?

Answer

x+F(LN(x))x+F(\mathrm{LN}(x))의 Jacobian에는 II가 별도로 남는다. Post-LN의 LN(x+F(x))\mathrm{LN}(x+F(x))는 identity 경로도 LN의 Jacobian을 통과한다. 이 차이가 인용된 initialization 분석의 gradient/warm-up 차이와 연결되며 모든 architecture의 warm-up 필요 여부를 정하는 법칙은 아니다.

Post-LN · Pre-LN를 축별로 비교하고 그 차이가 필요한 이유를 설명하세요.

Answer

Norm의 위치가 직접 identity 경로를 바꾼다.

식: Post-LN=LN(x+F(x)) / Pre-LN=x+F(LN(x))

Jacobian: Post-LN=JLN(I+JF)J_{\mathrm{LN}}(I+J_F) / Pre-LN=I+JFJLNI+J_FJ_{\mathrm{LN}}

Identity 경로: Post-LN=J_LN을 통과 / Pre-LN=직접 I term

인용된 setting: Post-LN=Output 근처 init gradient가 큼; warm-up 도움 / Pre-LN=초기 gradient가 더 안정적; warm-up 없이 비교 결과

Stack 끝: Post-LN=Sublayer output 정규화 / Pre-LN=Final Norm을 흔히 추가

조건: 해당 분석 결과를 모든 architecture의 warm-up 필요 여부로 일반화하지 않는다.

Contents