Da miscele sparse a miscele morbide di esperti

From Sparse to Soft Mixtures of Experts

August 2, 2023

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

cs.AI

Abstract

Le architetture sparse di mistura di esperti (MoE) aumentano la capacità del modello senza incrementi significativi dei costi di addestramento o inferenza. Nonostante il loro successo, i MoE presentano diverse problematiche: instabilità durante l'addestramento, perdita di token, incapacità di scalare il numero di esperti o fine-tuning inefficace. In questo lavoro, proponiamo Soft MoE, un Transformer sparso completamente differenziabile che affronta queste sfide, mantenendo i vantaggi dei MoE. Soft MoE esegue un'assegnazione implicita e soft passando diverse combinazioni ponderate di tutti i token di input a ciascun esperto. Come in altri lavori sui MoE, gli esperti in Soft MoE elaborano solo un sottoinsieme dei token (combinati), consentendo una maggiore capacità del modello a un costo inferiore di inferenza. Nel contesto del riconoscimento visivo, Soft MoE supera di gran lunga i Transformer standard (ViT) e le varianti popolari di MoE (Tokens Choice e Experts Choice). Ad esempio, Soft MoE-Base/16 richiede un costo di inferenza 10,5 volte inferiore (5,7 volte meno tempo di esecuzione) rispetto a ViT-Huge/14, pur eguagliandone le prestazioni dopo un addestramento simile. Soft MoE scala anche bene: Soft MoE Huge/14 con 128 esperti in 16 strati MoE ha oltre 40 volte più parametri di ViT Huge/14, mentre il costo del tempo di inferenza aumenta solo del 2%, e le prestazioni sono notevolmente migliori.

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.