Mixture-of-Supernets：アーキテクチャルーティング型Mixture-of-Expertsを用いた重み共有スーパーネット学習の改善

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

June 8, 2023

著者: 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

要旨

重み共有型スーパーネットは、最先端（SOTA）のニューラルアーキテクチャサーチ（NAS）フレームワークにおいて、性能推定の重要な要素となっています。スーパーネットは再学習なしで異なるサブネットワークを直接生成できますが、重み共有のため、これらのサブネットワークの品質が保証されるわけではありません。機械翻訳や事前学習済み言語モデリングなどのNLPタスクでは、同じモデルアーキテクチャであっても、スーパーネットとスクラッチからの学習との間に大きな性能差が観察されます。そのため、スーパーネットを直接使用することはできず、最適なアーキテクチャを見つけた後に再学習が必要です。 本研究では、スーパーネットの表現力を向上させるために、Mixture-of-Experts（MoE）を採用した一般化されたスーパーネットの定式化であるMixture-of-Supernetsを提案します。これにより、異なるサブネットワークはモデルの重みを直接共有するのではなく、アーキテクチャに基づくルーティングメカニズムを通じて共有されます。その結果、異なるサブネットワークのモデル重みは、それぞれの特定のアーキテクチャに合わせてカスタマイズされ、重み生成は勾配降下法によって学習されます。NLP向けの既存の重み共有型スーパーネットと比較して、本手法は再学習時間を最小化し、学習効率を大幅に向上させることができます。さらに、提案手法は、高速な機械翻訳モデルを構築するためのNASにおいてSOTA性能を達成し、HAT（機械翻訳向けのSOTA NAS）と比較して、レイテンシーとBLEUのトレードオフを改善します。また、メモリ効率の良いタスク非依存型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.