Des mélanges clairsemés aux mélanges doux d'experts

From Sparse to Soft Mixtures of Experts

August 2, 2023

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

cs.AI

Résumé

Les architectures de mélange d'experts parcimonieux (MoE) permettent d'augmenter la capacité des modèles sans entraîner de hausse significative des coûts d'entraînement ou d'inférence. Malgré leur succès, les MoE souffrent de plusieurs problèmes : instabilité lors de l'entraînement, perte de tokens, difficulté à augmenter le nombre d'experts ou réglage fin inefficace. Dans ce travail, nous proposons Soft MoE, un Transformer parcimonieux entièrement différentiable qui résout ces défis tout en conservant les avantages des MoE. Soft MoE effectue une assignation douce implicite en transmettant à chaque expert différentes combinaisons pondérées de tous les tokens d'entrée. Comme dans d'autres travaux sur les MoE, les experts dans Soft MoE ne traitent qu'un sous-ensemble des tokens (combinés), permettant une plus grande capacité de modèle à un coût d'inférence réduit. Dans le contexte de la reconnaissance visuelle, Soft MoE surpasse largement les Transformers standards (ViTs) et les variantes populaires de MoE (Tokens Choice et Experts Choice). Par exemple, Soft MoE-Base/16 nécessite un coût d'inférence 10,5 fois inférieur (et un temps d'exécution 5,7 fois plus court) que ViT-Huge/14 tout en atteignant des performances similaires après un entraînement comparable. Soft MoE montre également une bonne scalabilité : Soft MoE Huge/14 avec 128 experts répartis sur 16 couches MoE possède plus de 40 fois plus de paramètres que ViT Huge/14, tandis que le temps d'inférence n'augmente que de 2 %, et il offre des performances nettement supérieures.

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.