안녕하세요 Simon 입니다.
오늘 리뷰해볼 논문은 MoE 모델의 아키텍처인 Switch Transformers 입니다. 사실 제가 딥러닝을 배우기전에 나왔던 논문인데 신기하게 다 언젠가는 빛을 바라는 순간이 오는거 같습니다.
Abstrack
스위치 트랜스포머는 전통적인 딥러닝 모델의 한계를 극복하기 위해 개발된 새로운 접근 방식입니다. 기존 모델들이 모든 입력에 대해 동일한 매개변수를 사용하는 반면, 스위치 트랜스포머는 각 입력 예제마다 다른 매개변수를 선택하는 혼합(전문가 혼합) 모델을 채택합니다. 이 방식은 매개변수의 수를 크게 늘리면서도 계산 비용을 일정하게 유지합니다.
그러나 이러한 모델은 복잡성, 통신 비용, 그리고 훈련의 불안정성과 같은 문제들을 가지고 있었습니다. 스위치 트랜스포머는 이러한 문제들을 해결하기 위해 MoE(Mixture of Experts) 라우팅 알고리즘을 간소화하고 직관적으로 개선하였습니다. 이를 통해 통신 및 컴퓨팅 비용을 절감하고, 불안정성을 완화하는 동시에, 대규모 희소 모델을 처음으로 더 낮은 정밀도(bfloat16)로 훈련할 수 있게 되었습니다.
이 연구는 T5-Base와 T5-Large 모델을 기반으로 하여, 동일한 계산 리소스로 사전 훈련 속도를 최대 7배까지 향상시켰습니다. 또한, 다국어 설정에서도 이러한 개선 사항이 적용되어, 101개 언어에 대해 mT5-Base 버전보다 성능이 향상되었습니다. 마지막으로, 이 연구는 언어 모델의 규모를 현저히 확장하여, 최대 1조 개의 매개변수를 가진 모델을 대규모 크롤링 코퍼스에 사전 학습시키는 데 성공했으며, 이는 T5-XXL 모델보다 4배 빠른 속도를 달성했습니다.
Introduce
Switch Transformer는 다양한 자연어 처리 작업에서 우수한 확장성을 보여주며, 사전 훈련, 미세 조정, 다중 작업 훈련 등 NLP의 세 가지 영역에서 효과적입니다. 이 아키텍처는 슈퍼컴퓨터 영역뿐만 아니라 몇 개의 계산 코어만을 사용하는 경우에도 유용합니다. 또한, 큰 희소 모델은 작은 밀집 모델로 압축(distillation)할 수 있으며, 이때 희소 모델의 품질 향상의 30%를 유지합니다.
기여한점 :
- Switch Transformer 아키텍처: MoE를 단순화하고 개선함.
- T5 모델과 비교한 확장성 및 벤치마크: 같은 FLOPS를 사용하면서 7배 이상의 사전 훈련 속도 향상을 보임.
- 희소 사전 훈련 및 특화된 미세 조정 모델을 작은 밀집 모델로 성공적으로 압축.
- 개선된 사전 훈련 및 미세 조정 기술.
- 다국어 데이터에 대한 사전 훈련 혜택 측정: 모든 101개 언어에서 개선을 보이며, 91%의 언어가 mT5 기준 대비 4배 이상의 속도 향상을 경험함.
- 데이터, 모델, 전문가 병렬 처리를 효율적으로 결합하여 최대 1조 개의 파라미터를 가진 모델을 생성함으로써 신경 언어 모델의 규모를 확장함. 이러한 모델은 T5-XXL 기준의 사전 훈련 속도를 4배 향상시킴.
Switch Transformer
OpenAI의 "Scaling Law for Neural Language Models" 논문에 따르면, 언어 모델의 성능은 모델 크기, 컴퓨팅 예산, 데이터셋 크기라는 세 가지 주요 요인에 의존합니다. 모델 크기를 증가시키는 것은 비교적 간단한 작업입니다(파라미터 설정을 변경하기만 하면 됩니다). 그러나 이러한 확장은 컴퓨팅 예산의 증가 없이는 효과적이지 않습니다. 컴퓨팅 예산을 늘리는 것은 상당한 비용이 드는 작업이기 때문에, 모델을 확장하는 것은 비용적인 측면에서 중요한 고려사항이 되기에 MoE 아키텍처가 탄생하지 않았나 생각합니다. 이 모델은 GPU 및 TPU와 같은 밀집 행렬 곱셈에 최적화된 하드웨어를 효율적으로 사용하는 희소 활성화 모델을 설계함으로써 이를 달성합니다. 분산 훈련 설정에서, 이 희소 활성화 레이어는 다양한 장치에 고유한 가중치를 분할합니다. 따라서 모델의 가중치는 장치의 수와 함께 증가하지만, 각 장치에서는 관리 가능한 메모리 및 계산 부담을 유지합니다.
트랜스포머 모델에서 각 입력 토큰이 특정 전문가에 의해 독립적으로 처리되는 방식을 시각적으로 보여줍니다. 이는 모델의 파라미터 수를 크게 늘리면서도 각 장치에서의 계산 부담을 줄이는 방법입니다.
1. 스위치 FFN 레이어: 트랜스포머 모델의 일반적인 밀집 피드 포워드 네트워크(FFN) 대신, 희소 스위치 FFN 레이어를 사용합니다. 이 레이어는 연한 파란색으로 표시되어 있습니다.
2. 독립적 작동: 이 레이어는 입력된 시퀀스의 각 토큰에 독립적으로 작동합니다. 예를 들어, "More"와 "Parameters"라는 두 토큰이 각각 다른 처리 과정을 거칩니다.
3. 라우팅: 각 토큰은 라우터에 의해 네 개의 FFN 전문가 중 하나로 전달됩니다. 이는 실선으로 표시되어 있으며, 각 토큰이 어떤 전문가에게 할당되는지를 나타냅니다.
4. 게이트 값과의 곱셈: 선택된 FFN의 출력은 라우터 게이트 값(점선으로 표시)과 곱해져 최종 출력을 생성합니다. 이 과정은 각 토큰이 처리되는 방식을 최적화하여 효율성을 높입니다.
