[리뷰] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
OpenAI는 Scaling Law for Neural Language Models 논문(‘20)에서 언어 모델의 성능은 모델 사이즈, 컴퓨팅 budget, 데이터셋 사이즈 등 3가지 factor에 의존함을 밝혔다. Scaling law에 의해 모델 사이즈가 증가하는 만큼 컴퓨팅 budget이 패널티로 작용하지 않기 위해선, 컴퓨팅 budget 또한 증가해야 한다. 하지만 모델 사이즈를 증가시키는 것은 쉬운 일이지만(파라미터 config만 바꾸면 되는 일이므로..) 컴퓨팅 budget을 늘리는 것은 매우 많은 비용이 필요하다.
따라서 모델 사이즈가 늘어나면 증가하는 연산량을 효율적으로 제어할 수 있는 방법이 필요하다. 이를 해결하기 위해 Google Brain은 ‘21년 1월 Transformer 아키텍처의 FFNs 레이어를 MoE (Mixture of Experts)로 대체하여 연산량(FLOPS)는 유지하면 모델 사이즈를 크게 증가시킬 수 있는 Switch Transformer를 발표하였다.
Switch Transformer의 장점은 다음과 같다.
- 강력하게 튜닝된 T5 모델와 비교 벤치마크에서 토큰당 동일한 FLOPS를 사용하면서 7배 이상의 pre-training 속도 향상을 기록함. 제한된 연산 리소스(e.g., 2개의 experts를 사용)에서도 개선사항이 유지됨
- Sparse로 pre-training하고 특별하게 fine-tuned 모델을 작은 dense 모델로 distillation하는 것이 가능함. 큰 sparse 모델을 teacher로 사용하여 teacher 모델의 품질 이득 30%를 유지하면서 모델 크기를 최대 99%까지 감소시킴(기존 T5와 같은 파라미터로 줄임)
- Pre-training 시 더 낮은 bfloat16 정밀도로 학습이 가능하도록 선택적인 precision 학습
- 101개 모든 언어에서 보편적으로 개선을 발견하고 91%의 언어가 mT5 베이스라인보다 4 배 이상의 속도 향상을 보임
- 1 Trillion 이상의 파라미터를 갖는 모델을 생성하기 위해 Data, Model, Experts parallelism을 효율적으로 결합하여 자연어 모델의 규모를 증가시킴. 이 모델은 강력하게 튜닝된 T5-XXL 베이스라인의 pre-training 속도보다 4 배 이상 개선함
Mixure of Experts(MoE)
MoE는 91년도에 Geoffrey Hintton이 최초로 Adaptive Mixtures of Local Experts 논문에서 최초 제안하였다. MoE가 언어모델에 사용된 예는 다음과 같다.
Outrageously Large Neural Networks: The Sparsely-gated Mixture-of-Experts Layer(‘17)
Transformer가 발표되기 전인 ‘17년 초에 Transformer의 2 저자인 Noam Shazeer가 주저자이며 Geoffrey Hitton와 Jeff Dean이 참여한 논문[1]에서 LSTM 사이에 MoE를 내장한 언어 모델 논문을 발표하였다.
이때 MoE의 gating network는 입력 x에 맞는 K개의 experts를 선택할 수 있었다. Experts를 선택하기 위한 select 함수로 softmax를 사용하였으며 softmax에 의해 Top-K의 experts을 선택하여 weighted sum을 취한 후 다음 layer로 그 결과값을 보내게 된다.
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding(‘20)
앞서 소개했던 Sparsely-gated MOE 논문이 LSTM 사이에 MoE를 사용하였던 반면, Google이 ‘20년도에 발표한 GShard 논문에서는 Transformer Encoder의 FFN 레이어를 MoE로 대체하였다. 또한 experts들이 디바이스에 분산되어 있어 각 디바이스에 expert로 입력을 보낸 다음, 해당 expert의 연산 결과를 다시 가져오는 방식으로 동작한다.
GShard의 저자가 밝히는 GShard의 단점은 다음과 같다.
- Top-K experts를 고르는데 load balancing에 문제가 있다. 각 토큰에 맞는 expert를 찾는 게 중요하나 특정 experts에 몰리면서 계산이 느려진다.
- 디바이스에 experts들이 분산되어 있으므로 디바이스간 통신 비용이 매우 크다.
Switch Transformer의 아키텍처
Switch Transformer의 가장 큰 특징은 transformer의 FFN layer를 MoE로 바꾸고 MoE의 router(Gate function)를 통해 다수의 expert 중에서 token을 처리할 experts를 찾는다. Switch Transformer은 또한 GShard의 문제점을 해결하였다.
Switch Routing (Top-1)
Switch Transformer는 GShard가 Top-k의 experts를 뽑는 것을 Top-1의 expert만을 뽑는 것으로 바꿨다. (k=1 라우팅 전략) 즉 token마다 매칭되는 expert는 단 한개이다. 그리고 expert가 입력받을 수 있는 입력의 개수를 제한하며 한도를 넘는 입력에 대해서는 처리하지 않는다. 라우팅 방식을 Top-1 expert 라우팅 방식으로 바꾸면 다음과 같은 이점이 있다.
- (1) 라우팅을 위한 연산이 줄어든다.
- (2) 각 expert의 배치 사이즈(expert capacity)가 최소 절반이 된다.
- (3) 라우팅 구현이 간단해지고 Comm. 비용이 줄어든다.
Efficient Sparse Routing
Distributed Switch Implementation
모든 텐서 shape는 컴필레이션 시간에 정적으로 결정되지만, 계산은 학습과 추론에 라우팅이 결정되기 때문에 동적이다. 이 때문에 Switch Transformer에서는 expert capacity를 도입한다. Expert capacity가 커질수록 연산속도는 느려진다.
Capacity factor는 1.0 보다 더 큰 값(e.g, 1.0, 1.5 등)으로 token이 완벽히 experts에 균형있게 할당되지 못할 때 token을 수용하기 위한 추가 버퍼를 생성한다. 너무 많은 token이 expert에 라우팅되면, 계산을 건너뛰고(아래 그림에서 capacity factor:1.0일 때 빨간 색 라인은 expert capacity 이상의 token이 특정 expert에 전달되는 경우이다.) 그 token은 residual connection을 통해 다음 레이어로 직접 전달된다. Expert capacity를 늘리는 것은 computation와 메모리를 낭비하는 원인이 된다.
A Differentiable Load Balancing Loss
Switch Transformer는 학습하는 동안, 보조 loss를 전체 모델 loss 에 추가한다. 아래 보조 loss는 N개의 experts, T개의 Token을 갖는 Batch B와 같은 파라미터로 구성되어 있으며 벡터 f와 P간 scaled dot-product로 계산된다. 이 보조 loss는 token의 배치가 N개의 expert에 균등하게 라우팅되는 것을 목적으로 하기 때문에 uniform 분포를 갖을 때 최소화 된다.
- f_i: expert i에 디스패치되어진 token의 비율
- P_i: expert i에 할당된 라우터 확률의 비율
- α: multiplicative coefficient. 기본 cross-entropy 목적 함수를 압도하지 않으면서 로드 밸런싱을 보장할만큼 충분히 큰 값으로 α=0.01을 사용.
Switch Transformer와 MoE Transformer의 비교
MoE Transformer는 Top-2 라우팅을 사용하며, 동일한 HW상에서 동일 steps 동안 학습되었다. MoE 모델은 capacity factor를 2.0에서 1.25로 조정하였을 때 실제로 기대와 달리 학습 속도가 840 exams./s에서 790 exams./s로 느려졌다. 이 표에서는 다음과 같은 결과를 얻을 수 있었다.
- (1) Switch Transformer는 Speed-quality 측면에서 튜닝된 dense 모델과 MoE transformer를 압도한다. 한정된 computation과 wall-clock time에서 Switch Transformer가 가장 좋은 결과를 달성하였다.
- (2) Switch Transformer는 MoE transformer에 비해 더 작은 계산 footprint를 갖는다.
- (3) Switch Transformer는 더 낮은 capacity factor (1.0 ~ 1.25)에서 더 좋은 성능을 보인다. 더 작은 capacity factor은 모델 메모리가 매우 부족하고 capacity factor가 가능한한 더 작게 만들어지길 원하는 large model에 적합하다.
Training & Fine-tuning 기술
Sparse expert 모델은 vanilla transformer보다 학습에 어려움이 있다. 각 레이어에서 hard-switching 결정때문에 불안정성이 발생할 수 있다. 또한 bF16과 같은 낮은 precision 포맷은 router의 softmax 계산 시 이슈를 발생시킬 수 있다. 다음은 switch transformer는 안정적이고 확장 가능한 학습을 달성하는 방법이다.
Selective precision with large sparse models
Switch transformer는 router 입력을 FP32로 캐스트한다. FP32는 router 함수의 body 이내에서만 사용된다. 결과 디스패처와 조합 텐서는 router 함수의 끝에서 BF16으로 다시 캐스트된다. 이는 값 비싼 FP32 텐서가 all-to-all comm. 을 통해 브로드캐스트되지 않도록 하고, FP32의 향상된 학습 안정성으로부터 이득을 얻는다.
Smaller parameter initialization for stability
적절한 초기화는 딥러닝에서 성공적인 학습을 위해 매우 중요하다. Switch transformer는 초기화가 더욱더 중요한 것으로 관찰되었다. Switch의 초기화 기법은 다음과 같다.
- (1) Switch transformer는 평균 µ =0 과 표준 편차 σ = sqrt(s/n)인 잘린(truncated) 정규 분포로부터 요소를 뽑아내 weight matric를 초기화한다. (s: scale hyper-parameter, n: weight 텐서에 입력 유닛의 개수)
- (2) 기본 transformer의 초기 scale인 s=1.0을 10배로 줄이는 것을 추천한다.
이와 같은 방식을 통해 모델 품질의 개선 및 학습 초기에 variance를 줄이는 것으로 측정되었다.
Regularizing large sparse models
Fine-tuning 시 overfitting을 방지하기 위해 experts 안에 dropout을 증가시키는 소위 “expert dropout”을 사용한다. Fine-tuning 시 non-experts 레이어에 더 작은 dropout 비율(=0.1)을, experts 레이어에 더 많은 dropout 비율(=0.4)을 세팅하면 4개의 더 작은 downstream task에서 성능이 향상된다.
레퍼런스
[1] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
[2] Outrageously Large Neural Networks: The Sparsely-gated Mixture-of-Experts Layer
[3] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
