Mixture-of-Experts with Expert Choice Routing
Mixture of Experts
MoE라는 방식은 Mixture of Experts라는 의미로 여러 개의 서브 모델(전문가라고도 함)이 특정 작업을 수행하는 방식입니다. 기본적인 아이디어는 전체 작업을 여러 작은 부분으로 나누고 각 부분을 다른 전문가가 처리하게 함으로써 전체적인 처리 효율성과 효과를 높이는 것입니다. 성능적인 면도 중요하지만 최근 Green AI에 대한 관심이 높아지면서 MoE에 대한 관심도 같이 높아지고 있습니다.
이 방법에 대한 성능은 Mixtral 8x7B 오픈소스 모델이 증명했고, GPT-4도 MoE방식으로 동작한다고 알려져 있습니다. 이 모델은 토큰당 전체 매개변수 세트의 일부만 사용하므로 비용과 대기 시간을 제어하면서 모델의 매개변수 수를 늘립니다. 따라서 Mixtral은 총 46.7B 매개변수를 가지고 있지만 토큰당 12.9B 매개변수만 사용하고, 12.9B 모델과 동일한 속도와 비용으로 입력을 처리하고 출력을 생성합니다.
여기서 중요한 점은 공유하는 layer가 있기 때문에 8x7B = 56B가 아닌 46.7B라는 것입니다. 그렇다면 어떻게 3배나 더 작은 모델이 더 좋은 성능을 낼 수 있었을까요? 정답은 MoE!
Top-K Gating
우선 MoE의 기본적인 구조는 위와 같습니다. Softmax Gating(파란색)에서 K개의 Expert를 지정하고(나머지는 음의 무한대로 치환하여 무시한다), 각각 계산하고 Softmax에서 구한 가중치를 곱하여 Output vector(빨간색)으로 계산됩니다.
위와 같이 토큰 단위로 top-k를 계산해서 FFN(Expert)을 할당해주는 개념입니다. 위 그림에서는 k가 1인 모습입니다. 이처럼 기존의 연구에서는 상대적인 토큰의 중요도에 상관없이 각 토큰에 고정된 수의 전문가를 할당하는 top-k 함수를 사용했습니다.
Sparsely-activated Mixture-of-experts은 전체 모델의 매개변수 수는 크게 증가시키면서도, 주어진 토큰이나 샘플에 대한 계산량을 변화시키지 않는 방식으로 설계됩니다. 이러한 접근 방식은 모델의 크기와 능력을 확장할 수 있는 장점을 제공하지만, token-choice routing 전략이 적절하지 않으면 일부 전문가가 충분히 훈련되지 않아 과소 혹은 과도하게 전문화될 위험이 있습니다. (이전 연구들을 확인하려면 여기를 확인해주세요.)
기존에는 학생이 선생님을 고르는 방식이였지만, 이 연구에서는 선생님들이 학생 수를 골고루 고르는 방식입니다.
Expert Choice Routing
저자는 token-choice routing의 3가지 함정이 있다고 주장합니다.
- Load Imbalance: 일부 전문가들만 과도하게 학습되고, 나머지 전문가들은 잘 사용되지 않아 용량이 낭비되는 문제가 있습니다. 이는 latency에도 문제가 될 수 있습니다.
- Under Specialization: 최적화 되지 않은 load 전략은 불필요한 전문가나, 전문적이지 않은 전문가를 만들어 낼 수 있습니다.
- Same Compute for Every Token: 토큰 선택 전략에서는 각 토큰이 정확히 k개의 전문가를 받으므로 같은 양의 계산을 사용합니다. 계산의 복잡성에 따라 가변적으로 할당해야합니다.
이에 대한 해결책으로 토큰이 top-k 전문가를 선택하는 대신에, 전문가들이 top-k 토큰을 선택하는 방법을 제안합니다. 이로 인해 각 토큰은 가변적인 수의 전문가에게 라우팅될 수 있으며, 각 전문가는 고정된 버킷 크기(할당된 토큰 크기) 를 가지게 됩니다.
Heterogeneous MoE via Expert Choice
(이전 방식의 MoE와 조금 다른 구조로 동작해서 Heterogeneous MoE라고 하는 듯 합니다.)각 expert가 k개의 토큰을 선택합니다. k(expert가 가질 수 있는 고정된 크기)는 다음과 같이 정의할 수 있습니다.
n은 input batch에 있는 총 토큰 수, c는 평균적으로 사용하는 expert 수(capacity factor), e는 총 전문가의 수 입니다. input token은 다음과 같이 정의하며 d는 모델의 hidden state dimension 입니다.
Token-to-expert(gating function)는 세 개의 출력 행렬 I, G, P를 통해 나타냅니다.
- I: 인덱스 행렬로서, I번째 전문가에게 선택된 j번째 토큰을 지정합니다. 이 행렬은 각 전문가가 처리할 토큰을 결정하는 데 사용됩니다.
- G: 선택된 토큰에 대한 전문가의 가중치를 나타냅니다. 이 행렬은 각 전문가가 어느 토큰에 얼마나 많은 영향을 미칠지를 결정하는 데 사용됩니다.
- P: 인덱스 행렬 I의 one-hot 버전으로, 각 전문가에게 분배할 토큰을 수집하는 데 사용됩니다.
이 3가지 행렬은 다음과 같이 구할 수 있습니다.
S는 affinity scores(token과 expert의 관계)를 나타내고, W는 expert embedding을 나타냅니다. S를 transpose한 것중 top k를 골라 G, I를 결정합니다. 이때 Wg의 차원은 d와 e를 곱한 것입니다. 이 과정에서 가변적인 전문가의 수가 결정됩니다.
이후 Transformer 기반의 네트워크에서 가장 많은 연산량을 가지는 FFN에 대한 입력은 다음과 같이 정의합니다.
입력값과 P를 곱하여 다음과 같이 FFN에 대한 입력을 구한 후 i번째 Expoert에 대한 출력값을 구할 수 있습니다.
gating matrices P and G를 활용하여 Xout[토큰인덱스, 차원]을 다음과 같이 계산할 수 있습니다.
다시 정리하자면 다음과 같은 과정으로 expert-choice routing이 진행됩니다.
- 입력 토큰 X와 가중치행렬 Wg(각 expert의 embedding)을 곱한 후 Softmax로 각 토큰에 대한 affinity scores를 구합니다.
- affinity scores를 바탕으로 각 expert가 k개의 토큰을 선택합니다.
- I: 각 expert(i)에게 할당할 토큰(j)개에 대한 정보
- G: 각 토큰에 대한 가중치
- P: I에 대한 one-hot vector
- 입력 토큰 X와 P의 곱, 각 expert에게 할당할 토큰 정보
- FFN으로 각 expert에게 할당된 토큰 계산, Xe 생성
- Xe와 P, G와 곱해서 Xout[토큰 인덱스, 차원] 계산
Expert Choice with Additional Constraint
Expert choice routing에 대해 사전 학습(pre-training)과 미세 조정(fine-tuning) 결과를 향상시킬 수 있는지, 그리고 토큰 당 가변 수의 전문가를 사용하는 것이 모델 성능에 어떤 영향을 미치는지 분석하기 위해 몇가지 제약조건을 제안합니다.
제약 조건:
- ∀i: Σ(j’) A[i, j’] = k (각 전문가는 정확히 k개의 토큰을 선택해야 함)
- ∀j: Σ(i’) A[i’, j] ≤ b (각 토큰은 최대 b개의 전문가에 의해 선택될 수 있음)
- ∀i, j: 0 ≤ A[i, j] ≤ 1 (A의 원소는 0과 1 사이의 값, one-hot vector라서)
위 제약조건에 따라서 아래와 같은 경우의 수로 실험을 진행합니다.
Model Architecture
Transformer architecture를 사용하고, 최근 연구를 따라 번갈아가면서 transformer layer의 FFN가 다른 MoE 레이어로 대체합니다. 이렇게 일반 트랜스포머 레이어와 MoE 레이어를 교차 배치하면 서로 공유 하기 때문에 토큰을 skip하는 문제를 완화하여 모델 성능과 학습 효율성이 향상됩니다.
또한 최근 연구에 따라 몇가지 개선사항을 적용합니다.
- positional embedding을 per-layer relative positional bias로 바꿉니다.
- FFN의 구조를 GLU와 GeLU형태로 바꿉니다.
MoE layer는 다음과 같이 구성됩니다.
- “Expert”라고 불리는 독립적인 FFN으로 구성됩니다.
- Gating function은 소프트맥스 활성화 함수를 사용하여 전문가들에 대한 확률 분포를 모델링합니다. 이 분포는 각 입력 토큰이 전문가들에 대해 가지는 선호도를 나타냅니다.
- MoE Layer는 “셔플(shuffle)” 단계와 “언셔플(unshuffle)” 단계로 이루어집니다.
- shuffle 단계에서는 토큰을 지정된 전문가에게 모읍니다. Xe[i]: i번째 전문가가 계산한 토큰
- unshuffle 단계에서는 토큰을 입력 배치의 원래 순서로 되돌립니다. Xout[l, d]: 가중치를 곱하여 l번째 토큰 정보 d
- 각 토큰에 대해 활성화되는 전문가의 수가 다를 수 있지만, 실험을 위해 위 수식에서 c를 고정함으로써 전체 계산량을 기준 아키텍처(GShard의 top-2 토큰 선택)와 동일하게 유지합니다.
여기서 GShard는 구글이 2020년에 발표한 논문으로 experts들이 디바이스에 분산되어 있어 각 디바이스에 expert로 입력을 보낸 다음, 해당 expert의 연산 결과를 다시 가져오는 방식으로 동작합니다. 디바이스에 experts들이 분산돼 있어서 통신비용이 크다는 문제가 있습니다.
Experiments
setup
- GLaM에서 사용한 데이터셋을 활용했습니다. GLaM은 구글에서 개발한 token chice MoE모델입니다. GPT-3대비 전력 효율을 2배 이상 높이고, 성능도 더 개선시킨 모델입니다.
- GLaM의 설정을 따라 최대 시퀀스 길이를 1024 토큰으로 사용합니다.
- Adafactor 옵티마이저를 사용하고, 학습률은 10K 학습 스텝 동안 일정하게 유지한 후 inverse square root 스케줄에 따라 감소시킵니다.
- 부하 균형을 위한 보조 손실은 사용하지 않습니다.
- 256K 크기의 어휘를 가진 SentencePiece 서브워드 토크나이저를 사용합니다.
- 가장 큰 모델(8B/64E)은 512개의 TPU V4 칩에서 학습되었습니다.
- GLUE와 SuperGLUE 벤치마크에서 선택한 11개 작업에 대한 파인튜닝 성능을 주로 평가합니다.
Training Efficiency
Expert의 크기를 100M 매개변수로 고정하고 전문가 수를 변화시키는 실험을 진행했습니다.
(a) expert choice 방법을 사용하면 GShard의 top-2 게이팅에 비해 학습 수렴 속도가 2배 이상 빨라집니다.(그림에서는 더 늦게 수렴하고 있는데.. step은 더 많이 쓰지만 연산이 더 빠르다는 것?) (b) 전문가 크기를 고정하고 전문가 수를 늘리면 학습 perplexity가 크게 향상됩니다. EC-CF2가 GShard의 top-2 gating보다 일관되게 우수한 성능을 보입니다.
Fine-tuning on GLUE and SuperGLUE
향상된 perplexity가 실제로 downstream tasks에서 좋은 결과를 가져오는지 GLUE에서 선택한 11개의 작업에 대해 파인튜닝을 수행한 결과를 확인합니다.
Expert Choice with Capacity Factor of 2 (EC-CF2)가 Switch Transformer의 Top-1 gating과 GShard의 Top-2 gating보다 GLUE와 SuperGLUE 작업에서 더 우수한 성능을 보임을 보여줍니다.
GShard Top-2에서는 100M/32E 구조가 가장 잘 작동하는 반면, Switch Transformer Top-1에서는 100M/128E 구조가 더 좋은 성능을 보입니다. 두 방법 모두 일부 전문가가 원하는 용량보다 훨씬 더 많은 토큰을 받을 수 있기 때문에, step latency 병목 현상이 발생합니다. 하지만 EC-CF2는 전문가 간의 부하를 보다 균등하게 분산시킴으로써 이러한 병목 현상을 완화하고 더 나은 성능을 달성할 수 있음을 보여줍니다.
또한, EC-CF2 방법이 Dense(기존 FFN) 모델보다 더 좋은 파인튜닝 결과를 보였고, 평균 점수가 3.4점 향상됐습니다.
실험 결과에서 나온 흥미로운 사실
- 100M/32E 모델 설정이 GS Top-2와 EC-CF2 모두에서 가장 좋은 성능을 보였습니다.
- 이는 100M/64E나 100M/128E보다 모델 용량이 작음에도 불구하고 나타난 결과입니다.
- 이 결과는 좋은 학습 perplexity가 항상 다운스트림 작업에서의 더 나은 성능으로 이어지지는 않음을 보여줍니다.
Expert의 수가 많은 경우 학습할때는 perplexity가 좋았지만, 추론 과정에서는 오히려 좋지 않았다..(전문가 수가 많은 것 보다 크기가 큰게 좋다?)
Capped Expert Choice
EC-CAP2는 토큰당 최대 2개의 전문가를 허용하는 변형으로, 평균 정확도가 0.8점 감소합니다. EC-CAP3는 토큰당 최대 3개의 전문가를 허용하며, 기본 Expert Choice 방법과 유사한 결과를 달성합니다. 이 실험은 토큰당 가변적인 수의 전문가를 허용하는 것이 도움이 됨을 확인해줍니다.(기존 방법보다 좋다고 주장) 또한 부하를 평균적으로 주는 Hash Layer보다 좋은 성능을 보여주어, 부하 평균은 성능 향상에 도움이 안된다고 주장합니다.
Variable Experts per Token
Experts per Token 라우팅에 대한 통계를 계산하여, 특정 수의 전문가에 라우팅된 토큰의 비율을 분석합니다. 대부분의 토큰이 12개의 전문가에 라우팅되며, 23%는 34개의 전문가에, 3%는 4개 이상의 전문가에 라우팅됩니다. 이는 EC방법이 토큰을 여러 전문가에게 유연하게 할당할 수 있음을 보여줍니다.
Capacity Factor
Switch Transformer에서 사용된 기준 top-1 gating 방법과 학습 perplexity를 비교합니다. 예상대로 더 작은 capacity factor를 사용하면 perplexity가 증가하지만, 여전히 top-1 gating을 크게 능가합니다. 또한 0.5로 낮추어도 여전히 top-1 gating보다 우수한 성능을 보여줍니다.
Comparison with Dense Models on Pre-training
위 그림 (b)에서 expert 수를 늘리면 모델 성능이 향상되는 것과 별개로, expert capacity를 늘리는 것도 모델 성능을 크게 향상시킬 수 있음을 보여줍니다. EC 방법이 Dense 방법보다 일관되게 우수한 perplexity를 보입니다. 특히 전문가 크기가 작은 경우(100M 매개변수), MoE의 이점이 더욱 두드러집니다.
Conclusion & Limitations
Mixture-of-Experts (MoE) 모델을 위한 새로운 라우팅 방법을 제안합니다.(Expert choice routing) 이 방법은 기존 MoE 방법의 부하 불균형과 전문가 활용도 저하 문제를 해결하고, 각 토큰에 대해 서로 다른 수의 전문가를 선택할 수 있게 합니다. 이 모델은 최신 GShard 및 Switch Transformer 모델과 비교하여 2배 이상의 학습 효율성 향상을 보여주었으며, GLUE와 SuperGLUE 벤치마크의 11개 데이터셋에 대한 fine-tuning에서도 강력한 성능 향상을 달성했습니다.
하지만, 몇가지 문제점이 존재합니다.
현재 구현에서 과거와 미래 토큰을 모두 사용하여 top-k 선택을 수행하므로, auto-regressive text generation에 바로 적용하기 어려울 수 있습니다. 가능한 해결책으로는 입력 시퀀스의 대규모 배치를 수집하고, 동일한 시퀀스의 토큰을 별도 그룹으로 분배한 후, 각 그룹에 대해 Expert Choice 라우팅을 수행하는 것입니다.
또한, MoE의 오랜 문제 중 하나는 큰 메모리 공간 요구사항입니다. sparsely gated networks를 사용하여 계산 비용을 줄일 수 있지만, 전문가 수에 따라 총 매개변수 수는 선형적 또는 준선형적으로 증가합니다. 전문가 수를 늘리려면 많은 수의 하드웨어 장치를 예약해야 합니다. 따라서 동적(사용) 전력은 절약되지만 정적(예약) 전력은 절약되지 않습니다.
