‹ Lecture 3 11 / 21

Mixture of Experts

FFN은 많은 matrix weights를 가지면서도 token마다 독립적으로 계산할 수 있었다. 이 두 성질을 이용하면 FFN 하나 대신 여러 Experts를 두고, 각 token이 일부만 사용하게 할 수 있다. MoE의 목표는 저장된 전체 capacity는 늘리면서 한 token에 실제로 실행하는 계산은 제한하는 것이다. Attention 전체를 여러 모델로 갈라 놓는 구조와 구분한다.

  1. token x
  2. Router Softmax / top-kk
  3. selected Experts
  4. gate-weighted sum
  5. output

Token Routing

Router는 token state xx에 WrW_r을 곱하고 Softmax로 Experts의 선택 scores를 만든다. Top-k Experts만 실행한 뒤 그 outputs를 gate weights로 가중합한다. 슬라이드에서는 Expert 2와 Expert 8을 각각 0.7과 0.3으로 합하여 yy를 만든다.

중요한 것은 선택되지 않은 Expert의 output에 나중에 0을 곱하는 것이 아니라, 그 Expert의 FFN 계산 자체를 하지 않는다는 점이다. 모두 계산한 뒤 지우면 sparse compute를 절약하지 못한다. 또 샴 / 고양이라는 두 tokens가 각각 Experts 2·8과 2·7을 사용해도 된다. FFN이 token-wise이므로 같은 문장의 모든 tokens를 같은 Expert로 보낼 필요가 없다.

Capacity & Communication

Mixtral8×7B는 total≈47B, active≈13B/token이며 모든 expert weights의 저장이 필요하다. Load imbalance를 줄이는 auxiliary loss와 expert parallel token communication overhead가 남는다.

Mixtral 8×7B의 예시는 8 Experts 중 top-2를 사용하여 total 약 47B와 active 약 13B/token을 구분한다. Active 13B라는 것은 저장할 모델도 13B라는 뜻이 아니다. 다음 token이 다른 Experts를 선택할 수 있으므로 전체 Expert weights를 사용할 수 있게 저장해야 한다. 모델 카드의 total과 active 표기가 서로 다른 이유다.

Routing Is Also a Training Problem

초기 Router가 특정 Experts에만 tokens를 보내면 일부는 과부하가 되고 다른 Experts는 학습 기회를 적게 받는다. Auxiliary load-balancing loss는 사용이 과도하게 몰리지 않도록 돕는다. 교수님은 일부 설계에서 shared Experts로 공통 능력을 처리하고 routed Experts와 함께 쓰는 접근도 언급했다. 다만 어느 Expert가 수학·code 담당이라고 사람이 자동으로 확정할 수는 없다.

Experts를 여러 devices에 나누면 해당 token state를 선택된 device로 보내고 output을 모으는 communication도 생긴다. 따라서 더 적은 active parameters가 같은 비율의 latency 감소를 보장하지 않는다. Routing, load balance, 통신을 포함하여 total capacity와 실제 실행 비용을 따로 판단해야 한다.

g=Softmax⁡(xWr)g=\operatorname{Softmax}(xW_r) y=∑i∈Top-k⁡(g)giEi(x)y=\sum_{i\in\operatorname{Top-k}(g)}g_iE_i(x)

Active compute가 작아도 total expert weights는 저장해야 한다.

SiLU · SwiGLU · MoE

Scalar function, gated FFN, 여러 FFN의 routing을 서로 다른 층위로 읽는다.

AspectSiLUSwiGLUMoE
구조Scalar xσ(x)두 input branches + down projectionRouter + 선택된 expert FFNs
핵심 연산원소별 activationSiLU(xW_g) ⊙ (xW_u)Top-kk expert weighted sum
연결Gate branch에 사용 가능Dense expert의 FFN 형태 가능Token마다 다른 experts 선택
비용 관점단일 함수3dm matrix weightsTotal vs active parameters 구분

Check

Active parameters가 13B인 MoE는 weights도 13B만 저장하면 된다는 설명이 틀린 이유는?

Answer

Token마다 선택될 expert가 다르므로 전체 expert weights를 보관해야 한다. Active parameters는 그 token의 계산에 참여하는 부분이다.

MoE를 Attention보다 FFN에 적용하기 쉬운 구조적 이유는?

Answer

FFN은 token-wise 계산이므로 token별 독립 routing이 가능하고 block parameters 비중도 크다.

Mixtral의 total parameters가 active parameters보다 큰데도 두 숫자를 모두 알아야 하는 이유는?

Answer

Token마다 Router가 선택한 Experts만 계산하므로 active compute는 작아질 수 있다. 그러나 선택되지 않은 Experts도 다른 token을 위해 저장해야 하므로 total weights의 memory 비용은 남는다. Routing의 load balancing과 device 사이 communication 비용도 고려한다.

Contents