[논문 리뷰] A Closer Look at Codistillation for Distributed Training

Shagun Sodhani, Olivier Delalleau외 4명|arXiv (Cornell University)|2020. 10. 06.

Topic Modeling참고 문헌 20인용 수 4

한 줄 요약

이 논문은 표준 데이터 병렬 학습의 통신 효율적인 대안으로 codistillation을 제안한다. 여러 모델을 병렬로 학습시키되, 파라미터 평균화 대신 지식 정착 손실을 통해 동기화한다. 2-way codistillation은 2배의 장치를 사용하여 표준 데이터 병렬 학습과 유사한 정확도를 달성하면서도 통신량을 최대 1000배까지 줄일 수 있으나, n>2개 모델으로 확장할 경우 모델 및 데이터셋 특성과 관련된 도전 과제가 발생한다.

ABSTRACT

Codistillation has been proposed as a mechanism to share knowledge among concurrently trained models by encouraging them to represent the same function through an auxiliary loss. This contrasts with the more commonly used fully-synchronous data-parallel stochastic gradient descent methods, where different model replicas average their gradients (or parameters) at every iteration and thus maintain identical parameters. We investigate codistillation in a distributed training setup, complementing previous work which focused on extremely large batch sizes. Surprisingly, we find that even at moderate batch sizes, models trained with codistillation can perform as well as models trained with synchronous data-parallel methods, despite using a much weaker synchronization mechanism. These findings hold across a range of batch sizes and learning rate schedules, as well as different kinds of models and datasets. Obtaining this level of accuracy, however, requires properly accounting for the regularization effect of codistillation, which we highlight through several empirical observations. Overall, this work contributes to a better understanding of codistillation and how to best take advantage of it in a distributed computing environment.

연구 동기 및 목표

• Codistillation이 표준 데이터 병렬 학습의 확장성 있고 통신 효율적인 대안이 될 수 있는지 조사하기 위해.
• 기존의 all-reduce 기반 데이터 병렬 학습과 비교하여 2-way codistillation의 성능 및 통신 효율성을 평가하기 위해.
• 분산 학습에서 2개 모델을 초월해 codistillation을 확장할 때의 제한 사항과 과제를 이해하기 위해.
• Codistillation이 학습 동역학에 미치는 영향, 특히 정규화 효과 유사성에 대해 분석하기 위해.
• n-way codistillation이 2-way codistillation 또는 표준 학습을 능가할 수 있는 조건을 탐색하기 위해.

제안 방법

• Codistillation은 각각 데이터의 부분 집합에서 학습되는 여러 모델을 병렬로 훈련시키며, 모델 간 예측 불일치를 방지하기 위해 지식 정착 손실을 적용한다.
• All-reduce를 통한 파라미터 동기화 대신, 지식 정착을 통한 기능적 일치를 통해 모델을 동기화한다.
• 표준 교차 엔트로피 손실과 대칭적인 정착 손실을 조합한 훈련 목표 함수를 사용하며, 이는 모델 출력 간 일치를 장려한다.
• 전체 기울기나 파라미터가 아닌, 정착된 예측 또는 소프트 레이블만 전송함으로써 통신 오버헤드를 줄인다.
• ResNet50 (ImageNet) 및 트랜스포머 기반 번역 모델에서 평가하였으며, n-way codistillation에 대한 분석 실험을 수행하였다.
• 하이퍼파라미터 α는 정착 손실의 가중치를 제어하며 최적의 성능을 위해 튜닝된다.

실험 결과

연구 질문

• RQ1 2-way codistillation은 표준 데이터 병렬 학습과 유사한 정확도를 달성하면서도 반복마다 통신량을 수십만 배까지 줄일 수 있는가?
• RQ2 모델당 장치 수가 증가함에 따라 codistillation은 어떻게 확장되며, 선형적인 속도 향상이 유지되는가?
• RQ3 왜 2개 모델을 초월해 codistillation을 확장할 경우 다양한 워크로드에서 성능 향상이 일관되게 이루어지지 않는가?
• RQ4 모델 아키텍처와 데이터셋 특성이 n-way codistillation의 효과성에 어떤 영향을 미치는가?
• RQ5 Codistillation은 학습 동역학에 어떤 영향을 미치며, 어떤 정규화 효과를 유도하는가?

주요 결과

• 2-way codistillation은 2배의 장치를 사용할 경우 표준 데이터 병렬 학습과 유사한 정확도를 달성하면서도 반복당 통신량을 최대 1000배까지 줄일 수 있다.
• 이 방법은 모델당 장치 수에 대해 선형 확장이 가능하여 장치 수를 두 배로 늘일 경우 수렴에 필요한 학습 업데이트 수를 절반으로 줄일 수 있다.
• n>2개 모델로 codistillation을 확장할 경우 성능 향상 결과는 혼합된 편이며, 일부 사례에서는 향상이 관찰되지만 항상 일관되지는 않으며, 종종 2-way codistillation 수준 이하이거나 유사한 성능에 머무른다.
• 2개 모델을 초월해 확장하는 데서 발생하는 과제는 모델 아키텍처와 데이터셋 고유의 특성(일반화 행동 및 손실 곡면 등)과 강하게 연관되어 있다.
• Codistillation은 학습 동역학에 영향을 미치며 정규화 효과를 유도하므로, 최적의 성능를 달성하기 위해서는 이를 고려해야 한다.
• 통신량이 감소함에도 불구하고 월클록 속도 향상은 보장되지 않으며, 모델 크기와 통신 대역폭에 따라 달라지며, 더 큰 모델일수록 더 큰 이점을 얻는다.