Von spärlichen zu weichen Mischungen von Experten

From Sparse to Soft Mixtures of Experts

August 2, 2023

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

cs.AI

Zusammenfassung

Sparse Mixture of Expert-Architekturen (MoEs) skalieren die Modellkapazität ohne große Steigerungen der Trainings- oder Inferenzkosten. Trotz ihres Erfolgs leiden MoEs unter einer Reihe von Problemen: Trainingsinstabilität, Token-Dropping, Unfähigkeit, die Anzahl der Experten zu skalieren, oder ineffektives Fine-Tuning. In dieser Arbeit schlagen wir Soft MoE vor, einen vollständig differenzierbaren, sparsamen Transformer, der diese Herausforderungen bewältigt, während die Vorteile von MoEs erhalten bleiben. Soft MoE führt eine implizite Soft-Zuweisung durch, indem es unterschiedlich gewichtete Kombinationen aller Eingabe-Tokens an jeden Experten weitergibt. Wie in anderen MoE-Arbeiten verarbeiten Experten in Soft MoE nur eine Teilmenge der (kombinierten) Tokens, was eine größere Modellkapazität bei geringeren Inferenzkosten ermöglicht. Im Kontext der visuellen Erkennung übertrifft Soft MoE Standard-Transformer (ViTs) und beliebte MoE-Varianten (Tokens Choice und Experts Choice) deutlich. Beispielsweise benötigt Soft MoE-Base/16 10,5x geringere Inferenzkosten (5,7x geringere Echtzeit) als ViT-Huge/14, während es dessen Leistung nach ähnlichem Training erreicht. Soft MoE skaliert auch gut: Soft MoE Huge/14 mit 128 Experten in 16 MoE-Schichten hat über 40x mehr Parameter als ViT Huge/14, während die Inferenzzeitkosten nur um 2 % steigen, und es schneidet deutlich besser ab.

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.