[MoE] Mixture Of Experts

 

 

 

모델의 모든 파라미터를 RAM이나 VRAM에 올려서 추론할 때 사용하면 참 좋겠지만.. 컴퓨팅 리소스는 한정되어있다.

MoE는 적은 컴퓨팅 리소스로 모델을 효과적으로 돌릴 때 사용됨. 

 

모든 파라미터를 메모리에 올려서 사용하는 Dense 모델은 모든 데이터가 하나의 큰 Feed-Forward-Network를 통과한다.

MoE 방식에서는 FFN을 여러 개로 쪼개서 배치함. FFN을 Expert라고 생각하자.

 

Ensemble Model 방식으로 학습하는 모델이다.

여러 모델의 결과를 합치지는 않고, 입력 데이터에 따라 모델의 일부 경로만 선택하는 방식이니 완전히 똑같지는 않음.

 

 

 

 

 

 

앞에다 라우터를 하나 둬서 어떤 토큰을 어떤 전문가에게 보낼지를 결정한다.

그림을 보면 More 토큰은 두 번째 Expert에게 주고 Parameters 토큰은 첫 번째 Expert에게 줬음.

하나의 토큰을 여러 Expert에 보내는 것도 가능하다. 

 

라우터가 토큰과 일치하는 전문가를 찾아내는게 매우 중요함.

라우터 내부에서 학습 가능한 가중치가 있고, 학습할 때 특정 전문가만 찾는 Expert Collapse가 발생하지 않도록 주의해야 함.

 

Dense 모델은 모든 파라미터를 RAM에 올리지만, 사실 올려놓기만 하고 사용하지 않는 파라미터가 많다. 

그러니 필요한 부분만 골라 쓰자. 이러면 모델 크기를 1000배 키워도 계산량은 1000배가 되지 않음. 

 

Softmax 쓰면 모든 Expert에게 점수가 조금씩 분배되고 연산량이 많아진다.

그러니 라우터가 배정한 점수가 가장 높은 k개만 남기고 나머지는 -∞ 으로 밀어버린다. 

학습 초기에 라우터가 특정 Expert를 편애할 수 있으니 노이즈를 섞어주고, Loss Function을 적절하게 구성해준다.

 

 

 

 

 

구글은 GShard 구조를 사용해 파라미터가 6000억개가 넘는 거대 모델을 처리한다.

 

Transformer Model은 중첩된 Layer로 구성되는데, 모든 Layer가 MoE는 아님.

일반 FFN - MoE - 일반 FFN - MoE 

하나 건너서 MoE를 만나도록 구성한다. 모든 층을 MoE로 만들면 메모리 관리 비용이 너무 커지니 적당히 타협.

 

FFN은 너무 커서 GPU 하나에 다 들어갈 수 없다. GPU 마다 쪼개서 담아야 함.

데이터가 처리되다가 MoE 층을 만나면 라우터가 판단해서 n번 Expert에게 데이터를 보내는 방식을 사용한다.

 

추론 할 때는 라우터가 선택한 Expert만 활성화된다. 

Mixtral 8x7B 모델의 파라미터는 47B개로, 모두 RAM에 올라가 있어야 한다. (B는 10억)

다만, MoE 층에서는 8개의 전문가 중 딱 2개만 골라서 쓸 때 연산량이 12B 수준으로 줄어들게 됨. 

 

Dense 모델과 비교했을 때.

메모리에 모든 파라미터가 올라오는건 같고, 추론 할 때 연산량은 라우터가 선택한 Expert만 연산하게 되니 줄어든다.

 

 

 

 

---

 

 

 

 

이론은 여기까지고.. 실제로 돌려보니 생기는 문제들이 있고, 그걸 해결한 사례를 소개함.

 

Switch Transformer

Top-1 Routing으로 Expert를 한 명만 골라서 사용한다. 

Expert간 통신 오버헤드가 감소하고, 계산이 더 빨라짐 

 

 

ST-MoE 

모델이 너무 커서 학습이 불안정하다. 

Router Z-loss로 라우터가 뱉는 숫자가 지나치게 커지지 않도록 조정한다. 

 

MoE는 파라미터가 많아서 학습 데이터를 그대로 외워버리는 경향이 강해 Overfitting이 발생하기 쉽다.

그러니 Dropout을 강하게 걸어줘야 함. 

 

 

 

---

 

 

 

MoE는 모델의 설계 도면 그 자체로, 학습 단계에서부터 적용되는 아키텍처.

Mixtral 8x7B나 Switch Transformer 같은 모델들은 본투비 MoE 모델이라는거임. 

 

즉, 많은 파라미터를 가진 모델을 적은 연산량으로 굴리는 방법을 찾기 위해 MoE 방식을 생각해냈고, 그 방식대로 만들어진 모델을 MoE 모델이라고 부른다.

Dense 모델을 MoEfication 해서 배포에 초점을 맞추는 연구도 있는데. 우선 MoE의 본질은 아키텍처 라는 것.. 

 

Mixtral 8x7B, DeepSeek-V3, GPT-4, Qwen2-1.5B-MoE 등 본투비 MoE 모델들은 처음부터 대규모 추론 비용 문제를 해결하기 위해 설계된 아키텍처 계열이라는 것.. 

 

다만 모든 Expert를 메모리에 올려 둬야 하는건 어쩔 수 없다.

RAM이나 VRAM이 충분해야 MoE 모델을 돌릴 수 있음.

 

 

 

참고

1. Hugging Face 블로그 글
https://huggingface.co/blog/moe 

2. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
https://arxiv.org/pdf/1701.06538

3. MEGABLOCKS: EFFICIENT SPARSE TRAINING WITH MIXTURE-OF-EXPERTS

https://arxiv.org/pdf/2211.15841.pdf