Accelerated Optimization of Large Mixture-of-Experts Models by Density-Aware Multi-Stage Learning.

Abstract

This article aims to speed up the training of large neural networks with the Mixture-of-Experts (MoE) structure. Training MoE often needs a lot of computing resources due to its large scale. Traditional acceleration methods either degrade prediction performance or rely on dedicated hardware with additional resources, but the resources are usually limited in real applications.One solution is to resort to new optimization strategies, such as learning from easy to hard by multiple stages. However, existing strategies are designed mainly for networks with a serial structure, but MoE has multiple expert networks working in parallel. They employ an identical learning plan for all experts, ignoring that each expert's learning domain and speed differ, resulting in some experts being over-learned while others being under-learned. This mismatch will make it hard for experts to train together, harming training efficiency. To address this problem, we propose a new training acceleration framework. It can customize an effective learning plan for each expert by considering their training progress, avoiding blindly searching in a huge parameter space. In detail, we first design a multi-stage planner that starts with optimizing a subpart of the network and then scales it up to retrain until it expands to an entire network. It uses the density function to assess the knowledge gained by the expert in each stage, giving priority to the experts who learn faster to increase the training scale, so as to boost convergence. Afterward, we exploit the growth operator to add the expert training scale of the next stage. In each stage, the network would converge to some locally optimal values. That can provide a better initialization to train the next stage more easily, since the time and data required for training from scratch are greatly reduced. To alleviate the gradient vanishing problem caused by network growth, we develop a scheduler to dynamically adjust the learning rate. Extensive experiments are conducted to validate the effectiveness of our method. The results show that we can obtain more than 25% training acceleration on average.