Mixture-of-Supernets: Miglioramento dell'Addestramento delle Supernet con Condivisione dei Pesi attraverso un Approccio a Mistura di Esperti Basato su Architettura

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

June 8, 2023

Autori: Ganesh Jawahar, Haichuan Yang, Yunyang Xiong, Zechun Liu, Dilin Wang, Fei Sun, Meng Li, Aasish Pappu, Barlas Oguz, Muhammad Abdul-Mageed, Laks V. S. Lakshmanan, Raghuraman Krishnamoorthi, Vikas Chandra

cs.AI

Abstract

Il supernet con condivisione dei pesi è diventato un componente fondamentale per la stima delle prestazioni nei framework più avanzati (SOTA) di ricerca di architetture neurali (NAS). Sebbene il supernet possa generare direttamente diverse sottoreti senza bisogno di riaddestramento, non vi è alcuna garanzia sulla qualità di queste sottoreti a causa della condivisione dei pesi. Nei task di NLP come la traduzione automatica e la modellazione del linguaggio pre-addestrato, osserviamo che, data la stessa architettura del modello, esiste un ampio divario prestazionale tra il supernet e l'addestramento da zero. Pertanto, il supernet non può essere utilizzato direttamente ed è necessario un riaddestramento dopo aver individuato le architetture ottimali. In questo lavoro, proponiamo il mixture-of-supernets, una formulazione generalizzata del supernet in cui viene adottato il mixture-of-experts (MoE) per potenziare la capacità espressiva del modello supernet, con un sovraccarico di addestramento trascurabile. In questo modo, le diverse sottoreti non condividono direttamente i pesi del modello, ma attraverso un meccanismo di routing basato sull'architettura. Di conseguenza, i pesi del modello delle diverse sottoreti sono personalizzati in base alle loro specifiche architetture e la generazione dei pesi viene appresa tramite discesa del gradiente. Rispetto ai supernet con condivisione dei pesi esistenti per il NLP, il nostro metodo può minimizzare il tempo di riaddestramento, migliorando notevolmente l'efficienza dell'addestramento. Inoltre, il metodo proposto raggiunge le prestazioni SOTA nel NAS per la costruzione di modelli di traduzione automatica veloci, ottenendo un migliore compromesso tra latenza e BLEU rispetto a HAT, lo stato dell'arte del NAS per la MT. Raggiungiamo anche le prestazioni SOTA nel NAS per la costruzione di modelli BERT efficienti in termini di memoria e agnostici rispetto al task, superando NAS-BERT e AutoDistil in varie dimensioni del modello.

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.