Improving the Performance of Mixture-of-Expert by Feeding More Data to Hungry Expert

Abstract

In recent years, an increasing number of studies have used Large Language Models (LLMs) in the field of Natural Language Processing (NLP), which use large amounts of data and computing resources to allow the model to understand and think about human language. Many applications have already been in our lives, such as Chat-GPT, ClaudeAI, and MetaAI. Although LLMs achieves excellent performance in many tasks, it requires a large amount of space and performance for data storage and calculation, which has become a problem that LLMs must face. To address these challenges, the Mixture-of-Experts (MoE) model introduces a new chapter to LLMs, which can significantly improve performance and efficiency through dynamic routing and sparsity features. However, the MoE model is prone to overfitting, owing to its sparse characteristics.

In many past studies, the Dropout mechanism has been used to alleviate the problem of overfitting, but the accuracy of the model will be lost because experts are randomly discarded. To solve this problem, this study proposes a new dynamic allocation mechanism, G-Router, that adjusts the token allocation probability through the gradient value to solve the accuracy loss problem caused by Dropout. Experimental results show that G-Router outperforms traditional Dropout in various tasks, improving the overall MoE model, promising for future development. The superior performance of G-Router over Dropout highlights its potential as a valuable tool for enhancing the performance of MoE models, paving the way for further research and advancement in this field.

G-Router

Figure 1: System architecture of G-Router

We illustrate the G-Router in Figure 1. G-Router is a new router mechanism that recalculates each expert’s token dispatch probability. This mechanism optimizes the distribution of tasks, ensuring that each expert receives an appropriate number of tokens based on their calculated probability. By dynamically adjusting these probabilities, the G-Router enhances the overall accuracy of the model and alleviates the problem of overfitting. This chapter will explain the specific formula used for these calculations. Our goal is to demonstrate how the G-Router improves performance and facilitates better resource allocation within the system.

We denote the gradient of expert \( i \) as \( gi \). We normalize each expert gradient value by Algorithm 1 The new probability \( pi \) is calculated based on the expert’s gradient. The calculations are as follows: \[ p_i = \frac{1}{1 + g_i} \]

              Algorithm 1 Softmax Function

              Require: Input vector \( z = [z_1, z_2, \ldots, z_K] \)

              Ensure: Output vector \( y = [y_1, y_2, \ldots, y_K] \)

              1:  for \( i = 1 \) to \( K \) do

              2:    Compute \( y_i = \frac{\exp(z_i)}{\sum_{j=1}^{K} \exp(z_j)} \)

              3:  end for

              4:  return \( y \)

Figure 2 shows the method flowchart of the G-Router mechanism. During the training phase, the gradient value of each expert is recorded first. Then, forward propagation is performed to compute the gated network. The gated network operates on the input value and produces a tensor matrix with dimensions corresponding to the number of tokens and experts. Each token has a value associated with a different expert. The recorded gradient values are then input into Algorithm 1 for normalization processing to calculate the new distribution probability. The router score is recalculated, and finally, gating is performed to obtain the gate value.

Figure 2: The flowchart of G-Router mechanism.

Experiments Setup

We use Stanford Question Answering Dataset(SQuAD) and some of General Language Understanding Evaluation(GLUE) data in our experiment. Our experimental architecture uses the FastMoE framework to convert the BERT model into a Mixture-of-Experts (MoE) model and conduct experimental comparisons between Dropout and G-Router on Tiny BERT, Small BERT, and Medium BERT. It verifies our proposed method’s effect on different scales of BERT models.

Overall Performance

To ensure the reliability of our results, we averaged each experimental data point across multiple trials, instead of relying on specific instances of peak performance.

(a) Performance on SQuAD Dataset

	TinyBERT		SmallBERT		MediumBERT
	Dropout	G-Router	Dropout	G-Router	Dropout	G-Router
0	38.32	40.91	70.67	71.15	73.51	74.04
1	38.18	40.79	70.88	70.98	73.77	73.99
2	38.22	41.03	70.65	71.04	72.89	74.56
3	37.98	40.78	70.08	71.24	73.62	74.23
4	38.16	40.73	70.73	70.99	73.58	74.18
Accuracy	38.17	40.85	70.60	71.08	73.47	74.20

(b) Performance on CoLA Dataset

	TinyBERT		SmallBERT		MediumBERT
	Dropout	G-Router	Dropout	G-Router	Dropout	G-Router
0	69.15	69.06	74.11	73.83	70.95	73.15
1	69.60	69.41	74.20	74.40	72.48	73.15
2	69.22	69.51	73.83	74.68	72.77	73.34
3	69.51	69.42	73.98	74.57	73.04	73.22
4	69.36	69.60	74.02	74.72	72.50	73.27
Accuracy	69.37	69.40	74.02	74.44	72.20	73.23

(c) Performance on SST-2 Dataset

	TinyBERT		SmallBERT		MediumBERT
	Dropout	G-Router	Dropout	G-Router	Dropout	G-Router
0	81.88	81.88	87.95	87.84	90.59	89.56
1	82.22	82.56	88.18	88.42	89.79	90.37
2	82.30	82.11	87.84	88.65	89.45	89.67
3	81.96	82.68	87.95	89.11	89.67	90.13
4	82.15	82.34	88.64	87.73	89.63	89.94
Accuracy	82.10	82.31	88.12	88.35	89.81	89.93

(d) Performance on MRPC Dataset

	TinyBERT		SmallBERT		MediumBERT
	Dropout	G-Router	Dropout	G-Router	Dropout	G-Router
0	78.89	78.66	86.05	86.16	85.70	85.83
1	78.29	78.86	85.84	86.48	86.02	86.42
2	78.24	79.12	86.12	86.32	85.83	86.40
3	78.43	78.93	86.39	86.44	85.79	86.18
4	78.36	78.94	86.18	86.37	85.72	86.26
Accuracy	78.44	78.90	86.12	86.35	85.81	86.22

Convergence Analysis

Figure 3 presents our experimental results on the SQuAD dataset, showcasing the accuracy and loss values over the training period. The red curve corresponds to the experiment using G-Router, while the blue curve represents the experiment using Dropout.

These results collectively confirm the effectiveness of the dynamic allocation probability method we designed for model training. The G-Router not only improves model accuracy but also accelerates the convergence speed, providing a robust alternative to traditional Dropout techniques.

Figure 3: The figure shows the accuracy and loss value comparison between the experimental results of G-Router and Dropout on the SQuAD dataset.

Conclusion

In our work, we point out the critical issue by using Dropout technology. We propose G-Router, a gradient-informed token allocation mechanism designed to overcome the limitations of conventional Dropout in Mixture of Experts (MoE) models. By dynamically adjusting the routing probabilities based on expert-specific gradient magnitudes, our method effectively mitigates the stochastic nature of Dropout, leading to improved token assignment and enhanced utilization of expert networks. Empirical evaluations on benchmark datasets demonstrate that G-Router significantly improves convergence speed and model accuracy, while also reducing computational overhead. These findings underscore the potential of G-Router as a scalable and efficient solution for optimizing expert selection, thereby contributing to the advancement of overfitting mitigation techniques in deep learning.