[논문 리뷰] Scaling Vision with Sparse Mixture of Experts

Carlos Riquelme, Joan Puigcerver외 6명|arXiv (Cornell University)|2021. 06. 10.

Domain Adaptation and Few-Shot Learning참고 문헌 60인용 수 29

한 줄 요약

소개 Vision MoE (V-MoE), 일부 MLP 블록을 Mixture-of-Experts 계층으로 대체한 스파르스 변형의 Vision Transformer로, 대규모 비전 모델을 가능하게 하며 추론 비용을 낮추고 밀집 모델과의 성능을 맞추며 최대 15B 매개변수까지 확장 가능하다.

ABSTRACT

Sparsely-gated Mixture of Experts networks (MoEs) have demonstrated excellent scalability in Natural Language Processing. In Computer Vision, however, almost all performant networks are "dense", that is, every input is processed by every parameter. We present a Vision MoE (V-MoE), a sparse version of the Vision Transformer, that is scalable and competitive with the largest dense networks. When applied to image recognition, V-MoE matches the performance of state-of-the-art networks, while requiring as little as half of the compute at inference time. Further, we propose an extension to the routing algorithm that can prioritize subsets of each input across the entire batch, leading to adaptive per-image compute. This allows V-MoE to trade-off performance and compute smoothly at test-time. Finally, we demonstrate the potential of V-MoE to scale vision models, and train a 15B parameter model that attains 90.35% on ImageNet.

연구 동기 및 목표

• 희소 Mixture-of-Experts가 비전 모델을 효과적으로 확장할 수 있는지 조사한다.
• V-MoE가 감소된 추론 비용으로 밀집 ViT 성능에 필적하거나 이를 초과함을 입증한다.
• 학습을 안정화하고 전이 학습을 개선하기 위한 라우팅 및 용량 전략을 개발한다.
• 배치 우선 라우팅을 도입하여 이미지당 또는 배치당 컴퓨트를 조정한다.
• 최대 15B 매개변수의 비전 모델이 강력한 ImageNet 성능을 달성하는 것을 보인다.

제안 방법

• 각 토큰이 소수의 전문가에 라우팅되는 희소 MoE 계층으로 선택된 ViT MLP 블록을 대체한다.
• 토큰을 k명의 전문가에 할당하기 위해 softmax(Wx+ε)에 TOP_k를 적용하는 라우팅 함수 g(x)를 사용한다(일반적으로 k는 1 또는 2).
• 노이즈 ε를 추가하고 학습 중 전문가 부하를 균형 있게 조절하기 위해 용량 인식 버퍼링 B_e를 사용한다.
• 용량 비율 C로 전문가 버퍼 용량을 고정하고 부하 균형을 촉진하기 위해 보조 손실을 사용한다.
• 대용량의 노이즈 데이터(JFT-300M)로 학습하고 ImageNet 및 VTAB에서 선형 탐색과 전체 미세 조정을 통한 전이 성능을 평가한다.
• 배치 전체의 토큰에 우선 순위를 부여하고 추론 시 유용성이 낮은 토큰을 건너뛰도록 한다.

실험 결과

연구 질문

• RQ1 Vision Transformer에서 희소 MoE 계층이 감소된 계산으로도 경쟁력 있는 정확도를 달성할 수 있는가?
• RQ2 라우팅, 용량 제어, 노이즈가 V-MoE의 학습 안정성과 성능에 어떠한 영향을 미치는가?
• RQ3 배치 우선 라우팅과 조정 가능한 용량이 추론 시 컴퓨트 대 성능에서 어떤 이점을 제공하는가?
• RQ4 V-MoE 모델이 다운스트림 작업 및 소수 샷/미세 조정 상황에 얼마나 잘 전이되는가?
• RQ5 ImageNet에서 매개변수 수 및 정확도 측면에서 V-MoE의 확장 가능성은 어느 정도인가?

주요 결과

모델	매개변수	JFT prec@1	IN/1샷	IN/5샷	IN/10샷	IN/미세 조정	ExaFLOPs	TPUv3-일수
VIT-H/14	656M	56.68	62.34	76.95	79.02	88.08	4.27k	2.38k
V-MoE-L/16, Every-2	3.4B	57.65	62.41	77.10	79.01	87.41	2.17k	1.20k
V-MoE-H/14, Last-5	2.7B	60.12	62.95	78.08	80.10	88.23	4.75k	2.73k
V-MoE-H/14, Every-2	7.2B	60.62	63.38	78.21	80.33	88.36	5.79k	3.47k
V-MoE-15B, Every-2	14.7B	—	68.66	82.78	84.29	90.35	33.9k	16.8k

• V-MoE 변형은 대략 절반의 추론 계산으로도 밀집 ViT 성능에 필적하거나 이를 상회한다.
• 15B 매개변수의 V-MoE 모델(V-MoE-15B)은 완전히 미세 조정될 때 ImageNet에서 90.35%를 달성한다.
• 배치 우선 라우팅은 학습 FLOP를 약 20% 감소시키고 이미지당 컴퓨트 트레이드오프를 가능하게 한다.
• MoE 계층을 적용한 JFT-300M에서의 상위 단계 사전 학습은 소수샷 및 전체 미세 조정 설정에서 강한 전이 성능을 보인다.
• V-MoE 모델은 추론 시 k와 용량 C를 조정하는 유연성을 보여주며 최소한의 성능 손실로 상당한 컴퓨트 절감을 가능하게 한다.
• 가장 큰 V-MoE 모델(15B)은 경쟁력 있는 ImageNet 유사 벤치마크를 달성하고 확장 가능한 비전 모델 용량을 보여준다.