[논문 리뷰] Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints

Aran Komatsuzaki, Joan Puigcerver외 7명|arXiv (Cornell University)|2022. 12. 09.

Domain Adaptation and Few-Shot Learning인용 수 12

한 줄 요약

이 논문은 기존에 훈련된 밀집형 Transformer 모델을 더 강력한 전문가의 조합(MoE) 모델로 전환하는 '희박한 업싸이클링(sparse upcycling)'을 소개한다. 이 방법은 밀집형 체크포인트에서 MoE 레이어를 초기화하여, 원래의 전처리 계산 예산의 약 50%만으로 SuperGLUE와 ImageNet에서 최신 기술 수준의 성능을 달성한다. 이는 원래의 밀집형 모델과 완전히 새로 훈련된 MoE 모델을 모두 능가한다.

ABSTRACT

Training large, deep neural networks to convergence can be prohibitively expensive. As a result, often only a small selection of popular, dense models are reused across different contexts and tasks. Increasingly, sparsely activated models, which seek to decouple model size from computation costs, are becoming an attractive alternative to dense models. Although more efficient in terms of quality and computation cost, sparse models remain data-hungry and costly to train from scratch in the large scale regime. In this work, we propose sparse upcycling -- a simple way to reuse sunk training costs by initializing a sparsely activated Mixture-of-Experts model from a dense checkpoint. We show that sparsely upcycled T5 Base, Large, and XL language models and Vision Transformer Base and Large models, respectively, significantly outperform their dense counterparts on SuperGLUE and ImageNet, using only ~50% of the initial dense pretraining sunk cost. The upcycled models also outperform sparse models trained from scratch on 100% of the initial dense pretraining computation budget.

연구 동기 및 목표

• 큰 희박한 모델을 처음부터 훈련하는 데 드는 높은 계산 비용 문제를 해결하기 위해 기존의 밀집형 모델 체크포인트를 재사용한다.
• 새로운 데이터나 전체 재훈련 없이도 대규모 모델의 성능 향상을 가능하게 한다.
• 최소한의 추가 계산 자원으로 밀집형 모델을 더 효율적이고 고용량의 MoE 모델로 업그레이드할 수 있는 실용적인 방법을 개발한다.
• 동일한 계산 예산 내에서 밀집형 체크포인트에서 초기화된 MoE 모델이 원래의 밀집형 모델과 완전히 새로 훈련된 MoE 모델을 모두 능가할 수 있음을 입증한다.

제안 방법

• 사전에 훈련된 밀집형 Transformer 체크포인트를 활용하여 MoE 레이어의 온도를 높이는 방식으로 MoE 모델을 초기화한다.
• Transformer 블록 내의 MLP 구성 요소를 전문가 피드포워드 네트워크 집합과 라우터 네트워크로 대체하여 희박한 활성화를 가능하게 한다.
• 전문가 선택 라우팅을 사용하여 각 토큰에 대해 가장 관련성이 높은 전문가를 동적으로 선택함으로써 효율적인 계산을 구현한다.
• 성능 저하를 최소화하기 위해 MoE 구성 요소의 신중한 초기화 및 적응 조치를 포함하는 모델 수술 레시피를 적용한다.
• 원래의 밀집형 모델의 전처리 비용의 약 10%에서 60% 정도의 추가 계산 예산으로 업싸이클된 MoE 모델을 훈련한다.
• 원래의 밀집형 모델과 동일한 데이터와 훈련 스케줄을 사용하여 새로운 데이터를 도입하지 않는다.

실험 결과

연구 질문

• RQ1 기존에 사전 훈련된 밀집형 모델을 최소한의 추가 계산 자원으로 더 강력한 MoE 모델로 효과적으로 업싸이클링할 수 있는가?
• RQ2 밀집형 체크포인트에서 MoE 레이어를 초기화하면 원래의 밀집형 모델를 계속 훈련시키는 것보다 더 높은 성능을 낼 수 있는가?
• RQ3 동일한 총 계산 예산을 제공받을 때, 업싸이클된 MoE 모델이 처음부터 훈련된 MoE 모델보다 성능이 뛰어나게 될 수 있는가?
• RQ4 연속적인 밀집형 훈련에 비해 희박한 업싸이클링을 통해 성능 향상을 달성하기 위해 얼마나 많은 추가 계산 자원이 필요한가?
• RQ5 이 업싸이클링 방법은 자연어 처리와 비전 등 다양한 모델 아키텍처와 도메인에 일반화되는가?

주요 결과

• 업싸이클된 T5-Base, T5-Large, T5-XL 모델들은 원래의 밀집형 모델보다 SuperGLUE에서 1.5~2점의 절대 성능 향상을 달성했으며, 원래의 전처리 계산 예산의 46~55%만을 사용했다.
• 업싸이클된 ViT-Base와 ViT-Large 모델들은 ImageNet 10-shot에서 1%p 이상의 정확도 향상을 달성했으며, 연속적인 밀집형 훈련 대비 13%의 추가 훈련 시간만 필요로 했다.
• 동일한 총 계산 예산을 제공받을 때, 희박한 업싸이클링은 처음부터 훈련된 MoE 모델보다 더 뛰어난 성능을 내며, 사전 훈련된 가중치를 재사용하는 가치를 입증한다.
• ViT-Base의 경우, 13%의 추가 계산 자원으로 ImageNet 10-shot에서 1%p 이상의 성능 향상을 달성했고, 연속적인 밀집형 훈련은 동일한 성능 향상을 위해 58%의 추가 계산 자원이 필요했다.
• 이 방법은 다양한 모델 크기와 도메인에서 효과적이며, 제약 조건이 있는 계산 예산 하에서 언어 및 비전 작업 모두에서 밀집형 모델과 처음부터 훈련된 MoE 기준선을 모두 능가하는 업싸이클된 모델이 성능 향상을 이룬다.
• 업싸이클링 방법은 강건하고 일반화 가능하며, SuperGLUE와 ImageNet을 포함한 다양한 모델 아키텍처와 작업에서 일관된 성능 향상을 보였다.