Mixture of Experts(MoE)를 설명해 보았다 (Mixture of Experts Explained)

본 포스트는 huggingface의 moe 설명이 너무 훌륭하게 잘 되어있어서 혹시 도움이 될까 싶어 단순 번역한 글입니다. 모든 credit은 저자분들에게 있습니다. MoE에 관심이 있으신 분들이 보시면 좋을 것 같습니다. 또한 혹시 잘못된 정보나 수정 사항이 있으면 알려주시면 반영하겠습니다. 감사합니다.

Reference: https://huggingface.co/blog/moe 

Authors: Omar Sanseviero, Lewis Tunstall,Philipp Schmid,Sourab Mangrulkar,Younes Belkada,Pedro Cuenca 

이번에 발표된 Mixtral 8x7B (announcement, model card)로, transformer 관련해서는 Mixture of Experts 줄여서 MoE가 가장 open AI 커뮤니티에서는 가장 핫한 주제가 되었습니다. 본 blog 기고에서는, MoE의 구조가 어떠한지, 어떻게 학습하는지, 그리고 실제 serving시에 어떤 trade off를 고려해야한지 살펴보려합니다. 

한번 보시죠! 

Table of Contents 

Mixuture of Experts란 무엇인가? 

MoE에 대한 간략한 역사 

Sparsity란 무엇인가? 

MoE를 위한 Balancing token들을 불러오기 

MoE와 Transformers 

Switch Transformers 

router Z-loss를 통해 학습 안정화하기 

개별 Expert가 학습하는 것은 무엇일까? 

Experts 개수를 늘리는 것은 pre-training에 어떤 영향을 미칠까? 

Making MoEs go brrr 

Expert 병렬화 

Capacity Factor와 Communication 비용 

Serving 기법 

효율적인 학습 

Open Source MoEs 소개 

MoE의 기대되는 방향 

몇몇 Resources  

네줄 요약(TL;DR:Too long Didn’t Read) 

MoE는: 

• Dense 모델들에 비해 아주 빨리 학습 가능. 

• 같은 parameter수를 가진 모델에 비해 빠르게 inference 가능함. 

• 모든 experts들이 memory에 올라와야하므로 많은 VRAM을 요구. 

• Fine-tuning시 많은 어려움에 직시, 그러나 최근 작업(https://arxiv.org/pdf/2305.14705.pdf )인 instruction-tuning을 통한 방법이 희망적임. 

Mixture of Experts(MoE)란 무엇인가? 

모델 품질에 있어서 가장 중요한 축은 모델의 크기입니다. 정해진 예산 하에서, 더 큰 모델을 조금 학습하는 것이 작은 모델을 더 많이 학습하는 것 보다 더 좋습니다. 

MoE는 dense 모델을 같은 비용으로 학습시키는 것 보다, 더 적은 예산으로도 학습을 시킬 수 있고, 이는 모델의 크기를 늘리거나 데이터 크기를 늘릴 수 있다는 것을 의미합니다. 특히, MoE 모델은 반대급부에 있는 dense 모델과 같은 성능을 pre-training 시 빠르게 달성합니다. 

따라서, MoE는 대체 정확히 무엇일까요? Transformer 모델과 관련해서, MoE는 두가지 주요 요소로 이루어져있습니다. 

• Sparse MoE Layer들이 dense feed-forward network(FFN) layer대신에 사용됩니다. MoE layer들은 특정 수의 “experts”들을 가집니다, (예) 8개) 여기서 각 expert는 하나의 neural network을 뜻합니다. 보통은, experts들은 실제 사용 시에는 FFN들이지만, 더 복잡한 network들이 될 수 있으며, 심지어 MoE 스스로도 가능합니다-이를 통해 계층 구조의 MoE도 가능합니다. 

• Gate network, 또는 router는 expert중 누구에게 어떤 token들을 보낼지 결정하는 역할을 합니다. 예를 들어, 아래의 이미지에서, token “More”는 두번째 expert에게 보내지고, “Parameter”라는 token은 첫번째 network로 보내집니다. 나중에 더 살펴보겠지만, 하나의 token을 하나 보다 더 많은 expert에게 보내는 것도 가능합니다. 어디로 token을 보낼 것인지는 MoE를 다룰 때 고려해야할 큰 결정사항 중 하나입니다- router는 학습 가능한 변수들로 구성되어 있으며, 다른 network가 학습될 때 같이 학습됩니다. 

