[딥러닝] 논문 리뷰 - Normalization 꼭 필요없다. Transformers without Normalization

3일 전 Normalization과 관련된 굉장히 재밌는 논문이 메타에서 나왔다. 

 

https://arxiv.org/abs/2503.10622

Transformers without Normalization

 

나는 일전에도 Batch Normalization(BN) 과 Layer Normalization(LN) 에 관심이 많았어서 공부해서 포스팅을 정리하기도 했었고, 논문이 얀 르쿤이 관여했으며 메타에서 나온지라 흥미가 커서 한 번 읽어봤고, 내용이 직관적이고 많지 않아서 간단하게 리뷰 및 정리해보려고 한다. BN, LN 에 대한 내용은 아래의 이전 포스팅 참조.

 

https://tiabet0929.tistory.com/88

[딥러닝] Layer Normalization과 Batch Normalization 정리

Normalization in Transformer (Background)

논문의 저자들도 언급했듯이 Neural Network 에서 Normalization 이란 필수불가결한 요소로 여겨져왔다.

특히 Transformer의 Original Architecture를 확인해보면 LN이 매 헤드마다 사용되는 것을 볼 수 있는데, 이는 gradient vanishing 혹은 gradient exploding 같은 역전파 시 일어날 수 있는 문제들을 통제하기 위해 필수적인 방법으로 연구되었기 때문이다. 이후에 나온 여러 LLM들은 이 LN을 RMSNorm (GPT3, LLAMA 등등) 혹은 헤드 뒤가 아닌 앞에 붙이는 Pre-Normalization 등 다양하게 개선하여 사용해오고 있다.

What Do Normalization Layers Do?

저자들은 학습된 LLM 모델 (정확히 따지면 멀티모달 LLM) 들의 layer 몇 개를 뜯어봐서, Before LN 값과 After LN 값들을 1대1로 매칭시켜보았다. 아래 그림 참조.

그랬더니 층이 깊어질수록 After LN(LN Output) 값들은 S 자 커브 모양을 띄는 것을 알 수 있는데, 이는 마치 하이퍼볼릭탄젠트 함수 $tanh$ 와 유사하다고 볼 수 있다.

따라서 저자들은 LN의 역할이 결국 $tanh$ 와 별다를 것이 없다고 생각하고, 아예 LN (후의 실험에서 LLaMA를 사용하는데 거기선 RMSN) 자체를 $tanh$로 바꿔버리자는 아이디어를 냈다. 

 

물론 근거가 단순히 저 모양만 본 것은 아니다.

 

 

이 그림을 보면 LN input의 값들과 LN Output의 값들의 Plot이 상당히 선형적이면서도(0 주변에 몰려있는 값들을 확인해보면 선형성이 강하다), 일부 값들은 극단적으로 치우친 값을 갖기도 한다.

 

However, there are still many points that clearly fall out of this range, which are considered to have “extreme” values.
Normalization layers’ main effect for these values is to squash them into less extreme values, more in line with the majority of points.

 

저자들은 이런 값들을 그나마 선형적으로 바꿔주는 것이 LN의 역할이라고 봤다. 그리고 어차피 LN이 그걸 한다면, $tanh$로 가운데로 몰아줘도 (squash) 괜찮지 않나? 하는 생각에서 실험을 몇 가지 진행했다.

 

Dynamic Tanh (DyT)

저자들은 LN을 대체할 메소드로 Dynamic Tanh, DyT를 제시한다.

$$D_Y T(x) = \gamma \cdot \tanh(\alpha x) + \beta$$

 

정말 일반적인 $tanh$ 거의 똑같고, $\alpha, \beta, \gamma$ 라는 하이퍼파라미터들이 더해진 수식이다. 이 하이퍼파라미터들은 모두 학습되는 learnable parameter이며, $\beta$ 와 $\gamma$ 는 LN 에서와 똑같은 파라미터들로 층마다 다른 값을 가질 수 있다. 이렇게 학습되는 하이퍼파라미터들을 갖고 있기 때문에 Dynamic 이라는 이름을 붙였다고 한다.

 

Experiments and Analysis

실험 결과는 논문에 나온 그대로를 인용한다.

Training Loss DyT vs LN

 

 

Accuracy on ImageNet-1K

LLAMA Pretrain Loss RMSNorm vs DyT

LLAMA training loss and average performance with 15 zero-shot lm-eval tasks

 

놀랍게도 Loss 로 보건 성능으로 보건 LN, RMSNorm 과 차이가 거의 없다. 사실상 아예 똑같다고 보인다.

이 포스팅에는 ViT와 LLAMA 결과만 가져왔지만 논문을 보면 wav2vec 으로 음성인식모델, DiT 로 Diffusion 모델, DNA를 분류하는 DNA Sequence Model, MAE 같은 오토인코더에도 실험했지만 성능 차이가 아예 발생하지 않았다. 

 

Inference and Training Time

그럼 성능이 똑같다면 얼마나 효율적이냐를 봐야하는데, 아무래도 LN보다 단순한 계산의 $tanh$를 사용했다보니 (LN은 평균과 분산을 구하는 계산과정이 있으나 $tanh$에는 그런 게 없다.) Inference와 Training 둘 다에서 유의미한 시간 감소를 보였다.

 

 

물론 이 과정에는 섬세한 Ablation Study가 있었다. 이 포스팅에선 그건 생략하고, 저자들이 유심히 살펴본 영역이 $\alpha$의 initial value에 따른 성능 차이에 대해서만 살펴보려고 한다. 저자들은 다양한 $\alpha$ 값들로 학습률을 조정해가면서 ViT에 대해 실험을 해봤는데, 큰 $\alpha$ 값일수록 학습에 실패했다고 한다. 여기서 실패란 학습이 발산했다는 의미인데, 이는 결국 gradient exploding 혹은 vanishing 같은 문제 때문이라고 봐야 할 것 같다. 물론 LN도 학습률이 커질수록 학습에 실패하는 경향이 있긴 했는데, 작은 $\alpha$를 사용하면 이런 점도 어느정도 해결되는 모습을 보인다.

 

 

하지만 LLAMA로 실험해봤을 땐 적당히 큰 $\alpha$ 를 사용하는 것이 loss를 떨어트리는 데에 도움이 됐다고도 한다. 또한 attention layer에서의 optimal $\alpha$ 와 FFNN에서의 값이 달랐고, width 와 depth에 따라서도 최적값이 달랐다고 한다.

 

의의 및 한계점

저자들은 이렇게 분석을 했지만, 왜 이런 결과가 나오는지, 특히 LLM 내부에서 width 와 depth에 따라 왜 성능 차이가 발생하고 ViT와 LLM에서는 왜 $\alpha$ 의 값에 따른 경향성이 다르게 나오는지 이유까진 찾지 못 하여 추가적인 연구가 필요하다고 말했다. 그리고 LN이 아닌 BN에 대해서는 DyT가 대체하지 못 했다고 하는데, 이 이유 또한 명확히 밝히지 못하여 후속 연구의 필요성을 제기한다.

 

하지만 정말 필수불가결한 요소로 여겨진  LN이고, LN이 필요했던 이유를 값들을 정규분포에 근사해야 학습이 안정적으로 잘 이루어진다는 기존 상식을 깰 수도 있는 연구이기 때문에 굉장히 흥미로운 발견인 것 같다. 후속 연구가 나오면 그 논문도 리뷰해볼 생각이다.