De Misturas Especialistas Esparsas para Misturas Especialistas Suaves

From Sparse to Soft Mixtures of Experts

August 2, 2023

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

cs.AI

Resumo

Arquiteturas esparsas de mistura de especialistas (MoEs) escalam a capacidade do modelo sem grandes aumentos nos custos de treinamento ou inferência. Apesar do sucesso, os MoEs sofrem com uma série de problemas: instabilidade no treinamento, descarte de tokens, incapacidade de escalar o número de especialistas ou ajuste fino ineficaz. Neste trabalho, propomos o Soft MoE, um Transformer esparso totalmente diferenciável que aborda esses desafios, mantendo os benefícios dos MoEs. O Soft MoE realiza uma atribuição implícita suave, passando diferentes combinações ponderadas de todos os tokens de entrada para cada especialista. Como em outros trabalhos de MoE, os especialistas no Soft MoE processam apenas um subconjunto dos tokens (combinados), permitindo maior capacidade do modelo com menor custo de inferência. No contexto de reconhecimento visual, o Soft MoE supera amplamente os Transformers padrão (ViTs) e variantes populares de MoE (Tokens Choice e Experts Choice). Por exemplo, o Soft MoE-Base/16 requer 10,5 vezes menos custo de inferência (5,7 vezes menos tempo de execução) do que o ViT-Huge/14, enquanto iguala seu desempenho após treinamento semelhante. O Soft MoE também escala bem: o Soft MoE Huge/14 com 128 especialistas em 16 camadas MoE tem mais de 40 vezes mais parâmetros do que o ViT Huge/14, enquanto o custo de tempo de inferência aumenta apenas 2%, e ele apresenta um desempenho substancialmente melhor.

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.