Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

논문 개요

Mixture-of-Experts(MoE)는 모델의 파라미터 수를 늘리되 실제 활성화되는 파라미터(computation)는 고정하는 희소(sparse) 확장 방법이다. 각 입력 토큰이 전체 전문가(expert) 중 일부만 활성화하므로, 파라미터 수가 증가해도 추론·학습 FLOPs는 크게 늘지 않는다.

그러나 기존 MoE 방식(Shazeer et al., 2017)은 Top-2 라우팅, 복잡한 부하 분산, 학습 불안정성 등의 문제를 안고 있었다. Fedus et al.(2022)이 JMLR에 발표한 Switch Transformers는 이를 단순화하고 안정화하여 최초로 조 단위( $1 0^{12}$ ) 파라미터의 언어 모델을 학습하는 데 성공하였다.

핵심 기여

Switch(Top-1) 라우팅: 각 토큰을 단 하나의 전문가에게 라우팅하여 설계를 단순화하고 통신 비용을 절감한다.
보조 로드 밸런싱 손실: 전문가가 균등하게 활용되도록 유도하는 미분 가능한 보조 손실 함수를 제안한다.
학습 안정성 기법: 전문가별 가중치 초기화 스케일링, bf16 혼합 정밀도, 선택적 fp32 적용으로 학습 발산을 방지한다.
T5 대비 7× 속도 향상: 동일한 학습 FLOPs에서 Switch Transformer는 T5-Base보다 7배 빠르게 같은 품질에 도달한다.
조 파라미터 모델: Switch Transformer-1.6T (1.6조 파라미터, 2048개 전문가)를 처음으로 학습하였다.

방법론 상세

MoE 레이어 구조

Switch Transformer는 표준 트랜스포머의 FFN 레이어를 MoE 레이어로 교체한다. $N$ 개의 전문가 FFN ${E_{1}, E_{2}, \dots, E_{N}}$ 이 있고, 라우터 $W_{r}$ 이 각 토큰 $x$ 를 어느 전문가에게 보낼지 결정한다:

$h = \sum_{i = 1}^{N} g_{i} (x) \cdot E_{i} (x)$

여기서 게이팅 함수 $g_{i} (x)$ 는 소프트맥스 기반 확률: $g_{i} (x) = \frac{e x p ( x \cdot W _{r}^{(i)} )}{\sum _{j} e x p ( x \cdot W _{r}^{(j)} )}$

Switch(Top-1) 라우팅

기존 MoE는 $Top- k (k = 2)$ 라우팅을 사용하여 각 토큰이 두 전문가의 가중 합으로 처리된다. Switch Transformer는 $k = 1$ 로 단순화한다:

$h = E_{i^{*}} (x), i^{*} = ar g max_{i} (x \cdot W_{r}^{(i)})$

이 단순화의 장점은:

통신 감소: 분산 환경에서 토큰이 이동하는 디바이스 수가 절반으로 줄어든다.
연산 감소: FFN을 하나만 실행한다.
실증적 성능: $k = 1$ 이 $k = 2$ 보다 종종 더 좋거나 비슷한 성능을 보였다.

전문가 용량(Expert Capacity)

각 전문가는 한 배치 내에서 처리할 수 있는 토큰 수에 용량 제한을 둔다:

$C = ⌊ \frac{T}{N} \times c ⌋$

여기서 $T$ 는 배치 내 총 토큰 수, $N$ 은 전문가 수, $c$ 는 용량 계수(capacity factor, 보통 1.0~2.0)이다. 특정 전문가에 $C$ 를 초과하는 토큰이 할당되면 초과분은 오버플로우(overflow)되어 처리되지 않는다 (해당 토큰은 Residual Connection을 통해 전달됨).

로드 밸런싱 보조 손실

모든 토큰이 소수의 인기 전문가에게 집중되는 붕괴(collapse) 현상을 방지하기 위해, 보조 로드 밸런싱 손실을 메인 손실에 더한다:

$L_{aux} = α \cdot N \cdot \sum_{i = 1}^{N} f_{i} \cdot P_{i}$

여기서:

$f_{i} = \frac{전문가 i 에 라우팅된 토큰 수}{T}$ : 실제 라우팅 비율
$P_{i} = \frac{1}{T} \sum_{x} g_{i} (x)$ : 소프트 라우팅 확률의 평균
$α$ : 보조 손실 가중치 (보통 $1 0^{- 2}$ )

$f_{i}$ 와 $P_{i}$ 의 내적을 최소화하면 각 전문가에 균등하게 토큰이 분배된다. $f_{i}$ 는 미분 불가능하지만, $P_{i}$ 를 통해 그래디언트가 전달된다.

학습 안정성 기법

초기화 스케일링: 전문가 FFN의 가중치를 $N (0, σ^{2} / N)$ 으로 초기화하여 전문가 수가 증가해도 분산이 일정하게 유지된다.

혼합 정밀도: 대부분의 연산은 bf16으로 수행하되, 라우팅과 손실 계산은 fp32로 수행한다. bf16은 fp16과 같은 16비트지만 지수부가 더 넓어 오버플로우에 강하다.

소통 비용 최적화: 분산 학습(데이터 병렬 + 전문가 병렬)에서 전문가가 서로 다른 디바이스에 배치될 때, All-to-All 통신으로 토큰을 올바른 전문가에게 전달하는 통신 패턴을 최적화한다.

실험 결과

T5와의 학습 속도 비교

동일한 FLOPs(계산량)를 사용했을 때:

모델	파라미터	같은 품질 도달 시간
T5-Base	223M	기준
Switch-Base (128 전문가)	7.4B	7× 빠름
Switch-Large (128 전문가)	26.3B	4× 빠름

다운스트림 과제 (SuperGLUE)

Switch Transformer는 T5-11B와 비슷하거나 더 나은 SuperGLUE 점수를 훨씬 적은 FLOPs로 달성하였다.

전문가 수 스케일링

전문가 수를 늘릴수록 성능이 향상되지만, 256개 이상에서는 수익 체감(diminishing returns)이 관찰되었다. 전문가 수 대비 성능은 $lo g N$ 에 비례하는 경향이 있다.

의의 및 한계

의의

MoE의 민주화: 복잡했던 MoE 설계를 단순화하여 후속 연구(GLaM, Mixtral, DeepSeek-MoE 등)의 폭발적 성장을 촉발하였다.
스케일링 효율: 동일한 FLOPs로 더 많은 파라미터를 학습할 수 있어, 파라미터-효율적 스케일링의 길을 열었다.
T5와의 직접 비교: 동일한 프레임워크(T5/mesh-tensorflow)에서 비교하여 결과의 신뢰성이 높다.

한계

학습 불안정성: 논문에서도 인정하듯, 특히 대규모에서 학습이 발산하는 경우가 있다. 다양한 안정화 기법이 필요하다.
전문가 전문화 부족: 이상적으로는 각 전문가가 서로 다른 유형의 입력을 처리하길 바라지만, 실제로는 전문화가 불분명한 경우가 있다.
추론 시 메모리: 활성화되는 파라미터는 적지만, 전체 전문가 가중치를 메모리에 올려야 하므로 서빙 인프라 요구사항이 크다.
로드 불균형: 보조 손실에도 불구하고 학습 중 일부 전문가에 토큰이 집중되는 현상이 완전히 해결되지 않는다.

도형준 지식 볼트

탐색기