[논문 리뷰] Sparse MoE as the New Dropout: Scaling Dense and Self-Slimmable Transformers

Tianlong Chen, Zhenyu Zhang외 3명|arXiv (Cornell University)|2023. 03. 02.

Domain Adaptation and Few-Shot Learning인용 수 6

한 줄 요약

이 논문은 고정된 무작위로 초기화된 라우터를 사용하여 훈련 중에 활성화되는 전문가 수를 점진적으로 증가시키는 플러그 앤 플레이 훈련 프레임워크인 SMoE-Dropout을 제안한다. 이를 통해 트랜스포머 모델이 붕괴 없이 더 높은 정확도를 달성할 수 있다. 이 방법은 '자기 스러임메이블(self-slimmable)' 성질을 유도하여 사용 가능한 자원에 따라 추론 및 미세조정 시 부드러운 성능 향상을 가능하게 하며, 추론 작업에서 BERT 모델이 밀집형 대비 최대 1.09% 높은 성능을 기록한다.

ABSTRACT

Despite their remarkable achievement, gigantic transformers encounter significant drawbacks, including exorbitant computational and memory footprints during training, as well as severe collapse evidenced by a high degree of parameter redundancy. Sparsely-activated Mixture-of-Experts (SMoEs) have shown promise to mitigate the issue of training efficiency, yet they are prone to (1) redundant experts due to representational collapse; and (2) poor expert scalability for inference and downstream fine-tuning, primarily due to overfitting of the learned routing policy to the number of activated experts during training. As recent research efforts are predominantly focused on improving routing policies to encourage expert specializations, this work focuses on exploring the overlooked scalability bottleneck of SMoEs and leveraging it to effectively scale dense transformers. To this end, we propose a new plug-and-play training framework, SMoE-Dropout, to enable scaling transformers to better accuracy in their full capacity without collapse. Specifically, SMoE-Dropout consists of a randomly initialized and fixed router network to activate experts and gradually increases the activated expert number as training progresses over time. Transformers trained by SMoE-Dropout naturally exhibit a self-slimmable property subject to resource availability, offering smooth and consistent performance boosts with an increase in activated experts during inference or fine-tuning. Our extensive experiments demonstrate the superior performance and substantial computation savings of SMoE-Dropout, compared to dense training baselines with equivalent parameter counts. In particular, our trained BERT outperforms its densely trained counterpart with consistent improvements of {1.03%, 0.78%, 1.09%} on challenging reasoning tasks {ASDiv-A, MAWPS, SVAMP}, respectively.

연구 동기 및 목표

• 추론 및 미세조정 시 희소 전문가 혼합(SMoE) 모델에서의 표현 붕괴와 낮은 확장성 문제를 해결하기 위해.
• 고정된 전문가 수에 대한 라우팅 정책의 과적합으로 인해 발생하는 SMoE 훈련의 간과된 확장성 병목 현상을 탐색하기 위해.
• 붕괴 없이 전체 용량의 모델 확장을 가능하게 하는 플러그 앤 플레이 방법을 개발하여 전문가 사용 증가에 따라 일관된 성능 향상을 보장하기 위해.
• 사용 가능한 자원에 따라 동적으로 성능 조정이 가능한 자원 인식 추론 및 미세조정 능력을 제공하기 위해 '자기 스러임메이블' 트랜스포머 아키텍처를 구현하기 위해.

제안 방법

• 학습 중에 업데이트되지 않는 고정된 무작위로 초기화된 라우터 네트워크를 도입하여 학습 가능한 라우팅 메커니즘을 대체한다.
• 훈련 과정 전반에 걸쳐 활성화되는 전문가 수(k)를 점진적으로 증가시켜 점진적인 용량 확장을 가능하게 한다.
• SMoE 모듈을 트랜스포머 아키텍처의 후행 레이어에 삽입하여 전반적인 성능 향상이 전반적으로 더 우수하다는 것을 경험적으로 입증한다.
• 특정 트랜스포머 레이어에 SMoE 블록을 모듈식으로 삽입하며, 분석 결과 후행 레이어에 삽입할 경우가 가장 효과적임을 확인한다.
• 과적합을 방지하고 암묵적 정규화를 통해 전문가 전문화를 촉진하기 위해 무작위 라우팅 정책을 활용한다.
• 사용 가능한 자원에 따라 활성화되는 전문가 수를 조정할 수 있도록 하여 추론 및 미세조정 시 부드러운 성능 스케일링을 가능하게 한다.

실험 결과

연구 질문

• RQ1 고정된 무작위 라우팅 정책이 희소 전문가 혼합(SMoE) 모델에서 표현 붕괴를 완화하고 확장성을 향상시킬 수 있는가?
• RQ2 훈련 중에 점진적으로 활성화되는 전문가 수를 증가시키는 것이 추론 및 미세조정 시 더 나은 일반화 및 성능 향상에 기여하는가?
• RQ3 과적합 및 모델 용량 활용 측면에서 SMoE-Dropout은 학습 가능한 라우팅 및 밀집형 훈련과 비교해 어떻게 성능을 내는가?
• RQ4 트랜스포머에 SMoE 모듈을 삽입할 최적의 위치는 어디인가? 성능과 확장성 측면에서 최대화하기 위해.
• RQ5 자원 제약이 있는 환경에서 표준 SMoE 또는 밀집형 모델보다 더 잘 희석 가능한(SMooE-Dropout) 모델을 생성할 수 있는가?

주요 결과

• SMoE-Dropout로 훈련된 BERT 모델은 각각 ASDiv-A, MAWPS, SVAMP 추론 작업에서 밀집형 대비 1.03%, 0.78%, 1.09%의 일관된 성능 향상을 기록한다.
• 이 방법은 추론 및 미세조정 중 활성화되는 전문가 수에 따라 성능이 부드럽게 향상되는 '자기 스러임메이블' 성질을 유도하여 자원 기반 동적 효율-성능 트레이드오프를 가능하게 한다.
• enwik8 테스트 세트에서 SMoE-Dropout는 16개의 전문가를 사용하여 1자리당 비트 수(Bits-per-Character, BPC) 2.96×10⁻²를 기록하며, 완전히 무작위 및 결정적 해시 라우팅 기반선보다 뚜렷한 성능 우월성을 보였다.
• 학습 가능한 SMoE 모델이 과적합으로 인해 성능 저하를 보이는 반면, SMoE-Dropout 모델은 12층 또는 16개의 전문가에서도 과적합의 징후가 전혀 없이 안정적인 성능을 유지한다.
• 희석 실험 결과, SMoE-Dropout 모델은 SST-2 작업에서 단일 전문가 모델로 희석될 때 밀집형 또는 표준 SMoE 모델 대비 0.76%에서 1.89% 높은 정확도를 기록하였다.
• 분석 실험 결과 점진적으로 k를 증가시키는 것이 자기 스러임메이블 성질을 확보하는 데 필수적이며, 이 스케줄을 비활성화하면 모든 파rameter를 사용해도 성능이 열 劣화됨을 확인하였다.