Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

## Abstract

Abstract For Mixture-of-Experts (MoE) models, an unbalanced expert load will lead to routing collapse or increased computational overhead.
Existing methods commonly employ an auxiliary loss to encourage load balance, but a large auxiliary loss will introduce non-negligible interference gradients into training and thus impair the model performance.
In order to control load balance while not producing undesired gradients during training, we propose Loss-Free Balancing , featured by an auxiliary-loss-free load balancing strategy.
To be specific, before the top-K routing decision, Loss-Free Balancing will first apply an expert-wise bias to the routing scores of each expert.
By dynamically updating the bias of each expert according to its recent load, Loss-Free Balancing can consistently maintain a balanced distribution of expert load.
In addition, since Loss-Free Balancing does not produce any interference gradients, it also elevates the upper bound of model performance gained from MoE training.
We validate the performance of Loss-Free Balancing on MoE models with up to 3B parameters trained on up to 200B tokens.
Experimental results show that Loss-Free Balancing achieves both better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies.

## 1 Introduction

Mixture-of-Experts (MoE) architectures have emerged as a promising solution for managing computational costs when scaling up parameters in large language models (LLMs).
Recent applications of MoE in Transformer-based models (Vaswani et al., ) have led to successful attempts at scaling language models to substantial sizes (Shao et al., ; DeepSeek-AI et al., ; Dai et al., ; Fedus et al., ; Lepikhin et al., ), resulting in remarkable performance improvements.
However, training MoE models always face the circumstance of load imbalance, which may result in routing collapse (Shazeer et al., ) or increased computational overhead (Fedus et al., ; Lepikhin et al., ; Shazeer et al., ).
In order to avoid imbalanced routing, existing methods (Fedus et al., ; Lepikhin et al., ) commonly use an auxiliary loss to encourage balanced expert load.
Although the auxiliary loss can alleviate load imbalance during training, it also introduces undesired gradients that conflict with the language modeling objective.
These interference gradients will impair the model performance, so existing MoE methods always need to consider the trade-off between load balance and model performance.

In this paper, we propose Loss-Free Balancing, an auxiliary-loss-free load balancing strategy, aiming at maintaining control over expert load balance while not introducing interference gradients.
Loss-Free Balancing features an iterative process of token routing and bias updating.
As illustrated in Figure 1, before the top-K routing decision of MoE, Loss-Free Balancing will first apply expert-wise biases to the original routing scores to produce biased gating scores, which determine the actual routing targets of each token during training.
These expert-wise biases will keep updating according to the expert load observed on recent training tokens, where the biases of heavy-load experts will be depressed and those of lite-load experts will be elevated.
Through this dynamic updating strategy, Loss-Free Balancing ensures that the biased gating scores can consistently lead to balanced routing results.
Compared with the auxiliary-loss-controlled load balancing strategies, Loss-Free Balancing does not introduce undesired gradients that disrupt the primary language modeling objective, so its training process is more noise-free and friendly.

In order to validate the performance of Loss-Free Balancing, we train MoE language models with 1B parameters on 100B tokens and 3B parameters on 200B tokens from scratch.
Experimental results demonstrate that Loss-Free Balancing produces MoE models with better validation loss than traditional auxiliary-loss-controlled models.
Meanwhile, keeping the performance advantage, Loss-Free Balancing also achieves a significantly better load balance at the global and batch levels, and is naturally compatible with expert parallelism, which is usually employed for training extremely large MoE models.

Figure: Figure 1: Loss-Free Balancing selects experts according to a “biased gating score” in each training step and updates this expert-wise bias after each training step.
Refer to caption: extracted/5818453/figs/figure_815.png

## 2 Background

### 2.1 Mixture-of-Experts

Current dominant MoE architectures (Lepikhin et al., ; Fedus et al., ; Dai et al., ) replace the MLP layers in standard transformers with MoE layers. In an MoE layer, Top-K routing is employed to select the experts for each token. Let $\mathbf{u}_{t}$ denote the input of the $t$-th token to an $N$-expert MoE layer, the output $\mathbf{h}_{t}$ is computed as follows:

$$ $\displaystyle\mathbf{h}_{t}=\mathbf{u}_{t}+\sum_{i=1}^{N}g_{i,t}\operatorname{ FFN}_{i}\left(\mathbf{u}_{t}\right),$ (1) $\displaystyle g_{i,t}=\begin{cases}s_{i,t},&s_{i,t}\in\operatorname{Topk}\left (\left\{s_{j,t}\mid 1\leq j\leq N\right\},K\right),\\ 0,&\text{ otherwise, }\end{cases}$ $\displaystyle s_{i,t}=G\left(\mathbf{u}_{t}{}^{T}\mathbf{e}_{i}\right),$ $$

