Смесь суперсетей: Улучшение обучения суперсетей с совместным использованием весов с помощью архитектурно-направленной смеси экспертов

Аннотация

Суперсеть с общими весами стала важным компонентом для оценки производительности в современных (SOTA) фреймворках поиска нейронных архитектур (NAS). Хотя суперсеть может напрямую генерировать различные подсети без повторного обучения, нет гарантии качества этих подсетей из-за совместного использования весов. В задачах NLP, таких как машинный перевод и предварительное обучение языковых моделей, мы наблюдаем, что при одинаковой архитектуре модели существует значительный разрыв в производительности между суперсетью и обучением с нуля. Следовательно, суперсеть нельзя использовать напрямую, и необходимо повторное обучение после нахождения оптимальных архитектур. В данной работе мы предлагаем смесь суперсетей (mixture-of-supernets), обобщённую формулировку суперсети, в которой используется подход смеси экспертов (MoE) для повышения выразительной способности модели суперсети с минимальными накладными расходами на обучение. Таким образом, различные подсети не используют общие веса напрямую, а через механизм маршрутизации, основанный на архитектуре. В результате веса модели для разных подсетей адаптируются к их конкретным архитектурам, а генерация весов обучается с помощью градиентного спуска. По сравнению с существующими суперсетями с общими весами для NLP, наш метод позволяет минимизировать время повторного обучения, значительно повышая эффективность обучения. Кроме того, предложенный метод достигает SOTA-результатов в NAS для создания быстрых моделей машинного перевода, обеспечивая лучшее соотношение задержка-BLEU по сравнению с HAT, современным NAS для машинного перевода. Мы также достигаем SOTA-результатов в NAS для создания энергоэффективных моделей BERT, не зависящих от задачи, превосходя NAS-BERT и AutoDistil для различных размеров моделей.

English

Weight-sharing supernet has become a vital component for performance estimation in the state-of-the-art (SOTA) neural architecture search (NAS) frameworks. Although supernet can directly generate different subnetworks without retraining, there is no guarantee for the quality of these subnetworks because of weight sharing. In NLP tasks such as machine translation and pre-trained language modeling, we observe that given the same model architecture, there is a large performance gap between supernet and training from scratch. Hence, supernet cannot be directly used and retraining is necessary after finding the optimal architectures. In this work, we propose mixture-of-supernets, a generalized supernet formulation where mixture-of-experts (MoE) is adopted to enhance the expressive power of the supernet model, with negligible training overhead. In this way, different subnetworks do not share the model weights directly, but through an architecture-based routing mechanism. As a result, model weights of different subnetworks are customized towards their specific architectures and the weight generation is learned by gradient descent. Compared to existing weight-sharing supernet for NLP, our method can minimize the retraining time, greatly improving training efficiency. In addition, the proposed method achieves the SOTA performance in NAS for building fast machine translation models, yielding better latency-BLEU tradeoff compared to HAT, state-of-the-art NAS for MT. We also achieve the SOTA performance in NAS for building memory-efficient task-agnostic BERT models, outperforming NAS-BERT and AutoDistil in various model sizes.

Смесь суперсетей: Улучшение обучения суперсетей с совместным использованием весов с помощью архитектурно-направленной смеси экспертов

Mixture-of-Supernets: Improving Weight-Sharing Supernet Training with Architecture-Routed Mixture-of-Experts

Аннотация

Support