혼합 슈퍼넷: 아키텍처 기반 전문가 혼합을 통한 가중치 공유 슈퍼넷 학습 개선

초록

가중치 공유 슈퍼넷(weight-sharing supernet)은 최신(State-of-the-Art, SOTA) 신경망 구조 탐색(Neural Architecture Search, NAS) 프레임워크에서 성능 추정을 위한 핵심 구성 요소로 자리 잡았습니다. 슈퍼넷은 재학습 없이도 다양한 서브네트워크를 직접 생성할 수 있지만, 가중치 공유로 인해 이러한 서브네트워크의 품질이 보장되지 않는다는 문제가 있습니다. 기계 번역이나 사전 학습된 언어 모델링과 같은 자연어 처리(NLP) 작업에서, 동일한 모델 구조를 사용하더라도 슈퍼넷과 처음부터 학습한 모델 간에 큰 성능 차이가 관찰됩니다. 따라서 슈퍼넷은 직접 사용할 수 없으며, 최적의 구조를 찾은 후 재학습이 필수적입니다. 본 연구에서는 슈퍼넷의 표현력을 향상시키기 위해 전문가 혼합(Mixture-of-Experts, MoE) 방식을 도입한 일반화된 슈퍼넷 구조인 'mixture-of-supernets'를 제안합니다. 이 방법은 학습 오버헤드를 거의 증가시키지 않으면서도, 서로 다른 서브네트워크가 모델 가중치를 직접 공유하지 않고, 구조 기반 라우팅 메커니즘을 통해 간접적으로 공유하도록 합니다. 결과적으로, 각 서브네트워크의 모델 가중치는 해당 구조에 맞게 맞춤화되며, 가중치 생성은 경사 하강법을 통해 학습됩니다. 기존의 NLP용 가중치 공유 슈퍼넷과 비교했을 때, 본 방법은 재학습 시간을 최소화하여 학습 효율을 크게 개선합니다. 또한, 제안된 방법은 빠른 기계 번역 모델 구축을 위한 NAS에서 SOTA 성능을 달성하며, 최신 NAS 기법인 HAT보다 더 나은 지연 시간(latency)과 BLEU 점수 간의 균형을 제공합니다. 또한, 메모리 효율적인 작업 독립적(task-agnostic) BERT 모델 구축을 위한 NAS에서도 SOTA 성능을 달성하여, 다양한 모델 크기에서 NAS-BERT와 AutoDistil을 능가합니다.

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: Improving Weight-Sharing Supernet Training with Architecture-Routed Mixture-of-Experts

초록

Support