[Others] Transformers without Normalization

[Others] Transformers without Normalization

• 논문 링크: https://openaccess.thecvf.com/content/CVPR2025/papers/Zhu_Transformers_without_Normalization_CVPR_2025_paper.pdf

• 왜 대부분의 모델은 normalization layers 를 가져야 하는가에 대한, 기본 요구사항에 대한 질문을 제기하는 것에서 출발하는 논문이다. 그리고 나는 이런 질문을 좋아한다.

Introduction

지난 몇년간 normalization 은 거의 대부분의 modern neural networks 에서 fundamental componetns 로 여겨지며 지위를 공고히 했다. 몇가지 변주를 줬을 뿐이지 대부분의 모델들이 하나의 궤를 같이 했는데, 아래의 기본 normalization layers 를 보자.

• x는 기본 input, input 에서 mean 을 뺀 값을 square root of variance 로 나눠주는 게 기본이다.
• input x의 dimension 은 (B, T, C) 가 되는데, 여기서 B는 batch size, T는 number of tokens, C는 embedding dimension 이 된다.

normalization 은 이 기본 틀을 유지하면서 서로 다른 mean, variance 계산 방법을 도입하며 약간씩 변주를 줘 왔다.

• Batch Normalization 은 평균을 input 을 batch에 들어있는 number of tokens 만큼 계산한다음, 이 수들로 나누어 구한다. (그럼 BXT로 나누게 됨.) variance 도 마찬가지.

• Layer normalization 은 Transformer architectures 에서 사용되는 normalization 으로, 이번엔 평균을 각 샘플의 토큰에 걸쳐 내게 된다, 따라서 token dimension 으로 나누어 평균을 구한다. (C로 나누게 됨.) Variance 도 마찬가지.

• 비슷하게 RMSNorm (root mean square normalization)은 평균을 0으로 끌고 온 LN 버전이라 생각하면 됨.

오늘날 대부분의 modern nural networks 는 LN 을 사용하며, 특정 언어모델 (T5, LLaMA, Mistral, Qwen, InternLM, DeepSeek)의 경우 RMSNorm 도 사용한다.

What Do Normalization Layers Do?

ViT-B, wave2vec 2.0, LibriSpeech, Diffusion Transformer (DiT-XL)에 대하여, normalization layers 들이 하는 역할을 조사하기 위해 이들의 input-output 값을 조사해봤다. 특히 LN 은 dimension 을 보존하고 오직 값만 normalize 하기 때문에, one-to-one 으로 input and output elements 대응이 가능하다.

관찰 1. 초기 LN layers 에서의 input-output relationship은 linear 하다.

관찰 2. 더 깊은 layer에서는, 이 모양이 완전히 또는 부분적으로 S자 모양의 tanh function 의 모양을 가진다. 이는 굉장히 특징적인 부분인데, 이론상 LN layer의 계산은 input tensor에 대해 평균과 분산을 계산하는데, 이 계산은 본질적으로 linear 인데 반하여 tanh function 은 non-linear transformation 이기 때문이다. 저자들은 이러한 plot 의 이유가, LN 이 토큰마다 normalize 하기 때문에 여기까지는 linear 가 보존되지만, 각 토큰에서의 평균과 분산값은 모두 다르므로 결국은 non-linearity 가 추가될 수 밖에 없다는 주장을 한다.

관찰 3. Input 의 range 는 extream 하고, output 은 그렇지 않다.

결론적으로 Normalization 의 main effeect 는 input 의 값들을 squash (압축해서) them into less extreme values 하는 것과, non-linear 을 추가하는 것 두 가지로 보았다. 다음의 논문 에서는 왜 non-linearity 가 model's representation capacity 를 향상시키는가에 대한 이유를 보여주고 있는데, 분명히 당연한 질문이다. 읽어볼 가치가 있어보여 Keep. 지금은 그러면 이러한 역할을 하는 (특히 tanh 을 닮았으므로 저자들은 normalization layer 대신 tanh 함수 사용한 layer를 도입하는데) 대체 layer 로 저자들이 제안한 것을 보자.

Dynamic Tahn (DyT)

논문에서 normalization layers 의 역할로 제시한 layer는 다음과 같다.

논문에서 제시한 건 이게 다다! 아까 본 plot 에서 normalization layers 의 input-output function 이 tanh fucntion 을 닮았기 때문에 여기에 착안해서 tan함수를 도입했고, 알파는 learnerble scalar parameter 로, 다른 range 의 input 에 다르게 반응하도록 하기 위해 도입되었다. 감마와 베타는 learnable, per-chaneel vector paramter 로, output 이 scale 되도록 돕는다.

