[논문 리뷰] Stochastic Weight Averaging in Parallel: Large-Batch Training that Generalizes Well

Vipul Gupta, Santiago Akle Serrano외 1명|arXiv (Cornell University)|2020. 01. 07.

Advanced Neural Network Applications참고 문헌 26인용 수 15

한 줄 요약

이 논문은 깊이 신경망의 대용량 배치 학습을 가속화하기 위해 병렬에서 무작위 가중치 평균화(SWAP)를 제안한다. 이 방법은 먼저 대용량 미니배치를 사용해 신속하게 수렴하는 모델을 훈련한 후, 독립적으로 훈련된 소용량 미니배치 모델들의 가중치 평균을 통해 정밀도를 향상시킨다. SWAP는 소용량 배치 학습 수준의 일반화 성능를 달성하면서도 훨씬 더 짧은 시간에 학습을 완료하여, CIFAR10에서 최신 기술 대비 최대 68%의 학습 시간 단축을 이룬다.

ABSTRACT

We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm to accelerate DNN training. Our algorithm uses large mini-batches to compute an approximate solution quickly and then refines it by averaging the weights of multiple models computed independently and in parallel. The resulting models generalize equally well as those trained with small mini-batches but are produced in a substantially shorter time. We demonstrate the reduction in training time and the good generalization performance of the resulting models on the computer vision datasets CIFAR10, CIFAR100, and ImageNet.

연구 동기 및 목표

• 대용량 배치 딥 뉴럴 네트워크 학습에서 학습 속도와 일반화 성능 사이의 상충 관계를 해결하기 위해.
• 대용량 배치 초기화 이후 소용량 배치 모델의 병렬 학습을 활용해 일반화 성능를 희생시키지 않고 더 빠른 학습을 가능하게 하기 위해.
• 대규모 DNN 학습을 위한 계산 자원을 효율적으로 활용할 수 있는 실용적이고 최적화가 적은 방법을 개발하기 위해.
• 대용량 배치 학습 후 병렬 소용량 배치 보정을 통해 소용량 배치 학습 수준의 일반화 성능를 갖는 모델를 얻을 수 있음을 입증하기 위해.

제안 방법

• 방법은 먼저 대용량 미니배치를 사용해 낮은 훈련 손실 영역에 신속히 도달할 수 있도록 모델을 훈련한다.
• 그 후, 대용량 배치 모델에서 시작하여 소용량 미니배치를 사용해 병렬로 다수의 독립된 모델을 훈련한다.
• 최종 모델은 이러한 독립적으로 훈련된 소용량 배치 모델들의 가중치 평균을 통해 구성된다.
• 소용량 배치 단계에서 평균화된 최소값에 수렴하도록 유도하기 위해 순환 학습률 스케줄을 사용한다.
• 대용량 배치에서 소용량 배치로의 전환은 하이퍼파ram터에 의해 제어되며, 실험에서는 그리드 서치를 통해 선택된다.
• 이 방법은 다양한 최적화 기법과 호환되며, 통신 오버헤드를 줄이기 위해 로컬 SGD와도 조합할 수 있다.

실험 결과

연구 질문

• RQ1 일반화 성능를 훼손시키지 않고 대용량 배치 학습을 가속화할 수 있는가?
• RQ2 독립적으로 훈련된 소용량 배치 모델들의 가중치 평균화가 표준 대용량 배치 학습보다 일반화 성능를 향상시키는가?
• RQ3 대용량 배치 사전 훈련과 병렬 소용량 배치 미세조정의 조합이 소용량 배치 학습 수준의 일반화 성능를 달성할 수 있는가?
• RQ4 제안된 방법의 학습 시간은 표준 소용량 배치 및 대용량 배치 학습 전략과 비교해 어떻게 되는가?
• RQ5 성능와 효율성을 최적화하기 위해 대용량 배치와 소용량 배치 단계 사이에 임계 전이 지점이 존재하는가?

주요 결과

• SWAP는 CIFAR100에서 241.54초 만에 테스트 정확도 79.11%를 달성했으며, 이는 소용량 배치 SWA와 동일한 성능이지만 3.5배 빠른 속도로 달성된 것이다.
• CIFAR10에서 SWAP는 DAWNBench 경쟁 최종 우승자 대비 학습 시간을 68% 단축하면서도 유사한 일반화 성능를 확보했다.
• 대용량 배치 학습 후 소용량 배치 SWA를 수행한 경우, SWAP보다 3배 이상 오랜 시간이 소요되었으며, 이는 유사한 테스트 정확도(78.12% 대 78.18%)를 달성하기 위해 소요된 시간이었다.
• 40개의 소용량 배치 에포크를 거친 SWAP는 CIFAR100에서 241.54초 만에 79.11%의 테스트 정확도를 기록했으며, 이는 대용량 배치 SWA(76.00%)와 대용량 배치 후 소용량 배치 SWA(78.12%)보다 빠르고 정확도 면에서도 뛰어났다.
• SWAP의 최종 평균 모델은 표준 SGD 모델보다 손실 분지의 중심에 더 가까이 위치해 있어 더 높은 강건성을 가지는 것으로 나타났다.
• 시각화 결과는 대용량 배치 모델이 소용량 배치 모델들이 수렴하는 영역과 겹치는 분지에 도달한다는 것을 보여주며, 거의 볼록한 분지의 존재를 지지한다.