Mixture-of-Supernets: Melhorando o Treinamento de Supernets com Compartilhamento de Peso por meio de Mistura de Especialistas com Roteamento de Arquitetura

Resumo

A super-rede com compartilhamento de pesos tornou-se um componente vital para a estimativa de desempenho nos frameworks mais avançados (SOTA) de busca de arquitetura neural (NAS). Embora a super-rede possa gerar diretamente diferentes sub-redes sem necessidade de retreinamento, não há garantia da qualidade dessas sub-redes devido ao compartilhamento de pesos. Em tarefas de PLN, como tradução automática e modelagem de linguagem pré-treinada, observamos que, dada a mesma arquitetura de modelo, há uma grande lacuna de desempenho entre a super-rede e o treinamento a partir do zero. Portanto, a super-rede não pode ser usada diretamente, e o retreinamento é necessário após a descoberta das arquiteturas ótimas. Neste trabalho, propomos a mistura-de-super-redes, uma formulação generalizada de super-rede na qual a mistura-de-especialistas (MoE) é adotada para aumentar o poder expressivo do modelo de super-rede, com sobrecarga de treinamento insignificante. Dessa forma, diferentes sub-redes não compartilham os pesos do modelo diretamente, mas por meio de um mecanismo de roteamento baseado em arquitetura. Como resultado, os pesos do modelo de diferentes sub-redes são personalizados para suas arquiteturas específicas, e a geração de pesos é aprendida por descida de gradiente. Em comparação com a super-rede com compartilhamento de pesos existente para PLN, nosso método pode minimizar o tempo de retreinamento, melhorando significativamente a eficiência do treinamento. Além disso, o método proposto alcança o desempenho SOTA em NAS para a construção de modelos rápidos de tradução automática, proporcionando um melhor equilíbrio entre latência e BLEU em comparação com o HAT, o NAS SOTA para MT. Também alcançamos o desempenho SOTA em NAS para a construção de modelos BERT eficientes em memória e independentes de tarefa, superando o NAS-BERT e o AutoDistil em vários tamanhos de modelo.

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: Melhorando o Treinamento de Supernets com Compartilhamento de Peso por meio de Mistura de Especialistas com Roteamento de Arquitetura

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

Resumo

Support