where $G$ is a nonlinear gating function and $\mathbf{e}_{i}$ is the centroid of the $i$-th expert.

### 2.2 Auxiliary Loss for Load Balance

#### Auxiliary Loss

Uncontrolled routing strategies are likely to encounter load imbalance, which has two notable drawbacks.
Firstly, there is a risk of routing collapse (Shazeer et al., ), where the model consistently selects only a few experts, hindering sufficient training of the other experts.
Secondly, when experts are distributed across multiple devices, load imbalance can exacerbate computation bottlenecks.
To address these issues, an auxiliary loss (Fedus et al., ; Lepikhin et al., ) is commonly employed to control load balance.
For a sequence of length $T$, the auxiliary loss is defined as:

$$ $\displaystyle\mathcal{L}_{\text{Balance }}$ $\displaystyle=\alpha\sum_{i=1}^{N}f_{i}P_{i},$ (2) $\displaystyle f_{i}$ $\displaystyle=\frac{N}{KT}\sum_{t=1}^{T}\mathbbm{1}(\text{ Token }t\text{ selects Expert }i),$ $\displaystyle P_{i}$ $\displaystyle=\frac{1}{T}\sum_{t=1}^{T}s_{i,t},$ $$

where $N$ is the total number of experts, $K$ is the number of experts selected for each token, $s_{i,t}$ is the routing score of Expert $i$ for Token $t$, $f_{i}$ represents the fraction of tokens routed to Expert $i$, $P_{i}$ denotes the average gating scores of Expert $i$, and $\alpha$ is a hyper-parameter controlling the strength of the auxiliary loss.

Figure: Figure 2: The dilemma between load balance and model performance for auxiliary-loss-controlled training. A small auxiliary loss coefficient $\alpha$ leads to poor load balance, while a large $\alpha$ impairs the model performance. In contrast, our Loss-Free Balancing method breaks this dilemma.
Refer to caption: extracted/5818453/figs/balance_mv_perplexity_new_new.png

#### The Dilemma Between Load Balance and Model Performance

The auxiliary loss mentioned above can encourage load balance, but it also interferes with language modeling training as an additional regularization term. The absence of an auxiliary loss or a small auxiliary loss coefficient $\alpha$ can lead to poor balance, while a large $\alpha$ can impair training, resulting in suboptimal performance.
To illustrate this dilemma, we present the relationship between load balance and model performance in Figure 2.
We vary $\alpha$ among 1e-2, 1e-3, 1e-4, and 0, and present the corresponding $\text{MaxVio}_{\text{global}}$, which measures the degree of load balance and its computation details are described in § 4.1.
As shown in the figure, a small $\alpha$ causes routing collapse, affecting the model efficiency and potentially leading to some experts being insufficiently learned or exploited; while a large $\alpha$ keeps load balance under control but notably degrades the model performance.
In order to break this dilemma, we propose Loss-Free Balancing as a solution, which directly controls the expert load balance, but does not introduce unexpected gradients other than the gradients from the language modeling loss.

## 3 Auxiliary-Loss-Free Load Balancing Strategy

For a better load-balancing alternative that does not directly interfere with the main gradients from the training objective, we propose Loss-Free Balancing, which directly adjusts the gating scores of each expert according to their balance condition.
As illustrated in Figure 1, we add an expert-wise bias term $\bm{\{b_{i}\}_{i=1}^{N}}$ to the gating scores $s_{i,t}$ of each expert, and use the biased scores to determine the top-K selection:

$$ $\displaystyle g_{i,t}=\begin{cases}s_{i,t},&s_{i,t}+b_{i}\in\operatorname{Topk }\left(\left\{s_{j,t}+b_{j}\mid 1\leq j\leq N\right\},K\right),\\ 0,&\text{ otherwise. }\end{cases}$ (3) $$

Note that the expert bias term $b_{i}$ is only used to adjust the routing strategy by influencing the top-K selection.
It is not added to the $g_{i,t}$ that weights the output of the selected experts when computing the final output of the MoE layer.

In order to derive proper biases, we adjust each bias $b_{i}$ iteratively according to the following principle: decreasing it when the corresponding expert has a relatively heavy load, and vice versa.
To be specific, for each ${b}_{i}$, we keep monitoring its corresponding expert load on the previous batch.
If an expert has a heavy load on the previous batch, we will reduce its bias.
Otherwise, we will increase it.
Algorithm 1 describes the details of our update algorithm for the expert-wise biases.
It is worth noting that we update the biases based on the historical balance condition, since utilizing the load information of the current sequence will break the causal constraint of language modeling, leading to leakage of the information of future tokens.
Through the dynamic adjustment for the biases, we can achieve good expert load balance, but not directly introduce noisy gradients into the model like the auxiliary-loss-controlled method does.

