[논문 리뷰] Stop Wasting My Time! Saving Days of ImageNet and BERT Training with Latest Weight Averaging

Jean Kaddour|arXiv (Cornell University)|2022. 09. 29.

Topic Modeling인용 수 4

한 줄 요약

이 논문은 각 훈련 에포크의 끝에서 수집한 k개의 최신 모델 체크포인트의 가중치를 평균화하는 단순한, 모델 및 옵티마이저에 관계없는 방법인 Latest Weight Averaging (LAWA)을 제안한다. LAWA는 훈련 수렴 속도를 높여주며, 훈련 루프를 체크포인트 수집 및 평균화 외에 수정하지 않고도 ImageNet에서 최대 68 GPU 시간, WikiText-103에서 최대 30 GPU 시간의 훈련 시간을 단축시킨다. 이는 목표 손실 및 정확도 수준에 도달하는 데 수십 개의 에포크 전에 모델을 도달시켜 주며, 모델 아키텍처나 옵티마이저, 하이퍼파라미터를 수정하지 않는다.

ABSTRACT

Training vision or language models on large datasets can take days, if not weeks. We show that averaging the weights of the k latest checkpoints, each collected at the end of an epoch, can speed up the training progression in terms of loss and accuracy by dozens of epochs, corresponding to time savings up to ~68 and ~30 GPU hours when training a ResNet50 on ImageNet and RoBERTa-Base model on WikiText-103, respectively. We also provide the code and model checkpoint trajectory to reproduce the results and facilitate research on reusing historical weights for faster convergence.

연구 동기 및 목표

• 대규모 비전 및 언어 모델의 높은 계산 비용과 긴 훈련 시간 문제를 해결하기 위해, 고성능 하드웨어에 접근할 수 없는 연구자들에게도 유용하도록 하기 위해.
• 모델 아키텍처, 옵티마이저, 하이퍼파라미터를 수정하지 않고도 훈련 수렴 속도를 향상시키기 위해.
• 손실 경관이 여전히 변화하는 중인 훈련 중기 단계에서 최근 모델 가중치를 평균화하는 것이 수렴 속도를 가속화할 수 있는지 탐색하기 위해.
• 딥러닝 연구에서 더 빠른 반복적 실험을 가능하게 하는 실용적이고 즉시 사용 가능한 방법을 제공하기 위해.

제안 방법

• LAWA는 각 훈련 에포크의 끝에서 모델 가중치를 수집하고, 큐에 저장된 k개의 최신 체크포인트를 유지하는 슬라이딩 윈도우를 운영한다.
• 각 에포크마다, LAWA 앙상블의 모델 파라미터는 k개의 최신 체크포인트의 균일 평균을 통해 업데이트된다: θ^LAWA_E = (1/k) * Σ_{i=E-k+1}^{E} θ_i.
• 시간 복잡도를 줄이기 위해 원형 버퍼를 사용하여 평균 연산을 효율적으로 구현하며, 평균화된 모델은 필요할 때만 평가되어 메모리 및 계산 오버헤드를 최소화한다.
• 평균화된 모델의 배치 정규화 통계는 훈련 데이터를 전체적으로 통과시켜 추론하는 방식으로 재계산되어 일관된 배치 통계를 확보한다.
• 이 방법은 모델 아키텍처, 옵티마이저(예: SGD 또는 Adam), 작업에 관계없이 적용 가능하므로, 비전 및 NLP 벤치마크 전반에 걸쳐 널리 적용 가능하다.
• 하이퍼파라미터인 k(최신 체크포인트 수)는 경험적으로 튜닝되었으며, k=6이 ImageNet 및 WikiText-103 실험 전반에서 뛰어난 성능을 보였다.

실험 결과

연구 질문

• RQ1 훈련 중기 단계에서 k개의 최신 체크포인트를 평균화하는 것이 손실 및 정확도 측면에서 수렴 속도를 크게 향상시킬 수 있는가?
• RQ2 LAWA는 대규모 비전 및 언어 모델에서 표준 SGD 또는 Adam 훈련보다 시간-정확도 측면에서 우수한 성능을 보일 수 있는가?
• RQ3 LAWA의 성능는 k의 선택에 얼마나 민감한가? 그리고 다양한 훈련 단계에서 특정 k 값에서 최적의 성능를 나타내는가?
• RQ4 LAWA는 아키텍처나 옵티마이저 수정 없이도 컴퓨터 비전(ResNet50/ImageNet)과 NLP(RoBERTa-Base/WikiText-103) 양쪽에 효과적으로 적용될 수 있는가?
• RQ5 LAWA는 최종 모델 일반화 성능를 향상시키는가, 아니면 순수하게 수렴 속도를 가속화하는 기법인가?

주요 결과

• LAWA는 ResNet50/ImageNet에서 표준 SGD 훈련 대비 목표 검증 손실에 도달하는 데 최대 68 GPU 시간을 단축시켰다.
• RoBERTa-Base/WikiText-103에서 LAWA는 Adam의 최종 검증 손실에 약 45 에포크 전에 도달시켜 약 30 GPU 시간을 절약했다.
• LAWA 평균화 모델의 최종 일반화 성능는 최종 Adam 최적화 모델보다 항상 뛰어나게 나타나, 일반화 성능 향상이 확인되었다.
• LAWA는 k가 중간 정도일 때(예: k=6) 가장 효과적이었으며, k가 16을 초과하면 성능이 악화되어 과거 체크포인트를 너무 많이 평균화하는 것이 해로울 수 있음을 시사했다.
• 이 방법은 다양한 옵티마이저와 작업에 대해 강건하며, 비전 및 언어 모델링 벤치마크에서 일관된 수렴 속도 향상이 관찰되었다.
• 경험적 결과는 초기 훈련 단계에서의 평균화가 성능를 악화시킨다는 것을 확인하였으며, 이는 LAWA가 손실 경관이 여전히 변화하는 중기 단계에서 가장 효과적임을 뒷받침한다.