스위치 트랜스포머: 희소 전문가 혼합 (MoE)

스위치 트랜스포머: 희소 전문가 혼합 (MoE)

트랜스포머의 피드포워드 층을 여러 "전문가"로 늘리되, 토큰마다 전문가를 딱 하나만 골라 쓰게(top-1 라우팅) 단순화한 논문입니다. 토큰당 계산량은 그대로 두고 파라미터만 수십~수천 배로 늘려, 같은 연산으로 T5보다 최대 7배 빠르게 사전학습했고 1.6조 파라미터 모델까지 학습했습니다. 이후 Mixtral 같은 전문가 혼합 LLM의 기본 설계가 됐습니다.

들어가기 전에

이 해설은 아래 내용을 안다고 보고 씁니다. 처음이라면 먼저 읽고 오면 수식과 손계산을 훨씬 쉽게 따라올 수 있습니다.

• 모델 압축과 희소성: 지식 증류 · 가지치기 · 전문가 혼합 — 전문가 혼합의 라우팅을 손계산으로 먼저 보는 수업
• 어텐션만 있으면 된다 (트랜스포머) — 순환 없이 어텐션만으로 번역 모델을 만든 트랜스포머 — 스케일드 닷-프로덕트 · 멀티헤드 어텐션 · 위치 인코딩의 수식과 의미

한눈에 보기

• "파라미터 수"와 "토큰당 계산량"을 따로 떼어 키울 수 있다는 희소 활성화(sparse activation)의 생각을 알게 됩니다.
• 라우터가 토큰마다 전문가 하나를 고르는 스위치 라우팅과, 전문가 용량 · 버려지는 토큰의 관계를 손으로 계산해 보게 됩니다.
• 전문가들이 고르게 쓰이도록 만드는 부하 균형 보조 손실을 숫자로 확인합니다.
• 큰 희소 모델을 안정적으로 학습하는 세 가지 요령(선택적 정밀도 · 작은 초기화 · 전문가 드롭아웃)과, 실제로 얼마나 빨라졌는지 알게 됩니다.

초록 (번역)

딥러닝에서 모델은 보통 모든 입력에 같은 파라미터를 재사용합니다. 전문가 혼합(MoE)은 이를 거슬러 입력 예제마다 다른 파라미터를 고릅니다. 그 결과 엄청난 수의 파라미터를 가지면서도 계산 비용은 일정한, 희소하게 활성화되는 모델이 됩니다. 그러나 MoE의 몇몇 눈에 띄는 성공에도 불구하고, 복잡성 · 통신 비용 · 학습 불안정성 때문에 널리 쓰이지 못했습니다. 우리는 스위치 트랜스포머로 이 문제들을 다룹니다. MoE 라우팅 알고리즘을 단순화하고, 통신 · 계산 비용을 줄인 직관적인 개선 모델을 설계합니다. 우리가 제안하는 학습 기법들은 불안정성을 다스리는 데 도움이 되며, 큰 희소 모델을 처음으로 더 낮은 정밀도(bfloat16) 형식으로 학습할 수 있음을 보입니다. T5-Base와 T5-Large를 바탕으로 모델을 설계해, 같은 연산 자원으로 사전학습 속도를 최대 7배 높입니다. 이 개선은 다국어 설정으로도 이어져, 101개 모든 언어에서 mT5-Base보다 나은 결과를 측정합니다. 마지막으로 "Colossal Clean Crawled Corpus"(C4)로 최대 1조 파라미터 규모의 모델을 사전학습해 언어 모델의 규모를 한 단계 넓히고, T5-XXL 모델보다 4배 빨라집니다.

무엇을 풀었나

스케일링 법칙 연구는 "모델이 클수록 좋다"고 말했지만, 보통의 조밀한(dense) 트랜스포머는 파라미터가 늘면 토큰마다 하는 계산도 똑같이 늘어납니다. 모델을 10배 키우면 비용도 10배입니다.

전문가 혼합(MoE, Mixture of Experts)은 여러 하위 네트워크(전문가) 중 일부만 토큰마다 골라 쓰게 해서, 파라미터는 많지만 계산은 적게 합니다. 번역에서 성공 사례가 있었지만 널리 퍼지지 못했습니다. 논문이 꼽는 걸림돌은 세 가지입니다.

1. 복잡성: 기존 MoE는 토큰마다 상위 $k$개($k \ge 2$) 전문가를 골라 섞었고, 구현이 복잡했습니다.
2. 통신 비용: 전문가들이 여러 장치에 흩어져 있어 토큰을 주고받는 통신이 많습니다.
3. 학습 불안정: 라우팅의 딱딱한 선택과 bfloat16 같은 낮은 정밀도가 겹치면 학습이 발산하기 쉽습니다.