Figure: Algorithm 1 Adjusting the per-expert bias ${b}_{i}$ during training

#### Comparison with Other Load Balancing Methods.

In order to show the theoretical advantages of Loss-Free Balancing, we compare it with other two mainstream load balancing methods, i.e., the auxiliary-loss-controlled method (Lepikhin et al., ; Fedus et al., ) and the Expert Choice (EC) (Zhou et al., ) method.
As described in § 2.2, the auxiliary-loss-controlled method faces the dilemma between load balance and model performance, and a perfect trade-off may not exist.
As for the EC method, it will break the causal constraint of language modeling, since the target experts of each token are conditioned on the future tokens in the same sequence or batch.
This will result in the leakage of information about future tokens, thus destroying the generalization of the model.
Table 1 summarizes the properties of different load balancing methods.

**Table 1: Comparison among different load balancing methods. The good property is displayed in green and the bad property in red.**
| Load Balancing Methods | Balanced<br>Expert Load | Interference<br>Gradients | Future Token<br>Leakage |
| --- | --- | --- | --- |
| Loss-Controlled (strong auxiliary loss) | balanced | strong | no leakage |
| Loss-Controlled (weak auxiliary loss) | imbalanced | weak | no leakage |
| Expert Choice | balanced | none | with leakage |
| Loss-Free (Ours) | balanced | none | no leakage |

## 4 Experiments

### 4.1 Experimental Setups

#### Model Architecture.

We employ the DeepSeekMoE (Dai et al., ) architecture as the backbone since it outperforms conventional MoE architectures like GShard (Lepikhin et al., ) significantly.
Compared with GShard (Lepikhin et al., ), it segments experts into finer granularity and isolates some experts as shared ones.
Slightly different from DeepSeekMoE, in our main experiments, we choose sigmoid instead of softmax as the gating function $G$, since we find that the sigmoid baseline performs better than the softmax baseline.
Even so, we still provide the experimental results and discussion for the softmax gate in Appendix C.
Our experiments are based on two model sizes of 1B and 3B total parameters, and we tune the bias update rate under only the 1B scale.
Experiments under the 3B scale directly inherit the best configuration for the 1B scale.
Due to the page limit, we present more details about our architecture in Appendix A.

#### Training Settings

