[리뷰] 이젠 Vision에도 MoE를 사용하자: Google Brain의 Scaling Vision with Sparse Mixture of Experts
Google은 ‘20년 11월 NLP에 사용하던 Transformer 아키텍처를 vision에 적용한 ViT(Vision Transformer)[1]를 발표하며 Transformer가 NLP 외에도 Vision에서도 사용될 수 있음을 입증하였다.
이 논문의 실험 결과는 ViT는 CNN과 달리 image-specific inductive bias가 부족하므로 작은 데이터에서 CNN보다 뒤떨어지는 성능을 보여주지만 JFT-300M과 같은 거대 데이터셋을 적용하면 CNN보다 더 좋은 성능을 보여주었다.
하지만 ViT는 CNN에 비해 상대적으로 큰 모델과 데이터셋을 사용해야 하므로 높은 계산량을 요구한다. 따라서 최근에는 학습 리소스를 제한하면서 큰 모델을 학습하는 방법들이 활발하게 연구되고 있다. 이에 대해서는 이전에 작성한 페이지에서 이에 대한 연구 동향을 다룬 적이 있다.
구글의 Pathways
구글은 21년 10월말에 pathways[2]라는 차세대 AI 아키텍처를 소개하고 이 pathways를 통해 오늘날 AI 모델이 당면한 문제를 해결하고자 한다.
구슬은 현재 AI가 다음과 같은 3가지 문제를 가지고 있다고 정의하고 구글은 pathways를 이용하여 이 문제들을 해결한다. 특히 pathways의 sparse한 특징은 제한된 compute budget을 사용하면서 정확도는 잃지 않는 특징을 보여준다.
V-MoE(Vision Mixture of Experts)
구글은 위와 같은 Pathways를 ViT에 적용하기 위해 21년 “Sparse Vision with Sparse Mixture of Experts” 논문[3]을 발표하고 ViT에 MoE를 적용한 V-MoE를 제안하였다.
V-MoE는 Vision 모델을 대규모로 확장하기 위해 ViT 아키텍처에 MoE 레이어를 도입하였다. 좀더 정확하게 말하자면 아래 그림에서 보는 것과 같이 V-MoE는 ViT 블럭과 MoE 블럭이 서로 교차로 적층되며 MoE 블럭은 ViT 블럭의 FFN 부분을 MoE 레이어로 교체하였다.
이 논문의 주요 기여는 다음과 같다.
Vision models at scale
- 24 MoE 레이어(32 experts/layer)로 구성된 모델을 학습 (15B 파라미터)
- MoE를 안정적으로 학습시킴
- 1000개의 데이터 포인트로 성공적으로 fine-tuning함
- Fine-tuning 시, ImageNet에서 90.35% 테스트 정확도 달성함
Performance & Inference
- 컴퓨팅 양을 절반만 사용하여 가장 큰 dense 모델의 성능과 동일함
- 같은 비용에서 dense 모델보다 성능을 크게 높일 수 있음
Batch Prioritized Routing(BPR)
- 가장 유용하지 않은 patch들을 포기하는 새로운 우선순위기반 라우팅 알고리즘을 제안하여 각 이미지에 더 적은 연산을 할당함
- Dense 모델의 성능과 일치하면서 학습 FLOPs의 20%를 절감
실험결과
JFT-300M에서 ViT과 V-MoE를 pre-training한 결과(왼쪽 그림) 및 다시 새로운 downstream tasks (e.g. ImageNet. 오른쪽 그림)로 transfer learning한 결과이다. (오른쪽 그림은 JFT-300M으로 pre-training된 전체 ViT와 V-MoE 모델을 freezing한 후 새로운 head만을 소수의 예제로 fine-tuning한 few-shot transfer를 사용한 결과이다. )
두 결과 모두, 동일한 학습 컴퓨팅에서 V-MoE가 ViT을 능가하는 성능을 보이고 있다. 심지어 V-MoE 15B 모델은 ImageNet에서 90.35% 테스트 정확도 달성하였다.
Batch Priority Routing
MoE 모델을 추론할 경우, 각 expert용 pre-defined 버퍼를 사용하므로 제한된 버퍼 용량을 초과하여 할당된 토큰은 삭제되고 처리되지 않는다. 물론 버퍼 용량이 클수록 더 높은 정확도를 제공하지만 더 많은 컴퓨팅이 필요하다.
V-MoE는 추론 시간을 줄이기 위해 BPR(Batch Priority Routing)이라는 방식을 사용하여 처리해야 할 토큰의 개수를 총 버퍼 용량 이하로 줄여 네트워크의 expert 레이어에서 일부 토큰 처리를 건너뛸 수 있도록 강제화한다.
BPR은 정확도를 확보하기 위해 중요도 점수에 따라 토큰을 정렬하는 방법을 학습하여 컴퓨팅양을 줄이며 고품질의 예측을 유지할 수 있다.
아래 그림에서 볼 수 있는 것과 같이 Expert 용량(C)이 감소하면 vanilla 라우팅 메커니즘에서는 성능이 빠르게 줄어들지만, BPR 메커니즘을 사용하면 보다 견고하게 성능을 유지한다.
본 논문의 Insight
- OpenAI의 GPT-3 등장 후, 한동안 컴퓨팅 budget을 무시하고 모델 성능을 높이는데 주력해왔지만 점차 추론 서비스를 위한 컴퓨팅 budget을 고려하고 있다.
- MoE는 모델 성능을 유지하면서 컴퓨팅 budget을 줄이기 위한 대안으로 추론 서비스에 반드시 필요한 아키텍처이다.
- 컴퓨팅 budget을 줄이는 것은 에너지 소모를 줄이는 것과 연결되어 상용 서비스를 위해서는 반드시 해결해야 문제이다.
레퍼런스
[1] An Image is Worth 16x16 words: Transformers for Image Recognition at Scale
[2] Introducing pathways next generation AI architecture
[3] Scaling Vision with Sparse Mixture of Experts
[4] Google Research VMoE
