[논문 리뷰] On Layer Normalization in the Transformer Architecture

Transformer에서 layernorm을 사용할 때 post-LN을 사용하면 warm up stage에서 hyper-parameter tuning에 시간을 많이 써야하는데 pre-LN을 사용했을 때 그렇지 않다는 내용이다. 사실 내용이 별게 없고 지금은 사용하지 않는 방법인데 여기서는 왜 pre-LN을 사용하고자 하는지, 이론적인 근거가 뒷받침되는지 살펴보자. 

Summary

Abstracts를 읽어보면, 이 논문에서는 첫 번째로 이론적으로 왜 learning rate warm-up stage가 필수적인지와 layer normalization의 위치와의 관계를 보여준다. 구체적으로, post-LN의 initialization에서 mean field theory를 사용해 expected gradients near the output layer가 왜 큰지를 증명한다. 그리고 large learning rate를 이러한 gradients에 사용하는 것은 training을 unstable하게 만든다. warm-up stage는 이런 문제를 피하는데 도움이 되는데, 이 논문의 theory에 의하면 pre-LN 방법은 gradients가 initialization 때 well-behaved 되어있다는 것을 보여준다. 또한 warm-up stage를 제거하여도 comparable하다. 이게 무슨 말인지 뒤를 더 읽어보자. 

 

일단 layernorm이 residual block 사이에 있는 post-LN 방법의 단점에 대해 이야기 하기 시작한다. Transformer를 scratch부터 학습하기 시작할 때 어떤 optimization appraoch던 간에 learning rate warm-up stage가 필요하다. learning rate warm-up stage 라는 건 optimization을 시작할 때 extremely small learning rate로 시작해서 pre defined maximum value로 gradually increase하는 방법을 말한다. warm-up stage를 사용하지 않는 CNN같은 경우 large learning rate 부터 서서히 줄여나가는 반면, transformer는 warm-up stage가 필요하고, 이는 곧 optimization slow down과 sensitive hyper-parameter tuning이라는 단점을 가져온다. 

 

이 논문에서는  mean field theory를 통해 optimization behavior를 investigate한다. 일단 결과적으로 post-LN의 경우, expected gradients of parameters near ouput layer가 커지기 때문에 warm-up stage없이 large learning rate을 사용했을 때 optimization process가 unstable하다. gradients가 크다는 건 변화량이 높다는 뜻이고 large learning rate을 사용하면 더 커질테니 맞는 말이다. 또한 이 논문의 theory에 의하면 layernorm은 gradients를 control하는데 crucial role을 한다. 그것이 바로 well-behaved gradients를 위한 layernorm의 positioning에 대해 조사하게 되는 계기가 되었다고 한다. 이건 사실 조금 뜬금없다고 생각된다. 그리고 그 theory가 대체 어떻게 증명된건지 얼른 말해주면 좋겠다. 

 

아무튼 이 논문에서는 pre-LN transformer를 통해 initialization 때 gradients가 vanishing이나 exploding 없이 well-behave한다는 것을 이론적으로 뿐만아니라 실험적으로 보여줄 것이다. 그리고 gradients가 well-behave하므로 warm-up stage를 쓸 필요가 없어진다. English-German translation과 German-English translation, BERT pretraining task에서 warm-up stage는 문제 없이 제거되고 hyperparameter의 개수도 줄어든다. 또한 warm-up stage로 인한 slow down이 없기 때문에 loss는 더 빠르게 줄어든다. 따라서 pre-LN은 large-scale datasets의 large-scale models의 training에서 특별히 중요하다. 

 

그 다음에는 transformer with post normalization의 architecture에 대해 설명하는 내용인데, 크게 중요한 내용은 아닌 것 같은데 왜 넣었는지 모르겠다. 그리고 나서는 warm-up stage에 대한 설명도 하는데 이것도 위에서 언급한 내용 외에 더 말할 내용은 없는 것 같다. 그 다음에서야 understanding the transformer at initialization 파트로 넘어간다. 드디어 여기서 이론적으로 왜 post-LN의 gradients가 large learning rate를 initialization에 사용할 수 없는지 나온다. 

 

 

Strengths

vit와 그 뒤에 나온 vision transformer는 모두 pre-LN을 사용하고 있다. language는 내가 잘 모르지만.. 중요한 발견인 것 같은데 생각보다 인용수가 작다. 근데 Transformers without tears: Improving the normalization of self-attention 여기 논문에서 이미 empirically pre-LN이 warm-up stage를 없앨 수 있다는 말을 empirically 했다고 한다. 

Weaknesses

논문에 쓸데없는 내용이 많다.. 논문을 내기엔 별로 할 말이 없었나보다. 

My idea