다시 요약하면, MoE에서 transformer 모델의 FFN layer들은 gate network와 특정 개수의 expert들로  MoE layer들로 대체됩니다. 

MoE들이 dense모델에 비해 효율적인 pretraining이나 빠른 inference와 같은 이점들을 제공함에도, 많은 어려움이 존재합니다: 

학습: MoE들은 상당히 계산 효율적인 사전학습을 가능케 하지만, overfitting으로 이어지는 fine-tuning의 일반화에 있어서는 그동안 어려움을 겪어왔습니다. 

Inference: MoE가 많은 Parameter들을 가짐에도 그 중 몇몇만 inference 시 참여하게 됩니다. 이를 통해 같은 parameter수를 가지는 dense 모델에 비해 더 빠른 inference를 가능케 합니다. 그러나, 모든 parameter가 RAM에 올라가야하므로, 메모리 요구량이 높습니다. Mixtral 8x7B를 예로 들면, 47B parameter수의 dense모델을 들고 있을 수 있는 만큼의 VRAM을 요구 합니다. 왜 8x7B=56B가 아니라 47B parameter수 일까요?그것은 MoE모델들의 FFN layer들만 개별 expert로 취급되고, 나머지 parameter들은 공유하기 때문입니다. 또한, 두 expert가 한 token당 상용되었다고 가정했을 때, inference 속도(FLOPs)는 (14B 모델이 아니라) 12B 모델과 비슷하며, 이는 2x7B matrix 곱연산들이 수행되지만, 몇몇 layer들은 공유되고 있기 때문입니다.(더 자세한 내용은 나중에) 

자 이제, MoE에 대한 대략적인 아이디어들은 살펴본 것 같고, 이 발상을 이끌어온 연구들이 어떻게 발전해왔는지 살펴봅시다. 

간략한 MoE의 역사 