저자들의 가설은 "토큰당 연산(FLOPs)을 고정한 채 파라미터 수만 따로 늘리는 것도 독립적인 규모의 축"이라는 것입니다.

전체 흐름

1234567891011flowchart LR
  A[토큰 x] --> B[셀프 어텐션<br/>+ 더하기 · 정규화]
  B --> C[라우터 W_r<br/>전문가별 점수 h x]
  C --> D[소프트맥스 → 확률 p_i x]
  D --> E[확률 1등 전문가 하나 선택<br/>top-1 스위치]
  E --> F{전문가 용량 남았나}
  F -->|예| G[전문가 FFN_i 계산<br/>출력 × p_i x]
  F -->|아니요, 넘침| H[계산 건너뜀<br/>잔차 연결로 그대로 통과]
  G --> I[더하기 · 정규화 → 다음 층]
  H --> I
  D -.-> J[부하 균형 보조 손실<br/>α · N · Σ f_i · P_i]

핵심 아이디어

1. 기존 MoE 층 — 라우터와 가중 합

한 문장 직관: 큰 회사의 접수 창구(라우터)가 들어온 문의(토큰)를 보고 담당 부서(전문가)를 정해 넘기는 구조입니다.

단계별 설명

1. 토큰 표현 $x$에 라우터 가중치 $W_r$을 곱해 전문가 $N$명 각각의 점수 $h(x) = W_r \cdot x$를 만듭니다.
2. 소프트맥스로 점수를 확률 $p_i(x)$로 바꿉니다.
3. 상위 $k$개 전문가를 골라 각자의 출력 $E_i(x)$를 계산합니다.
4. 고른 전문가 출력을 확률로 가중 합해 층의 출력으로 씁니다.

게이트 값 $p_i(x)$가 출력에 곱해지므로 라우터도 역전파로 학습됩니다.

2. 스위치 라우팅 — 전문가는 딱 하나만 (k = 1)

한 문장 직관: 문의 하나를 두 부서에 동시에 보내지 말고, 가장 잘할 것 같은 한 부서에만 보냅니다.

단계별 설명

1. 이전 연구는 라우터가 학습되려면 최소 두 전문가를 비교해야 한다($k > 1$)고 추측했습니다.
2. 스위치 트랜스포머는 확률이 가장 높은 전문가 하나만 고릅니다. 이것을 스위치 층이라 부릅니다.
3. 출력은 $y = p_{i^*}(x)\, E_{i^*}(x)$, 곧 고른 전문가의 출력에 그 게이트 값을 곱한 것입니다.
4. 트랜스포머의 조밀한 FFN 층을 이 스위치 FFN 층으로 바꿉니다. 실험에서는 대개 두 층에 하나꼴로 바꿨습니다.

왜 하나만인가: 논문은 세 가지 이점을 듭니다. ① 라우터 계산이 줄어듭니다. ② 토큰이 전문가 하나에만 가니 각 전문가가 처리할 배치(용량)를 적어도 절반으로 줄일 수 있습니다. ③ 구현이 단순해지고 통신 비용이 줄어듭니다. 그러면서도 품질은 유지되고 오히려 나았다고 보고합니다.

숫자로 해 보기 (해설용, 전문가 4명)

라우터 점수가 $h(x) = (2.0, 1.0, 0.5, 0.0)$이라면

1. $e^{2.0} \approx 7.39$, $e^{1.0} \approx 2.72$, $e^{0.5} \approx 1.65$, $e^{0} = 1$, 합 $\approx 12.76$
2. 확률 $p \approx (0.579, 0.213, 0.129, 0.078)$
3. 기존 top-2 MoE라면 출력 $= 0.579\,E_1(x) + 0.213\,E_2(x)$ — 전문가 두 명분 계산
4. 스위치(top-1)라면 출력 $= 0.579\,E_1(x)$ — 전문가 한 명분 계산

3. 전문가 용량과 버려지는 토큰

한 문장 직관: 부서마다 하루에 처리할 수 있는 문의 수(용량)를 미리 정해 두고, 넘치는 문의는 그냥 다음 단계로 넘깁니다.

단계별 설명

1. TPU는 텐서 크기를 미리 정해야 하므로, 전문가마다 처리할 토큰 수(용량)를 고정합니다.
2. 용량 = (배치의 토큰 수 / 전문가 수) × 용량 계수
3. 용량 계수가 1보다 크면 쏠림을 대비한 여유가 생기지만, 남는 칸만큼 계산 · 메모리가 낭비됩니다.
4. 한 전문가에 너무 많은 토큰이 몰리면 넘친 토큰은 그 층 계산을 건너뛰고 잔차 연결로 그대로 다음 층에 갑니다(버려진 토큰).

