Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity - 논문 리뷰

https://arxiv.org/abs/2101.03961

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

 

 

1. 문제 정의

이 논문은 거대한 매개변수를 가진 언어 모델을 효율적으로 학습하기 위한 방법을 연구합니다. 전통적인 Transformer 모델은 모든 입력에 대해 동일한 매개변수를 사용해 전체 모델이 활성화되는 반면, 이 논문에서 다루는 Mixture of Experts (MoE) 모델은 각 입력에 맞는 전문가를 선택해 부분적으로만 모델을 활성화하는 방법을 사용합니다. 하지만, MoE 모델의 복잡한 라우팅 알고리즘과 높은 통신 비용, 훈련 불안정성 등으로 인해 널리 사용되지 못한 점이 문제로 제기되었습니다.

2. 접근 방법

논문은 이 문제를 해결하기 위해 Switch Transformer라는 새로운 모델 구조를 제안했습니다. 이 모델은 기존 MoE 모델에서 전문가 선택 과정을 단순화하고, 통신과 계산 비용을 줄여 모델의 안정성을 개선했습니다. 구체적으로, 논문에서는 각 입력을 한 명의 전문가에게만 라우팅하는 방식을 채택하여 라우팅 계산을 줄이고, bfloat16과 같은 낮은 정밀도를 사용하여 훈련 불안정성을 극복하는 방법을 사용했습니다.

3. 실험 방법

1. 모델 아키텍처 설계: Switch Transformer는 기존 T5 모델(T5-Base와 T5-Large)을 기반으로 전문가의 수를 증가시키고, 각 토큰을 선택된 전문가에게만 라우팅하여 계산 효율성을 높였습니다.
2. 훈련 및 미세 조정 기법: 낮은 정밀도의 훈련을 위한 선택적 정밀도 방식을 도입했으며, 불안정을 방지하기 위해 초기화와 정규화 기법을 적용했습니다. 또한, 각 전문가의 용량을 조절하는 보조 손실을 사용하여 균형적인 라우팅을 유도했습니다.
3. 확장성 테스트: 모델의 확장성을 평가하기 위해 전통적인 밀집 모델과 비교하여 동일한 계산 자원에서 7배 빠른 사전 훈련 속도를 기록하고, 이를 여러 언어 데이터셋과 다중 작업 학습 환경에서 평가했습니다.

4. 결과

Switch Transformer는 기존 밀집 모델(T5-Base) 대비 4배 이상의 속도 향상과 우수한 확장성을 보여주었으며, 101개 언어에 걸친 다중 언어 작업에서도 성능 향상을 기록했습니다. 또한, 모델 압축을 위해 대규모 희소 모델을 작은 밀집 모델로 증류하여 원본 성능의 30%를 유지하는 성과를 보였습니다. 다양한 실험을 통해 Switch Transformer가 단순히 대규모 모델로 확장될 수 있을 뿐만 아니라, 다중 전문가 시스템을 통해 효율적인 학습과 전이 학습이 가능함을 확인했습니다.

5. 한계점

논문에서 제안한 Switch Transformer는 대규모 모델에서의 불안정성과 고도의 통신 비용 문제를 어느 정도 해결했으나, 여전히 일부 극단적인 대규모 모델에서는 불안정성을 완전히 해결하지 못한 한계가 존재합니다. 또한, 훈련 불안정성은 높은 전송 비용과 결합되어 다중 전문가 모델을 확장하는 데 한계로 작용할 수 있습니다.

이 논문에서 제시된 방법과 결과는 향후 인공지능 해석 연구에서 희소성 기반의 라우팅 전략을 통한 모델의 효율적인 확장성을 고려하는 데 중요한 참고 자료가 될 수 있습니다.

 

 

 

기존 MoE 방식이 왜 복잡한 라우팅 알고리즘, 높은 통신 비용, 훈련 불안정성을 가졌는지 알려줘

기존 Mixture of Experts (MoE) 방식이 복잡한 라우팅 알고리즘, 높은 통신 비용, 훈련 불안정성을 가지게 된 이유는 다음과 같습니다.

1. 복잡한 라우팅 알고리즘

• MoE 방식에서는 각 입력 토큰마다 최적의 전문가(expert)를 선택하는 과정이 필요합니다. 일반적으로 여러 전문가 중 상위 k개를 선택하여 해당 입력 토큰을 라우팅하고, 라우팅을 위해 softmax로 모든 전문가에 대한 확률을 계산해 상위 전문가들을 선택합니다.
• 이 과정에서 각 토큰에 대해 적합한 전문가를 찾기 위해 복잡한 계산이 필요하고, 특히 많은 수의 전문가가 존재할 경우 라우팅 알고리즘이 매우 복잡해질 수 있습니다.
• 복잡한 라우팅은 전문가 간의 비대칭적인 로드 분배 문제를 야기할 수 있어, 일부 전문가에게만 데이터가 쏠리는 불균형 현상이 발생하고, 이는 성능 저하와 불안정성으로 이어질 수 있습니다.

2. 높은 통신 비용

• MoE 모델에서는 각 토큰을 최적의 전문가로 보내기 위해 분산 시스템 환경에서 장치 간 통신이 필수적입니다. 예를 들어, 수천 개 이상의 전문가가 여러 장치에 걸쳐 있을 때, 선택된 전문가에게 입력 토큰을 전달하고 그 결과를 다시 모으는 과정에서 막대한 통신 비용이 발생합니다.
• 또한, 이 과정에서 발생하는 통신 지연(latency) 문제는 특히 대규모 모델에서 훈련 속도를 크게 저하시킬 수 있습니다. 통신량이 많아질수록 처리 지연이 증가하고, 이는 학습 효율성과 속도에 악영향을 미칩니다.

3. 훈련 불안정성

• MoE에서는 라우팅 과정에서 비연속적인(스위칭) 선택이 이루어지므로, 각 입력이 특정 전문가에 선택될지 여부가 훈련 중에 동적으로 변화합니다. 이는 모델의 가중치 업데이트에 영향을 주어, 훈련 중 각 전문가의 활성화가 일관적이지 않게 됩니다.
• 이로 인해 그래디언트 흐름의 불안정성이 발생하고, 특히 bfloat16과 같은 낮은 정밀도 포맷에서 부동소수점 문제로 인해 오차가 더 커지게 되어 모델이 수렴하는 데 어려움을 겪을 수 있습니다.
• 또한, 불균형한 라우팅과 전문가의 비활성화(일부 토큰이 라우팅되지 않는 경우) 문제도 발생해 일부 전문가가 제대로 학습되지 않을 가능성이 있어, 결과적으로 모델 성능에 부정적인 영향을 미칩니다.

이러한 문제들은 MoE 방식이 큰 잠재력을 가지고 있음에도 불구하고 대규모 환경에서 적용하기 어려운 이유가 됩니다. Switch Transformer는 이를 해결하기 위해 라우팅 과정을 단순화하고 통신 비용을 줄여 효율적인 훈련을 가능하게 했습니다.