DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models

Abstract

conventional MoE 에서 N 개 expert 에서 top-K 를 사용하는 방식은 expert specialization 에서 어려움이 있다.

즉 각 expert 들이 overlapped 되지않은 지식을 학습할 수 있어야 한다.

그래서 이 논문에서는 이 기능을 강화한 DeepSeekMoE architecture를 제안한다. 2가지 주요 방법을 제안하는데,

 

1. expert 를 flexible combination 하게 사용할 수 있도록 한다.
2. Ks 개의 shared expert 를 분리하여 common knowledge 를 학습하도록 하고, routed experts 사이의 redundancy 를 줄인다.

 

Introduction

MoE 장점 : computational cost 를 유지하면서 parameter scaling 을 할 수 있다.

 

MoE 의 issue 에는 knowledge hybridity 와 knowledge redundancy 가 있다. 이 둘은 expert specialization 을 저해하는 요소이다.

 

conventional MoE 는 FFN 대신 FFN 과 똑같은 ffn 으로 구성된 expert 여러개로 대체하고 그 중 토큰은 1개 혹은 2개 expert 를 사용한다. 이 구조의 potential 문제점은

1. Knowledge Hybridity :

기존 MoE 는 제한된 수의 expert 를 (8-16) 사용하기 때문에 적은 expert 안에 다양한 지식을 커버해야 한다.

1. Knowledge Redundancy :

서로 다른 expert 에 할당된 토큰은 공통 지식을 필요로 할 수 있다. 그래서 학습을 통해 expert가 중복될 수 있다.

 

 

그래서 DeepSeekMoE 에선

1. Fine-Grained Expert Segmentation
2. Shared Expert Isolation 사용

DeepSeekMoE 2B를 12 zero-shot or few-shot benchmark 에 대해 평가한다.

 

1.5배 전문가 매개변수 및 계산을 가진 GShard 2.9B와 비슷한 성능을 달성한다는 것을 보여준다. 또한 같은 total 파라미터 사이즈의 dense 모델에 대해 성능이 거의 근접함을 알 수 있다.

 

그 후 ds MoE 를 16B 으로 확장하고 40% 의 computation 만으로 llama 2 7B 와 비슷한 성능을 달성함을 보인다. 또한 145B 로 확장하려 했을 때 GShard 와 비교해 이점이 있음을 검증하고 28.5%(혹은 18.2%)의 computation 만으로 DeepSeek 67B와 비슷한 성능을 보인다.

 

2. Preliminaries: Mixture-of-Experts for Transformers

Transformer 에 사용되는 기존 MoE 구조

DeepseekMoE paper

 

위와 같이 Transformer attn block 을 나타낼 수 있다.

 

𝑇 : sequence length

 

Self-Att(·) : self-attention module, FFN(·) : Feed-Forward Network, layer norm 은 생략한다.

 

$u^𝑙_{1:𝑇}∈ R^{𝑇×d}$ 는 𝑙번째 attn module 에 대한 모든 token 의 hidden state 이고,

 

$h^𝑙_𝑡 ∈ R^𝑑$ 은 𝑙번째 트랜스포머 블록 이후 𝑡번째 토큰의 출력 hidden state이다.

 

보통 MoE layer 는 특정 interval 단위로 FFN 을 대체한다.

 

𝑙 번째 output hidden state $h_{𝑙𝑡}$ 는 아래와 같다.

 

DeepseekMoE paper

 

𝑁 : 전체 expert 수

 

$𝑔_{𝑖,𝑡}$ : 각 i 번째 gate value = softmax 한 score(의 Top k 값)

$𝑠_{𝑖,𝑡}$ : token-to-expert affinity score (행렬곱으로 구하는 게 맞다)

 

3. DeepSeekMoE Architecture

DeepseekMoE paper

 

3.1. Fine-Grained Expert Segmentation

expert 개수가 제한되어있으면 특정 전문가는 다양한 knowledge 를 커버하게 된다. 각 토큰을 더 많은 expert 에 할당할 수 있으면 knowledge 가 각각 expert 에 분해되어 학습될 수 있다.

 

expert 를 finer grain 하게 만들기 위해 : 각 expert FFN 의 hidden dimension 을 1/m 으로 줄여 더 작은 ffn 이 되도록 한다. 그대신 activated expert 수를 m 만큼 늘린다.

 

DeepseekMoE paper

 

expert 의 총 파라미터수는 기존 LLM 의 ffn 매개변수 개수 x N(expert 개수) 이고, fine-grained expert 총 개수는 𝑚 𝑁 이다. 그리고 activate 되는 expert 수도 𝑚 𝐾 로 늘어난다.

 

fine-grain 으로 expert combination 가짓수를 늘려 flexibility 를 높인다.

 

3.2. Shared Expert Isolation

expert 가 같은 knowledge 를 저장하여 생길 수 있는 weight 중복을 막기 위해 𝐾 𝑠 개의 shared expert 를 도입한다. 모든 토큰에 대해 shared expert 가 사용된다.

그래서 activated routed expert 는 𝐾 𝑠 개 만큼 감소시킨다. computation cost 를 유지하기 위해서이다.

 

DeepseekMoE paper

 

routed experts 의 총 개수는 𝑚𝑁 − 𝐾𝑠, activated routed expert 개수는 𝑚𝐾 − 𝐾𝑠.

 

3.3. Load Balance Consideration

routing 방식은 load imbalance문제를 가져올 수 있다.

 

첫 번째 문제는 모델이 항상 몇몇 특정 expert 만 선택하여 다른 expert 가 학습이 제대로 되지 않는 routing collapse 이고, 두 번째는 expert 가 여러 device 에 분산되어 있을 때 computation bottlenecks을 야기할 수 있다.

 

Expert-Level Balance Loss

routing collapse 를 막기 위해 사용.

DeepseekMoE paper

 

N’ = mN-K_s (실제 라우팅하는 expert 개수)

K’ = mK-K_s (실제 activate 되는 expert 개수)

 

1(.) : indicator function (= 해당하는 expert 에 대한 값은 1이고 나머지는 0으로 해주는 함수(라고 추정))

→ 즉 모든 토큰에 대해 expert 에 대한 affinity score 를 얻은 후 각 activate 되는(=token 이 선택하는) expert 에 대한 값을 줄이도록 하는 로스이다.

 

Device-Level Balance Loss

device level 에서 load 균형을 위해 도입.

DeepseekMoE paper

 

${\varepsilon_1, \varepsilon_2, ..., \varepsilon_𝐷}$ : routed expert 가 할당된 D개의 device

각 P 와 f 값은 특정 디바이스에 대한 expert 에 대한 위 loss 에서 계산한 값.

 

을 모든 디바이스에 대해 줄이도록 한다.