숫자로 해 보기 (해설용): 배치에 토큰 1,024개, 전문가 8명, 용량 계수 1.25라면 전문가당 용량은 $1{,}024 / 8 \times 1.25 = 160$개입니다. 어떤 전문가에 200개가 몰리면 40개는 그 층에서 계산되지 않습니다. 용량 계수를 2.0으로 올리면 256개까지 받지만, 고르게 128개씩 왔을 때는 절반이 빈칸으로 낭비됩니다.

논문 실험에서 버려지는 토큰은 전문가 수와 상관없이 대개 1% 미만이었고, 스위치는 낮은 용량 계수(1.0, 1.25)에서 특히 잘했습니다.

4. 부하 균형 보조 손실

한 문장 직관: 한 부서에만 일이 몰리지 않도록, "고르게 나눌수록 벌점이 작아지는" 항을 손실에 더합니다.

단계별 설명

1. 배치 $\mathcal{B}$(토큰 $T$개)에서 전문가 $i$에 실제로 보내진 토큰 비율 $f_i$를 셉니다.
2. 전문가 $i$에 준 라우터 확률의 평균 $P_i$를 구합니다.
3. 두 벡터의 내적에 전문가 수 $N$과 계수 $\alpha$를 곱해 보조 손실로 더합니다.
4. 완전히 고르면 $f_i = P_i = 1/N$이고 이때 손실이 최소입니다. $f$는 미분이 안 되지만 $P$는 미분됩니다.

숫자로 해 보기 (해설용, 전문가 2명)

• 고르게 나뉜 경우: $f = (0.5, 0.5)$, $P = (0.5, 0.5)$ → $2 \times (0.25 + 0.25) = 1.0$
• 한쪽에 쏠린 경우: $f = (1, 0)$, $P = (0.9, 0.1)$ → $2 \times (0.9 + 0) = 1.8$

$\alpha = 0.01$이면 보조 손실은 각각 0.010, 0.018입니다. 쏠릴수록 벌점이 커지므로, 라우터가 $P$를 고르게 만드는 쪽으로 학습됩니다. 계수가 작아서 주 손실(교차 엔트로피)을 압도하지는 않습니다.

5. 안정적으로 학습하는 세 가지 요령

한 문장 직관: 덩치 큰 희소 모델은 쉽게 흔들리니, 위험한 부분만 정밀하게 계산하고, 작게 시작하고, 미세조정 때는 전문가 쪽을 더 세게 규제합니다.

(품질은 음의 로그 퍼플렉서티, 높을수록 좋음)

왜 라우터만 float32인가: 라우터의 소프트맥스 · 지수 계산이 낮은 정밀도에서 특히 불안정합니다. 라우터 계산은 한 장치 안에서만 일어나고, 결과를 다시 bfloat16으로 바꿔 보내므로 비싼 float32 통신 없이 안정성만 얻습니다.

6. 파라미터는 늘리고 연산은 그대로

한 문장 직관: 같은 계산량의 모델끼리 비교하면, 전문가를 늘린 쪽이 파라미터가 수십 배 많습니다.

숫자로 해 보기 (해설용 계산): Switch-Base는 T5-Base와 시퀀스당 FLOPs가 같지만 파라미터는 $70 / 2 = 35$배입니다. 토큰 하나는 여전히 전문가 하나만 거치므로 계산은 늘지 않고, 저장해 둔 "지식"의 양만 늘어납니다.

결과

스위치 vs MoE vs 조밀 모델(10만 단계 후, 전문가 128명, TPUv3 32코어)

같은 연산 · 같은 시간 기준으로 스위치가 가장 좋았고, 낮은 용량 계수에서 더 유리했습니다.

사전학습 속도

• 전문가 수를 2명에서 256명(147억 파라미터)까지 두 배씩 늘리자, 토큰당 연산이 같은데도 손실이 꾸준히 좋아졌습니다.
• Switch-Base(전문가 64명)는 T5-Base가 45만 단계에 낸 성능을 6만 단계에 냈습니다(단계 기준 7.5배). 실제 시간 기준으로도 7분의 1 시간에 같은 품질에 도달했습니다.
• 토큰당 연산을 3.5배 더 쓰는 T5-Large와 비교해도 Switch-Base가 시간 기준 2.5배 빨랐습니다.

미세조정(검증 세트, 높을수록 좋음)

ARC를 제외한 대부분 과제에서 스위치가 나았습니다(SuperGLUE에서 Base 4.4%p, Large 2%p 향상).

