Pre-LN Transformer: On Layer Normalization in the Transformer Architecture 논문 리뷰

The overview of this paper

 Transformer는 NLP task에서 널리 사용된다. 하지만 Transformer를 학습시키기 위해, 대게 신중하게 디자인된 learning rate warm-up stage를 사용한다. 이 learning rate warm-up stage는 최종 성능에 많은 영향을 끼치지만, optimization의 속도를 저하시키고 더 많은 hyper-parameter tuning을 필요로 한다. 이 논문에서는 learning rate warm-up stage가 왜 필수적인지와 layer normalization$($LN$)$의 위치에 대한 연구를 진행하였다. 구체적으로, 논문에서는 초기화 시 residual block 사이에 layer normalization을 배치하는 원래 설계된 Post-LN Transformer의 경우 출력 레이어 근처의 매개변수의 예상 기울기가 크다는 것을 mean field 이론으로 증명하였다. 따라서 그 기울기에서 큰 learning rate를 사용하는 것은 학습을 불안정하게 만든다. warm-up stage는 이러한 문제를 피하는데 매우 실용적이다. 반대로, 논문의 이론은 layer normalization이 residual block 사이에 들어가게 되면, 초기화시에 기울기가 잘 적용되는 모습을 보여준다. 논문의 실험에서는 warm-up stage가 없는 Pre-LN Transformer이 상당히 적은 학습 시간과 hyper-parameter tuning에도 불구하고, baseline에 필적하는 결과를 얻는 것을 보여줬다.

 

 

Table of Contents

1. Introduction

2. Optimization for the Transformer

   2-1. Transformer with Post-Layer Normalization

   2-2. The learning rate warm-up stage

   2-3. Understanding the Transformer at initialization

3. Experiment Results

 

 

1. Introduction

 Transformer는 NLP에서 가장 흔하게 사용되는 신경망 architecture이다. LN은 Transformer의 성공에 중요한 역할을 하였다. 기존의 Transformer는 residual block 사이에 LN을 두는데, 이를 Transformer with Post-Layer Normalization이라 한다. 이 architecture는 여러 분야의 NLP task에서 SOTA를 달성하였다. 그리고 Post-LN Transformer에 기반을 둔 비지도학습 pre-trained model은 다양한 downstream task에서 인상적인 성능을 보여줬다.

 

 이러한 성공에도 불구하고, 사람들은 보통 Post-LN Transformer는 CNN이나 seq2seq model보다 더욱 조심히 optimization을 다룬다. 특히, model을 처음부터 학습시키기 위해서는, 어떠한 기울기 기반의 최적화 방식이든 간에 learning rate warm-up stage를 필요로 한다. 이 learning rate warm-up stage는 매우 작은 learning rate에서 시작해서 이를 pre-defined 반복수로 pre-defined maximum value까지 점진적으로 증가시킨다. 이러한 warm-up stage는 최적화의 속도를 늦출 뿐만 아니라 더 많은 parameter tuning을 필요로 한다. 

 

 이 논문에서는 learning rate warm-up stage를 안전하게 제거하는 방법을 찾음으로써 문제를 완화하였다. 첫 번째 반복에서 warm-up stage가 발생하는 것처럼, mean field 이론을 사용해서 초기화 상태에서 optimization 행동을 조사하였다. 논문의 이론적 분석에 의하면, residual block 사이에 LN을 배치할 때, output layer 근처 파라미터의 예상되는 기울기는 크다. 그래서, warm-up stage 없이 large learning rate를 이 파라미터에 사용하는 것은 optimization 프로세스를 불안정하게 만든다. warm-up stage와 small learning rate로 model을 학습시키는 것은 이 문제를 실용적으로 피할 수 있게 도와준다. 광범위한 실험은 논문의 이론적 발견을 지지하는 내용을 제공해준다.

 

 논문의 이론은 LN이 기울기 scale을 조절하는 데 중요한 역할을 한다는 것을 보여줬다. 이는 연구자들에게 LN을 어떤 위치에 놓았을 때 잘 적응된 기울기를 얻게 하는지에 대한 다른 방법을 조사하고자 하게 하였다. 특히, 논문에서는 다른 변형인 Pre-LN을 사용하는 Transformer를 연구하였다. Pre-LN Transformer는 residual connection 안에 LN을 넣고, 예측 전에 추가적인 final-layer normalization을 추가하였다. 논문에서는 초기화 시 기울기가 이론적으로나 경험적으로 Pre-LN Transformer에 대해 폭발하거나 사라지지 않고 잘 동작함을 보여줬다.

 

 논문의 실험 결과를 살펴보면, 모든 task에서 learning rate warm-up stage가 안전하게 제거되었고, hyper-parameter의 수 또한 줄어들었다. 게다가, 논문에서는 Pre-LN Transformer에 대해 loss decay가 더 빨랐다. Pre-LN Transformer는 비슷한 최종 성능을 더욱 적은 학습 시간을 사용해서 달성하였다. 이는 large-scale model을 large-scale dataset에 대해 학습시키는 것에 대해 특히 중요하다.

 

 논문의 Contribution은 다음과 같다.

 

