Switch Transformer

2021년 1월, Google Brain 연구팀이 공개하며 전 세계 AI 커뮤니티를 놀라게 했던 Switch Transformer에 대해 알아보겠습니다.

이 모델은 당시 "역대 최대 규모"라는 수식어와 함께, 거대 모델의 고질적인 문제였던 '계산 비용' 문제를 우아하게 해결한 모델로 평가받습니다.

 

Switch Transformer란?

Switch Transformer는 Google의 T5(Text-to-Text Transfer Transformer)를 백본(Backbone)으로 사용하는 모델입니다.

주요 특징:

• 압도적인 스케일: 무려 1조 6천억 개(1.6 Trillion)의 파라미터를 가졌습니다. 당시 GPT-3(1,750억 개)보다도 약 9배나 큰 규모였죠.
• 압도적인 속도: 파라미터 수는 훨씬 많지만, 고정된 계산 예산(Computational Budget) 내에서 T5 대비 약 4배에서 7배까지 빠른 학습 속도를 보여줍니다.

일반적으로 파라미터가 늘어나면 성능은 좋아지지만 계산 비용과 시간이 기하급수적으로 증가합니다. Switch Transformer는 어떻게 이 모순을 해결했을까요?

 

 

MoE (Mixture of Experts): "전문가들의 모임"

Switch Transformer를 이해하기 위한 핵심 개념은 MoE(Mixture of Experts)입니다.

기존 Transformer의 FFN(Feed Forward Network) 레이어는 모든 입력 데이터가 모든 노드를 통과해야 했습니다. 하지만 MoE는 다릅니다.

• Expert(전문가): 특정 데이터 패턴을 처리하는 작은 FFN들입니다.
• Router(게이트): 입력된 토큰을 어떤 전문가(Expert)에게 보낼지 결정하는 안내 데스크 역할을 합니다.

예를 들어 '고양이'라는 클래스가 있을 때 고양이 클래스를 나타내는 이미지는 다양할 것입니다.
고양이의 일부만 보일 수 있고, 다른 객체들이 고양이와 함께 있는 이미지도 있을 수 있겠죠.
그런 경우 하나의 모델이 다양한 입력 이미지를 학습하기 어려울 수 있기 때문에 'expert'라는 개념을 도입했습니다.

고양이의 일부만 보는 expert, 객체 탐지 expert, 배경 분리 expert 등의 다양한 expert를 두고, gate 또는 router라고 불리는 곳에서 각각의 입력에 대해 적합한 expert를 할당하는 거죠.

gate에서는 하나의 입력과 가중치를 곱해 softmax를 통과시켜, 높은 확률 순으로 적당한 expert 후보를 top K개 뽑고, top K개의 expert와 입력을 연산한 결과를 가중합하여 출력하게 됩니다.

 

그런데 MoE에는 몇 가지 단점이 있었습니다. 구조가 복잡하고 expert간 커뮤니케이션 비용이 높으며 학습이 불안정하다는 것이었죠.

왜 "Sparsely-activated model"이라 부를까요?

모든 노드가 활발히 움직이는 기존 모델과 달리, MoE는 특정 입력에 대해 필요한 전문가(일부 노드)만 활성화됩니다. 즉, 전체 파라미터는 엄청나게 많지만 실제로 연산에 참여하는 파라미터는 적기 때문에 '희소하게 활성화된다(Sparsely-activated)'고 표현합니다.

 

Switch Transformer의 혁신 기술

기존 MoE 방식(예: Shazeer et al., 2017)은 보통 하나의 토큰을 2개 이상의 전문가(Top-K)에게 보냈습니다. 하지만 Switch Transformer는 이를 더 단순화하고 최적화했습니다.

① Switch Routing (Top-1 Strategy)

Switch Transformer는 과감하게 단 하나의 전문가(Top-1)에게만 토큰을 보냅니다.

• 장점:
  • 1라우터 연산량이 절반으로 줄어듭니다.
  • 구현이 단순해지며 성능 저하가 거의 없습니다.
  • 전문가 간의 커뮤니케이션 비용(Network Bandwidth)이 획기적으로 감소합니다.

② Distributed Switch Implementation & Capacity Factor

전문가가 1개로 제한되더라도, 특정 전문가에게만 토큰이 몰리는 병목 현상이 발생할 수 있습니다. 이를 해결하기 위해 Expert Capacity라는 개념을 도입했습니다.

• Capacity Factor: 각 전문가가 처리할 수 있는 토큰의 양을 제한합니다.
• 만약 특정 전문가에게 토큰이 너무 많이 몰려 용량을 초과하면, 해당 토큰은 'Dropped' 처리되거나 다음 레이어로 바로 넘겨집니다(Residual connection). 이를 통해 연산의 균형을 맞춥니다.

③ Differentiable Load Balancing Loss

특정 전문가만 학습되고 나머지는 노는 현상을 막기 위해, 모든 전문가에게 토큰이 골고루 배분되도록 유도하는 손실 함수를 추가했습니다.

$$Loss = \alpha N \cdot \sum_{i=1}^{N} f_i \cdot P_i$$

• $f_i$: 전문가 $i$에 전달된 토큰의 비율
• $P_i$: 라우터가 전문가 $i$를 선택할 확률
• 이 두 값의 곱을 최소화하도록 학습하여, 결과적으로 균등한 분배를 만들어냅니다. (논문에서는 $\alpha = 0.01$ 사용)

 

학습의 안정성을 위한 추가 기법

거대 모델은 학습 중 수치가 발산(Explode)하기 쉽습니다. 이를 막기 위해 Google은 몇 가지 트릭을 더 썼습니다.

1. Selective Precision (선택적 정밀도): 전체 모델은 효율을 위해 bfloat16을 사용하지만, 불안정한 Router 연산에서만 float32를 사용합니다. 이를 통해 성능과 안정성을 모두 잡았습니다.
2. Smaller Initialization: 가중치 초기화 값을 일반적인 경우보다 훨씬 작게 설정하여 학습 초기 단계의 불안정성을 줄였습니다.
3. Expert Dropout: 과적합(Overfitting)을 방지하기 위해 각 전문가 레이어에 개별적으로 드롭아웃을 적용했습니다.

 

결론: Switch Transformer가 남긴 것

Switch Transformer는 "모델이 커진다고 해서 반드시 느려질 필요는 없다"는 것을 증명했습니다.

성과 요약:

• C4 데이터셋 학습: 동일한 계산 자원으로 T5-Base보다 7배 이상 빠른 수렴 속도를 보임.
• 효율성: 파라미터는 1조 개가 넘지만, 실제 추론 시 필요한 연산량은 훨씬 작은 모델 수준으로 유지.

이 모델의 등장 이후, 최근 등장하는 GPT-4나 Mixtral 같은 최신 모델들도 MoE 구조를 적극적으로 채택하게 되었습니다. 거대 언어 모델(LLM)의 효율화를 이끈 선구자적인 모델이라고 할 수 있겠네요!