[논문 리뷰] GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training

Chen Zhu, Renkun Ni외 4명|arXiv (Cornell University)|2021. 02. 16.

Advanced Neural Network Applications인용 수 15

한 줄 요약

GradInit는 신경망 레이어의 최적 초기화 스케일링 계수를 한 번의 SGD 또는 Adam 최적화 단계 후 손실을 최소화함으로써 학습하는 아키텍처에 종속되지 않는 방법이다. 이는 수렴 속도를 가속화하고 테스트 정확도를 향상시키며, 트랜스포머의 학습률 웜업 없이도 학습이 가능하게 한다. ImageNet 및 IWSLT 번역 벤치마크에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Innovations in neural architectures have fostered significant breakthroughs in language modeling and computer vision. Unfortunately, novel architectures often result in challenging hyper-parameter choices and training instability if the network parameters are not properly initialized. A number of architecture-specific initialization schemes have been proposed, but these schemes are not always portable to new architectures. This paper presents GradInit, an automated and architecture agnostic method for initializing neural networks. GradInit is based on a simple heuristic; the norm of each network layer is adjusted so that a single step of SGD or Adam with prescribed hyperparameters results in the smallest possible loss value. This adjustment is done by introducing a scalar multiplier variable in front of each parameter block, and then optimizing these variables using a simple numerical scheme. GradInit accelerates the convergence and test performance of many convolutional architectures, both with or without skip connections, and even without normalization layers. It also improves the stability of the original Transformer architecture for machine translation, enabling training it without learning rate warmup using either Adam or SGD under a wide range of learning rates and momentum coefficients. Code is available at https://github.com/zhuchen03/gradinit.

연구 동기 및 목표

• 나이드한 초기화로 인한 깊은 신경망의 학습 불안정성 문제, 특히 복잡하거나 새로운 아키텍처에서의 문제를 해결한다.
• 새로운 또는 이질적인 아키텍처 간에 이식성이 떨어지는 아키텍처에 특화된 초기화 규칙의 한계를 극복한다.
• 아키텍처 수정 없이도 학습 안정성과 수렴 속도를 향상시키는 일반적인 목적의 최적화기 인식 초기화 방법을 개발한다.
• SGD 또는 Adam을 사용할 때도 학습률 웜업 없이 1202층의 ResNet과 포스트 레이어 노멀라이제이션(Pre-LN) 트랜스포머를 학습할 수 있도록 한다.
• 전체 학습 시간의 1% 미만을 소비하여 최적의 스케일링 계수를 학습하는 확장성 있고 효율적인 초기화 기법을 제공한다.

제안 방법

• 각 파라미터 블록(예: 컨볼루션 커널 또는 선형 레이어)에 스케일링 계수를 도입하여 무작위로 초기화된 가중치를 재스케일링한다.
• 지정된 최적화기(SGD 또는 Adam)의 한 단계 후 손실을 최소화하도록 이러한 스케일링 계수를 수치적 방법으로 최적화한다.
• 최적화 단계를 미분 가능한 과정으로 간주하여 첫 번째 업데이트를 통해 역전파를 통해 최적의 스케일링 계수를 학습한다.
• 미니배치 샘플링의 확률성, 학습률, 최적화기 동역학(예: 모멘타ム 및 적응형 학습률)을 고려한다.
• 표준 학습을 시작하기 전에 학습된 스케일링 계수를 사용해 원래의 무작위 가중치를 재스케일링한다.
• Adam 및 SGD 최적화기를 모두 지원하며, 각각의 최적화기 동역학에 맞는 별도의 스케일링 패턴을 학습한다.

실험 결과

연구 질문

• RQ1 단일 자동화된 방법이 다양한 아키텍처에서 학습 안정성과 수렴 속도를 향상시키는 초기화 스케일링 계수를 학습할 수 있는가?
• RQ2 GradInit는 Adam 또는 SGD로 학습하는 트랜스포머에서 학습률 웜업의 필요성을 어느 정도 제거할 수 있는가?
• RQ3 기존 초기화 기법(Xavier, He, Admin 등)과 비교해 GradInit는 최종 테스트 정확도와 학습 안정성 측면에서 어떻게 성능을 내는가?
• RQ4 GradInit는 표준 초기화로 실패하는 1202층의 매우 깊은 네트워크(예: ResNet)를 안정적으로 학습시킬 수 있는가?
• RQ5 학습된 스케일링 계수의 패턴은 기울기 분산과 초기화에 기인한 불안정성에 대한 통찰을 어떻게 드러내는가?

주요 결과

• GradInit를 사용하면 IWSLT-14 DE-EN에서 Post-LN 트랜스포머를 학습률 웜업 없이 Adam 또는 SGD로 학습할 수 있으며, 극단적인 초모수 설정에서도 평균 BLEU 점수 36.0을 달성한다.
• GradInit를 사용하면 1202층의 ResNet이 ResNet-110보다 훨씬 높은 테스트 정확도를 달성하여, 표준 초기화로는 실패하는 매우 깊은 네트워크의 학습에 성공함을 보여준다.
• ImageNet에서 GradInit는 배치 정규화 없이도 ResNet-50의 최종 테스트 정확도를 향상시켜 더 나은 초기화를 학습함을 입증한다.
• GradInit는 대부분의 레이어에서 초기 기울기 분산을 감소시키며, 특히 인코더와 잔차 브랜치에서 두드러지게 작용하고, 스키프 커넥션의 균형을 맞추기 위해 출력 프로젝션 및 피드포워드 네트워크(FFN) 레이어의 스케일링을 낮춘다.
• GradInit는 Adam과 SGD에 대해 별개의 스케일링 패턴을 학습함으로써 최적화기 특화 초기화가 기울기 분산 제어에 매우 중요하다는 것을 시사한다.
• 시각화 결과는 GradInit가 초기화 단계에서 네트워크를 더 매끄러운 손실 영역로 재스케일링하여, 샘플 간 기울기 분산을 감소시키고 수렴 속도를 가속화함을 보여준다.