• 논문에서는 두 개의 Transformer 변형인, Post-LN Transformer과 Pre-LN Transformer에 대해 mean field 이론을 사용해서 조사하였다. 초기화 시에 기울기를 연구하여, 논문에서는 왜 learning rate warm-up stage가 Post-LN Transformer을 학습시킬 때 필수적인지에 대한 증거를 제공하였다.
• 논문에서는 처음으로 learning rate warm-up stage가 hyperparameter tuning을 쉽게 만들어주는 Pre-LN Transformer를 위해 제거될 수 있음을 보여줬다. 논문에서는 추가적으로 적정한 learning rate scheduler를 사용함으로써 training time이 광범위한 응용에서 크게 감소될 수 있음을 보여주었다.

 

2. Optimization for the Transformer

2-1. Transformer with Post-LN Normalization

 

 이 섹션에서는 vanilla Transformer에 대해서 설명하고 있다. Transformer에 대한 자세한 설명이 궁금하다면 여기를 참고하길 바란다. 

 

 Transformer layer에서 sub layer, reidual connection, LN의 서로 다른 순서는 Transformer architecture의 변형을 야기한다. Transformer와 BERT를 위한 가장 기본적이고 가장 유명하게 사용되는 architecture는 "self-attention$($FFN$)$ sub-layer → residual connection → layer normalization"을 따른다. 이것은 Post-LN with Transformer이라고 부르고, 다음의 그림 1에 나타나 있다.

 

그림 1. a는 Post-LN Transformer, b는 Pre-LN Transformer

 

Post-LN Transformer  $x_{l, i}$를 위치 $i$에서 $l$번째 Transformer layer의 입력으로 표기하고, 여기서 $x_{l, i}$는 차원이 $d, i = 1, 2, ..., n, l = 1, 2, ..., L$의 실제값 벡터이다. $n$은 sequence의 길이이고, $L$은 layer의 수이다. 완성도를 위해, 논문에서는 일반적으로 word embedding과 positional embedding의 조합인 위치 $i$에서 input embedding을 $x_{0, i}$로 정의한다. $l$번째 레이어의 안에 들어있는 계산은 여러 단계로 이루어져 있고, 표 1의 왼쪽처럼 다른 스텝의 input과 output을 표현하기 위해 $x$에 super-script를 사용한다. 여기서 $W^{1,l}, W^{2,l}, b^{1,l}, b^{2,l}$은 $l$번째 레이어에서 FFN sub-layer의 파라미터들이다.

 

표 1. Post-LN Transformer vs. Pre-LN Transformer

 

 

2-2. The learning rate warm-up stage$($Post-LN Transformer$)$

 

논문에서는 Post-LN Transformer의 optimization에서 learning rate wam-up stage에 관심을 가졌다. learning rate가 상대적 큰 값에서 시작해서 감소하는 다른 많은 architecture의 optimization과 달리 Post-LN Transformer의 learning rate warm-up stage는 매우 중요해 보인다. 논문에서는 $t$번째 반복의 learning rate를 $lr(t)$로 표기하고, 학습 중에 maximum learning rate를 $lr_{max}$로 표기하였다. pre-defined 타임 프레임 $T_{warmup}$이 주어지면, 첫 번째 $T_{warmup}$ 반복을 위한 learning rate scheduler는 다음과 같이 정의된다.

 

$lr(t) = \frac {t}{T_{warmup}}lr_{max}, t \leq T_{warmup}$

 

 이 warm-up stage 이후에, learning rate는 일반적인 learning rate scheduler로 설정된다. 논문에서는 Post-LN Transformer를 학습시키기 위해 learning rate warmup-stage가 필수적이라는 것을 실험을 통해 보여줬다.

 