MoE의 뿌리는 1991년 논문인 Adaptive Mixture of Local Experts(https://www.cs.toronto.edu/~hinton/absps/jjnh91.pdf)에서 나왔습니다.  ensemble 방법과 비슷한 이 발상은 newtork들(학습 케이스의 서로 다른 부분 집합을 다루게되는)로 구성된 시스템을 관리하는 방법이었습니다. 개별 분리된 시스템, 또는 expert는 입력 공간의 서로 다른 영역에 특화되어 있습니다. 어떻게 expert가 선택되냐구요? Gating network가 각 expert의 가중치를 결정합니다. 학습 중에, expert와 gating 부분이 학습되구요. 

2010-2015년 사이에, 두 서로 다른 연구 영역이 훗날의 MoE의 진전에 기여합니다: 

• 구성 요소로써의 Experts: 전통적인 MoE setup에서는 전체 시스템은 gating nework와 여러 experts로 구성됩니다. 전체 모델로써 MoE는 SVM들과 Gaussian Process들 그리고 다른 방법들에서 탐색됐습니다. Eigen,Ranzato, 그리고 Ilya의 작업(https://arxiv.org/abs/1312.4314)에서 MoE를 더 깊은 신경망의 구성 요소로써 살펴봤습니다. 이는 다층 신경망의 layer들로써의 MoE 개념을 제시했고, 더 크고 동시에 효율적인 모델로 만들었습니다. 

• 조건부 연산: 전통적인 network들은 모든 input 데이터를 layer 전부를 거쳐 처리합니다. 이 때, Yoshua Bengio는 입력 token에 기반하여 구성 요소들을 동적으로 활성화 또는 비활성화하는 방안을 내놓았습니다. 

이러한 연구들은 NLP의 문맥에서 MoE를 다룰 수 있게 했습니다. 더욱 공고하게, Shazeer et al.(2017, 여기서 “et al.”은 Geoffrey Hinton과 구글의 척노리스 Jeff Dean을 포함합니다.)은 이 아이디어를, 희소성(Sparsity)를 도입하여,큰 모델임에도 불구하고 빠른 inference가 가능했던 137B의 (당시 Schmidhuber에 의해 만들어진 NLP 아키텍처 그 자체였던) LSTM에 적용하여 크기를 키웠습니다. 이 연구는 번역에 중점을 둔 연구였고, 높은 통신 비용과 학습 불안정성 등 많은 어려움을 겪었습니다. 

MoE는 예를 들어 open-source로 풀린 1.6T개의 변수를 가진 Switch Transformers 외에 많은 모델처럼, 수조개의 변수를 가진 모델도 학습 가능하게 했습니다. MoE는 Computer Vision에서도 연구됐지만 본 post에서는 NLP domain만 집중할 예정입니다. 

희소성(Sparsity)이란 무엇인가? 

희소성은 조건부 연산에 대한 아이디어라고 할 수 있습니다. Dense 모델은 모든 변수들이 모든 input들에 대해 사용됨과 달리, 희소성은 전체 시스템의 일부만 작업을 수행합니다. 

Shazeer의 번역을 위한 MoE 모델을 자세히 살펴보도록 합시다. (네트워크의 일부가 예제 한개를 기반으로 활성화 되는) 조건부 연산의 아이디어는 계산량을 증가시키지 않고 모델을 키우는 것을 가능하게 했고, 따라서 수천개의 experts가 각 MoE layer에서 사용될 수 있었습니다.  

이러한 setup은 몇몇 어려움을 맞닥뜨리게 했습니다. 예를 들어, batch 크기가 커지면 보통은 더 성능이 증가하는데, MoE의 batch크기는 데이터 흐름에 따라 활성화된 experts들에 퍼지면서 효과적으로 줄어들어 버렸습니다. 예를 들어, 만약 batch로 묶인 input이 10개 토큰들로 구성되어 있다면, 5개의 토큰은 한 expert에, 다른 5개 토큰은 서로 다른 5개 experts에 도달해, batch 크기 불균등, 그리고 활용률 저하로 이어졌습니다. 아래의 “MoE의 시동을 걸어보자”(Making MoEs go brrr) 섹션에서는  도전점들과 해결책에 대해 이야기해볼 것입니다. 

어떻게 이 문제를 풀 수 있을까요? 학습된 gating network (G) 는 expert들 (E) 이 입력의 한 부분을 어디로 보내야할지 정합니다. 

이 setup에서는, 가중 곱셈이므로, 모든 expert들가 모든 입력에 대해서 작업을 수행합니다. 그러나 만약 G가 0이면 어떻게 될까요? 만약 그렇다면, 각 expert에 대한 계산을 수행하지 않아도 되므로, 계산량을 줄일 수 있습니다. 전형적인 Gating 기능은 어떤가요? 이 전통적인 setup에서는, softmax 함수를 이용한 간단한 신경망을 사용하게 됩니다. 이 신경망에서는 입력을 어떤 expert로 보낼지 학습하게 됩니다. 

Shazeer의 작업에서는 Noisy Top-k Gating과 같은 다른 gating 메커니즘도 연구했습니다. 이 gating 방식은 어느정도 (조정가능한) noise를 도입하고, top k개의 값을 유지하는 것입니다. 이는: 

1. Noise를 더합니다. 

2. Top k개 만을 고릅니다. 

3. Softmax를 적용합니다. 

이러한 희소성은 재미있는 특성들을 가지고 왔습니다. 이러한 충분히 작은 k 값을 사용함으로써 (예) 1이나 2), 많은 expert들이 활성화 되어있을 때 보다, 학습하거나 inference를 수행할 때 더 빠르게 할 수 있습니다. 왜 그냥 가장 높은 expert를 골라 버리지 않는 걸까요? 초기의 추측은 둘 이상의 expert를 선택해야, 서로 다른 experts들에 대해 어떻게 선택하는 법을 선택할지 배울 수 있기 때문에 최소 둘 이상의 expert들을 골라야한다 였습니다. Switch Transformers(huggingface.co/blog/moe#switch-transformers)에서 이 결정에 대해 다시 생각해볼 것입니다. 

왜 noise를 더할까요? 그것은 부하 균형 조정(laod balancing) 때문입니다! 

MoE를 위한 균등 부하 조정(load balancing) 토큰들 

앞서 논의한 바와 같이, 만약 모든 토큰들이 몇몇 인기많은 expert들에만 보내지게 된다면, 학습이 비효율적이게 됩니다. 일반적인 MoE 학습에서는, gating 신경망은 거의 동일한 몇개의 활성화된 expert들로 수렴합니다. 이러한 선호되는 expert들로의 자가 강화는 학습을 빠르게 하며, 따라서 더 선택받기 쉽습니다. 이 문제를 해결하기 위해, 보조 loss를 추가해 모든 expert들이 동등하게 중요하게 다뤄지도록 유도합니다. 이 loss는 모든 expert들이 거의 동등한 개수의 학습 예제들을 보도록 보장합니다. 이어지는 섹션에서는 expert 용량이라는, 한 expert에 의해 얼마나 많은 토큰들이 다뤄질 수 있는지에 대한 임계 값을 도입하는, 컨셉에 대해 알아볼 것입니다. Transformers에서는 보조 loss는 aux_loss 변수로 확인할 수 있습니다. 

참조: auxiliary loss in ST-MoE 

MoE와 Transformers 

Transformers는 변수의 개수를 늘릴 수록 성능이 향상된다는 명확한 사례이며, 따라서 Google이  GShard(https://arxiv.org/abs/2006.16668)에서 600B 이상의 transformers 모델로 확장하여 탐색한 것은 놀라운 일은 아닙니다. 

GShard에서는 모든 FFN layer를 top-2 gating MoE layer로 encoder, decoder를 바꾸었습니다. 다음 이미지는 encoder부가 어떻게 생겼는지 보여줍니다. 이 setup은 대규모 계산에서 상당한 이득을 볼 수 있습니다. 우리가 여러 기기로 모델을 확장할 때, 다른 layer들이 복제하여 사용하는데 반해, MoE layer는 모든 기기에서 공유하게 됩니다. 이는 “MoE의 시동을 걸어보자”에서 다뤄질 예정입니다. 

균등한 부하를 유지하고, 대규모에서 효율적으로 작동하기 위해, GShard 저자들은 보조 loss 외에 이전 섹션에서 논의한 것과 비슷한 몇 가지 변화를 도입했습니다. 

Random routing: top-2 setup에서 우리는 언제나 top expert는 선택하고, 두 번째 expert는 가중치 분의 1에 해당하는 확률로 선택됩니다. 

Expert Capacity: 한 expert가 처리할 수 있는 토큰의 개수에 대한 임계값을 설정합니다. 두 expert들이 용량 한계에 이르면 토큰이 넘쳤다고(overflowed) 가정하고, residual connection들을 통해 다음 layer로 보냅니다. (또는 다른 프로젝트에서는 전부 버리는 경우도 있습니다.) 이 컨셉은 MoE에서 가장 중요한 개념이 되었습니다. 왜 expert 용량 개념이 필요할까요? 모든 tensor의 모양은 컴파일 시 정적으로 정해지고, 우리는 얼마나 많은 토큰들이 expert에 보내지게 될지 미리 알 수 없기 때문에, 이를 용량 계수(expert capacity)를 도입해 보완하는 것입니다. 

GShard 논문은 MoE와 어울려 잘 동작하는 병렬 계산 패턴으로 표현하는데에도 기여했지만, 이에 대해 논의하는 것은 본 blog post의 범위를 벗어납니다. 

주: inference를 수행할 때, 몇몇 expert들이 동작하게 됩니다. 동시에, 모든 토큰들이 연관된 self-attention과 같은 공유된 연산들도 있습니다. 이 때문에 8개 expert로 구성된 47B 모델을 12B dense 모델과 같은 연산량으로 돌릴 수 있다고 하는 것입니다. 만약 top-2, 14B 변수가 사용됐다면 말이죠. 그러나 attention 연산이 (다른 연산들처럼) 공유되기 때문에, 실제 변수 개수는 12B에 해당합니다. 

Switch Transformers 

MoE가 희망적인 모습을 많이 보여줬음에도 불구하고, 학습이나 fine-tuning 시에 불안정한 모습을 보이며 어려움을 겪었습니다. Switch Transformers(https://arxiv.org/abs/2101.03961)는 이러한 주제에 대해 깊이 탐구한 아주 흥미로운 연구입니다. 저자는 2048개 experts로 구성된 1.6조개 변수로 이루어진 MoE를, 당신도 transformers library로 돌려볼 수 있도록, huggingface에 공개했습니다. Switch Transformers는 T5-XXL에 비해 4배 빠른 pre-train 속도를 달성했습니다. 

GShard와 마찬가지로, 저자는 FFN layer를 MoE layer로 대체했습니다. Switch Transformers 논문은 두 개의 입력들을 받고 (두 개의 서로 다른 token들을 받는) 그리고 4개의 expert들을 가지는 Switch Transformer layer를 제안했습니다. 

초기의 적어도 두개의 expert들을 사용한다는 아이디어와는 다르게, 간략화된 단일-expert 전략을 사용합니다. 여기서 오는 효과는: 

• Router 연산량이 줄어듦.  

• 각 expert의 Batch 크기가 최소 절반으로.  

• 통신 비용이 줄어듦. The router computation is reduced 

• 품질은 유지됨.  

Switch Transformers는 expert 용량 개념에 대해서도 연구했습니다. 

용량은 한 batch 안의 token개수를 expert 개수로 나눈 것이 됩니다. 만약 용량 계수(capacity factor)가 1 보다 크면, 토큰들이 완벽히 균형을 이루지 않았을 때, 버퍼로 작용하게 됩니다. 용량을 키우는 것은 기기 내부 통신 비용을 늘리것이 되어 트레이드 오프가 있음을 명심해야 합니다. 특히, Switch Transformers는 낮은 용량 계수(1-1.25)일 때 잘 작동합니다. 

Switch Transformer의 저자는 이 섹션에서 다뤘던 부하 균등 조정을 다시 방문하여 단순화했습니다. 각 Switch layer에서는, 보조 loss는 학습 중 전체 모델의 loss에 더해집니다. 이 loss는 균등한 routing을 하도록 유도하고, 하이퍼파라미터를 통해 가중치를 정할 수 있습니다.  

저자는 또한 expert는 bfloat16으로 나머지는 full precision으로 사용하는 등 선택적 precision 정책을 실험했습니다. Precision이 낮을 수록 processor에서 처리하는 통신 비용, 계산 비용 그리고 tensor를 저장하는 메모리 사용량이 줄게 됩니다. 초기 실험에서는, bflaot16으로 expert들과 gate network를 학습했고, 결과적으로 학습이 불안정해졌습니다. 이는 특히 router의 계산 특성 때문이었습니다: router는 지수 함수이므로, 높은 precision을 유지하는 것이 중요했습니다. 이러한 불안정성을 해결하기 위해, routing은 full precision을 사용했습니다. 

이 노트북(https://colab.research.google.com/drive/1aGGVHZmtKmcNBbAwa9hbu58DDpIuB5O4?usp=sharing)은 Swith Transformers를 요약해 사용할 수 있도록 fine-tuning하는 쇼케이스이지만, 뒤에 나올 fine-tuning 섹션을 먼저 살펴보고 보시기를 권장합니다. 

Switch Transformers는 T5도 MoE에 대응하는 encoder-decoder setup도 사용했습니다. GLaM 논문은 스케일을 올려서 GPT-3의 품질을 1/3의 에너지만을 사용하고도(https://arxiv.org/abs/2112.06905) 달성했습니다. (맞습니다, MoE를 학습시킬 때 필요한 계산량이 더 적어짐으로써, 탄소배출량을 한 차원 줄일 수 있었습니다.) 저자들은 fine-tuning 접근보다는 decoder-only 모델들과 few-shot, one-shot 평가에 집중했습니다. Top-2 routing과 더 큰 용량 계수를 사용했습니다. 게다가, 용량 계수를 얼마만큼 계산량을 쓰고 싶은지에 따라서 학습과 평가 중에 변경하여 사용할 수 있는 척도로써 연구했습니다. 

router Z-loss를 활용한 학습 안정화 

이전에 살펴본 균형 loss를 도입하면 불안정성 이슈와 맞닥뜨릴 수 있습니다. 우리는 품질을 희생시키면서 희소 모델들을 안정화시키는 많은 방법들을 사용할 수 있습니다 예를 들어, dropout은 모델 품질을 희생시켜 안정성을 올리는 방법입니다. 반면에, 곱셈 성분을 더 늘리면, 품질은 올라가지만, 안정성은 내려갑니다.  

ST-MoE(https://arxiv.org/abs/2202.08906)에서 소개된 Router z-loss는 gating 신경망에 들어가는 큰 logits에 페널티를 주어 품질의 저하 없이 학습 안정성을 크게 향상시키는 방법입니다. 이 loss는 값들의 절대값 크기를 더 작아지도록 유도하여, gating처럼 지수함수에 큰 영향을 주는 반올림 에러들을 줄입니다. 더 자세한 사항은 논문을 살펴보는 것을 권장합니다. 

그렇다면 expert는 무엇을 배우는가? 

ST-MoE의 저자들은 encoder expert가 한 그룹이나 얕은 개념들에 특화되었음을 알아냈습니다. 예를 들어, 우리는 학습결과로 구두점, 고유 명사 등에 대한 expert를 얻을 수도 있습니다. 반면에, decoder expert들은 특화되는 모습을 덜 보입니다. 저자들은 다국어 setup도 학습했습니다. 누군가는 한 언어에 대한 특화된 expert가 생겼을 것이라고 상상하겠지만, 정반대의 일이 일어났습니다: 토큰 routing과 균등 부하 조정의 영향으로, 하나의 언어에 특화된 expert는 등장하지 않았습니다. 

Expert의 개수를 늘리는 것이 pretraining에 어떤 영향을 미치는가? 

처음에는 더 많은 expert는 학습에 필요한 데이터 효율성(sample efficiency)과 더 빠른 속도로 이어지지지만, (특히 256 또는 512개 이후에는) 그 이득이 줄어드는 모습을 보입니다. 대규모에서 Switch Transformers의 확인된 특성은 소규모, 심지어 layer 당 2,4 또는 8개의 expert를 가질 때에도 일관되게 나타났습니다. 

MoE Fine-tuning하기 

| “Mixtral은 transformers 4.36.0에서 지원됩니다. Pip install transformers==4.36.0 --upgrade로 설치할 수 있습니다. 

과적합의 역학은 dense 모델과 희소 모델에서 매우 다르게 나타납니다. 희소 모델들은 과적합하기 쉽고, 때문에 더 높은 단계의 정규화(예를 들어 dropout)과 같은 방법을 expert들에 적용합니다.(예) dense layer에 dropout비율을 하나 상정하고, 희소 layer들에서는 더 높은 값을 적용합니다.) 

다른 질문은 fine-tuning시에 보조(auxiliary) loss를 적용하느냐입니다. ST-MoE저자들은 보조 loss 끄는 실험을 했는데, 11%의 토큰이 버려지는 상황에서도 품질은 크게 변함이 없었습니다. 토큰이 버려지는 것이 과적합을 막는데 도움이 되는 정규화 방법의 형태일 수 있겠습니다. 

Switch Transformers에서 pretrain perplexity를 고정시킨 상황에서 희소 모델이 dense모델보다, 특히 SuperGLUE와 추론 기반 태스크와 같은, downstream 태스크에서 성능이 떨어지는 모습을 관찰했습니다. 반면에, TriviaQA와 같은 지식 기반 태스크에서는, 희소 모델이 훨씬 나은 모습을 보였습니다. 저자들은 더 적은 expert를 가질 때 fine-tuning에 유리함을 확인했습니다. 다른 발견은 모델이 더 큰태스크에서 잘 하던 것을 작은 태스크에서 못하는 일반화에 대한 이슈를 보였다는 점입니다. 

작은 태스크 (좌측)에서는 과적합하여 희소 모델이 validation set에서 더 좋지 않음이 명확히 보입니다. 더 큰 태스크 (우측)에서는 MoE가 잘 작동합니다. 이 이미지는 ST-MoE 논문에서 가져왔습니다. 

One could experiment with freezing all non-expert weights. That is, we'll only update the MoE layers. This leads to a huge performance drop. We could try the opposite: freezing only the parameters in MoE layers, which worked almost as well as updating all parameters. This can help speed up and reduce memory for fine-tuning. This can be somewhat counter-intuitive as 80% of the parameters are in the MoE layers (in the ST-MoE project). Their hypothesis for that architecture is that, as expert layers only occur every 1/4 layers, and each token sees at most two experts per layer, updating the MoE parameters affects much fewer layers than updating other parameters. 

By only freezing the MoE layers, we can speed up the training while preserving the quality. This image is from the ST-MoE paper. 

One last part to consider when fine-tuning sparse MoEs is that they have different fine-tuning hyperparameter setups - e.g., sparse models tend to benefit more from smaller batch sizes and higher learning rates. 

Sparse models fine-tuned quality improves with higher learning rates and smaller batch sizes. This image is from the ST-MoE paper. 

At this point, you might be a bit sad that people have struggled to fine-tune MoEs. Excitingly, a recent paper, MoEs Meets Instruction Tuning (July 2023), performs experiments doing: 

• Single task fine-tuning 

• Multi-task instruction-tuning 

• Multi-task instruction-tuning followed by single-task fine-tuning 

When the authors fine-tuned the MoE and the T5 equivalent, the T5 equivalent was better. When the authors fine-tuned the Flan T5 (T5 instruct equivalent) MoE, the MoE performed significantly better. Not only this, the improvement of the Flan-MoE over the MoE was larger than Flan T5 over T5, indicating that MoEs might benefit much more from instruction tuning than dense models. MoEs benefit more from a higher number of tasks. Unlike the previous discussion suggesting to turn off the auxiliary loss function, the loss actually prevents overfitting. 

When to use sparse MoEs vs dense models? 

Experts are useful for high throughput scenarios with many machines. Given a fixed compute budget for pretraining, a sparse model will be more optimal. For low throughput scenarios with little VRAM, a dense model will be better. 

Note: one cannot directly compare the number of parameters between sparse and dense models, as both represent significantly different things. 

Making MoEs go brrr 

The initial MoE work presented MoE layers as a branching setup, leading to slow computation as GPUs are not designed for it and leading to network bandwidth becoming a bottleneck as the devices need to send info to others. This section will discuss some existing work to make pretraining and inference with these models more practical. MoEs go brrrrr. 

Parallelism 

Let’s do a brief review of parallelism: 

• Data parallelism: the same weights are replicated across all cores, and the data is partitioned across cores. 

• Model parallelism: the model is partitioned across cores, and the data is replicated across cores. 

• Model and data parallelism: we can partition the model and the data across cores. Note that different cores process different batches of data. 

• Expert parallelism: experts are placed on different workers. If combined with data parallelism, each core has a different expert and the data is partitioned across all cores 

With expert parallelism, experts are placed on different workers, and each worker takes a different batch of training samples. For non-MoE layers, expert parallelism behaves the same as data parallelism. For MoE layers, tokens in the sequence are sent to workers where the desired experts reside. 

Illustration from the Switch Transformers paper showing how data and models are split over cores with different parallelism techniques. 

Capacity Factor and communication costs 

Increasing the capacity factor (CF) increases the quality but increases communication costs and memory of activations. If all-to-all communications are slow, using a smaller capacity factor is better. A good starting point is using top-2 routing with 1.25 capacity factor and having one expert per core. During evaluation, the capacity factor can be changed to reduce compute. 

Serving 기법들 

| “ mistralai/Mixtral-8x7B-Instruct-v0.1 를 Inference Endpoint들로 배포해 볼 수 있습니다” 

MoE의 가장 큰 단점은 Parameter의 수가 많다는 것입니다. 로컬 사용 시에는, 더 작은 모델을 테스트해보고 싶은 분도 있을 겁니다. Serving을 도와줄 몇몇 기법에 대해 이야기해봅시다. 

• Switch Transformers 저자는 초기 증류(distillation) 기법들을 실험했습니다. MoE를 해당하는 dense 모델에 증류했을 때, 30-40% 희소성에서 이득이 있었습니다. 증류는 따라서 더 빠른 pretraining과 실서비스에서 더 작은 모델을 사용할 수 있다는 이점이 있습니다. 

• 최신 방식들에서는 serving에 유리하도록 network의 일부만 추출하여, 전체 문장 또는 태스크를 한 expert에서 수행하도록 변경하기도 합니다. 

• Experts(MoE)의 병합: 이 기법은 expert들의 weight들을 병합해서, inference 시의 parameter 수를 줄입니다. 

더 효율적인 학습 

FasterMoE (2022년 3월) 에서는 아주 효율적인 분산 시스템에서 MoE의 성능을 분석했고, expert의 선호도 왜곡에 대한 기법, 지연시간을 줄이기 위한 통신 시 세분화된 스케쥴링, 최소의 지연시간을 지니는 expert를 고르는 위상을 고려한 g조정된 gate도 포함하여, 서로 다른 병렬화 전략의 이론적 한계를 분석해 17배 속도향상을 이끌어 냈습니다. 

Megablocks(2022년 11월) 은 MoE에 등장하는 역동성을 다룰 수 있는 새로운 GPU 커널을 제공하여 효율적인 희소 모델 사전학습을 모색했습니다. 그들의 제안은 토큰을 탈락시키지 않으면서 최신 하드웨어에 효율적으로 연동하여, 큰 속도향상을 이뤄냈습니다. 어떻게 했냐구요? 전통적인 MoE는 expert들이 같은 형태 같은 개수의 토큰을 가지고 있다고 가정하고, batch로 묶인 행렬 곱연산을 수행합니다. 반면에, Megablocks는 MoE layer를 블록화된-희소 작업들로 표현하여 불균등 할당을 해소했습니다. 

Block-sparse matrix multiplication for differently sized experts and number of tokens (from [MegaBlocks](https://arxiv.org/abs/2211.15841)). 

Open Source MoEs 

현재는 MoE를 학습하는 몇몇 오픈 소스 프로젝트들이 존재합니다. 

• Megablocks: https://github.com/stanford-futuredata/megablocks 

• Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm 

• OpenMoE: https://github.com/XueFuzhao/OpenMoE 

공개되어 접근 가능한 MoE의 경우, 아래를 살펴보시면 됩니다: 

• Switch Transformers (Google): Collection of T5-based MoEs going from 8 to 2048 experts. The largest model has 1.6 trillion parameters. 

• NLLB MoE (Meta): A MoE variant of the NLLB translation model. 

• OpenMoE: A community effort that has released Llama-based MoEs. 

• Mixtral 8x7B (Mistral): A high-quality MoE that outperforms Llama 2 70B and has much faster inference. A instruct-tuned model is also released. Read more about it in the announcement blog post. 

흥미로워 보이는 연구 방향들 

Further experiments on distilling a sparse MoE back to a dense model with less parameters but similar number of parameters. 

Another area will be quantization of MoEs. QMoE (Oct. 2023) is a good step in this direction by quantizing the MoEs to less than 1 bit per parameter, hence compressing the 1.6T Switch Transformer which uses 3.2TB accelerator to just 160GB. 

So, TL;DR, some interesting areas to explore: 

• Distilling Mixtral into a dense model 

• Explore model merging techniques of the experts and their impact in inference time 

• Perform extreme quantization techniques of Mixtral 

Some resources 

• Adaptive Mixture of Local Experts (1991) 

• Learning Factored Representations in a Deep Mixture of Experts (2013) 

• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (2017) 

• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (Jun 2020) 

• GLaM: Efficient Scaling of Language Models with Mixture-of-Experts (Dec 2021) 

• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (Jan 2022) 

• ST-MoE: Designing Stable and Transferable Sparse Expert Models (Feb 2022) 

• FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models(April 2022) 

• MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (Nov 2022) 

• Mixture-of-Experts Meets Instruction Tuning:A Winning Combination for Large Language Models (May 2023) 

• Mixtral-8x7B-v0.1, Mixtral-8x7B-Instruct-v0.1. 

Citation 

@misc {sanseviero2023moe, 
author = { Omar Sanseviero and 
Lewis Tunstall and 
Philipp Schmid and 
Sourab Mangrulkar and 
Younes Belkada and 
Pedro Cuenca 
}, 
title = { Mixture of Experts Explained }, 
year = 2023, 
url = { https://huggingface.co/blog/moe }, 
publisher = { Hugging Face Blog } 
} 
 

Sanseviero, et al., "Mixture of Experts Explained", Hugging Face Blog, 2023.