안녕하세요. 이번에 소개할 논문도 MoE 분야의 연구입니다. 이 논문은 기존 dense 모델을 MoE 모델로 확장하는 과정에서, MoE 레이어의 전문가 네트워크(expert)를 단순히 랜덤으로 초기화해 학습하는 것이 아니라, 기존 dense 레이어의 가중치를 그대로 복사해 사용함으로써 학습의 효율성과 성능을 동시에 향상시키는 방법을 제안하고 있습니다. 이 논문은 새로운 방법론보다는 다양한 실험을 통해 아이디어를 검증하는 데 초점을 맞추고 있습니다. 따라서 이후 제 실험에서 참고할 만한 부분이 있는지 살펴보기 위해 이번 리뷰를 진행하게 되었습니다. 그럼 바로 리뷰를 시작하겠습니다.
1. Introduction
일반적으로 딥러닝에서는 규모(scale)가 커질수록 성능이 향상됩니다. 예를 들어 NLP 분야에서는 BERT, GPT-3, PaLM 같은 대형 모델들이 등장하면서 성능이 크게 발전했고 Vision 분야에서도 초기의 AlexNet에서 Vision Transformer 계열의 ViT-G까지, 더 큰 규모의 모델들이 계속해서 성능을 향상시켰습니다. 이런 규모의 중요성은 NLP나 Vision에만 적용되는 이야기가 아니고 음성인식이나 강화학습, 멀티모달 학습에도 이어져오면서 더 큰 데이터셋이나 더 복잡하고 큰 모델의 아키텍처로 발전하게 되었습니다. 하지만 대부분의 SOTA 모델들은 랜덤으로 초기화된 가중치로 시작하여 처음부터 학습(scratch train)을 진행하기 때문에, 이러한 모델들을 학습하는데 드는 비용 또한 빠르게 증가하고 있습니다. 예를 들어 BERT-Large(2018년, 3억 4,500만 개의 파라미터)는 학습하는데 약 0.5 ZFLOPS가 필요하며, GPT-3(2020년, 1,750억 개의 파라미터)는 314 ZFLOPS가 필요했고, PaLM(2022년, 5,400억 개의 파라미터)은 2527 ZFLOPS를 필요로 합니다. (FLOPS는 FLoating-point Operations Per Second을 뜻하며, 초당 부동소수점 연산 횟수를 나타내는 단위입니다. 1 ZFLOPS는 10²¹ FLOPS을 의미합니다.) 이처럼 모델의 규모가 커질수록 계산량(cost)가 엄청나게 증가하기 때문에 큰 자원을 가진 몇몇 회사나 연구기관에서만 대규모 모델을 연구할 수 있게되었습니다. 그래서 최근에는 더 적은 자원으로도 큰 모델을 효과적으로 학습하거나 개선하는 방법에 대한 연구가 늘어나고 있습니다.
이번 논문에서도 적은 자원으로 기존 모델의 성능을 개선할 수 있는 model upcycling에 대해 다루고 있습니다. Model upcycling이란 이미 학습된 모델을 재활용하는 것을 의미합니다. 저자는 이것을 Mixture-of-Experts(MOE)에 적용하여 더 크고, spares하게 활성화되는 모델을 설계하였습니다. 결과적으로 기존의 dense 모델보다 더 좋은 성능을 얻을 수 있었고 새로운 dense 모델을 처음부터 다시 훈련시키는 것보다 훨씬 적은 연산 비용을 사용했습니다. 이러한 Sparse upcycling은 두 가지 상황에서 매우 유용하게 사용될 수 있습니다. 첫 번째는 사전학습된 Transformer 모델을 가지고 있고, 이 모델을 가지고 적은 계산 cost로 성능을 끌어올리고 싶을 때 Sparse upcycling이 유용할 수 있습니다. 두 번째는 새로운 대형 모델을 아예 처음부터 학습하려고 하지만 dense 모델 이나 MoE 모델 중 어떤 것이 더 효과적인지 모르는 경우입니다. 왜냐하면 MoE는 dense 모델보다 보통 성능이 좋지만 학습 과정이 더 복잡하기 때문입니다. 이럴 경우 먼저 dense 모델을 먼저 학습하고, 그 다음 성능이 더 안 오를 때 Sparse upcycling을 통해 MoE로 바꿔주는 방법을 사용하여 성능을 더 향상시킬 수 있습니다.
결과적으로 저자는 model upcycling 학습 시 이미 학습된 네트워크의 구조를 바꿀 때 발생하는 초기 성능 하락 문제를 해결하는 방법을 제시하였고, 모델 학습에 중요한 핵심 구성 요소들에 대해 다양한 실험을 수행하였습니다. 실험 결과에 따르면,Vision Transformer와 T5 모델에서 원래의 dense 모델을 훈련하는 데 들었던 계산 비용 대비 +10%에서 +60% 정도의 추가 비용으로 저자의 방법론이 효과적이라는 것을 입증하였습니다.
2. THE UPCYCLING ALGORITHM
2.1 CLIP with Mixture-of-Experts upcycling
위 그림 1이 저자가 제안하는 알고리즘으로 매우 간단한 구조를 가지고 있습니다. 우선 학습된 dense 모델의 파라미터(체크포인트)가 필요합니다. 새로 만들 모델의 Transformer block의 수는 기존의 dense 모델과 일치하지만 MLP 부분에서만 차이가 있습니다. 여기서는 MLP 레이어를 E개의 MLP를 가진 MoE 계층으로 확장하고 MLP , layer-norm, attention, embedding, output layers의 파라미터는 원래 모델에서 그대로 복사해 사용합니다. 그리고 여기에 무작위로 초기화된 router를 추가해주면 새로운 upcycled model을 설계할 수 있습니다.
DESIGN DECISIONS
Upcycling을 할 때 모델의 성능은 MoE의 구성에 크게 좌우됩니다. 보통 MoE layer의 수를 늘리거나 Expert의 수를 늘리거나, 각 Expert가 처리하는 양(용량)을 늘리면 모델의 용량이 커져서 성능이 좋아집니다. 하지만 이렇게 하면 계산 비용도 늘어나고, 이미 학습된 dense 모델에서 갑자기 구조를 크게 바꾸기 때문에 처음에는 성능이 떨어질 수 있습니다.
그럼 MoE layer를 얼마나 추가하느냐가 중요한데, MoE layer를 많이 넣을수록 성능 향상은 커지지만 계산 비용도 훨씬 늘어나고, dense 모델보다 초기 성능 하락폭이 더 커지는 문제가 있습니다. 이전 MoE 논문들에서도 그렇고 이 논문에서도 기본적으로는 MLP layer의 절반만 MoE로 바꿔서 Upcycling합니다.
MoE layer에 몇 개의 Expert를 추가할지도 중요한데, Expert가 많을수록 학습할 수 있는 파라미터가 늘어나서 모델의 성능 향상 가능성은 커집니다. 하지만 너무 많이 추가하면 처음에는 성능이 많이 떨어지기 때문에, 적당한 수준으로 설정을 해야합니다. 저자는 실험을 통해 32개의 Expert를 사용하는 것이 성능과 cost 측면에서 가장 적절했다 합니다. 또한 각 Expert가 처리하는 양(Expert Capacity)도 중요한데, 이 계수도 실험 결과를 통해 C=2로 두면 성능과 계산 효율 측면에서 가장 좋았다고 합니다
3. Experiment
- 이제 실험 파트 살펴보겠습니다. 모든 Upcycling 실험은 사전 학습된 dense 모델 체크포인트에서 시작. 기존 learning rate schedule로 학습을 이어가고, dense 모델 또한 연속으로 학습을시켜 비교를 진행하였습니다. 아래에 Vision, Language의 실험 셋업 환경을 정리해 놓았으니 참고해주세요.
- Vision experiments :
- V-MoE를 JFT300M 데이터셋에서 사전학습.
- 이미지넷 10-shot 전이학습과 full finetuning 실험.
- Language experiments:
- English C4에서 사전학습.
- SuperGLUE 태스크에서 동시에 파인튜닝.
- Base 모델의 경우, 저자가 직접 새로운 dense 모델을 처음부터 사전학습하여 사용
- 반면 Large와 XL 모델의 경우 T5 1.1 체크포인트를 사용
3.1 Recipe Study
그림 2에서는 Vision(왼쪽)과 Language(오른쪽)에서, Upcycling 모델과 기존 dense 모델의 성능을 다양한 모델 크기에서 비교한 결과를 보여줍니다. 모델 크기나 태스크가 달라도, 추가 학습 cost가 아주 작을 때는 dense 모델과 Upcycling 모델 성능이 거의 같습니다. 하지만 추가 연산량을 학습을 이어가면, Upcycling한 모델이더 높은 성능을 내는 모습을 보입니다.
그림 3은 그림 2에서 훈련한 모델들을 추가로 파인튜닝한 후의 성능을 보여줍니다. Vision 모델(왼쪽)의 경우, 사전 학습에서 얻은 성능 향상이 대체로 다운스트림으로 전이되는 경향이 있습니다. 반면, Language 모델(오른쪽)에서는 파인튜닝 후 성능이 더 다양하게 나오긴 했지만, 전반적으로 Upcycling된 Language 모델이 더 좋은 성능을 보였습니다.
그림 4에서는, sparse Upcycling과 처음부터 sparse MoE 모델을 학습한 경우를 비교했습니다. 처음부터 sparse 모델을 학습하면 이미 학습된 dense 체크포인트의 연산비용을 재사용하지 못하므로, Upcycling 모델을 따라잡는데 더 긴 추가 학습 시간이 필요합니다. 특히 Language 모델의 경우, 처음부터 sparse로 학습하면 Upcycling 모델과 비슷한 성능에 도달하는 데 원래 dense 모델 연산량의 약 120%를 더 써야 합니다. 다만, 처음부터 sparse로 학습한 MoE 모델은 상대적으로 큰 학습률을 쓸 수 있고, Expert들이 처음부터 독립적으로 학습되고 다양화될 수 있기 때문에, 한 스텝당 성능 향상은 더 빠르게 나타날 수 있습니다. 그림 4의 결과로, 만약 아주 큰 연산 자원을 쓸 수 있다면, 처음부터 sparse 모델을 학습하는 게 결국엔 Upcycling 모델을 따라잡을 수 있다는 걸 보여줍니다. 하지만 연산 자원이 제한되거나 충분하지 않을 때에는, sparse Upcycling이 훨씬 더 효율적인 접근법이 될 수 있습니다.
Upcycling 효율성은 원래 dense 모델이 얼마나 잘 수렴했는지에 따라 달라질 수 있습니다. 이를 알아보기 위해, 그림 6에서는 dense 체크포인트의 사전학습 정도가 서로 다른 상태에서 B/16 비전 모델을 Upcycling했습니다. 같은 dense 체크포인트에서 출발해, Upcycling과 dense continuation을 각각 20만 스텝씩 추가 학습시켜 성능을 비교한 결고어느 단계에서 Upcycling을 시작하든, Upcycling으로 인한 성능 향상 폭은 일관되게 나타났습니다.
마지막으로 Expert의 수와 MoE Layer 수에 대한 실험 결과를 살펴보고 마무리 하겠습니다. 먼저 Expert 수를 늘리면 모델의 파라미터 수가 증가하고, 어느 정도까지는 모델의 품질도 좋아집니다. 각 Expert가 처리하는 토큰 수는 Expert 수가 많아질수록 줄어들기 때문에, Expert 수를 늘려도 계산량에는 큰 계산 비용이 추가되지 않습니다. 다만, Expert 수가 매우 많아지면 Upcycling된 모델은 원래의 dense 모델과 비교해 초반 품질이 더 크게 떨어질 수 있습니다.
Vision 태스크에서 6개의 MoE 레이어를 사용해 Expert 수를 2에서 128까지 바꿔가며 실험한 결과가 그림 10의 왼쪽 2개의 차트에 나와 있습니다. x축에서 계산량이 일정할 때, B/16 모델에서는 Expert 수가 많을수록 대체로 더 좋은 성능을 보였습니다. 그리고 그림 10의 오른쪽 2개의 차트를 보면, MoE 레이어를 더 많이 추가한다고 해서 항상 더 좋은 성능이 나오는 것은 아님을 확인할 수 있습니다. 결론적으로 대략 마지막 5~6개의 레이어를 사용했을때 모델이 가장 좋은 효율-성능 균형을 보인다는 결론을 내릴 수 있습니다.