• 감마는 1로, 베타는 0, 알파는 0.5로 초깃값을 준다. 후에는 초기 알파값으로 어떤 것이 적절한지에 대한 experiments 가 조금 있다.

• DyT는 normalization layers 를 대체한다. (attention blocks, FFN blocks, normalization layer에 각각 존재.)

결과 1. DyT는 Vision, LLM, Diffusion, Speech, DNA sequence modeling 에 대해 잘 작동한다.

논문은 다양한 모델, vision, LLM, Diffusion, Speech, 심지어는 DNA sequence modeling (normalization layer 자체가 전 neural networks 의 기본이 되고 있었기 때문에) 에 대해 DyT를 테스트하는데, 아래와 같이 comparable, 조금 더 높은 혹은 낮은 성능을 기록함을 확인했다.

(Supervised learning in vision)

(Self-supervised learning in vision)

(Diffusion models)

(Large Language Model)

(Self-supervised learning in Speech)

(DNA sequence modeling)

결과 2. tanh and $\alpha$ 효과

그런데 이정도의 결과로는 충분치 않다고 느껴지는 게, 그냥 input 을 output 으로 매핑하는 함수 (f(x) = x) 같은 것과, 또는 다른 함수 (논문에서는 hardtanh 와 sigmoid) 와의 비교도 필요해 보이기 때문이다. 이에 대한 결과가 다음과 같다.

결론적으로 다른 함수들보다 tanh fucntion 에서의 성능이 가장 좋았으며, 이는 함수 자체의 smothness 와 zero-centered properties 때문이라고 저자들은 분석하고 있다. 함수 자체의 그림은 다음과 같다.

다음으로는 알파의 효과를 보는 것인데, 아래는 알파가 learnable parameter 로 전반적인 성능 향상에 기여하고 있음을 확인할 수 있다.

결과 3. 알파 초깃값에 따른 변화

다음은 서로 다른 알파를 초깃값으로 줬을 경우 모델별로 성능을 비교한 것이다. 결과는 다음과 같다.

관찰 1. Non-LLM 모델의 경우 알파의 초깃값에 덜 민감했다.: 서로 다른 알파 값을 도입하는 건 훈련의 초기 단계에서 큰 영향을 미쳤고, 훈련이 점점 진행되며 insensitive 한 쪽으로 수렴했다.

관찰 2. 작은 알파값이 stable training 에 기여한다. 알파값을 높이는 건 훈련을 불안정하게 만들었으며, divergence 하는 경우들을 만들어냈다.

관찰 3. LLM의 경우 알파값 세팅이 중요했다. 더 큰 모델일수록, 훈련 시간이 길수록 작은 알파값이 필요했고, 상대적으로 다른 곳 (FFN, linear projection) 에서의 DyT보다 attention blocks 에서의 DyT를 높게 줘서, 서로 다르게 튜닝하는 것이 성능 향상에 기여했다.

Insights

• normalizaiton 에 대한 근본적인 물음, 반드시 필요한 요소인가? 에 대한 질문을 던짐. 비슷한 함수를 찾아 바꿈.

• 그런데 몇가지 아쉬운 점은, 이것이 근본적인 물음을 던졌다는 것 외에, 반드시 왜 필요한가? 에 대한 질문에는 답하지 못한다는 것. 성능을 크게 높인 것 X, computation cost 를 낮춘 것 X. 비슷한 성능을 낼 수 있었던 이유를 theoretical 하게 분석한 것은 아니라는 점, 등이 있다.

• 그럼에도 불구하고 이런 연구를 하고 싶고, 중요한 가치가 있다 생각하는 게, Transformer architecture 에서는 덜 잘 알려진 이런 요소들이 많을 것이라는 생각이 듦. 예컨대 1. 지금 포함된 components 가 실험해봤더니 완전히 필요 없는 것으로 판명나거나. 2. 그 components 가 다른 components 로 대체될 수 있는 경우. (이 논문처럼 비슷한 기능을 하는 다른 요소의 도입. 성능과 computational cost 를 고려했을 때) 3. 필요 없지도 않고, 다른 components 로 대체할 필요도 없는데, 왜 잘 되는지 theoretic 하게 모르는 경우 4. Transformer architecture 외에 아예 다른 형식의 training architecture 도 나올 것이라 (적어도 나는) 생각하기에.. 연구할 거리가 굉장히 많다 생각 든다. 이런 류의 fundamental question 을 던지는 논문들을 꾸준히 추적하면 좋을 듯.