Sparcity in Transformers
(1) 모델의 파라미터 크기
(2) 데이터 셋의 크기
(3) 계산량(Computational Budget)
Switch Transformer의 저자들은 여기에 한 가지 차원을 더 추가하는데, 그것이 example당 FLOPs를 고정한 채로 늘어나는 parameter count를 늘릴 수록 성능이 좋아진다는 것이다. 이 개념과 Switch Transformer를 설명하기 위해서는 먼저 Mixture of Experts라는 개념을 살펴봐야 한다.
Mixture of Experts(MoE)
MoE는
Shazeer et al. 에서 소개된 개념으로 Neural Model(간단하게는 FFN)을 expert라고 부를 때 example 별로 적절한 expert에 routing을 해줘서 계산량을 줄이면서도 큰 모델을 만드는 방식이다. 위의 이미지는 RNN 같은 recurrent 구조에 MoE를 적용한 형태인데, Gating Network에서 주어진 샘플 \( x \)가 주어지면 \(n\)개의 expert에 대해서 어디에 할당할지에 대한 확률 분포를 계산하게 된다. 가장 단순한 Softmax Gating 방식은 다음과 같다.
$$ G(x) = \textrm{Softmax}(x \cdot W_g) \\ y = \sum_{i=1}^{n}G(x)_i E_i(x)\quad\textrm{where}\quad E_i = (i\textrm{-th expert})$$
softmax output을 그대로 사용하면 expert가 densely activate 되므로 sparse model의 이점을 누릴 수 없다. 따라서 여기에서 top-k개의 Expert만 사용하는 것이 일반적이다.
Shazeer et al. 에서는 적어도 두 개의 expert간 비교가 필요하기 때문에 k 값은 1보다는 커야한다고 주장했지만 Switch Transformer의 저자들은 각 token은 하나의 expert로 routing 되어도 충분하다는 것을 실험적으로 밝혔다. 1개의 expert로 routing 함으로서 얻는 이점은 굉장히 크다. 먼저 단적으로 1개의 token이 k개의 expert에서 계산될 것을 1개로 줄였기 때문에 1/k배로 계산량이 줄어든다. 라우팅 구현도 단순해지고 communication cost도 줄어든다.
위의 이미지는 각 token들이 expert로 분배되는 것을 나타낸 것인데 Capacity Factor(CF)라는 개념이 등장한다. 기본적으로 token 개수가 L개가 expert의 개수가 n개면 L/n개씩 각 expert에서 분담하면 된다. 그런데 CF=1.0 그림처럼 용량이 넘쳐서 가지 못하는 token들까지 수용해주기 위해서 용량을 몇배만큼 늘릴 것인지를 결정하는 값이 CF이다. CF가 늘어나면 계산량과 Communication Cost가 늘어나게 된다. Padding이 많이 늘어나는건 덤! CF=1.0의 케이스처럼 Expert 용량이 꽉 차서 더 이상 보내지 못하는 경우에는 계산을 하지 않고 skip connection으로 다음 layer로 더해주게 된다. (저자들이 TPU를 사용해서 실험했기 때문에 Compilition 시점에 이미 각 expert들의 용량을 정해야 했기 때문에 CF 라는 개념이 나온 부분도 있다.)
Load Balancing Loss
Routing 과정에서 특정 expert로 token이 몰리는 것은 학습을 비효율적으로 만든다. 따라서 각 expert로 uniform하게 token들이 분산될 수 있도록 하는 Loss를 추가한다. 여기에서 \(P_i\)가 미분가능하기 때문에 backpropagation이 가능하고 Expert의 개수에 관계 없이 최소값을 유지하기 위해 scaling factor로 \(N\)을 곱해줬다.
$$ \textrm{loss} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i \\ f_i = \frac{1}{T} \sum_{x \in \mathcal{B}} \mathbb{I}\{\textrm{argmax}p(x) = i \} \\ P_i = \frac{1}{T} \sum_{x \in \mathcal{B}} p_i(x)$$
Switch Transformer outperforms MoE and Dense T5 in speed-quality basis
Training Details
Selective Precision Training
기존에 제시되던 mixed-precision training과 유사하게 router 부분의 계산(router function의 input 부터 확률 곱 부분까지)만 32bit로 처리하는 Selective Precision Training을 제안하였고 full precision training과 같은 Quality를 가지는 것을 확인했다.
Smaller Parameter Initialization
보통 input-unit의 개수 \(n\)에 대해서 \(\mathcal{N}(0, 1/n)\)에서 weight를 초기화한다. (fan-in 방식) 여기에서 2\(\sigma\) 이상인 경우 다시 뽑게 되어 truncated normal이라고 한다. Switch Transformer에서는 standard deviation을 0.1배 하여 더 학습이 안정화되는 것을 확인할 수 있었다.
Expert Dropout
Switch Transformer의 parameter 수는 다른 FLOP matched dense model에 비해 월등히 많으므로 downstream task에 fine-tuning할 때 overfit이 일어나기 쉽다. 저자들은 단순히 expert 내에 중간 FFN에 강한 dropout(0.4)를 적용하는 것으로 이 이슈를 해결했다. 전체적으로 강한 dropout을 적용하는 것은 오히려 성능을 감소시켰다.
full-precision training and Router-Z Loss (st-moe-32B)
Switch Transformer를 아주 크게 scaling 할 경우에 학습 불안정의 문제가 종종 발생했는데 후속 논문인
Zoph et al. 에서 해결했다. 우선 Selective Precision 대신에 full Precision을 사용하였다. Sparse model 같은 경우 Router 때문에 Exponential 연산이 추가되기 때문에 round-off error에 dense model 보다 민감해진다. 위의 그림처럼 bfloat16과 float32는 roundoff error의 차이가 수가 커질 경우에 최대 65536배까지 날 수 있다. 마찬가지 관점에서 라우터에서 logit 값이 매우 커지는 경우 round-off error 값이 커지기 때문에 logit값 자체를 작게 유지하기 위해서 Router-Z Loss를 추가하였다.
$$ L_z(x) = \frac{1}{B} \sum_{i=1}^{B} \left( \log\sum_{j=1}^{N} e^{x_j^{(i)}} \right)^2 $$
Is Sparse Model better than Dense one?
먼저 왼쪽 그림은 parameter count를 2배씩 늘리면서 test loss를 확인한 결과인데 log-linear하게 개선되는 것을 확인할 수 있다. 오른쪽 그림은 FLOP-matched dense model인 T5-Base 모델과의 training sample efficiency를 비교한 것이다. Expert가 많아질수록 효율성이 좋아지는 것을 확인할 수 있다.
이번 그림은 Step별 비교가 아니라 Wall Clock으로 비교한 부분인데, Step 비교에서 효율적이라도 한 Step에 드는 시간이 아주 길어서 실제로는 더 느릴 경우도 있기 때문이다. Wall Clock Time 비교에서 T5-Base의 수렴성능에 도달하는데 7배 정도 시간이 감소된 것을 확인할 수 있다.
이번에는 scale을 키워서 switch-base 모델과 T5-Large 모델을 비교해봤다. Dense model에서도 파라미터 수가 많아지면 학습 효율성이 증가하는 것은 잘 알려져있는데, Switch-base는 Step별 효율성에서는 T5-Large랑 동일하고, Wall Clock Time으로 봤을 때는 T5-Large의 최고성능에 도달하는데 2.5배 빠른 것을 확인할 수 있다. 그러니까.. 결론은 Storage 이슈가 없다면 Sparse Model을 안 쓸 이유가 없는 것 같다.
Finetuning on 32B model achieves SOTA for many benchmarks
Sentinel Token Routing Analysis
- Expert들이 Specialization을 한다. 어떤 Expert는 sentinel token만 받고 어떤 Expert는 구두점만 받고, 어떤 것은 대명사만 받는 것을 확인할 수 있다. multilingual 모델에서도 신기하게 비슷한 역할을 하는 단어들끼리 묶이는 것을 확인할 수 있었다.
- 이런 Specialization은 Encoder에서만 일어나고 Decoder에서는 거의 일어나지 않는다. Router의 Entropy 분석을 통해 알아볼 수 있다. 그래서 저자들은 후속 연구에서는 디코더에 Sparsity를 빼는 방향을 고려하고 있다고 한다.
Discussions
- 하이퍼파라미터 세팅이 Dense Model과 다르다. 최적 Learning Rate(1e-3)는 더 높고, Batch size(65k)는 더 작다.
- non-MoE weight만 finetuning 하는 것이 전체 parameter의 20% 정도만 업데이트 해서 효율적이면서 전체를 tuning하는 것과 비슷한 성능을 보인다. non-MoE 안에서도 FFN layer만 학습하면 심지어 약간의 성능 향상이 나타났다. 저자들은 MoE layer의 라우팅 때문에 토큰별로 최대 2개의 expert에 대응될 수밖에 없고, 전체 expert가 골고루 학습되지 못하기 때문에 이런 현상이 나타난다고 보고 있다.
Data, Model, Expert Parallelism
대형 모델을 학습하기 위해 data, model parallelism을 사용하였다. data parallelism은 batch 내에 있는 instance들을 core별로 나누는 것으로 model weight는 모든 core가 가지고 있다고 생각할 수 있다. model parallelism은 일반적으로 tensor parallelism을 말하는데 horizontal하게 모델 weight를 잘라서 코어별로 나누는 것이다. 이 경우 데이터는 모든 코어에 복사되어 있고 각 코어에서 연산한 결과를 AllReduce를 통해 합친 뒤 다음 레이어로 이동하게 된다. 이 논문에서는 추가적으로 Expert Parallelism을 제안한다. 코어별로 expert를 할당하는 것으로 일종의 model parallel이라고 볼 수도 있지만 expert 내에 있는 weight를 분할하지 않는다는 특징이 있다.
가장 큰 모델 학습을 위해서는 Data, Model, Expert parallelism을 사용하였다. 이 방식은 expert 또한 weight를 쪼개서 내부적으로 AllReduce 연산을 사용하는 대신 더 큰 모델을 학습할 수 있게 한다.
댓글
댓글 쓰기