Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
Switch Transformers: 단순하고 효율적인 희소성으로 조 단위 파라미터 모델로 스케일링
William Fedus, Barret Zoph, Noam Shazeer (2022)
Mixture of Experts(MoE) 라우팅을 단순화하여 각 토큰이 하나의 전문가만 선택하는 Switch Transformer를 제안했다. 이를 통해 통신 비용과 학습 불안정성을 줄이면서 조 단위 파라미터 규모로 효율적으로 확장할 수 있음을 보였다.
배경
Transformer 모델의 성능은 파라미터 수에 따라 향상되지만, 밀집(dense) 모델은 모든 입력에 대해 전체 파라미터를 활성화하므로 계산 비용이 파라미터 수에 비례하여 증가한다. Mixture of Experts(MoE)는 입력에 따라 일부 파라미터만 활성화하여 이 문제를 해결할 수 있지만, 기존 MoE(Shazeer et al., 2017)는 top-k(보통 k=2) 전문가를 선택하는 라우팅의 복잡성, 전문가 간 부하 불균형, 학습 불안정성 등의 문제가 있어 대규모 적용이 어려웠다. 특히 분산 학습 환경에서 다수 전문가 간 통신 비용이 병목이었다.
핵심 아이디어
Switch Transformer의 핵심 아이디어는 MoE 라우팅을 극단적으로 단순화하는 것이다. 기존의 top-2 이상 전문가 선택 대신 각 토큰이 정확히 하나의 전문가(top-1)만 선택하는 'Switch 라우팅'을 도입한다. 이는 직관에 반하지만, 라우터 계산량 감소, 전문가 용량(capacity factor) 절반화, 통신 비용 감소라는 세 가지 이점을 동시에 달성한다. 또한 전문가 부하 균형을 위한 보조 손실(auxiliary load-balancing loss)을 단순화하고, bfloat16 정밀도로 학습 안정성을 확보하며, 소수 전문가에서 다수 전문가로 점진적으로 확장하는 전략을 제안한다.
방법론
Transformer의 각 피드포워드(FFN) 레이어를 N개의 독립적인 전문가 FFN으로 대체하고, 학습 가능한 라우터 네트워크가 각 토큰을 하나의 전문가에 할당한다. 라우터는 토큰 임베딩에 대한 선형 변환 후 소프트맥스를 적용하여 확률 분포를 생성하고, 가장 높은 확률의 전문가를 선택한다. 전문가 용량 팩터(capacity factor)를 설정하여 한 전문가에 과도한 토큰이 할당되는 것을 방지하며, 초과 토큰은 잔차 연결을 통해 다음 레이어로 전달된다. 부하 균형 손실 alpha * N * sum(f_i * P_i)를 추가하여 전문가 활용의 균형을 유도한다. 1.6조 파라미터 모델까지 확장하여 C4 데이터셋에서 학습했다.
주요 결과
Switch Transformer는 동일한 계산 예산(FLOPS)에서 밀집 T5 모델 대비 최대 7배 빠른 사전학습 속도를 달성했다. T5-Base 규모에서 Switch-Base(128 전문가)는 동일 학습 시간 대비 밀집 모델보다 현저히 낮은 perplexity를 보였다. T5-XXL과 비교하여 Switch-XXL은 4배 적은 학습 스텝으로 동등한 성능에 도달했다. 다운스트림 태스크(SuperGLUE, ARC, Winogrande 등)에서도 밀집 모델을 일관되게 능가했으며, 1.6T 파라미터 모델의 안정적 학습을 시연했다.
임팩트
Switch Transformer는 MoE 아키텍처를 대규모 언어 모델에 실용적으로 적용할 수 있는 경로를 열었다. top-1 라우팅의 성공은 이후 GShard, GLaM, Mixtral 등 MoE 기반 모델 설계에 직접적 영향을 미쳤으며, 특히 Mixtral 8x7B의 상업적 성공에 이론적 기반을 제공했다. 희소 활성화(sparse activation)를 통해 추론 효율성을 유지하면서 모델 용량을 확장하는 패러다임은 현재 LLM 아키텍처 연구의 핵심 방향 중 하나가 되었다.