Mezcla-de-Superredes: Mejorando el Entrenamiento de Superredes con Compartición de Peso mediante Mezcla de Expertos con Enrutamiento Arquitectónico

Resumen

La superred con compartición de pesos se ha convertido en un componente esencial para la estimación de rendimiento en los marcos de búsqueda de arquitecturas neuronales (NAS) más avanzados (SOTA). Aunque la superred puede generar directamente diferentes subredes sin necesidad de reentrenamiento, no hay garantía de la calidad de estas subredes debido a la compartición de pesos. En tareas de PLN como la traducción automática y el modelado de lenguaje preentrenado, observamos que, dada la misma arquitectura de modelo, existe una gran brecha de rendimiento entre la superred y el entrenamiento desde cero. Por lo tanto, la superred no puede usarse directamente y es necesario reentrenar después de encontrar las arquitecturas óptimas. En este trabajo, proponemos una mezcla de superredes, una formulación generalizada de superred en la que se adopta la mezcla de expertos (MoE) para mejorar el poder expresivo del modelo de superred, con un costo de entrenamiento insignificante. De esta manera, las diferentes subredes no comparten directamente los pesos del modelo, sino a través de un mecanismo de enrutamiento basado en la arquitectura. Como resultado, los pesos del modelo de diferentes subredes se personalizan según sus arquitecturas específicas y la generación de pesos se aprende mediante descenso de gradiente. En comparación con las superredes con compartición de pesos existentes para PLN, nuestro método puede minimizar el tiempo de reentrenamiento, mejorando significativamente la eficiencia del entrenamiento. Además, el método propuesto logra el rendimiento SOTA en NAS para construir modelos de traducción automática rápidos, obteniendo un mejor equilibrio entre latencia y BLEU en comparación con HAT, el NAS SOTA para MT. También alcanzamos el rendimiento SOTA en NAS para construir modelos BERT eficientes en memoria y agnósticos a la tarea, superando a NAS-BERT y AutoDistil en varios tamaños 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.

Mezcla-de-Superredes: Mejorando el Entrenamiento de Superredes con Compartición de Peso mediante Mezcla de Expertos con Enrutamiento Arquitectónico

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

Resumen

Support