[논문 리뷰] Scaling Laws for Fine-Grained Mixture of Experts

Jakub Krajewski, Jan Ludziejewski외 10명|arXiv (Cornell University)|2024. 02. 12.

Advanced Statistical Process MonitoringDecision Sciences인용 수 3

한 줄 요약

이 논문은 Mixture of Experts (MoE) 모델을 최적화하기 위한 새로운 초모수로 ' granularity '을 도입하여 전문가 크기의 미세 조절을 가능하게 한다. 모델 크기, 훈련 토큰 수, 그리고 granularity 를 포함하는 스케일링 법칙을 유도함으로써, 저자들은 MoE 모델이 계산 비용 절감 측면에서 밀도 있는 Transformer 보다 일관되게 뛰어난 효율성을 보임을 입증한다. 특히 대규모 예산에서 계산 비용 절감이 40배 이상 초과됨을 보이며, 이는 MoE의 효율성이 규모가 증가함에 따라 감소한다는 이전의 주장과 정면으로 배치된다.

ABSTRACT

Mixture of Experts (MoE) models have emerged as a primary solution for reducing the computational cost of Large Language Models. In this work, we analyze their scaling properties, incorporating an expanded range of variables. Specifically, we introduce a new hyperparameter, granularity, whose adjustment enables precise control over the size of the experts. Building on this, we establish scaling laws for fine-grained MoE, taking into account the number of training tokens, model size, and granularity. Leveraging these laws, we derive the optimal training configuration for a given computational budget. Our findings not only show that MoE models consistently outperform dense Transformers but also highlight that the efficiency gap between dense and MoE models widens as we scale up the model size and training budget. Furthermore, we demonstrate that the common practice of setting the size of experts in MoE to mirror the feed-forward layer is not optimal at almost any computational budget.

연구 동기 및 목표

• 대규모 언어 모델(Large Language Models, LLMs) 훈련의 높은 계산 비용 문제를 Mixture of Experts (MoE) 아키텍처 최적화를 통해 해결하기 위해.
• MoE 효율성이 규모가 증가함에 따라 감소한다는 가정을 도전하기 위해, 특히 고정된 훈련 기간과 전문가 크기 조건 하에서.
• 다양한 계산 예산 하에서 MoE 모델의 최적 초모수—특히 granularity —을 규명하기 위해.
• 변동하는 훈련 기간과 전문가 granularity 를 포함하는 스케일링 법칙을 활용해 MoE 모델의 계산 최적 훈련을 위한 통합 프레임워크를 제공하기 위해.

제안 방법

• 전문가 크기 조절을 가능하게 하기 위해 granularity (G) 를 새로운 초모수로 도입하여 MoE 모델의 전문가 크기 조절을 미세하게 조정할 수 있도록 한다.
• 모델 크기, 훈련 토큰 수, granularity, 훈련 기간 간의 관계를 모델링하는 새로운 스케일링 법칙을 유도하며, 주로 FLOPs 를 비용 측정 기준으로 사용한다.
• 다양한 granularity 수준에서 효율적인 훈련을 가능하게 하기 위해 공유 전문가와 세그먼트 전문가를 포함한 수정된 MoE 아키텍처를 활용한다.
• 다양한 모델 크기와 계산 예산을 대상으로 광범위한 실험을 수행하여 스케일링 법칙의 타당성과 최적 설정을 검증한다.
• 훈련 효율성을 평가하기 위해 월클럭 시간과 퍼플렉서티를 사용하며, 높은 granularity 가 동일 성능을 달성하는 데 더 짧은 훈련 시간을 요구함을 보여준다.
• 확장률(E) 에 대한 분석 실험을 수행하여, E=16 및 E=64 를 포함한 다양한 E 값에서도 결과가 유지됨을 입증한다.

실험 결과

연구 질문

• RQ1 MoE 모델의 granularity 를 조정할 경우, 다양한 계산 예산에서 훈련 효율성과 성능에 어떤 영향을 미치는가?
• RQ2 모델 크기와 훈련 예산이 증가함에 따라 MoE 모델이 밀도 있는 Transformer 보다 효율성을 유지하거나 향상시키는가?
• RQ3 전문가 크기를 피드포워드 레이어 크기와 동일하게 설정하는 것이 MoE 모델에 최적의 전략인가?
• RQ4 변동하는 훈련 기간과 스케일링 법칙이 결합되어 계산 최적의 MoE 설정을 어떻게 결정하는가?
• RQ5 극도의 계산 예산(예: 10^25 FLOPs)에서도 미세 조절된 MoE 모델이 밀도 있는 Transformer 보다 뛰어난 효율성을 달성할 수 있는가?

주요 결과

• 최적의 granularity 를 가진 MoE 모델은 계산 효율성 측면에서 밀도 있는 Transformer 를 일관되게 능가하며, 10^25 FLOP 예산에서 계산 비용 절감이 40배 이상 초과된다.
• 표준적인 전문가 크기 설정 방식(즉, G=1)은 거의 모든 계산 예산에서 최적의 성능을 내지 못한다.
• 모델 크기와 훈련 예산이 증가함에 따라 MoE와 밀도 있는 모델 간의 효율성 격차가 더욱 벌어지며, 이는 이전에 보고된 수익 감소 현상에 대한 주장과 정면으로 배치된다.
• 최적의 훈련 기간은 모델 크기와 granularity 에 따라 달라지며, 스케일링 법칙에 변동하는 훈련 시간을 통합하는 것이 계산 최적 설정을 도출하는 데 필수적이다.
• 미세 조절된 MoE 모델은 더 낮은 퍼플렉서티를 더 빠르게 달성하여, 표준 MoE 및 밀도 있는 기준 모델 대비 월클럭 훈련 시간에서 뚜렷한 성과 향상을 보였다.
• granularity, 모델 크기, 훈련 토큰 수, 그리고 변동하는 훈련 기간을 포함하는 스케일링 법칙은 주어진 FLOP 예산에 대한 최적 초모수를 정확하게 예측할 수 있다.