[논문 리뷰] From Sparse to Soft Mixtures of Experts

Joan Puigcerver, Carlos Riquelme외 2명|arXiv (Cornell University)|2023. 08. 02.

Domain Adaptation and Few-Shot Learning인용 수 23

한 줄 요약

Soft MoE는 모든 토큰에 걸쳐 입력을 각 전문가에게 부드럽게 혼합하는 완전한 미분 가능 희소 트랜스포머를 도입하여 더 낮은 추론 비용과 안정성으로 확장 가능한 용량을 달성하고, 비전 태스크에서 ViT 및 기존 MoE를 능가합니다.

ABSTRACT

Sparse mixture of expert architectures (MoEs) scale model capacity without significant increases in training or inference costs. Despite their success, MoEs suffer from a number of issues: training instability, token dropping, inability to scale the number of experts, or ineffective finetuning. In this work, we propose Soft MoE, a fully-differentiable sparse Transformer that addresses these challenges, while maintaining the benefits of MoEs. Soft MoE performs an implicit soft assignment by passing different weighted combinations of all input tokens to each expert. As in other MoEs, experts in Soft MoE only process a subset of the (combined) tokens, enabling larger model capacity (and performance) at lower inference cost. In the context of visual recognition, Soft MoE greatly outperforms dense Transformers (ViTs) and popular MoEs (Tokens Choice and Experts Choice). Furthermore, Soft MoE scales well: Soft MoE Huge/14 with 128 experts in 16 MoE layers has over 40x more parameters than ViT Huge/14, with only 2% increased inference time, and substantially better quality.

연구 동기 및 목표

• 트랜스포머 모델의 확장을 실행 비용이나 메모리 비용 없이 추진한다.
• 전통적인 희소 MoE의 불안정성과 토큰 드랍 문제를 전문가 전문화의 이점을 유지하면서 해결한다.
• 수천 개의 전문가를 가능하게 하는 완전 미분 가능 소프트 라우팅 메커니즘을 제안한다.
• 이미지 분류에 대한 접근 방식을 시연하고 ViT 및 기존 MoE와 비교한다.

제안 방법

• Soft MoE를 소프트 할당을 통해 라우팅이 달성되는 완전 미분 가능 계층으로 정의한다.
• 토큰-슬롯 축 및 슬롯-토큰 축에 대한 소프트맥스를 사용하여 디스패치 및 결합 가중치를 계산한다(논문의 식 (1)–(3에 비견되는 식).
• 각 입력 슬롯을 해당 전문가(일반적으로 MLP)와 함께 처리한다.
• 학습을 안정화하기 위해 토큰당 및 슬롯당 L2 정규화를 사용하여 입력 및 라우터 매개변수를 정규화한다.
• 트랜스포머의 조밀한 MLP 블록을 Soft MoE 블록으로 대체하고, 총 슬롯 수를 통해 컴퓨트를 제어한다.
• 간단한 JAX 구현을 제공하고 전체 코드를 Google Research GitHub(vmoe)에서 참조한다.

Figure 1 : Main differences between Sparse and Soft MoE layers. While the router in Sparse MoE layers (left) learns to assign individual input tokens to each of the available slots, in Soft MoE layers (right) each slot is the result of a (different) weighted average of all the input tokens. Learning

실험 결과

연구 질문

• RQ1 Soft MoE가 Dense ViT 및 기존 희소 MoE와 학습 및 추론 예산 전반에서 동등하거나 우수한 정확도를 달성할 수 있는가?
• RQ2 소프트 라우팅이 토큰 드랍 및 규모에 따른 전문가 불균형과 같은 전형적인 MoE 이슈를 완화하는가?
• RQ3 전문가 수와 전문가당 슬롯 수가 늘어날 때 Soft MoE의 확장성은 어떠하며 비전 태스크에 대한 최적 구성이 무엇인가?
• RQ4 Soft MoE가 이미지-텍스트 대조 학습과 같은 다운스트림 태스크에 이점을 확장할 수 있는가?
• RQ5 Dense 및 Sparse 기준선과 비교했을 때 학습 시간, FLOPs, 실사용 시간의 트레이드오프는 어떠한가?

주요 결과

• Soft MoE가 다양한 모델 크기에서 학습 비용과 성능의 파레토 프런트에서 Dense ViTs 및 인기 있는 Sparse MoE 모두를 지배합니다.
• Soft MoE Base/16은 ViT-Huge/14와 유사한 학습 후 성능을 유지하면서 추론 비용을 10.5배 낮추고 실사용 시간을 5.7배 빠르게 달성합니다.
• Soft MoE가 128개의 전문가와 16개의 레이어를 갖는 경우 ViT Huge/14에 비해 추론 비용이 약 2% 증가하는 수준에서 파라미터 수를 40배 이상 늘릴 수 있으며 성능은 현저히 향상됩니다.
• 장기 규모 실험에서 Soft MoE 모델은 비전 트랜스포머를 유사한 컴퓨트 예산에서 능가하는 성능을 보여주며, 작은 백본에서 특히 큰 이점이 있고 더 큰 규모에서도 경쟁력 있거나 우수한 결과를 제공합니다.
• Soft MoE B/16 및 L/16 변형은 강력한 업스트림 및 파인튜닝 결과를 달성하며, ViT 기준선 대비 큰 추론 속도 향상을 제공합니다(예: Soft MoE L/16은 Dense H/14를 능가하면서도 더 빠릅니다).
• Soft MoE는 전문가 수가 늘어나도 잘 확장되며, 전문가마다 하나의 슬롯을 두고 수백에서 수천 개의 전문가가 성능을 향상시키지만 비용은 지나치게 증가하지 않으며, 이는 일부 Sparse MoE 변형과 달리 비용 효용이 좋습니다.

Figure 2 : The Soft MoE routing algorithm. Soft MoE first computes scores or logits for every pair of input token and slot, based on some learnable per-slot parameters. These logits are then normalized per slot (columns) and every slot computes a linear combination of all the input tokens based on t