We use a multilingual training corpus created by DeepSeek-AI, sourced from a diverse range of textual materials including web text, mathematical material, coding scripts, and published
literature.
We employ the HuggingFace Tokenizer(^2^22[https://github.com/huggingface/tokenizers](https://github.com/huggingface/tokenizers)) to train a byte pair encoding (BPE) (Sennrich et al., ) tokenizer with a vocabulary size of 32K.
In order to draw solid conclusions, we train the 1B model on 100B tokens and the 3B model on 200B tokens to ensure sufficient training.
We apply the cosine learning rate scheduler (Loshchilov & Hutter, ) and multi-step learning rate scheduler (Dai et al., ) for the 1B and 3B models, respectively.
Due to the page limit, we list more details about our training settings and hyper-parameters in Appendix B).

#### Baseline.

We compare our Loss-Free Balancing method with the conventional auxiliary-loss-controlled method.
For the baseline, we set the auxiliary loss coefficient $\alpha$ to 0.001 to achieve a reasonable trade-off between model performance and load balance (see Figure 2).
We do not take the EC method into comparison due to its issue of future token leakage, which we will discuss in depth in § 5.2.

#### Metrics.

We reserve a validation set from the training corpus to evaluate model performance and load balance.
For model performance, we take perplexity as the metric.
For load balance, we introduce a metric called maximal violation (MaxVio) to quantify the degree of load balance of an MoE layer:

$$ $\text{MaxVio}=\frac{\max_{i}\text{Load}_{i}-\overline{\text{Load}_{i}}}{ \overline{\text{Load}_{i}}},$ (4) $$

where $\text{Load}_{i}$ represents the number of tokens assigned to the $i$-th expert, and $\overline{\text{Load}_{i}}$ denotes the expected expert load under perfect load balance.

MaxVio has two variants: $\textbf{MaxVio}_{\textbf{global}}$ and $\textbf{MaxVio}_{\textbf{batch}}$.
For $\textbf{MaxVio}_{\textbf{global}}$, we count $\text{Load}_{i}$ on the whole validation set, so it reflects the degree of balanced expert utilization and efficiency upper bound when the batch size approaches the limitation.
For $\textbf{MaxVio}_{\textbf{batch}}$, we count $\text{Load}_{i}$ on each training batch, so it is more related to the training efficiency.
For simplicity, in the rest of this paper, we report the MaxVio averaged across all layers as a load balance measurement of the whole model.

### 4.2 Main Results

Table 2 shows the validation perplexity and $\text{MaxVio}_{\text{global}}$ for the 1B and 3B MoE models trained with auxiliary loss or our auxiliary-loss-free load balancing strategy.
As shown in the table, compared with the auxiliary-loss-controlled method, our Loss-Free Balancing achieves better perplexity and much better global load balance for both 1B and 3B models.
In addition, to present the load balance condition during training, we provide a load balancing curve depicting $\text{MaxVio}_{\text{batch}}$ over training steps in Figure 3, which demonstrates the persistent advantage of Loss-Free Balancing on load balance.
In summary, our Loss-Free Balancing method avoids interfering gradients during training and effectively controls the load balance, breaking the dilemma between load balance and model performance in MoE training.

**Table 2: Loss-Free Balancing achieves lower perplexity and better load balance on both 1B and 3B models. A validation set is used to calculate these metrics (see details in Appendix B).**
| Model Size | Load Balancing Methods | Validation Perplexity | $\textbf{MaxVio}_{\textbf{global}}$ |
| --- | --- | --- | --- |
| 1B | Loss-Controlled | 9.56 | 0.72 |
| Loss-Free | 9.50 | 0.04 |  |
| 3B | Loss-Controlled | 7.97 | 0.52 |
| Loss-Free | 7.92 | 0.04 |  |

Figure: Figure 3: Loss-Free Balancing maintains a better load balance throughout most of the training time. Here, $\text{MaxVio}_{\text{batch}}$ is averaged over 100 neighboring steps for visibility purposes.
Refer to caption: x1.png

### 4.3 Empirical Studies on Bias Update Algorithm

We conduct empirical studies on the update rate and variants of the bias update algorithm to validate the optimal configuration used in our main experiments.

#### Update rate.

The update rate $u$ in Algorithm 1 controls the speed at which the expert bias $\{{b}_{i}\}_{i=1}^{N}$ converges to the “suitable bias”. Figure 4 illustrates that an overly low update rate $u=0.0001$ may lead to slow convergence, while an unnecessarily high update rate $u=0.01$ can cause undesirable fluctuations of the expert bias ${b}_{i}$ during the later stage of training, deteriorating load balance in this stage. Both situations can impair performance. An appropriate choice is $u=0.001$, which shows good training balance and validation perplexity.

Figure: Figure 4: The impact of update rate on training load balance. A low update rate shows poor load balance in the early stage of training, while a high update rate deteriorates load balance in the later stage. Validation PPL denotes the validation perplexity.
Refer to caption: x2.png

#### Update rule.

We investigate a different update rule of the expert-wise biases.
To be specific, we attempt to change the update rule of $b_{i}=b_{i}+u*\mathrm{sign}(e_{i})$ to $b_{i}=b_{i}+u*e_{i}$, which encourages the bias of experts with high violation errors to change faster. Although this variant slightly improves load balance, it does not lead to better performance, as shown in Table 3.
Therefore, we maintain the $\mathrm{sign}$ version.

**Table 3: The variant $b_{i}=b_{i}+u*e_{i}$ slightly improves load balance but does not show improvement in model performance.**
| Method | Perplexity | $\textbf{MaxVio}_{\textbf{global}}$ |
| --- | --- | --- |
| $b_{i}=b_{i}+u*\text{sign}(e_{i})$, $u=0.001$ | 9.50 | 0.044 |
| $b_{i}=b_{i}+u*e_{i}$, $u=0.01$ | 9.53 | 0.028 |
| $b_{i}=b_{i}+u*e_{i}$, $u=0.001$ | 9.51 | 0.036 |
| $b_{i}=b_{i}+u*e_{i}$, $u=0.0001$ | 9.51 | 0.040 |

#### Multiplicative bias.

In addition to adding the expert-wise biases to the gating scores, using multiplicative biases is also a potential variant:

$$ $\displaystyle g_{i,t}=\begin{cases}s_{i,t},&s_{i,t}*{b}_{i}\in\operatorname{ Topk}\left(\left\{s_{j,t}*{b}_{j}\mid 1\leq j\leq N\right\},K\right),\\ 0,&\text{ otherwise, }\end{cases}$ (5) $$

These $\{{b}_{i}\}_{i=1}^{N}$ can be updated using a similar procedure to Algorithm 1, except that they should be initialized as 1 instead of 0.
Table 4 shows that using multiplicative biases results in slightly worse model performance compared to using additive biases, without significant improvements in load balance.
Based on these findings, we conclude that additive biases are a more suitable choice for our method.

**Table 4: Multiplicative bias shows similar load balance but slightly worse performance compared to additive bias.**
| Method | Perplexity | $\textbf{MaxVio}_{\textbf{global}}$ |
| --- | --- | --- |
| Addative Bias, $u=0.001$ | 9.50 | 0.044 |
| Multiplicative Bias, $u=0.01$ | 9.52 | 0.041 |
| Multiplicative Bias, $u=0.001$ | 9.52 | 0.036 |
| Multiplicative Bias, $u=0.0001$ | 9.54 | 0.048 |

## 5 Discussion

### 5.1 Loss-Free Balancing Is Compatible with Expert Parallelism

Extremely large-scale MoE models often employ expert parallelism (Lepikhin et al., ) for training or inference, which distributes experts across different devices to reduce memory requirements.
In such scenarios, load balance on the data in a single computation step is crucial for efficiency.
Due to expert parallelism, each computation step involves micro_batch_size * ep_data_parallel_size samples, which we refer to as a computation batch.
Here, micro_batch_size denotes the number of samples processed in one gradient accumulation step on a single device.

Loss-Free Balancing can achieve nearly optimal global load balance, and the load balance in each computation step will get closer to the global load balance as the computation batch size increases.
In Figure 5, we examine the computation-batch-level load balance with the $\text{MaxVio}_{\text{computation-batch}}$ metric.
The results show that the load balance of our Loss-Free Balancing always keeps improving as the computation batch size increases, but the load balance of the auxiliary-loss-controlled method approximately maintains a constant level when the computation batch is large.
Since expert parallelism will significantly increase the computation batch size by ep_data_parallel_size times, Loss-Free Balancing is naturally compatible with large-scale MoE training, and its advantage on the load balance will be further enhanced as the size of expert parallelism increases.

Figure: Figure 5: Loss-Free Balancing achieves improved balance compared to auxiliary-loss training as the computation-batch size increases, demonstrating its superiority when a moderately sized computation-batch is utilized.
Refer to caption: x3.png

### 5.2 Load Balancing and Future Token Leakage

For casual language models, load balancing methods must adhere to the causal constraint of language modeling to avoid future token leakage. While conventional auxiliary-controlled balancing and our Loss-Free Balancing obey this constraint, Expert Choice (EC) (Zhou et al., ) violates it. EC ensures perfect load balance by assigning exactly the same number of tokens to each expert. However, this approach inherently leads to a severe issue of future token leakage.

In EC, future tokens can influence the expert assignment of previous tokens. Figure 6 illustrates how information can be easily transmitted within a sequence via such influence. Theoretically, the token assignment of an MoE layer with sparse ratio $R$ (average activated experts per token $K$ divided by total expert number $N$) can leak more than $K\log_{2}\frac{1-R}{R}$ bits per token (proof in Appendix D.1). For a 9-layer MoE model with 16 experts and an average of 2 experts per token, this amounts to 50 bits, sufficient for each token to determine its successor’s identity.

Figure: Figure 6: An example of future token leakage in EC. Future tokens can influence the expert assignment of previous tokens. Such an assignment can help previous tokens to infer the identity of their successors.
Refer to caption: extracted/5818453/figs/EC_fig_new.png

We designed experiments to demonstrate the existence of future token leakage in realistic model training. (1) We reduced the chunk size, within which top-K selection is performed, from 8192 tokens (4 sentences) to 512 (1/4 sentence), with the expectation of exposing such leakage. We observed an abnormal loss drop (about 10%), confirming the presence of leakage. (2) We made leakage more difficult by shuffling tokens across chunks in the top-K selection step, and observed that the abnormal loss drop was mitigated. Detailed experimental results on EC’s information leakage are provided in Appendix D.2.

Future token leakage is fatal since it destroys the generalization of a model and prevents reliable evaluation of the model performance.
Therefore, compared with EC, scaling up an MoE model with our Loss-Free Balancing is safer.

## 6 Conclusion

In this work, we introduced Loss-Free Balancing, a novel MoE load balance control method without introducing auxiliary-loss gradients. Loss-Free Balancing addresses the issue of traditional auxiliary-loss load balance control, which introduces additional gradients during training and potentially impairs model performance when enforcing load balance. Experiments conducted on 1B and 3B MoE models, trained on 100B and 300B tokens respectively, demonstrate that Loss-Free Balancing achieves better model performance and load balance compared to the traditional auxiliary-loss training.

## Appendix A Model Architecture

We employ the DeepSeekMoE (Dai et al., ) architecture as the backbone, which introduces shared experts to mitigate knowledge redundancy among routed experts:

$$ $\displaystyle\mathbf{h}_{t}=\mathbf{u}_{t}+\sum_{i=1}^{N_{s}}\operatorname{FFN }_{i}^{(s)}\left(\mathbf{u}_{t}\right)+\sum_{i=1}^{N_{r}}g_{i,t}\operatorname{ FFN}_{i}^{(r)}\left(\mathbf{u}_{t}\right),$ (6) $$

where $r$ denotes the routed experts, while $s$ the shared experts. DeepSeekMoE replaces all FFN layers with MoE layers, except the dense FFN layer just after the input embedding layer.

The detailed architecture hyper-parameters are listed in Table 5.

**Table 5: Model architecture.**
| hyper-parameters | 1B | 3B |
| --- | --- | --- |
| Vocab size | 32064 | 32064 |
| Hidden size | 1024 | 1280 |
| Attention heads | 8 | 10 |
| MoE layers | 9 | 11 |
| Granularity ($\frac{d_{\text{ff}}}{d_{\text{expert}}}$) | $\frac{\text{16}}{\text{3}}$ | 4 |
| Shared experts | 2 | 2 |
| Routed experts | 64 | 64 |
| Activated routed experts | 6 | 6 |

## Appendix B Training Settings

Following the work of Dai et al. (), we initialize all learnable parameters with a standard deviation of 0.006, and set the maximum training sequence length to 2048.

For the 1B model, we employ a cosine learning rate scheduler with warmup, setting the learning rate to 1e-3, the minimum learning rate to 1e-4, and the warmup steps to 1000. The training batch size for the 1B model is set to 1152, resulting in a total of 40000 training steps (100B tokens).

For the 3B model, we use a multistep learning rate scheduler with stage steps = [45211, 50862, 56514] and corresponding stage learning rates of [7.8e-4, 2.47e-4, 7.8e-5]. The warmup steps for the 3B model are set to 2000. We use a training batch size of 1728 for the 3B model, resulting in a total of 56514 training steps (200B tokens).

For validation, we leave around 70M tokens from the training corpus as the validation set (30 * 1B_batch_size * max_seq_len = 20 * 3B_batch_size * max_seq_len = 71M tokens).

## Appendix C Experiments with Softmax Gate

### C.1 Comparison of Sigmoid Gate Baseline and Softmax Gate Baseline

We compare the sigmoid gate baseline and the softmax gate baseline with varying auxiliary loss coefficients $\alpha$ on a 1B-sized model. As shown in Figure 7, the softmax gate exhibits higher perplexity under similar load balance conditions, and its performance is more sensitive to load imbalance compared to the sigmoid gate.

Figure: Figure 7: Comparison of the sigmoid gate baseline and the softmax gate baseline. The softmax gate exhibits higher perplexity under similar load balance conditions and is more sensitive to load imbalance compared to the sigmoid gate.
Refer to caption: extracted/5818453/figs/balance_mv_perplexity_softmax_new.png

### C.2 Loss-Free Load Balancing with Softmax Gate

Adjusting the per-expert bias for the softmax gate is more challenging due to the normalization property of softmax, which makes the score gap between two experts sensitive to the scores of other experts. In such a situation, we choose the $\mathbf{b}_{i}=\mathbf{b}_{i}+u*e_{i}$ variant to maintain load balance, where $u$ is set to 1e-3. For the baseline, we choose $\alpha$ = 0.0003, which yields the lowest perplexity for the softmax gate. The results are presented in Table 6, showing that Loss-Free Balancing achieves a slightly lower perplexity while maintaining significantly better load balance compared to the auxiliary-loss training method. Figure 8 confirms that Loss-Free Balancing maintains a superior load balance throughout most of the training process.

**Table 6: For softmax gate, Loss-Free Balancing achieves a slightly lower perplexity while reaching a significantly better load balance compared to the auxiliary-loss training method.**
| Load Balancing | Perplexity | $\textbf{MaxVio}_{\text{global}}$ |
| --- | --- | --- |
| Loss-Controlled | 9.604 | 0.937 |
| Loss-Free | 9.599 | 0.027 |

Figure: Figure 8: For softmax gate, Loss-Free Balancing maintains a superior load balance throughout most of the training process.
Refer to caption: x4.png

## Appendix D Future Token Leakage in Expert Choice

### D.1 Proof for Theoretical Leakage Amount

Let $R=\frac{K}{N}$ denote the MoE sparsity. Here $K$ denotes the average number of experts activated per token, and $N$ is the total number of experts.
For an MoE layer in Expert Choice, the maximum information leakage $I$ (in bits per token), i.e., the information that the combinations of routing allocation can carry is:

$$ $\displaystyle I=$ $\displaystyle\left.\log_{2}\binom{\frac{KT}{N}}{T}^{N}\right/T$ (7) $\displaystyle>$ $\displaystyle N\left.\log_{2}\frac{((1-\frac{K}{N})T)^{\frac{K}{N}T}}{(\frac{K }{N}T)^{\frac{K}{N}T}}\right/T$ $\displaystyle=$ $\displaystyle K\log_{2}\frac{1-R}{R}.$ $$

For a model with a sparse ratio $R=\frac{2}{16}=0.125$ and 9 MoE layers, the total leakage information is more than 50 bits per token.

### D.2 Experimental Evidence

We investigate the potential future token leakage of the Expert Choice by varying the chunk size used for experts’ top-$k$ selection, ranging from 512 tokens to 8192 tokens.(^3^33A chunk size of 2048 tokens means performing top-$k$ selection inside a sentence, while 512 tokens correspond to a quarter of a sentence and 8192 tokens to four sentences.) We train a 2B MoE model on 100B tokens. The results, shown in Table 9, reveal two key findings:

- 1.
Using a small chunk size of 512 leads to an abnormal loss drop, which can be attributed to significant future token leakage. A smaller chunk size allows the model to more easily exploit information from future tokens within the chunk during training.
- 2.
Shuffling tokens within a batch before chunking and selecting mitigates the observed loss drop. Such shuffling makes it more challenging for the model to utilize information leakage, as the future tokens are no longer in their original context. This finding supports the hypothesis that the loss drop originates from the model’s accessing and exploiting future token information.

Figure: Figure 9: Comparison of Expert Choice with different chunk sizes and shuffling. Expert Choice with a chunk size of 512 exhibits a significant loss drop compared to chunk sizes of 8192 or 2048. Shuffling tokens eliminates this loss drop, indicating the presence of future token leakage.
Refer to caption: x5.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: The load-balancing loss brings the advantage of load balancing to the inference stage by directly affecting the routing weights? How does the proposed method behave during inference.
soundness: 2
strengths: - The paper is well-written and the method is clearly explained
- Good visualizations
- Simple approach that should be very easy to test and cheaper to compute than the conventionally used load-balancing loss
summary: The authors propose an alternative loss-free method for load balancing of experts during MoE training. Load imbalance is a critical issue in MoE training as it can lead to expert collapse or increased utilization of some experts over the others. The proposed method achieves load balancing by dynamically applying expert-wise biases on routing scores according to their recent load, avoiding interfering gradients. The added bias is designed to only affect the top-k selection without changing the routing weights for combining the selected experts.

The loss-free load balancing approach is applied to DeepSeekMoE models with sizes 1B and 3B. The authors report perplexity on the validation set vs Maximal Violation scores.
weaknesses: - I am concerned over the validity of the claims. The empirical evaluations are very limited constrained to two DeepSeekMoE models and perplexity differences among the models and the baselines are at the level of 0.05 difference. Is this difference in perplexity significant?

- The evaluation is limited to language modelling and perplexity values. It would be better to see the actual effect of the loss-free load balancing on other downstream tasks such as MMLU or GLUE.

- The proposed Max Violation (MaxVio) score subtracts the mean load from the maximum load an expert receives, which highlights the worst-case scenario of load imbalance. Since MaxVio focuses on the maximum load, it can be highly sensitive to outliers. A single batch with unusually high token routing to one expert can disproportionately affect the MaxVio calculation, making it appear as though there is a significant imbalance even if most batches are well-balanced.

- To make the comparison fair, it would be useful to also see the load balancing (loss) values (without back propagating) and see how the loss-free variant behaves on that score as the training continues (e.g., a plot of load balancing loss over training steps for both methods).

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: Please see my questions in Weakness column.
soundness: 2
strengths: This paper has the following strengths:

1. The proposed Loss-Free Balancing method eliminates the need for auxiliary loss, which traditionally adds undesirable interference gradients. This results in a cleaner training signal focused solely on the primary language modeling objective, potentially enhancing overall model performance.

2. By dynamically adjusting biases for each expert based on recent load data, the method ensures a balanced expert load without compromising model efficiency.

3. The strategy is compatible with expert parallelism, a key feature for training extremely large MoE models across multiple devices. This makes it highly suitable for scaling up model sizes while maintaining efficient load balancing.

4. Unlike the Expert Choice (EC) method, which risks future token leakage, Loss-Free Balancing maintains causal constraints in language modeling, thus preserving the integrity of model generalization.
summary: The paper introduces Loss-Free Balancing, a new load balancing strategy for MoEs that avoids auxiliary loss, which traditionally introduces interference gradients and hampers model performance. The motivation stems from the need to balance expert loads in MoE models to prevent routing collapse and computational overhead, issues that current auxiliary-loss methods attempt to address but with performance trade-offs. The major research question is whether a balancing strategy can maintain expert load distribution without harming model performance. The proposed method adjusts each expert's routing score using a dynamically updated bias based on recent loads, promoting balanced expert utilization without adding interference gradients. Experiments on 1B and 3B parameter MoE models trained on abundant datasets show that Loss-Free Balancing achieves better load balance and improved validation perplexity compared to traditional methods, making it a promising approach for scaling large language models while preserving efficiency and performance.
weaknesses: This paper has the following weaknesses:

1. I am at first astonished by the short reference list of this paper, as the authors only cited 10 papers. Clearly, this paper did a very bad job on surveying the related work, including the various auxilliary-loss-based balancing methods, the major improvement of MoEs, the current MoE-based LLMs. Normally, I would list a few of the works for your reference, but the authors missed too many so I do not know where to start. I would strongly suggest the authors to check other MoE-LLM papers to improve the related work part.

2. The authors seemed to enlarge the figures in this paper in order to make the paper length reach 9 pages, which result in a disproportionate paper layout, and the figures look abrupt. 

3. The architecture design of the MoE-based LLMs are versatile. This paper only demonstrates its effectiveness on a small DeepSeek MoE model, while leaving other MoE models unvisited, such as Mistral, QWen, LLaMA-MoE, OpenMoE.

4. The improvement in performance seem trivial, which challenges the motivation of this study: why do we need a loss-free balancing at all.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. On Page 4, in Formula 3, the selection and initialization process for b_i remains unclear. It would be helpful to clarify whether b_i  is influenced by K or independent of it. Additionally, in the adjustment process—where each bias b_i is iteratively modified based on the load of the corresponding expert—it is unclear by what amount b_i  should be increased or decreased and according to which theoretical basis this adjustment is made.
2. On Page 8, it is stated that 'the load balance of our Loss-Free Balancing consistently improves as the computation batch size increases.' Is there an upper bound to this improvement? 
3. Is there a specific reason for selecting 9 MoE layers for testing?
soundness: 3
strengths: This is an interesting research problem and the author aims to develop an efficient solution approach
summary: This paper proposes a Logg-free balancing strategy to solve the mixture of expert models' unbalanced expert load problem. The advantage of the approach is that it does not produce any interference gradients.
this paper proved that Loss-Free Balancing achieves better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies
weaknesses: 1. The motivation and underlying intuition for the proposed approach could be clarified further to enhance understanding.
2. Additional experiments are recommended to demonstrate the robustness of this approach when applied across varying numbers of expert mixtures. A scalable analysis would also be beneficial.
3. The approach would be strengthened with theoretical justification to substantiate its effectiveness.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: 1. Theoretically, can it be shown that an auxiliary loss produces interference gradients that impact model performance?

2. Please clarify further how expert choice leads to future token leakage. 

3. Can this method be applied to a more diverse range of MoE models and for computer vision as well. The improvement over the baseline seems marginal. I have doubts about the impact of the method in improving performance.
soundness: 1
strengths: 1. There is some originality in viewing the problem of load balancing from the perspective of interfering gradients. 

2. The problem and solution proposed are simple and easy to understand.
summary: This paper aims to replace the standard auxiliary load-balancing loss in mixture of experts (MoE) models with a new load balancing strategy that does not involve additional losses. This is to ensure tokens are being routed evenly to each expert without introducing gradients that do not directly contribute to achieving the language model objective. The authors also introduce a metric to quantify the degree of load balance in a MoE model.
weaknesses: 1. It is unclear if or how the gradients due to the auxiliary load balancing loss results in the model not achieving its training objective. 

2. It is unconvincing that there is a future token leakage when using expert choice. The explanation was brief and relied on Fig. 6 which does not explain how there is a break in the causal constraint of the language modeling. While some empirical evidence is reported, the evidence assumes that "an abnormal loss drop" must be due to a future token leakage. 

3. The experimental results are severely lacking. Only one model is used, DeepSeekMoE, instead of the more commonly used Mixtral or traditional Switch transformer. The improvements are marginal even if there is a significant improvement in load balancing (according to their proposed metric). 

Minor: in the definition for MaxVio, the $\max_i$ is included in the numerator of fraction which does not make sense as the index i is also involved in $\bar{Load_i}$

Meta-review: After reviewing the authors' rebuttal and considering the discussion phase, all reviewers maintained reservations about this submission, resulting in final ratings of 3 (Reviewer tLSf), 5 (Reviewer TZLQ), 5 (Reviewer 8VSt), and 3 (Reviewer QrX9). 

While the rebuttal partially addressed some concerns, notably leading Reviewer TZLQ to raise their score to 5, the reviewers collectively concluded that the remaining weaknesses outweigh the strengths of the work. Important issues, as identified during the review process, limit the overall contribution and impact of the paper. Regrettably, given the current state of the submission and the lack of a compelling consensus in favor of its acceptance, I recommend rejection.

Decision: Reject