Results and discussions  논문에서는 validation loss와 BLEU score를 학습하고 계산할 때 매 epoch에 대해 model checkpoint를 기록하였다. model의 성능은 그림 2$($a$)$와 그림 2$($b$)$에 나타나있다. $x$축은 epoch의 수이고, $y$축은 BLEU score/validation loss이다. "w/o warm-up"은 "warm-up stage가 없음"을 의미하는 반면에 "w/ warm-up"은 "warm-up stage와 함께"를 의미한다.

 

그림 2. Adam과 SGD에 의해 최적화된 model의 성능

 

 첫 번째로, 두 optimizer에 대해 learning rate warm-up stage가 필수적이라는 것을 알아냈다. warm-up stage 없이 Adam optimizer와 함께 학습된 모델은 BLEU score에서 8.45를 달성하였다. 비교에 따르면, warm-up stage를 사용해서 학습한 모델은 BLEU score에서 34를 달성하였다. 이러한 경향은 validation loss curve에서도 확인할 수 있었다. SGD를 이용하여 학습된 모델은 Adam에 비해 떨어지는 성능을 보여주긴 했지만, 이러한 경향은 똑같이 나타났다. warm-up stage를 사용하지 않은 모델은 15 epoch의 학습 뒤에도 BLEU score가 0을 조금 넘는 모습을 보여줬다.

 

 두 번째로, optimization 프로세스는 $T_{warmup}$의 값에 대해서 민감하다는 것을 알 수 있었다. 이것은 Post-LN Transformer을 학습시킬 때 $T_{warmup}$이 중요한 hyper-parameter라는 의미이다. 예를 들어, $T_{warmup} = 500$으로 설정했을 때, Adam으로 학습된 모델은 각각 $lr_{max}=5e^{-4}$와 $1e^{-3}$에 대해 BLEU score에서 겨우 31.16과 2.77의 점수를 기록하였다.

 

 이러한 warm-up stage는 여러 단점을 가지고 있다. 첫 번째로, warm-up stage의 구성은 최종 성능에 상당한 영향을 끼친다는 것이다. 그래서 개발자들은 large-scale NLP task에 대해 계산적으로 비싼 신중한 hyper-parameter tuning을 필요로 한다. 두 번째로, warm-up stage는 optimization의 속도를 늦춘다는 것이다. 기존의 optimization 알고리즘은 보통 빠른 수렴을 위해 커다란 learning rate에서 시작한다. 하지만, warm-up stage를 사용할 때, learning rate는 0에서부터 점차 증가해 나가는데, 이는 학습을 비효율적으로 만든다. 이전의 연구에 따르면, warm-up stage는 모델 훈련의 초기 단계에서 Adam의 바람직하지 않은 상당한 분산을 줄이는 데 역할을 한다는 것을 알 수 있었다. 하지만, 논문의 결과에 따르면, warm-up stage는 SGD의 학습 또한 도와주는 것을 알 수 있다. 이는 warm-up stage의 이익은 특정 optimizer에 국한되지 않음을 보여준다.

 

2-3. Understanding the Transformer at initialization$($Pre-LN Transformer$)$

 

 논문에서 제안된 Pre-LN Transformer는 다음과 같다.

 

그림 3. Pre-LN Transformer

 

 이 그림 3을 살펴보면 Post-LN Transformer와 달리 LN의 위치가 변해있는 것을 알 수 있다. 이것이 Post-LN Transformer와 Pre-LN Transformer의 가장 큰 차이점이다.

 

Gradient of Weight Prameter  직감적으로 랜덤 변수 $Z$가 $(\varepsilon /, \delta)$-bounded이면, 높은 확률로 그 realization이 기대에서 너무 멀리 떨어져있지 않을 것이다. $L$번째 레이어에서 Post-LN Transformer에 대한 마지막 레이어의 파라미터의 기울기는 다음을 만족한다.

 

$\begin{Vmatrix}
\frac {\partial \tilde{\mathfrak{L}}}{\partial W^{2,L}}
\end{Vmatrix}_{F} \leq O(d \sqrt{ln d})$

 

 반면에 $L$번째 레이어에서 Pre-LN Transformer의 기울기는 다음과 같다.

 

$\begin{Vmatrix}
\frac {\partial \tilde{\mathfrak{L}}}{\partial W^{2,L}}
\end{Vmatrix}_{F} \leq O(d \sqrt{\frac {ln d}{L}})$

 

 Post-LN Transformer의 경우 마지막 FFN layer에 대한 기울기의 스케일이 $L$과 무관한 $O(d \sqrt{ln d})$ 차수임을 알 수 있다. Pre-LN Transformer의 경우에는, 기울기의 스케일이 훨씬 작은 $O(d \sqrt{\frac {ln d}{L}})$이다.

 

