Mixture-of-Supernets: Verbetering van Weight-Sharing Supernet Training met Architecture-Routed Mixture-of-Experts

Samenvatting

Het delen van gewichten in een supernet is een cruciaal onderdeel geworden voor prestatie-inschatting in de meest geavanceerde (state-of-the-art, SOTA) neurale architectuurzoekmethoden (neural architecture search, NAS). Hoewel een supernet direct verschillende subnetwerken kan genereren zonder hertraining, is er geen garantie voor de kwaliteit van deze subnetwerken vanwege het delen van gewichten. Bij NLP-taken zoals machinaal vertalen en vooraf getraind taalmodeleren, observeren we dat bij dezelfde modelarchitectuur er een groot prestatieverschil bestaat tussen het supernet en training vanaf nul. Daarom kan het supernet niet direct worden gebruikt en is hertraining noodzakelijk na het vinden van de optimale architecturen. In dit werk stellen we mixture-of-supernets voor, een gegeneraliseerde supernetformulering waarbij mixture-of-experts (MoE) wordt toegepast om de expressieve kracht van het supernetmodel te vergroten, met verwaarloosbare trainingsoverhead. Op deze manier delen verschillende subnetwerken de modelgewichten niet direct, maar via een op architectuur gebaseerd routeringsmechanisme. Hierdoor worden de modelgewichten van verschillende subnetwerken afgestemd op hun specifieke architecturen en wordt de gewichtsgeneratie geleerd door gradient descent. In vergelijking met bestaande supernetten voor NLP die gewichten delen, kan onze methode de hertrainingstijd minimaliseren, waardoor de trainingsefficiëntie aanzienlijk wordt verbeterd. Bovendien behaalt de voorgestelde methode de SOTA-prestaties in NAS voor het bouwen van snelle machinale vertaalmodellen, wat een betere afweging tussen latentie en BLEU oplevert vergeleken met HAT, de state-of-the-art NAS voor machinaal vertalen. We behalen ook de SOTA-prestaties in NAS voor het bouwen van geheugenefficiënte, taakonafhankelijke BERT-modellen, waarbij we NAS-BERT en AutoDistil overtreffen bij verschillende modelgroottes.

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: Verbetering van Weight-Sharing Supernet Training met Architecture-Routed Mixture-of-Experts

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

Samenvatting

Support