De Mezclas Dispersas a Mezclas Suaves de Expertos

From Sparse to Soft Mixtures of Experts

August 2, 2023

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

cs.AI

Resumen

Las arquitecturas de mezcla dispersa de expertos (MoEs, por sus siglas en inglés) escalan la capacidad del modelo sin aumentar significativamente los costos de entrenamiento o inferencia. A pesar de su éxito, los MoEs presentan varios problemas: inestabilidad en el entrenamiento, descarte de tokens, incapacidad para escalar el número de expertos o ajuste fino ineficaz. En este trabajo, proponemos Soft MoE, un Transformer disperso completamente diferenciable que aborda estos desafíos, manteniendo los beneficios de los MoEs. Soft MoE realiza una asignación suave implícita al pasar diferentes combinaciones ponderadas de todos los tokens de entrada a cada experto. Como en otros trabajos de MoE, los expertos en Soft MoE solo procesan un subconjunto de los tokens (combinados), lo que permite una mayor capacidad del modelo con un menor costo de inferencia. En el contexto de reconocimiento visual, Soft MoE supera ampliamente a los Transformers estándar (ViTs) y a las variantes populares de MoE (Tokens Choice y Experts Choice). Por ejemplo, Soft MoE-Base/16 requiere 10.5 veces menos costo de inferencia (5.7 veces menos tiempo de ejecución) que ViT-Huge/14, mientras que iguala su rendimiento después de un entrenamiento similar. Soft MoE también escala bien: Soft MoE Huge/14 con 128 expertos en 16 capas MoE tiene más de 40 veces más parámetros que ViT Huge/14, mientras que el costo de tiempo de inferencia aumenta solo un 2%, y su rendimiento es sustancialmente mejor.

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.