Van Sparse naar Zachte Mengsels van Experts

From Sparse to Soft Mixtures of Experts

August 2, 2023

Auteurs: Joan Puigcerver, Carlos Riquelme, Basil Mustafa, Neil Houlsby

cs.AI

Samenvatting

Sparse mixture of expert architecturen (MoEs) schalen de modelcapaciteit zonder grote toename in trainings- of inferentiekosten. Ondanks hun succes kampen MoEs met een aantal problemen: trainingsinstabiliteit, token dropping, onvermogen om het aantal experts te schalen, of ineffectieve finetuning. In dit werk stellen we Soft MoE voor, een volledig differentieerbare sparse Transformer die deze uitdagingen aanpakt, terwijl de voordelen van MoEs behouden blijven. Soft MoE voert een impliciete soft assignment uit door verschillende gewogen combinaties van alle invoertokens door te geven aan elke expert. Net als in andere MoE-werken verwerken experts in Soft MoE slechts een subset van de (gecombineerde) tokens, waardoor een grotere modelcapaciteit mogelijk is tegen lagere inferentiekosten. In de context van visuele herkenning presteert Soft MoE aanzienlijk beter dan standaard Transformers (ViTs) en populaire MoE-varianten (Tokens Choice en Experts Choice). Zo vereist Soft MoE-Base/16 10,5x lagere inferentiekosten (5,7x kortere wall-clock tijd) dan ViT-Huge/14, terwijl het dezelfde prestaties levert na vergelijkbare training. Soft MoE schaalt ook goed: Soft MoE Huge/14 met 128 experts in 16 MoE-lagen heeft meer dan 40x zoveel parameters als ViT Huge/14, terwijl de inferentietijdkosten slechts met 2% toenemen en het aanzienlijk beter presteert.

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.