Mixture-of-Supernets : Amélioration de l'entraînement des super-réseaux à partage de poids grâce à un mélange d'experts basé sur l'architecture

Résumé

Le supernet à partage de poids est devenu un composant essentiel pour l'estimation des performances dans les frameworks de recherche d'architecture neuronale (NAS) de pointe. Bien que le supernet puisse générer directement différents sous-réseaux sans réentraînement, il n'y a aucune garantie quant à la qualité de ces sous-réseaux en raison du partage de poids. Dans les tâches de traitement du langage naturel (NLP) telles que la traduction automatique et la modélisation de langage pré-entraînée, nous observons que, pour une même architecture de modèle, il existe un écart de performance important entre le supernet et un entraînement à partir de zéro. Par conséquent, le supernet ne peut pas être utilisé directement et un réentraînement est nécessaire après avoir trouvé les architectures optimales. Dans ce travail, nous proposons le mixture-of-supernets, une formulation généralisée du supernet où le mixture-of-experts (MoE) est adopté pour améliorer la puissance expressive du modèle supernet, avec une surcharge d'entraînement négligeable. De cette manière, les différents sous-réseaux ne partagent pas directement les poids du modèle, mais à travers un mécanisme de routage basé sur l'architecture. En conséquence, les poids du modèle pour les différents sous-réseaux sont personnalisés en fonction de leurs architectures spécifiques et la génération des poids est apprise par descente de gradient. Par rapport aux supernets à partage de poids existants pour le NLP, notre méthode permet de minimiser le temps de réentraînement, améliorant ainsi considérablement l'efficacité de l'entraînement. De plus, la méthode proposée atteint les performances de pointe en NAS pour la construction de modèles de traduction automatique rapides, offrant un meilleur compromis latence-BLEU par rapport à HAT, le NAS de pointe pour la traduction automatique. Nous atteignons également les performances de pointe en NAS pour la construction de modèles BERT agnostiques à la tâche et économes en mémoire, surpassant NAS-BERT et AutoDistil pour différentes tailles de modèles.

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 : Amélioration de l'entraînement des super-réseaux à partage de poids grâce à un mélange d'experts basé sur l'architecture

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

Résumé

Support