TITLE: Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
ARXIV: 2408.15664
AUTHORS: Lean Wang , Huazuo Gao , Chenggang Zhao , Xu Sun , Damai Dai
DATE: [Submitted on 28 Aug 2024]

=== SECTION HEADINGS ===
## Report GitHub Issue
## Auxiliary-Loss-Free Load Balancing
## Abstract
## 1 Introduction
## 2 Background
## 2.1 Mixture-of-Experts
## 2.2 Auxiliary Loss for Load Balance
## Auxiliary Loss
## The Dilemma Between Load Balance and Model Performance
## 3 Auxiliary-Loss-Free Load Balancing Strategy
## Comparison with Other Load Balancing Methods.
## 4 Experiments
## 4.1 Experimental Setups
## Model Architecture.
## Training Settings
## Baseline.
## Metrics.
## 4.2 Main Results
## 4.3 Empirical Studies on Bias Update Algorithm
## Update rate.
## Update rule.
## Multiplicative bias.
## 5 Discussion
## 5.1 Loss-Free Balancing Is Compatible with Expert Parallelism
## 5.2 Load Balancing and Future Token Leakage
## 6 Conclusion

=== BODY (head) ===
Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
## Report GitHub Issue
×
Title:
Content selection saved. Describe the issue below:
Description:
Submit without GitHub
Submit in GitHub
arXiv is now an independent nonprofit!
Learn more
×
Back to arXiv
Why HTML?
Report Issue
Back to Abstract
Download PDF
Abstract
1 Introduction
2 Background
2.1 Mixture-of-Experts
2.2 Auxiliary Loss for Load Balance
Auxiliary Loss
The Dilemma Between Load Balance and Model Performance
3 Auxiliary-Loss-Free Load Balancing Strategy
Comparison with Other Load Balancing Methods.
4 Experiments
4.1 Experimental Setups
Model Architecture.
Training Settings
Baseline.
Metrics.
4.2 Main Results
4.3 Empirical Studies on Bias Update Algorithm
Update rate.
Update rule.
Multiplicative bias.
5 Discussion
5.1 Loss-Free Balancing Is Compatible with Expert Parallelism
5.2 Load Balancing and Future Token Leakage
6 Conclusion
References
A Model Architecture
B Training Settings
C Experiments with Softmax Gate
C.1 Comparison of Sigmoid Gate Baseline and Softmax Gate Baseline
C.2 Loss-Free Load Balancing with Softmax Gate
D Future Token Leakage in Expert Choice
D.1 Proof for Theoretical Leakage Amount
D.2 Experimental Evidence
License: arXiv.org perpetual non-exclusive license
arXiv:2408.15664v1 [cs.LG] 28 Aug 2024
## Auxiliary-Loss-Free Load Balancing
Strategy for Mixture-of-Experts
Lean Wang 1 1 footnotemark:
1
Affiliation: DeepSeek-AI
Affiliation: State Key Laboratory of Multimedia Information Processing,School of Computer Science, Peking University
Email: lean@pku.edu.cn
Huazuo Gao
Affiliation: DeepSeek-AI
Email: xusun@pku.edu.cn
Chenggang Zhao
Affiliation: DeepSeek-AI
Email: damai.dai@deepseek.com
Xu Sun 2 2 footnotemark:
2
Affiliation: State Key Laboratory of Multimedia Information Processing,School of Computer Science, Peking University
Damai Dai 2 2 footnotemark:
2
Affiliation: DeepSeek-AI
## 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., 2017 ) have led to successful attempts at scaling language models to substantial sizes ( Shao et al., 2024 ; DeepSeek-AI et al., 2024 ; Dai et al., 2024 ; Fedus et al., 2021 ; Lepikhin et al., 2020 ) , 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., 2017 ) or increased computational overhead ( Fedus et al., 2021 ; Lepikhin et al., 2020 ; Shazeer et al., 2017 ) .
In order to avoid imbalanced routing, existing methods ( Fedus et al., 2021 ; Lepikhin et al., 2020 ) 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 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.
## 2 Background
## 2.1 Mixture-of-Experts
Current dominant MoE architectures ( Lepikhin et al., 2020 ; Fedus et al., 2021 ; Dai et al., 2024 ) 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 𝐮 t \mathbf{u}_{t} denote the input of the t t -th token to an N N -expert MoE layer, the output 𝐡 t \mathbf{h}_{t} is computed as follows:
𝐡 t = 𝐮 t + ∑ i = 1 N g i , t ​ FFN i ⁡ ( 𝐮 t ) , \displaystyle\mathbf{h}_{t}=\mathbf{u}_{t}+\sum_{i=1}^{N}g_{i,t}\operatorname{FFN}_{i}\left(\mathbf{u}_{t}\right),
(1)
g i , t = { s i , t , s i , t ∈ Topk ⁡ ( { s j , t ∣ 1 ≤ j ≤ N } , K ) , 0 , otherwise, \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}
s i , t = G ⁡ ( 𝐮 t ​ 𝐞 i T ) , \displaystyle s_{i,t}=G\left(\mathbf{u}_{t}{}^{T}\mathbf{e}_{i}\right),
where G G is a nonlinear gating function and 𝐞 i \mathbf{e}_{i} is the centroid of the i 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., 2017 ) , 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., 2021 ; Lepikhin et al., 2020 ) is commonly employed to control load balance.
For a sequence of length T T , the auxiliary loss is defined as:
ℒ Balance \displaystyle\mathcal{L}_{\text{Balance }}
= α ​ ∑ i = 1 N f i ​ P i , \displaystyle=\alpha\sum_{i=1}^{N}f_{i}P_{i},
(2)
f i \displaystyle f_{i}
= N K ​ T ​ ∑ t = 1 T 𝟙 ​ ( Token ​ t ​ selects Expert ​ i ) , \displaystyle=\frac{N}{KT}\sum_{t=1}^{T}\mathbbm{1}(\text{ Token }t\text{ selects Expert }i),
P i \displaystyle P_{i}
= 1 T ​ ∑ t = 1 T s i , t , \displaystyle=\frac{1}{T}\sum_{t=1}^{T}s_{i,t},
where N N is the total number of experts, K K is the number of experts selected for each token, s i , t s_{i,t} is the routing score of Expert i i for Token t t , f i f_{i} represents the fraction of tokens routed to Expert i i , P i P_{i} denotes the average gating scores of Expert i i , and α \alpha is a hyper-parameter controlling the strength of the auxiliary loss.
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.
## 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 MaxVio global \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 s_{i,t} of each expert, and use the biased scores to determine the top-K selection:
g i , t = { s i , t , s i , t + b i ∈ Topk ⁡ ( { s j , t + b j ∣ 1 ≤ j ≤ N } , K ) , 0 , otherwise. \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 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 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 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 {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.
Algorithm 1 Adjusting the per-expert bias b i {b}_{i} during training
Input: MoE model θ \theta , training batch iterator B B , bias update rate u u .
1. Initialize b i = 0 b_{i}=0 for each expert;
for a batch { ( 𝐱 k , 𝐲 k ) } k \{(\mathbf{x}_{k},\mathbf{y}_{k})\}_{k} in B B do
2. Train MoE model θ \theta on the batch data { ( 𝐱 k , 𝐲 k ) } k \{(\mathbf{x}_{k},\mathbf{y}_{k})\}_{k} , with gating scores calculated according to Eq. ( 3 );
3. Count the number of assigned tokens c i c_{i} for each expert, and the average number c i ¯ \overline{c_{i}} ;
4. Calculate the load violation error e i = c i ¯ − c i e_{i}=\overline{c_{i}}-c_{i} ;
4. Update 𝐛 i \mathbf{b}_{i} by b i = b i + u ∗ sign ⁡ ( e i ) b_{i}=b_{i}+u*\mathrm{sign}(e_{i}) ;
end for
Output: trained model θ \theta , corresponding bias 𝐛 i \mathbf{b}_{i}
## 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., 2020 ; Fedus et al., 2021 ) and the Expert Choice (EC) ( Zhou et al., 2022 ) 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
Expert Load
Interference
Gradients
Future Token
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., 2024 ) architecture as the backbone since it outperforms conventional MoE architectures like GShard ( Lepikhin et al., 2020 ) significantly.
Compared with GShard ( Lepikhin et al., 2020 ) , 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 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 2
2
https://github.com/huggingface/tokenizers to train a byte pair encoding (BPE) ( Sennrich et al., 2015 ) 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, 2016 ) and multi-step learning rate scheduler ( Dai et al., 2024 ) 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:
MaxVio = max i ⁡ Load i − Load i ¯ Load i ¯ , \text{MaxVio}=\frac{\max_{i}\text{Load}_{i}-\overline{\text{Load}_{i}}}{\overline{\text{Load}_{i}}},
(4)
where Load i \text{Load}_{i} represents the number of tokens assigned to the i i -th expert, and Load i ¯ \overline{\text{Load}_{i}} denotes the expected expert load under perfect load balance.
MaxVio has two variants: MaxVio global \textbf{MaxVio}_{\textbf{global}} and MaxVio batch \textbf{MaxVio}_{\textbf{batch}} .
For MaxVio global \textbf{MaxVio}_{\textbf{global}} , we count Load i \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 MaxVio batch \textbf{MaxVio}_{\textbf{batch}} , we count Load i \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 MaxVio global \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 MaxVio batch \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
MaxVio global \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 3: Loss-Free Balancing maintains a better load balance throughout most of the training time. Here, MaxVio batch \text{MaxVio}_{\text{batch}} is averaged over 100 neighboring steps for visibility purposes.
## 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 u in Algorithm 1 controls the speed at which the expert bias { b i } i = 1 N \{{b}_{i}\}_{i=1}^{N} converges to the “suitable bias”. Figure 4 illustrates that an overly low update rate u = 0.0001 u=0.0001 may lead to slow convergence, while an unnecessarily high update rate u = 0.01 u=0.01 can cause undesirable fluctuations of the expert bias b i {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 u=0.001 , which shows good training balance and validation perplexity.
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.
## 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 ∗ sign ⁡ ( e i ) b_{i}=b_{i}+u*\mathrm{sign}(e_{i}) to b i = b i + u ∗ e i 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 sign \mathrm{sign} version.
Table 3: The variant b i = b i + u ∗ e i b_{i}=b_{i}+u*e_{i} slightly improves load balance but does not show improvement in model performance.
Method
Perplexity
MaxVio global \textbf{MaxVio}_{\textbf{global}}
b i = b i + u ∗ sign ​ ( e i ) b_{i}=b_{i}+u*\text{sign}(e_{i}) , u = 0.001 u=0.001
9.50
0.044
b i = b i + u ∗ e i b_{i}=b_{i}+u*e_{i} , u = 0.01 u=0.01
9.53
0.028
b i = b i + u ∗ e i b_{i}=b_{i}+u*e_{i} , u = 0.001 u=0.001
9.51
0.036
b i = b i + u ∗ e i b_{i}=b_{i}+u*e_{i} , u = 0.0001 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:
g i , t = { s i , t , s i , t ∗ b i ∈ Topk ⁡ ( { s j , t ∗ b j ∣ 1 ≤ j ≤ N } , K ) , 0 , otherwise, \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 \{{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
MaxVio global \textbf{MaxVio}_{\textbf{global}}
Addative Bias, u = 0.001 u=0.001
9.50
0.044
Multiplicative Bias, u = 0.01 u=0.01
9.52
0.041
Multiplicative Bias, u = 0.001 u=0.001
9.52
0.036
Multiplicative Bias, u = 0.0001 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., 2020 ) 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 MaxVio computation-batch \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 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.
## 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., 2022 ) 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 R (average activated experts per token K K divided by total expert number N N ) can leak more than K ​ log 2 ​ 1 − R R 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 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.
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.

=== BODY (tail) ===
