ReZero is All You Need: Fast Convergence at Large Depth

딥러닝의 발전에 있어서 중요한 발견 중 한가지는 ResNet의 발명이라고 할 수 있습니다. 매우 deep한 뉴럴넷을 학습시키기위해서 input의 정보를 그대로 output에 더해줌으로써 모델은 input과 output의 "차이"에 해당하는 부분만 학습할 수 있도록 해주는 것 입니다. 실제로 Residual Connection을 사용했을 대 vanishing gradient problem도 어느정도 해결이 된다고 알려져 있습니다. 하지만 시간이 지날수록 점점 더 크고 깊은 모델을 사용하게 되면서, Residual Connection만으로는 부족하다고 느껴지고 이를 개선한 여러가지 모델들과 방법론이 등장하게 됩니다. 예를 들어서, Batch Normalization이나 Layer Normalization과 같은 적절한 정규화를 통해 covariance shift를 줄임으로써 매우 깊은 모델의 훈련 속도를 향상시키기도 하고, 요즘 자연어처리에서 빼놓을 수 없는 존재인 Transformer에도 Residual connection이 사용됩니다. 

 

이 논문에서는, 아주 간단한 방법을 통해서 ResNet을 효과적으로 개선할 수 있는 방법인 ReZero(Residual with zero initialization)를 제안합니다. (뭔가를 떠올리신 분이 있을수도 있는데, 아마 아닐겁니다..아마) 이는 Residual Connection에서 layer의 ouput에 residual parameter라는 숫자 하나를 곱해주는것인데, 이를 0으로 initialize하고 모델을 training하는 동안 이 역시 학습될 수 있도록 합니다. 식으로 나타내면 아래와 같습니다.

기존의 Residual Connection은 alpha = 1이면서 학습을 시키지 않고 고정하는 경우에 해당한다고 볼 수 있습니다. ReZero의 경우 처음에는 identity function으로써 Residual Connection Layer가 initialize되고, 학습이 진행되면서 alpha값에 따라 Layer F의 output이 반영되는 것 입니다. 이게 끝입니다! 결과적으로는 Residual Connection이 사용되는 곳 마다 parameter가 딱 하나씩 늘어나기 때문에 전체 paramter 숫자도 거의 늘어나지 않습니다. 기존의 Residual connection들과 variation들을 비교하면 아래와 같습니다.

ReZero와 다른 Residual Connection들과의 비교, Bachlechner et al. 2020

저자들이 주장하는 ReZero의 장점은 다음과 같습니다.

 

1. Widely Applicable: 정말 간단하고 구현 역시 어렵지 않습니다. 그리고 앞에서 언급했듯이 파라미터의 숫자도 거의 늘어나지 않습니다.

2. Deeper Learning: ReZero는 매우 deep한 network를 훈련시킬 수 있다는 장점이 있습니다. 뒤의 실험에서 더 자세히 설명하겠지만, 10000개짜리 layer의 fully-connected network를 훈련시키는 것이 가능하다고 합니다. 이는 기존의 Vanilla ResNet을 포함한 다른 variation으로는 불가능합니다.

3. Faster Convergence: 수렴 속도가 매우 빠른 편입니다. 예를 들어서, enwik8 benchmark에서 vanilla transformer에 비해 두배정도 빠른 수렴 속도를 보입니다. (성능은 1.2BPT로 vanilla에 크게 뒤지지 않습니다.)

 

실험은 toy example과 CIFAR-10 image classification, 그리고 enwik-8에서의 character-level language modeling이 이루어졌는데, 여기서는 뒤의 둘에 대해서만 설명하도록 하겠습니다. CIFAR-10에서의 실험의 경우, 먼저 32개의 layer를 가진 fully connected network를 CIFAR-10에 훈련시켰습니다. 아래 그래프에서 볼 수 있듯이, ReZero를 사용했을 때의 수렴 속도가 다른 baseline (그냥 FC network, ResNet, LayerNorm)에 비해서 눈에 띄게 빠른 것을 알 수 있습니다. 또한, GPU 한개만을 사용해서 10000 layer짜리 FC network를 훈련시키는것도 가능하다고 언급하고 있습니다. (아마 다른 baseline들은 너무 느리거나 수렴 자체를 안하는 것일텐데, 이에 대해서는 자세히 나와있지 않습니다.)

 

