От разреженных к мягким смесям экспертов

From Sparse to Soft Mixtures of Experts

August 2, 2023

Авторы: Joan Puigcerver, Carlos Riquelme, Basil Mustafa, Neil Houlsby

cs.AI

Аннотация

Разреженные архитектуры смеси экспертов (MoE) позволяют масштабировать емкость модели без значительного увеличения затрат на обучение или вывод. Несмотря на их успех, MoE сталкиваются с рядом проблем: нестабильность обучения, потеря токенов, невозможность масштабирования числа экспертов или неэффективная тонкая настройка. В данной работе мы предлагаем Soft MoE — полностью дифференцируемую разреженную архитектуру Transformer, которая решает эти проблемы, сохраняя при этом преимущества MoE. Soft MoE выполняет неявное мягкое назначение, передавая каждому эксперту различные взвешенные комбинации всех входных токенов. Как и в других работах по MoE, эксперты в Soft MoE обрабатывают только подмножество (комбинированных) токенов, что позволяет увеличить емкость модели при меньших затратах на вывод. В контексте задач визуального распознавания Soft MoE значительно превосходит стандартные Transformer (ViT) и популярные варианты MoE (Tokens Choice и Experts Choice). Например, Soft MoE-Base/16 требует в 10,5 раз меньше затрат на вывод (в 5,7 раз меньше времени выполнения) по сравнению с ViT-Huge/14, при этом демонстрируя сопоставимую производительность после аналогичного обучения. Soft MoE также хорошо масштабируется: Soft MoE Huge/14 с 128 экспертами в 16 слоях MoE имеет более чем в 40 раз больше параметров, чем ViT Huge/14, при этом время вывода увеличивается всего на 2%, а производительность значительно улучшается.

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.