Scale of Hidden States  서로 다른 레이어의 hidden state의 스케일을 측정하였다. 입력과 초기화의 무작위성에 대한 기대가 있다. 만약 $X \in R^{d}$가 가우시안 벡터 $X ~ N(0, \sigma^{2}I_d)$이면, $\mathbb{E}(||ReLU(X)||_{2}^{2})=\frac {1}{2}\sigma^{2}d$이다. 초기화 시점에서 Post-LN Transformer에 대해서는 $\mathbb{E}(||x_{l,i}^{post, 5}||_{2}^{2})=\frac {3}{2}d$이고, Pre-LN Transformer에 대해서는 $(1 + \frac {l}{2})d \leq \mathbb{E}(||x_{l,i}^{pre}||_{2}^{2}) \leq (1 + \frac {3l}{2})d$이다.

 

Advantage

 

그림 3. 1. 6-6 Transformer에서 서로 다른 레이어의 기울기 a,b  /  Transformer의 서로 다른 크기에서 W^{2,L} c,d

 

 위의 그림에서 나와있는 것처럼, 예상되는 기울기의 scale은 Post-LN Transformer에 대한 layer index와 함께 성장한다. 이와는 반대로, Pre-LN Transformer에서 서로 다른 레이어에 대해 scale은 거의 똑같은 값을 유지한다. 여기서 main idea는 LN이 기울기를 정규화할 것이라는 것이다.

 

 Post-LN Transformer에서 LN에 대한 입력의 스케일은 $L$과 무관하므로 마지막 계층의 매개변수 기울기는 $L$과 무관하다. Pre-LN Transformer에 있는 동안 최종 LN에 대한 입력의 스케일은 $L$에서 선형이므로 모든 매개변수의 기울기는 $\sqrt{L}$로 정규화된다. 

 

 

3. Experiment Results

 이전 섹션에서 Pre-LN Transformer를 위한 초기화 시에 기울기가 잘 작동한다는 것을 알 수 있었다. 이러한 사실을 기반으로, Pre-LN Transformer를 학습시킬 때 learning rate warm-up stage가 안전하게 제거될 수 있음을 보여줬다. 이번 섹션에서는 이를 2개의 NLP task에 대해 실험적으로 입증하도록 하였다. 이 두 개의 NLP task는 Machine Translation과 Unsupervised Pre-Training$($BERT$)$이다.

 

Machine Translation  이 task에 대한 모델의 성능은 다음의 그림 4에 a부터 d까지 나와있다. 결과를 분석해보면 다음과 같다.

 

1. 더 이상 learning rate warm-up stage가 강력하지 않음. 따라서 Pre-LN Transformer도 경쟁력을 갖추게 되었음.
2. Pre-LN Transformer가 Post-LN Transformer보다 똑같은 $lr_{max}$에 대해 더 빨리 수렴함.
3. LN 위치의 변화가 optimizer의 변화를 '지배'한다는 사실 발견

 

그림 4. Machine Translation task에 대한 모델의 성능

 

Unsupervised Pre-training$($BERT$)$  model의 validation loss를 다음의 그림 5의 a와 같이 기록하였다. machine translation task와 비슷하게 Pre-LN Transformer를 위해 learning rate warm-up stage는 제거될 수 있다. 이렇게 하면 Pre-LN Transformer는 더욱 빨리 학습할 수 있다. 그리고 Pre-LN Transformer는 큰 learning rate를 사용해서 더욱 쉽게 최적화될 수 있다. 실험에서는 서로 다른 model의 checkpoint를 MRPC와 RTE downstream task에 대해 진행하였다. 실험에 대한 결과는 그림 5의 b와 c에 나와있다. 결과를 보면 Pre-LN Transformer가 downstream task에 더욱 빠르게 수렴한다는 것을 알 수 있다.

 

그림 5. 비지도 pre-training BERT와 downstream task에 대한 모델의 성능

 

 요약하면, 서로 다른 task에 대한 모든 실험은 Pre-LN Transformer을 학습시키는 것이 learning rate warm-up stage에 의존하지 않고 Post-LN Transformer에 비해 더욱 빠르게 학습될 수 있다는 것을 보여준다.

 

 

 

 

출처

https://sh-tsang.medium.com/review-pre-ln-transformer-on-layer-normalization-in-the-transformer-architecture-b6c91a89e9ab

Review — Pre-LN Transformer: On Layer Normalization in the Transformer Architecture

https://arxiv.org/pdf/2002.04745.pdf