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

Area of Science:

• Artificial Intelligence
• Machine Learning
• Deep Learning

Background:

• Training large neural networks with Mixture-of-Experts (MoE) architecture demands significant computational resources.
• Existing acceleration techniques often compromise prediction performance or require limited dedicated hardware.
• Current MoE training strategies apply uniform learning plans, neglecting individual expert differences and leading to inefficient training.

Purpose of the Study:

• To develop a novel training acceleration framework for Mixture-of-Experts (MoE) neural networks.
• To address the challenge of divergent expert learning speeds and domains within MoE architectures.
• To improve the overall training efficiency and convergence of large-scale MoE models.

Main Methods:

• Proposed a multi-stage training planner that optimizes network subparts sequentially, scaling up progressively.
• Utilized a density function to assess expert knowledge and prioritize faster-learning experts for increased training scale.
• Implemented a growth operator to manage expert training scale across stages and a scheduler for dynamic learning rate adjustment to mitigate gradient vanishing.

Main Results:

• The proposed framework customizes learning plans for individual experts based on their training progress.
• Achieved an average of over 25% training acceleration in extensive experimental validations.
• Demonstrated improved training efficiency by avoiding uniform learning plans and addressing expert-specific needs.

Conclusions:

• The novel framework effectively accelerates MoE training by personalizing expert learning strategies.
• The multi-stage approach with expert-aware planning enhances convergence and reduces training time.
• This method offers a practical solution for resource-intensive MoE model training without sacrificing performance.