[논문 리뷰] Task-Specific Expert Pruning for Sparse Mixture-of-Experts

Tianyu Chen, Shaohan Huang외 6명|arXiv (Cornell University)|2022. 06. 01.

Indoor and Outdoor Localization Technologies인용 수 9

한 줄 요약

이 논문은 다운스트림 작업을 위한 효율적인 단일 전문가 밀도 모델로 희박한 믹스처 오브 전문가(MoE) 모델을 변환하기 위해 작업별 전문가 프루닝을 제안한다. 기여도 점수를 바탕으로 피나이팅 중 비전문가 전문가를 점진적으로 프루닝함으로써, 이 방법은 MoE 성능의 99.3%를 유지하면서도 인퍼런스 속도를 2배 빠르게 하고, 상호 장치 간 통신 비용을 제거한다.

ABSTRACT

The sparse Mixture-of-Experts (MoE) model is powerful for large-scale pre-training and has achieved promising results due to its model capacity. However, with trillions of parameters, MoE is hard to be deployed on cloud or mobile environment. The inference of MoE requires expert parallelism, which is not hardware-friendly and communication expensive. Especially for resource-limited downstream tasks, such sparse structure has to sacrifice a lot of computing efficiency for limited performance gains. In this work, we observe most experts contribute scarcely little to the MoE fine-tuning and inference. We further propose a general method to progressively drop the non-professional experts for the target downstream task, which preserves the benefits of MoE while reducing the MoE model into one single-expert dense model. Our experiments reveal that the fine-tuned single-expert model could preserve 99.3% benefits from MoE across six different types of tasks while enjoying 2x inference speed with free communication cost.

연구 동기 및 목표

• 전문가 병렬 처리와 통신 오버헤드로 인해 자원이 제한된 다운스트림 배포 환경에서 희박한 MoE 모델의 비효율성을 해결하기 위해.
• 사전에 트레이닝된 MoE 모델에서 특정 다운스트림 작업에 대해 가장 유능한 전문가만이 실제로 기여하는지 조사하기 위해.
• 가장 작업에 유능한 전문가를 식별하고 유지함으로써 효율적인 인퍼런스를 가능하게 하는 프루닝 전략을 개발하기 위해.
• 프루닝된 단일 전문가 모델이 밀도 기반 베이스라인을 능가하고, 전체 MoE 피나이팅 성능과 동일하거나 이를 초월함을 보여주기 위해.
• 최종 모델 성능과 인퍼런스 효율성을 극대화하기 위해 전문가 프루닝의 시점과 기준을 최적화하기 위해.

제안 방법

• 고정 길이의 슬라이딩 트레이닝 윈도우를 기반으로 계산된 전문가 전문성 점수에 따라 전문가가 프루닝된다.
• 정기적인 간격으로 낙제 기준이 적용되며, 이 기준 이하의 전문가는 제거된다.
• 모든 MoE 레이어에 대해 한 명의 전문가만 남거나, 피나이팅의 절반까지 프로세스가 계속된다.
• ‘즉각적’ 프루닝 전략이 제안되며, 이는 중간 지점에서 비전문가 전문가를 제거하여 두 번째 반기 동안 선택된 전문가의 최적화에 집중할 수 있도록 한다.
• 최종 선택된 전문가는 나머지 트레이닝 스케줄 동안 독립적으로 피나이팅된다.
• 이 방법은 지식 정렬을 피하고, MoE 사전 트레이닝 모델의 지식을 직접 단일 밀도 전문가로 이전한다.

실험 결과

연구 질문

• RQ1 사전 트레이닝된 MoE 모델에서 특정 다운스트림 작업에 대해 전체 MoE 아키텍처의 대부분의 성능 이점을 유지할 수 있는 단일 전문가를 식별할 수 있는가?
• RQ2 피나이팅 중 비전문가 전문가를 프루닝할 최적의 시점과 기준은 무엇인가, 이를 통해 최종 정확도를 극대화할 수 있는가?
• RQ3 트레이닝 윈도우의 길이가 선택된 전문가의 성능에 어떤 영향을 미치는가?
• RQ4 비전문가 전문가를 프루닝함으로써 성능을 손상시키지 않고도 인퍼런스 속도를 높이고 통신 비용을 줄일 수 있는가?
• RQ5 프루닝된 단일 전문가 모델이 처음부터 트레이닝된 밀도 모델을 능가할 수 있는가?

주요 결과

• 프루닝된 단일 전문가 모델은 여섯 가지 다양한 다운스트림 작업에서 전체 MoE 모델의 성능 향상의 99.3%를 유지한다.
• MNLI 작업에서 프루닝된 모델은 전체 MoE 모델 대비 2배 빠른 인퍼런스 속도를 기록했으며, SQuAD에서는 1.28배의 속도 향상을 달성했다.
• 선택된 전문가에 하위-MoE 레이어가 포함되어 있음에도 불구하고, 프루닝된 모델의 인퍼런스 효율성은 밀도 모델의 80%에 도달한다.
• 중간 지점에서 비전문가 전문가를 제거하는 ‘즉각적’ 프루닝 전략은 단계적 또는 이중 통과 방법보다 더 뛰어난 인퍼런스 성능을 제공한다.
• 이 방법은 상호 장치 간 통신 비용을 완전히 제거하여 클라우드 및 모바일 배포에 매우 적합하다.
• 선택된 전문가는 트레이닝 초반에 기여도가 높아지며, 이는 두 번째 반기 동안 집중 최적화를 위해 조기에 프루닝하는 것이 타당함을 정당화한다.