증류: 큰 희소 모델을 작은 조밀 모델(T5-Base, 2억 2,300만)로 증류하면, 파라미터를 95% 넘게 줄이면서 품질 향상분의 약 30%를 지켰습니다(99% 압축에서도 28%).

다국어(mC4, 101개 언어): 101개 모든 언어에서 mT5-Base보다 좋아졌고, 평균 5배 빨랐으며 91%의 언어에서 4배 이상 빨랐습니다.

조 단위 모델: Switch-C(1.6조, 전문가 2,048명)는 같은 연산에서 T5-XXL보다 고정 퍼플렉서티 도달이 4배 빨랐습니다. Switch-C는 학습 불안정이 전혀 없었지만, FLOPs가 거의 10배인 Switch-XXL은 가끔 불안정했습니다.

왜 중요한가 · 이후 영향

• "파라미터 ≠ 연산"의 분리: 모델의 지식 용량과 토큰당 비용을 따로 조절할 수 있다는 생각이 대중화됐습니다. 이후 Mixtral 같은 공개 MoE LLM이 같은 원리(토큰마다 일부 전문가만 사용)를 씁니다.
• 단순함의 힘: top-1이라는 가장 단순한 선택이 오히려 더 좋았다는 결과는 MoE 설계를 크게 쉽게 만들었습니다.
• 실무 요령: 라우터만 float32로 계산하기, 부하 균형 손실, 작은 초기화, 전문가 드롭아웃은 이후 희소 모델 학습의 기본 도구가 됐습니다.
• 작은 환경에서도: 전문가 2명만으로도, 코어가 적은 환경에서도 이득이 있었다고 보고합니다.

한계 · 주의할 점

• 학습 불안정은 완전히 풀리지 않았습니다. 가장 큰 Switch-XXL은 가끔 불안정해서 계획한 100만 단계까지 학습하지 못했습니다.
• 사전학습 이득이 미세조정으로 다 이어지지는 않았습니다. Switch-XXL은 사전학습 퍼플렉서티가 최고였지만 SuperGLUE 87.5로 T5 쪽(89.3)보다 낮았습니다. 저자들도 이 차이를 열린 문제로 남깁니다.
• 메모리와 통신: 계산은 그대로여도 파라미터가 많으니 저장 메모리가 크고, 전문가가 여러 장치에 흩어져 통신 비용이 듭니다. 배포하려면 증류 등이 필요합니다.
• 버려지는 토큰: 용량을 넘친 토큰은 그 층을 건너뜁니다. 용량 계수를 키우면 줄지만 낭비가 늘어나는 맞교환이 있습니다.
• TPU · Mesh TensorFlow 중심: 구현과 속도 수치는 이 환경 기준입니다. 저자들도 속도는 알고리즘뿐 아니라 저수준 구현에도 좌우된다고 밝힙니다.

이어서 읽을 것

• 앞: 어텐션만 있으면 된다 (트랜스포머) · 신경망 언어 모델의 스케일링 법칙
• 뒤: Mixtral: 전문가 혼합 모델 · 연산 최적 대형 언어 모델 학습 (Chinchilla)

용어 정리

스스로 점검

1. 스위치 라우팅은 기존 MoE와 무엇이 다르고, 왜 그렇게 했나요?

   기존 MoE는 토큰마다 상위 $k \ge 2$개 전문가 출력을 섞었고, 스위치는 확률 1등 전문가 하나만 씁니다. 라우터 계산 · 전문가 용량 · 통신이 줄고 구현이 단순해지면서 품질은 유지되거나 더 좋았기 때문입니다.

2. 배치 토큰 2,048개, 전문가 16명, 용량 계수 1.0이면 전문가당 용량은 얼마이고, 한 전문가에 150개가 몰리면 어떻게 되나요?

   용량은 $2{,}048 / 16 \times 1.0 = 128$개입니다. 넘친 22개 토큰은 그 층 계산을 건너뛰고 잔차 연결로 다음 층에 그대로 넘어갑니다.

3. 부하 균형 보조 손실은 언제 가장 작아지나요?

   실제 배정 비율 $f_i$와 라우터 확률 평균 $P_i$가 모두 $1/N$으로 고를 때입니다. 이때 $N \sum f_i P_i = 1$이 됩니다.

다음에 볼 것

이 글을 이해했다면 이어서 보면 좋은 항목입니다.

• Mixtral: 전문가 혼합 모델 — 전문가 8명 중 2명만 쓰는 MoE로 활성 130억 파라미터에 Llama 2 70B급 성능을 낸 Mixtral

출처

• 원문: arxiv.org/abs/2101.03961 (2021 · William Fedus 외)