[논문 리뷰] ReZero is All You Need: Fast Convergence at Large Depth

Thomas Bachlechner, Bodhisattwa Prasad Majumder외 3명|arXiv (Cornell University)|2020. 03. 10.

Domain Adaptation and Few-Shot Learning참고 문헌 39인용 수 150

한 줄 요약

이 논문은 ReZero를 제안한다. 0으로 초기화된 잔차 게이트로 초기 다이나믹 아이소메트리를 가능하게 하여 매우 깊은 네트워크의 학습과 수렴 속도를 크게 개선한다. FC 네트, ResNet, Transformer에서의 학습에 적용 가능하고 수렴이 빨라진다.

ABSTRACT

Deep networks often suffer from vanishing or exploding gradients due to inefficient signal propagation, leading to long training times or convergence difficulties. Various architecture designs, sophisticated residual-style networks, and initialization schemes have been shown to improve deep signal propagation. Recently, Pennington et al. used free probability theory to show that dynamical isometry plays an integral role in efficient deep learning. We show that the simplest architecture change of gating each residual connection using a single zero-initialized parameter satisfies initial dynamical isometry and outperforms more complex approaches. Although much simpler than its predecessors, this gate enables training thousands of fully connected layers with fast convergence and better test performance for ResNets trained on CIFAR-10. We apply this technique to language modeling and find that we can easily train 120-layer Transformers. When applied to 12 layer Transformers, it converges 56% faster on enwiki8.

연구 동기 및 목표

• 경사 소실/발산으로 인해 매우 깊은 네트워크를 학습하는 어려움을 동기부여한다.
• 초기 다이나믹 아이소메트리를 강제하기 위한 간단하고 아키텍처에 구애받지 않는 수정 제안.
• 다수의 아키텍처(FCN, CNN/ResNet, Transformers)에서 큰 깊이의 학습 능력과 더 빠른 수렴을 시연.
• CIFAR-10 및 enwiki8과 같은 언어 모델링 벤치마크에서 실험적 이점을 보인다.

제안 방법

• 각 층에 잔차 게이트 alpha를 추가하고 alpha를 0으로 초기화한다: x_{i+1}=x_i+alpha_i F(x_i).
• 이 초기화가 시작 시 네트워크를 항등 함수처럼 작동하게 만들어 다이나믹 아이소메트리를 달성함을 시연한다.
• ReZero를 Fully connected networks, convolutional ResNets, 및 Transformer 아키텍처에 적용한다.
• 표준 잔차 및 정규화 기반 접근법과 수렴 속도 및 최종 성능을 비교한다.
• 학습 중 신호 전파가 유지되는 것을 보여주기 위해 Jacobian 스펙트럼을 분석한다.

실험 결과

연구 질문

• RQ1 최소한의 제로 초기화 잔차 게이트가 정규화 없이도 매우 깊은 네트워크의 안정적인 학습을 가능하게 하까?
• RQ2 ReZero가 FCN, CNN/ResNets, Transformer 전반에서 수렴 속도와 테스트 성능에 어떻게 영향을 미치는가?
• RQ3 alpha를 0으로 초기화하는 것이 1로 초기화하는 것 대비 학습 역학에 어느 정도 영향을 미치는가?
• RQ4 현대 아키텍처에 ReZero를 적용할 때 깊이의 실용적 한계와 깊이 관련 성능 향상은 무엇인가?
• RQ5 깊은 모델에서 warm-up이나 강한 정규화의 필요성을 줄이거나 제거하는가?

주요 결과

• ReZero는 수천 층의 네트워크 학습을 가능하게 한다(예: 10,000-layer FC networks).
• 100층이 넘는 Transformer는 ReZero 하에서 LayerNorm이나 워밍업 없이 학습할 수 있다.
• enwiki8에서 ReZero가 적용된 12-layer Transformer는 일반 Transformer에 비해 56% 빠르게 1.2 BPB에 수렴한다.
• CIFAR-10에서 ReZero를 적용한 ResNet 변형은 더 빠른 수렴과 더 나은 테스트 정확도를 보인다(예: ResNet-110 with ReZero는 향상된 지표를 보임).
• ReZero는 여러 ResNet 깊이에서 목표 정확도에 도달하는 반복 수를 최대 약 32%까지 단축한다.
• 더 깊은 Transformer(64, 128 layers)에서 ReZero는 Post-Norm/GPT2-Norm 배치보다 다이나믹 아이소메트리를 더 잘 유지하고 일반 Transformer가 불안정해지는 곳에서도 학습을 가능하게 한다.