32-layer FC network를 통한 수렴 속도 비교 실험, Bachlechner et al. 2020

또한, 기존의 ResNet(CNN이 사용된 그것)과 그 variation들(Gated ResNet, Zero gamma, FixUp)과의 수렴 속도 및 성능 비교도 진행하였고, 아래 표에서 볼 수 있듯이 성능과 속도 면에서 모두 ReZero가 다른 baseline들에 비해서 뛰어나다는 것을 알 수 있습니다. 또한, 이 표에 나와있지는 않지만 Super-Convergence와도 비교를 했는데 94%의 정확도에 도달하는데에 걸리는 시간이 60epoch에서 45epoch으로 줄어들었다고 합니다.

 

ResNet과 그 variation들, 그리고 ReZero와의 성능 및 수렴 속도 비교. Bachlechner et al. 2020

Transformer의 경우, Vanilla Transformer에서 LayerNorm을 제거하고 기존의 Residual connection layer대신 ReZero를 이용한, ReZero Transformer를 바탕으로 language modeling 실험을 하였습니다. 먼저, 기존의 여러가지 Normalization방법(Pre-Norm, Post-Norm, GPT2-Norm 등, 기본적으로 LayerNorm과 SubLayer, Residual Connection의 순서를 바꾼 variation들이고 vanilla transformer는 Post-Norm에 해당합니다.)과 비교했을 때 1.2BPT에 도달하는 속도가 50%정도 빨라지는 것을 볼 수 있습니다. 또한, 초기 alpha값을 0으로 두는 것이 1로 두는 것에 비해서 훨씬 좋다는 것도 확인했습니다(Table 3).

 

또한, 앞에서 FC network의 layer를 10000개까지 늘리는데에 성공했으니, transformer역시 비슷한걸 시도해볼 수 있습니다. 이 논문에서는 중간 layer에서의 output에 대한 예측이나 바로 다음이 아닌 더 이후의 character에 대한 예측까지 포함하는 auxiliary loss를 도입함으로써 기존의 transformer가 character-level language modeling에서 좋은 성능을 보이지 못한 점을 보완하고 더 많은 layer를 사용할 수 있도록 했는데, 이렇게 만들어진 모델(논문에서는 Char-TX라고 표기하고 있습니다)과 ReZero Transformer(auxiliary loss가 없는)를 비교했을 때, 추가적인 regularization 없이도 Character Transformer보다 더 적은 파라미터를 이용해서 비슷한 성능을 보여줍니다(Table 4).

ReZero Transformer의 구조, Bachlechner et al. 2020

enwik8 dataset에서의 실험 결과, Bachlechner et al. 2020

여기에 추가적으로, ReZero Transformer를 학습시킬 때 alpha(의 절댓)값의 변화를 layer별로, 그리고 시간별로 그 변화를 관찰하였는데, heatmap(L=64인 경우)에서 볼 수 있듯이 초반에는 alpha가 빠르게 증가했다가 서서히 줄어들고, 특히 위쪽 layer로 갈수록 그 값이 더 컸다고 합니다. 

 

요약하자면, 매우 간단한 방법으로 Residual Connection의 수렴 속도를 빠르게 할 수 있는 아이디어를 제시한 논문이라고 할 수 있습니다. 논문에 꽤 자주 등장하는 dynamical isometry에 관한 내용은 제 지식 부족으로 잘 이해를 하지 못해서 언급하지 않았는데, 혹 알게 되면 나중에 추가하도록 하겠습니다.