スパースなエキスパート混合からソフトなエキスパート混合へ

要旨

スパースなエキスパート混合アーキテクチャ（MoE）は、トレーニングや推論コストを大幅に増加させることなくモデルの容量を拡張します。しかし、MoEにはいくつかの課題があります：トレーニングの不安定性、トークンのドロップ、エキスパート数のスケーリングの困難さ、または効果的なファインチューニングの欠如などです。本研究では、これらの課題に対処しつつ、MoEの利点を維持する完全微分可能なスパースTransformerであるSoft MoEを提案します。Soft MoEは、すべての入力トークンの異なる重み付けされた組み合わせを各エキスパートに渡すことで、暗黙的なソフト割り当てを行います。他のMoE研究と同様に、Soft MoEのエキスパートは（組み合わされた）トークンのサブセットのみを処理し、低い推論コストでより大きなモデル容量を実現します。視覚認識の文脈において、Soft MoEは標準的なTransformer（ViT）や人気のあるMoEの変種（Tokens ChoiceやExperts Choice）を大きく上回ります。例えば、Soft MoE-Base/16は、ViT-Huge/14と同等の性能を達成するために、推論コストを10.5倍低く（ウォールクロック時間を5.7倍短く）必要とします。Soft MoEはまた、スケーラビリティにも優れています：16のMoE層に128のエキスパートを持つSoft MoE Huge/14は、ViT Huge/14よりも40倍以上のパラメータを持ちながら、推論時間コストはわずか2％増加し、大幅に優れた性能を発揮します。

English

Sparse mixture of expert architectures (MoEs) scale model capacity without large 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 proposeSoft 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 MoE works, experts in Soft MoE only process a subset of the (combined) tokens, enabling larger model capacity at lower inference cost. In the context of visual recognition, Soft MoE greatly outperforms standard Transformers (ViTs) and popular MoE variants (Tokens Choice and Experts Choice). For example, Soft MoE-Base/16 requires 10.5x lower inference cost (5.7x lower wall-clock time) than ViT-Huge/14 while matching its performance after similar training. Soft MoE also scales well: Soft MoE Huge/14 with 128 experts in 16 MoE layers has over 40x more parameters than ViT Huge/14, while inference time cost grows by only 2%, and it performs substantially better.

スパースなエキスパート混合からソフトなエキスパート混合へ

From Sparse to Soft Mixtures of Experts

要旨

Support