Maximum Score Routing For Mixture-of-Experts
Authors: Bowen Dong, Yilong Fan, Yutao Sun, Zhenyu Li, Tengyu Pan, zhou Xun, Jianyong Wang
Accepted at: acl25

Maximum Score Routing For Mixture-of-Experts

Bowen Dong1, Yilong Fan2, Yutao Sun1, Zhenyu Li1,
Tengyu Pan1, Xun Zhou3†, Jianyong Wang1†

1Tsinghua University, 2Tianjin University,
3Seed-Foundation-Model Team, ByteDance
dbw22@mails.tsinghua.edu.cn

Abstract

Routing networks in sparsely activated mixture-of-experts (MoE) dynamically allocate input tokens to top-k experts through differentiable sparse transformations, enabling scalable model capacity while preserving computational efficiency. Traditional MoE networks impose an expert capacity constraint to ensure GPU-friendly computation. However, this leads to token dropping when capacity is saturated and results in low hardware efficiency due to padding in underutilized experts. Removing the capacity constraint, in turn, compromises load balancing and computational efficiency. To address these issues, we propose Maximum Score Routing (MaxScore), a novel MoE routing paradigm that models routing as a minimum-cost maximum-flow problem and integrates a SoftTopk operator. MaxScore resolves the fundamental limitations of iterative rerouting and optimal transport formulations, achieving lower training losses and higher evaluation scores at equivalent FLOPs compared to both constrained and unconstrained baselines. Implementation details and experimental configurations can be obtained from https://github.com/dongbw18/MaxScore.git.

1 INTRODUCTION

The Mixture of Experts (MoE) paradigm has emerged as a compelling architectural strategy for scaling neural networks while maintaining computational efficiency. This approach dynamically combines multiple subsets of parameters (experts) by a learnable routing network, aiming to improve model capacity and computational efficiency. The routing network of sparsely activated MoE (Shazeer et al., 2017) dynamically allocates input tokens to top-k experts through differentiable sparse transformations, enabling conditional computation that scales model parameters without proportionally increasing FLOPs.

Softmax is conventionally employed to compute token-expert affinity coefficients in MoE routing networks, which promotes inter-expert competition. To mitigate winner-takes-all and preserve load balance, both hard constraints using expert capacity (Eigen et al., 2014), and soft constraints using auxiliary losses (Bengio et al., 2016), are incorporated into the routing network (Shazeer et al., 2017). GShard (Lepikhin et al., 2020) pioneers the integration of MoE with Transformer architectures (Vaswani et al., 2017), where expert capacity constraints enable GPU-friendly computation patterns. ExpertChoice (Zhou et al., 2022) directly enables experts to select tokens based on capacity constraints. However, token dropping occurs when inputs are routed to capacity-saturated experts, while padding operations in underutilized experts create hardware inefficiencies. Empirical analysis reveals that approaches such as expanding capacity (Hwang et al., 2023) or removing capacity constraints altogether (Gale et al., 2022; Muenninghoff et al., 2024) effectively eliminate token dropping, but inevitably introduce a trade-off between computational efficiency and load balancing performance. Efforts to prevent token dropping via refined routing strategies (Fedus et al., 2022; Clark et al., 2022) have not yielded performance improvements, highlighting unresolved challenges in dynamic resource allocation.

This work introduces Maximum Score Routing (MaxScore), a novel MoE routing paradigm that formulates token-expert routing as a minimum-cost maximum-flow problem (Waissi, 1994), integrated with a SoftTopk operator. To the best of our knowledge, this is the first successful integration of network flow modeling and SoftTopk in MoE routing.

MaxScore preserves GPU-compatible expert capacity constraints and achieves better load balanc-
ing. Under the same FLOPs, MaxScore exhibits lower training loss and higher evaluation scores compared to both constrained and unconstrained baselines. Ablation studies demonstrate the necessity of both network flow modeling and the SoftTopk operator, revealing fundamental limitations in the iterative rerouting mechanism of Fedus et al. (2022) and the optimal transport-based routing of Clark et al. (2022). The synergistic combination of two methodological enhancements yields superadditive performance gains, with empirical results demonstrating that their integrated efficacy surpasses the linear summation of individual improvements. Scaling experiments show that MaxScore delivers consistent performance improvements with larger activated parameter budgets, and achieves more gains when increasing the number of experts, compared with standard MoE approaches.

2 PRELIMINARIES

2.1 Top-k Sparsely Activated MoE
The top-k routing mechanism is a cornerstone of sparsely activated MoE architectures, enabling efficient scaling of model capacity while maintaining computational tractability. Originally popularized in language modeling (Shazeer et al., 2017), this paradigm dynamically routes each input token to a subset of \( k \) expert networks (where \( k \ll e \), for \( e \) total experts). Unlike dense models that activate all parameters per input, top-k routing induces conditional computation by selecting experts based on learned gating scores, typically computed via softmax over a trainable projection of input embeddings (Lepikhin et al., 2020).

For a given input \( x \), the output \( y \) of the MoE module can be written as follows:

\[
y = \sum_{i=1}^{E} R(x)_i \cdot E_i(x),
\]

\[
R(x) = \text{KeepTopk}(\text{Softmax}(x \cdot W_g)),
\]

where \( R(x) \) is the sparsely activated routing function, KeepTopk(\( \cdot \)) retains the top-k largest values while setting others to zero, \( W_g \) is the weight matrix of the routing function, \( E_i(x) \) is the output of the \( i \)-th expert network and the computation is performed only when \( R(x)_i > 0 \).

By leveraging sparse activation, MoE decouples total capacity \( \mathcal{O}(e) \) from per-step computational cost, activating only \( \mathcal{O}(k) \) parameters during both training and inference.

2.2 Operators in Top-k MoE Routing
Routings in MoE commonly use Softmax(\( \cdot \)) to calculate the token-expert affinity coefficients, which encourages competition between experts. However, Softmax(\( \cdot \)) serves as a smooth approximation to the one-hot Argmax(\( \cdot \)) function, which can lead to inefficiencies in top-\( k \) routing, as the top-1 expert often receives a disproportionately large affinity score compared to the remaining \( k-1 \) experts.

Alternative routing operators have also been investigated. DeepSeek-AI et al. (2024b) replaces Softmax(\( \cdot \)) with Sigmoid(\( \cdot \)) to align with its auxiliary-loss-free load balancing strategy, while ReMoE (Wang et al., 2025) explores the feasibility of using ReLU(\( \cdot \)) for routing decisions.

We define SoftTopk(\( \cdot \)) as a smooth approximation to ArgTopk(\( \cdot \)), which represents the top-k selection in a one-hot form, formally given by:

\[
\text{ArgTopk}(\mathbf{a})_i = \begin{cases}
1, a_i \in \text{Topk}(\mathbf{a}) \\
0, \text{otherwise},
\end{cases}
\]

where \( \mathbf{a} = (a_1, a_2, ..., a_e) \) represents the affinity coefficients between the token and e experts.

Martins and Astudillo (2016) and Peters et al. (2019) proposed Sparsemax(\( \cdot \)) and Entmax(\( \cdot \)) as differentiable approximations for top-k probability truncation. Su (2024) further introduced a broader family of SoftTopk(\( \cdot \)) operators. However, their integration into MoE routing has not been investigated, leaving a promising direction underexplored.

2.3 Expert Capacity Constrained
To counteract the winner-takes-all phenomenon and maintain load balancing in the routing network, traditional routing architectures integrate dual constraint mechanisms: (i) hard limits through expert capacity (Eigen et al., 2014), and (ii) soft regularization via differentiable auxiliary losses (Bengio et al., 2016; Shazeer et al., 2017; Zoph et al., 2022).

GShard (Lepikhin et al., 2020) strategically harmonizes capacity-constrained MoE design with Transformer architectures (Vaswani et al., 2017). For a batch of \( n \) tokens, GShard fixes per-expert capacity with \( c = \frac{k \cdot n}{e} \) to enable parallel-friendly computation patterns. This routing mechanism, however, poses optimization challenges due to imbalanced expert utilization. While underloaded experts incur computational overhead through padding (mathematically sound but hardware-inefficient), overloaded experts lead to token dropping. Increasing expert capacity \( c' = c_f * \frac{k \cdot n}{e} \) by a
Figure 1: Different top-2 routing paradigms for 3 experts and 6 tokens. (a) sets capacity-factor \( c_f = 1.0 \), and token dropping occurs; (b) sets capacity-factor \( c_f = 1.5 \), there is no more token dropping, but more computation is wasted; (c) uses iterative rerouting mechanism, the dropped token is reassigned to expert with remaining capacity.

capacity-factor \( c_f \) can alleviate token dropping. Tutel (Hwang et al., 2023) uses a highly scalable stack design and sets the \( c_f \) dynamically, but it would lead to additional computational costs and reduced load balancing. Figure 1(a) and 1(b) shows the trade-off between token dropping and additional computation by increasing expert capacity. Figure 2(a) shows the token dropping proportion in the MoE routing of each layer in a GShard model with \( e = 16 \) and \( k = 2 \), and approximately 35% of tokens routed to the second experts experience dropping.

ExpertChoice (Zhou et al., 2022) inverts the conventional routing paradigm by allowing experts to select their top-c tokens, thereby achieving optimal load balancing. However, this strategy allows each token to be assigned to an arbitrary number of experts, including zero, which exacerbates token dropping. More importantly, it introduces a data leakage issue: determining whether a token belongs to the top-c set of a given expert requires comparisons not only with preceding tokens but also with subsequent ones, thereby violating the causal structure required by autoregressive models.

Another class of approaches, referred to as Drop-Less MoE, eliminates capacity constraints entirely to prevent token dropping. Those methods allocate an indefinite number of tokens to experts via direct indexing (e.g., DeepSeekMoE (Dai et al., 2024; DeepSeek-AI et al., 2024a,b), OLMoE (Muenninghoff et al., 2024; Gale et al., 2022)).

Switch Transformers (Fedus et al., 2022) explored an iterative rerouting mechanism for dropped tokens as shown in Figure 1(c): in the first stage, tokens are assigned to experts using the top-k strategy; in the second stage, any dropped tokens are greedily reassigned to the highest-affinity expert among those with remaining capacity. However, empirical results show that this approach does not lead to improvements in model quality.

SBASE (Clark et al., 2022) formulates MoE routing as an optimal transport problem: \( c = (c_1, c_2, ..., c_e) \) denotes the capacity of each expert, and \( \mathbf{k} = (k_1, k_2, ..., k_n) \) specifies the number of experts each token should be assigned to. The matrix \( \mathbf{A} \in \mathbb{R}^{n \times e} \) represents token-expert affinity coefficients. The feasible solution space is defined as

\[
U(c, \mathbf{k}) = \{ \mathbf{P} \in \mathbb{R}_{\geq 0}^{n \times e} | \mathbf{P}^T \mathbf{1}_n = c, \mathbf{P} \mathbf{1}_e = \mathbf{k} \},
\]

and the optimization objective is

\[
d_{\mathbf{A}}(c, \mathbf{k}) = \max_{\mathbf{P} \in U(c, \mathbf{k})} \sum_{ij} \mathbf{P}_{ij} \mathbf{A}_{ij}.
\]

To efficiently approximate the solution, SBASE employs the parallelizable Sinkhorn algorithm (Cuturi, 2013). Nonetheless, this formulation primarily contributes to improved training stability, offering limited gains beyond this benefit.

![Bar charts showing token dropping and expert affinities for different routing paradigms](page_1012_186_377_232.png)

Figure 2: The proportion of tokens not dropped and the mean token-expert affinities in top-2 routing are analyzed separately. The data is derived from the GShard MoE with \( e = 16 \) after training on 65 billion tokens. (a) shows that tokens assigned to the top-1 experts are rarely dropped, whereas approximately 35% of tokens routed to the second experts experience dropping. (b) illustrates that the top-1 token-expert affinities are typically much higher than those of other experts.
transport formulation (Sinkhorn) fail to improve model quality, and propose Maximum Score Routing (MaxScore), a novel mixture-of-experts routing strategy that integrates network flow modeling and a differentiable SoftTopk(\cdot) operator.

3.1 Limitations of Iter and Sinkhorn

Softmax operator. Both the iterative rerouting mechanism and the optimal transport formulation aim to achieve a globally improved allocation by replacing locally optimal assignment strategies. However, as discussed in Section 2.2, using the conventional Softmax(\cdot) to compute token-expert affinity scores results in the top-1 affinity being significantly higher than those of other token-expert pairs. We statistically analyze the probability distribution in a top-2 GShard MoE, as shown in Figure 2(b), where the top-1 token-expert affinities markedly exceeds that of the second-ranked expert. For example, if a token’s top-2 affinities are 0.8 and 0.05 respectively, then when the first expert is saturated, substituting with any expert outside the top-2 (with affinity below 0.05) yields no meaningful benefit; similarly, if the second expert is saturated, replacing it has negligible impact on the model’s gradient.

Limitation of optimal transport formulation. Modeling MoE routing using Equations (4) and (5) has inherent limitations: in MoE routing strategies, the actual gain of a token-expert pair appearing multiple times is equivalent to that of a single occurrence. This constraint cannot be enforced in the optimal transport formulation. As illustrated in Figure 3, high-probability token-expert pairs may be matched repeatedly, causing redundant reward accumulation and effectively degenerating to a top-1 routing scheme, which results in wasted computational resources.

<table>
  <tr>
    <th>Name</th>
    <th>Expression</th>
  </tr>
  <tr>
    <td>Softmax(x)</td>
    <td>\( y = e^x / \sum_j^N e^{x_j} \)</td>
  </tr>
  <tr>
    <td>Sigmoid(x)</td>
    <td>\( y = 1/(1 + e^{-x}) \)</td>
  </tr>
  <tr>
    <td>SoftKmax(x)<sup>(k)</sup></td>
    <td>\( y^{(k)} = y^{(k-1)} + \text{Softmax}(g^{(k-1)}) \)</td>
  </tr>
  <tr>
    <td></td>
    <td>\( g^{(k-1)} = (1 - y^{(k-1)}) \otimes x \)</td>
  </tr>
  <tr>
    <td>IterTopk(x)<sup>(k)</sup></td>
    <td>\( y^{(k)} = y^{(k-1)} + g(x; 1 - y^{(k-1)}) \)</td>
  </tr>
  <tr>
    <td></td>
    <td>\( g(x; w) = w \cdot e^x / \sum_j^N w_j \cdot e^{x_j} \)</td>
  </tr>
  <tr>
    <td></td>
    <td>\( y^{(k)} = e^{g^{(k)}} - z^{(k)} \)</td>
  </tr>
  <tr>
    <td>GradTopk(x)<sup>(k)</sup></td>
    <td>\( g^{(k)} = x + \log(e^{g^{(k-1)}} - e^{g^{(k-1)}}) \)</td>
  </tr>
  <tr>
    <td></td>
    <td>\( z^{(k)} = \log(\sum_j^N e^{g^{(k)}}) - \log k \)</td>
  </tr>
</table>

Table 1: Operators can be used for MoE routing. SoftKmax, IterTopk and GradTopk are mentioned in Su (2024).

complexity, we did not achieve better results than Softmax(\cdot). We propose a simple but highly effective SoftTopk(\cdot) operator for MoE routing:

\[
\text{SoftTopk(a)}^{(k)} = \text{SoftTopk(a)}^{(k-1)} + \text{SE(a)},
\]
\[
\text{SE(a)}_i = \begin{cases}
0, a_i \in \text{Topk(a)} \\
t \cdot \text{Softmax(a)}_i, \text{otherwise},
\end{cases}
\]

where \( t \) is a constant that gradually decays from the initialization value \( t_0 \) to 0.

Network flow modeling. To better capture the characteristics of MoE routing, Equations (4) and (5) are revised as follows:

\[
U'(c, k) = \{ P \in \mathbb{F}_2^{n \times e} | P^T 1_n = c, P 1_e = k \},
\]
\[
d'_{A}(c, k) = \max_{P \in U'(c, k)} \sum_{ij} P_{ij} A_{ij},
\]

where \( \mathbb{F}_2 \) denotes the finite field of \( \{0, 1\} \) equipped with addition and multiplication operations. To address this problem, MoE routing can be formulated as a minimum-cost maximum-flow problem as shown in Figure 4. We model tokens and experts as nodes in a flow network graph. Edges from the super source to tokens have capacities representing that each token must be assigned to k experts, while edges from experts to the super sink enforce capacity constraints of c per expert. These edges carry zero cost. Edges between tokens and experts have unit capacity, allowing at most one match per token-expert pair, with costs defined as the negation

![Limitation of optimal transport formulation. The fifth token and the second expert matched twice.](page_109_1047_627_181.png)

Figure 3: Limitation of optimal transport formulation. The fifth token and the second expert matched twice.

3.2 Maximum Score Routing

SoftTopk operator. We first tried different operators as shown in Table 1, but due to the potential damage caused by the increased computational
of their affinity coefficients. A detailed summary of the graph edge properties is provided in Table 2.

![The minimum-cost maximum-flow modeling for MoE routing.](page_184_370_563_246.png)

Figure 4: The minimum-cost maximum-flow modeling for MoE routing.

<table>
  <tr>
    <th>From</th>
    <th>To</th>
    <th>Capacity</th>
    <th>Cost</th>
    <th>Count</th>
  </tr>
  <tr>
    <td>Source</td>
    <td>Token<sub>i</sub></td>
    <td>k<sub>i</sub></td>
    <td>0</td>
    <td>n</td>
  </tr>
  <tr>
    <td>Expert<sub>j</sub></td>
    <td>Sink</td>
    <td>c<sub>j</sub></td>
    <td>0</td>
    <td>e</td>
  </tr>
  <tr>
    <td>Token<sub>i</sub></td>
    <td>Expert<sub>j</sub></td>
    <td>1</td>
    <td>-A<sub>ij</sub></td>
    <td>n * e</td>
  </tr>
</table>

Table 2: Edges in the graph of Figure 4. Source is the super source, Sink is the super sink, \( A_{ij} \) represents the affinity coefficient between Token<sub>i</sub> and Expert<sub>j</sub>.

Algorithm complexity optimization. TA commonly used and effective approach to solving the minimum-cost maximum-flow problem is the Shortest Path Faster Algorithm (SPFA) (Bellman, 1958; Ford, 1956), which iteratively searches for the lowest-cost augmenting path until no such path remains. However, this method is computationally expensive and inherently sequential, limiting its parallelizability. In top-2 MoE routing, given that the token drop rate in top-1 routing is relatively low (approximately 0 as shown in Figure 2) and that the Sinkhorn algorithm corresponds to the minimum-cost maximum-flow formulation under top-1 routing, we propose a two-stage strategy: first allocate tokens using top-1 routing, followed by applying the Sinkhorn algorithm to handle the residual routing problem. The complete algorithm process is shown in Algorithm 1. For top-k MoE routing with \( k > 2 \), a trade-off needs to be made between quality (SPFA) and speed (Iter).

4 EVALUATION

4.1 Experimental Setup

Model Architecture. We conduct our experiments using the Llama architecture (Touvron et al., 2023a,b; Grattafiori et al., 2024), incorporating grouped query attention (GQA) (Ainslie et al., 2023), SwiGLU activation function (Shazeer, 2020), RoPE position embedding (Su et al., 2023), and RMSNorm (Zhang and Sennrich, 2019). Our sparsely activated models are constructed by substituting the MLP layers of the dense baseline with MoE layers. We explore three different backbone sizes, as detailed in Table 8.

Baselines. We compared the dense model, GShard MoE (Lepikhin et al., 2020) and GShard-I MoE, the variant with iterative routing strategy (Fedus et al., 2022), SBASE MoE (Clark et al., 2022), ExpertChoice MoE (Zhou et al., 2022), DropLess MoE (Gale et al., 2022), DeepSeek-V2 MoE (Dai et al., 2024; DeepSeek-AI et al., 2024a) along with our proposed MaxScore MoE and MaxScore-I MoE, which replaces network flow modeling with the iterative rerouting mechanism. All MoEs except DeepSeek use the base configuration with \( k = 2 \) and \( e = 16 \), while DeepSeek MoE employs fine-grained experts with \( k = 6 \) and \( e = 64 \) and a double-sized shared expert.

Load Balance Loss. All MoE models employ the same auxiliary loss function, defined as

\[
\mathcal{L}_{aux} = \lambda \cdot \frac{1}{e} \sum_{j=1}^{e} \left( \frac{1}{n} \sum_{i=1}^{n} \mathbf{A}_{i,j} \right) \left( \frac{1}{n} \sum_{i=1}^{N} \mathbf{P}_{i,j} \right),
\]

where the \( \mathbf{A}_{i,j} \) and \( \mathbf{P}_{i,j} \) correspond to the terms defined in Equations (7) and (8).

Training Settings. We adopt the tokenizer from

Algorithm 1 Maximum Score Routing For Top-2 Mixture-of-Experts

Input: Weight matrix \( W_g \) in the routing function, the number of experts \( e \), temperature \( t \gets t_0 \), a batch of \( n \) tokens \( \{ x_i \} \)
1: Initialization expert capacity: \( c_j \gets 2 * n / e \)
2: Calculate the token-expert affinity coefficients: \( a_{i,j} \gets \mathrm{SoftTopk}(x_i \cdot W_g)_j \)
3: Update temperature: \( t \)
4: Calculate the mask matrix of top-1: \( \mathrm{mask}_{i,j} \gets \mathrm{onehot}(\mathrm{Argmax}(a_i), e)_j \)
5: Remove top-1: \( a_{i,j} \gets a_{i,j} - \mathrm{mask}_{i,j} \)
6: Update expert capacity: \( c_j \gets \max(0, c_j - \sum_i \mathrm{mask}_{i,j}) \)
7: Set \( k_i: k_i \gets 1 \)
8: The feasible solution space: \( U'(c, k) = \{ \mathbf{P} \in \mathbb{R}_+^{n \times e} | \mathbf{P}^T \mathbf{1}_n = c, \mathbf{P} \mathbf{1}_e = k \} \),
9: Use Sinkhorn for an approximate solution: \( d'_{\mathbf{A}}(c, k) = \max_{\mathbf{P} \in U'(c, k)} \sum_{ij} \mathbf{P}_{ij} \mathbf{A}_{ij} \)
Output: \( \{ \mathbf{P}_{ij} \} \)
<table>
  <tr>
    <th>Model</th>
    <th>ARC challenge</th>
    <th>ARC easy</th>
    <th>BoolQ</th>
    <th>HellaSwag</th>
    <th>LAMBADA</th>
    <th>PIQA</th>
    <th>RACE</th>
    <th>SciQ</th>
    <th>Record</th>
    <th>OBQA</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td>Dense</td>
    <td>18.69</td>
    <td>40.19</td>
    <td>57.06</td>
    <td>28.91</td>
    <td>16.28</td>
    <td>63.71</td>
    <td>25.65</td>
    <td>64.2</td>
    <td>56.05</td>
    <td>15.0</td>
    <td>38.57</td>
  </tr>
  <tr>
    <td>GShard</td>
    <td>18.86</td>
    <td><b>44.49</b></td>
    <td>61.90</td>
    <td>31.74</td>
    <td>21.54</td>
    <td>66.38</td>
    <td><b>28.52</b></td>
    <td>69.4</td>
    <td>62.08</td>
    <td>16.2</td>
    <td>42.11</td>
  </tr>
  <tr>
    <td>GShard-I</td>
    <td>19.80</td>
    <td>44.36</td>
    <td>59.94</td>
    <td>32.54</td>
    <td>21.52</td>
    <td>67.03</td>
    <td>28.23</td>
    <td>68.7</td>
    <td>62.84</td>
    <td>16.0</td>
    <td>42.10</td>
  </tr>
  <tr>
    <td>SBASE</td>
    <td>18.34</td>
    <td>43.73</td>
    <td>57.61</td>
    <td>30.96</td>
    <td>19.70</td>
    <td>65.18</td>
    <td>27.37</td>
    <td>68.3</td>
    <td>60.06</td>
    <td>16.2</td>
    <td>40.75</td>
  </tr>
  <tr>
    <td>ExpertChoice</td>
    <td>19.37</td>
    <td>42.00</td>
    <td>61.74</td>
    <td>32.10</td>
    <td>21.19</td>
    <td>66.16</td>
    <td>27.18</td>
    <td>68.4</td>
    <td>62.26</td>
    <td>17.6</td>
    <td>41.80</td>
  </tr>
  <tr>
    <td>DropLess</td>
    <td>19.28</td>
    <td>44.07</td>
    <td>61.16</td>
    <td>32.03</td>
    <td>21.35</td>
    <td>67.14</td>
    <td>27.08</td>
    <td>67.9</td>
    <td>61.55</td>
    <td>16.0</td>
    <td>41.76</td>
  </tr>
  <tr>
    <td>DeepSeek</td>
    <td>19.88</td>
    <td>44.28</td>
    <td>60.55</td>
    <td>32.23</td>
    <td>21.93</td>
    <td>66.97</td>
    <td>27.94</td>
    <td>70.9</td>
    <td>62.57</td>
    <td>17.6</td>
    <td>42.49</td>
  </tr>
  <tr>
    <td>MaxScore-I</td>
    <td><b>20.90</b></td>
    <td>43.22</td>
    <td>61.71</td>
    <td>32.51</td>
    <td>21.66</td>
    <td><b>67.41</b></td>
    <td>28.42</td>
    <td>69.9</td>
    <td>63.61</td>
    <td><b>18.4</b></td>
    <td>42.77</td>
  </tr>
  <tr>
    <td>MaxScore</td>
    <td>20.73</td>
    <td><b>44.49</b></td>
    <td><b>62.23</b></td>
    <td><b>32.85</b></td>
    <td><b>23.27</b></td>
    <td><b>67.41</b></td>
    <td><b>28.52</b></td>
    <td><b>72.5</b></td>
    <td><b>64.00</b></td>
    <td><b>18.4</b></td>
    <td><b>43.44</b></td>
  </tr>
</table>

Table 3: Results for the base-sized models.

LLama (Touvron et al., 2023a,b; Grattafiori et al., 2024) and set the context length to 512. The batch size is 688, which is the largest setting that allows all baseline models to be trained on 8 NVIDIA A800 GPUs (this constraint arises primarily from the DeepSeek, as shown in Table 10). We can train all baselines with 8 NVIDIA A800 GPUs. All models are trained for 180k steps (approximately 65B tokens) on C4 dataset (Raffel et al., 2019). This exceeds the compute-optimal dataset size identified by Krajewski et al. (2024), ensuring convergence. For training, we leverage the HuggingFace Trainer (Wolf et al., 2020) integrated with DeepSpeed optimizations, including Zero Redundancy Optimizer (ZeRO) (Rajbhandari et al., 2020) and activation checkpointing (Chen et al., 2016), and we employ bfloat16 for numerical precision and efficiency. We adopt AdamW (Loshchilov and Hutter, 2019) as the optimizer with weight decay \( wd \), adam betas \( (\beta_1, \beta_2) \) and adam epsilon \( \epsilon \). The learning rate is set to be \( lr \) following a WSD scheduler (Hu et al., 2024) with a warmup for 2k steps and decay over the last 6k steps.

Hyperparameters. We perform grid searchs over learning rate \( lr \), weight decay \( wd \), adam betas \( (\beta_1, \beta_2) \), and adam epsilon \( \epsilon \) on the GShard baseline, and apply the selected hyperparameters uniformly across all other baselines, as summarized in Table 5. For the scaling factor \( \lambda \) of the auxiliary loss in Equation (9), we perform a grid search over the set \( \{10^{-1}, 10^{-2}, 10^{-3}, 10^{-4}\} \) for each baseline. The final selected values are \( 10^{-3} \) for DeepSeek and \( 10^{-2} \) for all other baselines.

Evaluation Settings. We leverage the open source lm-evaluation-harness (Gao et al., 2024) for standardized evaluation on various types of tasks: ARC challenge, ARC easy (Clark et al., 2018), BoolQ (Clark et al., 2019), HellaSwag (Zellers et al., 2019), LAMBADA (Paperno et al., 2016), PIQA (Bisk et al., 2019), RACE (Lai et al., 2017), SciQ (Welbl et al., 2017), Record (Zhang et al., 2018) and OpenBookQA (OBQA) (Mihaylov et al., 2018).

![Training loss curve.](page_1012_1042_410_312.png)

Figure 5: Training loss curve.

ARC challenge, ARC easy (Clark et al., 2018), BoolQ (Clark et al., 2019), HellaSwag (Zellers et al., 2019), LAMBADA (Paperno et al., 2016), PIQA (Bisk et al., 2019), RACE (Lai et al., 2017), SciQ (Welbl et al., 2017), Record (Zhang et al., 2018) and OpenBookQA (OBQA) (Mihaylov et al., 2018).

4.2 Main Results

Figure 5 presents the training loss curves for all evaluated base-sized models, and Table 3 summarizes the evaluation results of models after training on about 65B tokens.

Our proposed MaxScore and MaxScore-I consistently achieve lower training loss compared to all baseline methods throughout the training process and outperform existing baselines on the evaluation datasets. Notably, MaxScore attains the lowest final training loss of approximately 2.62, indicating more effective optimization and improved convergence behavior, and achieves the highest average accuracy of 43.44%, surpassing the best baseline
<table>
  <tr>
    <th>Model</th>
    <th>ARC challenge</th>
    <th>ARC easy</th>
    <th>BoolQ</th>
    <th>HellaSwag</th>
    <th>LAMBADA</th>
    <th>PIQA</th>
    <th>RACE</th>
    <th>SciQ</th>
    <th>Record</th>
    <th>OBQA</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td>GShard</td>
    <td>18.86</td>
    <td><b>44.49</b></td>
    <td>61.90</td>
    <td>31.74</td>
    <td>21.54</td>
    <td>66.38</td>
    <td><b>28.52</b></td>
    <td>69.4</td>
    <td>62.08</td>
    <td>16.2</td>
    <td>42.11</td>
  </tr>
  <tr>
    <td>GShard-I</td>
    <td>19.80</td>
    <td>44.36</td>
    <td>59.94</td>
    <td>32.54</td>
    <td>21.52</td>
    <td>67.03</td>
    <td>28.23</td>
    <td>68.7</td>
    <td>62.84</td>
    <td>16.0</td>
    <td>42.10</td>
  </tr>
  <tr>
    <td>GShard-M</td>
    <td>20.14</td>
    <td>43.74</td>
    <td>59.38</td>
    <td>32.27</td>
    <td>22.30</td>
    <td>66.63</td>
    <td>27.61</td>
    <td>68.7</td>
    <td>62.59</td>
    <td>18.2</td>
    <td>42.16</td>
  </tr>
  <tr>
    <td>GShard-S</td>
    <td>20.52</td>
    <td>44.30</td>
    <td>59.13</td>
    <td>32.34</td>
    <td>22.54</td>
    <td>66.74</td>
    <td>28.19</td>
    <td>69.4</td>
    <td>63.94</td>
    <td><b>18.4</b></td>
    <td>42.55</td>
  </tr>
  <tr>
    <td>GShard-SI (MaxScore-I)</td>
    <td><b>20.90</b></td>
    <td>43.22</td>
    <td>61.71</td>
    <td>32.51</td>
    <td>21.66</td>
    <td><b>67.41</b></td>
    <td>28.42</td>
    <td>69.9</td>
    <td>63.61</td>
    <td><b>18.4</b></td>
    <td>42.77</td>
  </tr>
  <tr>
    <td>GShard-SM (MaxScore)</td>
    <td>20.73</td>
    <td><b>44.49</b></td>
    <td><b>62.23</b></td>
    <td><b>32.85</b></td>
    <td><b>23.27</b></td>
    <td><b>67.41</b></td>
    <td><b>28.52</b></td>
    <td><b>72.5</b></td>
    <td><b>64.00</b></td>
    <td><b>18.4</b></td>
    <td><b>43.44</b></td>
  </tr>
</table>

Table 4: Ablation study results. We validate the contributions of the SoftTopk(\( \cdot \)) Operator (\( S \)), the Minimum-cost Maximum Flow Modeling (M), and the Iterative Routing Strategy (I).

(DeepSeek) by approximately 0.95%. It also attains state-of-the-art performance on almost all individual tasks. The iterative variant MaxScore-I demonstrates competitive results, particularly excelling on ARC challenge and PIQA.

These findings validate the superiority of our routing mechanisms in integrating the SoftTopk(\( \cdot \)) operator and the minimum cost maximum flow modeling in improving MoE routing quality.

<table>
  <tr>
    <th>Name</th>
    <th>Gird Search</th>
    <th>Result</th>
  </tr>
  <tr>
    <td>\( lr \)</td>
    <td>\( \{1,3\} * \{10^{-4}, 10^{-5}, 10^{-6}\} \)</td>
    <td>\( 3 * 10^{-5} \)</td>
  </tr>
  <tr>
    <td>\( wd \)</td>
    <td>\( \{0, 1, 2, 3, 4\} * 0.05 \)</td>
    <td>0.1</td>
  </tr>
  <tr>
    <td>\( (\beta_1, \beta_2) \)</td>
    <td>\( (0.9, \{0.999, 0.99, 0.95, 0.9\}) \)</td>
    <td>(0.9, 0.95)</td>
  </tr>
  <tr>
    <td>\( \epsilon \)</td>
    <td>\( \{10^{-5}, 10^{-6}, 10^{-7}, 10^{-8}\} \)</td>
    <td>\( 10^{-6} \)</td>
  </tr>
</table>

Table 5: Gird search and results for hyperparameters.

4.3 Ablation Evaluation

Table 4 presents the ablation study results, validating the individual contributions of the SoftTopk(\( \cdot \)) operator (\( S \)), the minimum-cost maximum flow modeling (M), and the iterative routing strategy (I). The variants GShard-S, GShard-M, and GShard-I correspond to incorporating SoftTopk, network flow modeling, and iterative routing respectively, while GShard-SI (MaxScore-I) and GShard-SM (MaxScore) combine these components.

GShard exhibits negligible improvements when employing either network flow modeling or the iterative strategy alone, consistent with observations reported in SwitchTransformer. However, incorporating the SoftTopk(\( \cdot \)) operator individually yields noticeable gains. Furthermore, combining the iterative strategy or network flow modeling with the SoftTopk(\( \cdot \)) operator results in substantial performance improvements. This demonstrates the necessity of the SoftTopk(\( \cdot \)) operator, revealing fundamental limitations in the iterative rerouting mechanism of Fedus et al. (2022) and the optimal transport-based routing of Clark et al. (2022).

By comparing Figure 2 and Figure 6, we observe that network flow modeling effectively eliminates token dropping, and the SoftTopk(\( \cdot \)) operator significantly improves the distribution of token-expert affinities.

Our full model, GShard-SM (MaxScore), consistently achieves the best average performance of 43.44%, outperforming all ablated variants. The synergistic combination of two methodological enhancements yields superadditive performance gains, with empirical results demonstrating that their integrated efficacy surpasses the linear summation of individual improvements.

![Bar charts showing proportion of not dropping and mean token-expert affinities in top-2 routing](page_1012_1012_388_234.png)

Figure 6: The proportion of not dropping and the mean token-expert affinities in top-2 routing are analyzed separately. The data is derived from our MaxScore MoE with \( e = 16 \) after training on 65 billion tokens.

4.4 Scalability

We perform scaling experiments along two dimensions: model size and sparsity. Detailed configurations are provided in Table 8 and Table 9.

As shown in Figure 7 and Tables 6 and 7, our MaxScore MoE consistently achieves a more significant reduction in training loss and superior evaluation performance compared to traditional MoE baselines such as GShard and DropLess across varying scales. In contrast, DropLess MoE suffers from increased expert load imbalance as sparsity increases, adversely affecting its scalability and overall performance. These results underscore
<table>
  <tr>
    <th>Size</th>
    <th>Model</th>
    <th>ARC challenge</th>
    <th>ARC easy</th>
    <th>BoolQ</th>
    <th>Hella-Swag</th>
    <th>LAM-BADA</th>
    <th>PIQA</th>
    <th>RACE</th>
    <th>SciQ</th>
    <th>Record</th>
    <th>OBQA</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td rowspan="3">Base</td>
    <td>GShard</td>
    <td>18.86</td>
    <td><b>44.49</b></td>
    <td>61.90</td>
    <td>31.74</td>
    <td>21.54</td>
    <td>66.38</td>
    <td><b>28.52</b></td>
    <td>69.4</td>
    <td>62.08</td>
    <td>16.2</td>
    <td>42.11</td>
  </tr>
  <tr>
    <td>DropLess</td>
    <td>19.28</td>
    <td>44.07</td>
    <td>61.16</td>
    <td>32.03</td>
    <td>21.35</td>
    <td>67.14</td>
    <td>27.08</td>
    <td>67.9</td>
    <td>61.55</td>
    <td>16.0</td>
    <td>41.76</td>
  </tr>
  <tr>
    <td>MaxScore</td>
    <td><b>20.73</b></td>
    <td><b>44.49</b></td>
    <td><b>62.23</b></td>
    <td><b>32.85</b></td>
    <td><b>23.27</b></td>
    <td><b>67.41</b></td>
    <td><b>28.52</b></td>
    <td><b>72.5</b></td>
    <td><b>64.00</b></td>
    <td><b>18.4</b></td>
    <td><b>43.44</b></td>
  </tr>
  <tr>
    <td rowspan="3">Large</td>
    <td>GShard</td>
    <td>19.88</td>
    <td>45.58</td>
    <td>62.16</td>
    <td>33.34</td>
    <td>23.69</td>
    <td>67.74</td>
    <td>29.28</td>
    <td>70.0</td>
    <td>64.99</td>
    <td>19.2</td>
    <td>43.59</td>
  </tr>
  <tr>
    <td>DropLess</td>
    <td>20.05</td>
    <td>45.24</td>
    <td>61.19</td>
    <td>33.97</td>
    <td>23.60</td>
    <td>67.63</td>
    <td>27.66</td>
    <td>69.7</td>
    <td>63.95</td>
    <td>17.2</td>
    <td>43.02</td>
  </tr>
  <tr>
    <td>MaxScore</td>
    <td><b>20.90</b></td>
    <td><b>45.92</b></td>
    <td><b>62.39</b></td>
    <td><b>34.00</b></td>
    <td><b>24.96</b></td>
    <td><b>68.28</b></td>
    <td><b>29.67</b></td>
    <td><b>74.3</b></td>
    <td><b>66.12</b></td>
    <td><b>19.8</b></td>
    <td><b>44.63</b></td>
  </tr>
  <tr>
    <td rowspan="3">XL</td>
    <td>GShard</td>
    <td>20.05</td>
    <td>46.68</td>
    <td>63.09</td>
    <td>35.14</td>
    <td>25.05</td>
    <td>69.31</td>
    <td>29.19</td>
    <td>72.7</td>
    <td>67.50</td>
    <td>20.0</td>
    <td>44.87</td>
  </tr>
  <tr>
    <td>DropLess</td>
    <td>20.14</td>
    <td>46.60</td>
    <td>61.69</td>
    <td>35.11</td>
    <td>24.34</td>
    <td>68.34</td>
    <td>29.19</td>
    <td>72.4</td>
    <td>67.55</td>
    <td>20.2</td>
    <td>44.56</td>
  </tr>
  <tr>
    <td>MaxScore</td>
    <td><b>21.22</b></td>
    <td><b>47.60</b></td>
    <td><b>63.60</b></td>
    <td><b>35.57</b></td>
    <td><b>25.93</b></td>
    <td><b>69.95</b></td>
    <td><b>29.90</b></td>
    <td><b>75.2</b></td>
    <td><b>67.90</b></td>
    <td><b>21.6</b></td>
    <td><b>45.85</b></td>
  </tr>
</table>

Table 6: Results of scaling in model size.

<table>
  <tr>
    <th>Sparsity</th>
    <th>Model</th>
    <th>ARC challenge</th>
    <th>ARC easy</th>
    <th>BoolQ</th>
    <th>Hella-Swag</th>
    <th>LAM-BADA</th>
    <th>PIQA</th>
    <th>RACE</th>
    <th>SciQ</th>
    <th>Record</th>
    <th>OBQA</th>
    <th>Avg.</th>
  </tr>
  <tr>
    <td rowspan="3">2 : 16</td>
    <td>GShard</td>
    <td>18.86</td>
    <td><b>44.49</b></td>
    <td>61.90</td>
    <td>31.74</td>
    <td>21.54</td>
    <td>66.38</td>
    <td><b>28.52</b></td>
    <td>69.4</td>
    <td>62.08</td>
    <td>16.2</td>
    <td>42.11</td>
  </tr>
  <tr>
    <td>DropLess</td>
    <td>19.28</td>
    <td>44.07</td>
    <td>61.16</td>
    <td>32.03</td>
    <td>21.35</td>
    <td>67.14</td>
    <td>27.08</td>
    <td>67.9</td>
    <td>61.55</td>
    <td>16.0</td>
    <td>41.76</td>
  </tr>
  <tr>
    <td>MaxScore</td>
    <td><b>20.73</b></td>
    <td><b>44.49</b></td>
    <td><b>62.23</b></td>
    <td><b>32.85</b></td>
    <td><b>23.27</b></td>
    <td><b>67.41</b></td>
    <td><b>28.52</b></td>
    <td><b>72.5</b></td>
    <td><b>64.00</b></td>
    <td><b>18.4</b></td>
    <td><b>43.44</b></td>
  </tr>
  <tr>
    <td rowspan="3">2 : 32</td>
    <td>GShard</td>
    <td>19.62</td>
    <td>44.57</td>
    <td>62.28</td>
    <td>32.63</td>
    <td>21.99</td>
    <td>67.19</td>
    <td><b>29.04</b></td>
    <td>69.6</td>
    <td>62.77</td>
    <td>18.2</td>
    <td>42.79</td>
  </tr>
  <tr>
    <td>DropLess</td>
    <td>19.62</td>
    <td>44.51</td>
    <td>62.23</td>
    <td>32.36</td>
    <td>21.79</td>
    <td>67.10</td>
    <td>27.46</td>
    <td>68.3</td>
    <td>62.20</td>
    <td>16.8</td>
    <td>42.24</td>
  </tr>
  <tr>
    <td>MaxScore</td>
    <td><b>20.90</b></td>
    <td><b>44.60</b></td>
    <td><b>63.73</b></td>
    <td><b>33.24</b></td>
    <td><b>23.76</b></td>
    <td><b>67.63</b></td>
    <td><b>29.04</b></td>
    <td><b>73.5</b></td>
    <td><b>64.41</b></td>
    <td><b>18.8</b></td>
    <td><b>43.96</b></td>
  </tr>
  <tr>
    <td rowspan="3">2 : 64</td>
    <td>GShard</td>
    <td>19.80</td>
    <td>44.86</td>
    <td>62.26</td>
    <td>33.05</td>
    <td>22.20</td>
    <td>67.30</td>
    <td>28.46</td>
    <td>69.4</td>
    <td>63.17</td>
    <td>17.6</td>
    <td>42.81</td>
  </tr>
  <tr>
    <td>DropLess</td>
    <td>19.60</td>
    <td>44.69</td>
    <td>62.40</td>
    <td>32.79</td>
    <td>21.65</td>
    <td>67.27</td>
    <td>27.56</td>
    <td>69.5</td>
    <td>63.05</td>
    <td>17.6</td>
    <td>42.61</td>
  </tr>
  <tr>
    <td>MaxScore</td>
    <td><b>21.11</b></td>
    <td><b>46.17</b></td>
    <td><b>64.24</b></td>
    <td><b>33.38</b></td>
    <td><b>23.41</b></td>
    <td><b>67.95</b></td>
    <td><b>28.90</b></td>
    <td><b>73.3</b></td>
    <td><b>64.60</b></td>
    <td><b>19.0</b></td>
    <td><b>44.21</b></td>
  </tr>
</table>

Table 7: Results of scaling in sparsity.

![Scalability with respect to model size and sparsity. The Y-axis represents the training loss of each model after training on approximately 65 billion tokens.](page_101_1042_627_246.png)

Figure 7: Scalability with respect to model size and sparsity. The Y-axis represents the training loss of each model after training on approximately 65 billion tokens.

MaxScore’s effectiveness in harnessing both model capacity and sparsity to improve MoE routing and model accuracy.

4.5 Load Balancing Analysis

Figure 8 illustrates the sorted ratio between the number of tokens assigned to each expert and the capacity \( c = \frac{k * n}{e} \) in the first MoE layer with \( k = 2 \) and \( e = 16 \) after training on about 65 billion tokens. For ExpertChoice MoE, this ratio remains strictly equal to 1, indicating perfect load balancing by design. MaxScore MoE achieves near-ideal load balance with a mean ratio of 0.9996, closely approximating ExpertChoice. In contrast, GShard exhibits notable load imbalance caused by token dropping, resulting in a lower mean ratio of 0.8237 and uneven token distribution across experts. Drop-Less displays extreme variability, with ratio values ranging from 0.55 to 1.74, indicating significant disparity in expert loads. These findings demonstrate MaxScore’s superior capability in mitigating load imbalance relative to traditional approaches.

![The sorted ratio between the number of tokens each expert allocated and Capacity c = k * n / e in the first layer of MoE with k = 2 and e = 16. For ExpertChoice MoE, the ratio is always equal to 1. The mean ratios of GShard MoE, MaxScore MoE, and Drop-Less MoE are 0.8237, 0.9996, and 1, respectively.](page_101_1342_627_246.png)

Figure 8: The sorted ratio between the number of tokens each expert allocated and Capacity \( c = k * n / e \) in the first layer of MoE with \( k = 2 \) and \( e = 16 \). For ExpertChoice MoE, the ratio is always equal to 1. The mean ratios of GShard MoE, MaxScore MoE, and Drop-Less MoE are 0.8237, 0.9996, and 1, respectively.

4.6 Different SoftTopk Operators

We evaluate various SoftTopk(\( \cdot \)) operators listed in Table 1. As illustrated in Figure 9, none yield performance improvements except for our proposed operator defined in Equation (6). We hypothesize that the increased complexity of alternative opera-
tors may hinder effective model learning.

![Results of different operators.](page_186_120_495_312.png)
Figure 9: Results of different operators.

4.7 Hyperparameter t in SoftTopk Operator

We perform hyperparameter tuning for the parameter \( t \) in our SoftTopk(\( \cdot \)) operator defined in Equation (6), exploring two strategies: maintaining a constant value or decaying \( t \) to 1 over training on 10b tokens. As shown in Figure 10, the optimal approach initializes \( t_0 = 4 \) and gradually decays it to 1.

![Hyperparameter tuning experiment.](page_186_573_495_312.png)
Figure 10: Hyperparameter tuning experiment.

5 Conclusion and Future Work

This work introduces MaxScore MoE, a novel mixture-of-experts routing paradigm formulated via minimum-cost maximum flow modeling and the integration of a differentiable SoftTopk(\( \cdot \)) operator. To our knowledge, this is the first successful integration of network flow modeling and SoftTopk within MoE routing. The synergistic combination of these components yields superadditive performance gains, with empirical evidence showing that their joint effect surpasses the linear sum of individual contributions. Future work will focus on evaluating the method at larger model scales and across more diverse benchmarks to validate its generality and robustness.

Limitations

Due to limited computational resources, our experiments are restricted to smaller-scale models, precluding direct comparison with larger, state-of-the-art models. Additionally, the training data volume is relatively modest; further experiments with substantially larger token budgets are necessary to fully assess the ultimate benefits and convergence properties of our approach.

<table>
  <tr>
    <th>Model</th>
    <th>Base</th>
    <th>Large</th>
    <th>XL</th>
  </tr>
  <tr>
    <td><b>Activated Params</b></td>
    <td>162\( M \)</td>
    <td>317\( M \)</td>
    <td>600\( M \)</td>
  </tr>
  <tr>
    <td><b>Total Params</b></td>
    <td>757\( M \)</td>
    <td>1.6\( B \)</td>
    <td>3.2\( B \)</td>
  </tr>
  <tr>
    <td><b>FLOPs</b></td>
    <td>302\( G \)</td>
    <td>603\( G \)</td>
    <td>1.2\( T \)</td>
  </tr>
  <tr>
    <td>hidden_size</td>
    <td>768</td>
    <td>1128</td>
    <td>1608</td>
  </tr>
  <tr>
    <td>num_heads</td>
    <td>12</td>
    <td>12</td>
    <td>12</td>
  </tr>
  <tr>
    <td>num_layers</td>
    <td>12</td>
    <td>12</td>
    <td>12</td>
  </tr>
</table>

Table 8: Configurations for different dense backbones. FLOPs are calculated with a single sequence of 512 tokens. The intermediate\_size of the MLP layer in the dense model is four times that of the hidden\_size, while for the top-\( k \) MoE, the intermediate\_size in each expert is reduced to \( 1/k \), compared with the dense model.

<table>
  <tr>
    <th>Sparsity</th>
    <th>2 : 16</th>
    <th>2 : 32</th>
    <th>2 : 64</th>
  </tr>
  <tr>
    <td><b>Activated Params</b></td>
    <td>162\( M \)</td>
    <td>162\( M \)</td>
    <td>162\( M \)</td>
  </tr>
  <tr>
    <td><b>Total Params</b></td>
    <td>757\( M \)</td>
    <td>1475\( M \)</td>
    <td>2867\( M \)</td>
  </tr>
  <tr>
    <td><b>FLOPs</b></td>
    <td>302\( G \)</td>
    <td>302\( G \)</td>
    <td>302\( G \)</td>
  </tr>
</table>

Table 9: Configurations of different sparsity.

<table>
  <tr>
    <th>Models</th>
    <th>Peak GPU Memory Usage</th>
    <th>Tokens processed Per Hour</th>
  </tr>
  <tr>
    <td>GShard</td>
    <td>71.7\( GB \)</td>
    <td>0.308\( b \)</td>
  </tr>
  <tr>
    <td>ExpertChoice</td>
    <td>71.7\( GB \)</td>
    <td>0.301\( b \)</td>
  </tr>
  <tr>
    <td>DropLess</td>
    <td>73.4\( GB \)</td>
    <td>0.296\( b \)</td>
  </tr>
  <tr>
    <td>DeepSeek</td>
    <td>78.8 GB</td>
    <td>0.277\( b \)</td>
  </tr>
  <tr>
    <td>MaxScore-I</td>
    <td>71.7\( GB \)</td>
    <td>0.305\( b \)</td>
  </tr>
  <tr>
    <td>GShard</td>
    <td>71.7\( GB \)</td>
    <td>0.299\( b \)</td>
  </tr>
</table>

Table 10: The peak GPU memory usage and the speed of processing of MoE models during training. DeepSeek MoE’s use of fine-grained experts leads to larger GPU memory and slower speed (Dai et al., 2024; DeepSeek-AI et al., 2024a).

Acknowledgments

This work was supported in part by National Key Research and Development Program of China under Grant No. 2020YFA0804503, National Natural Science Foundation of China under Grant No. 62272264, and ByteDance Doubao Large Model Fund Project under Grant No. CT20240909109354.
References

Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghavi. 2023. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. Preprint, arXiv:2305.13245.

Richard Bellman. 1958. On a routing problem. Quarterly of applied mathematics, 16(1):87–90.

Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. 2016. Conditional computation in neural networks for faster models. Preprint, arXiv:1511.06297.

Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. 2019. Piqa: Reasoning about physical commonsense in natural language. Preprint, arXiv:1911.11641.

Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. 2016. Training deep nets with sublinear memory cost. Preprint, arXiv:1604.06174.

Aidan Clark, Diego de las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew Johnson, Katie Millican, Albin Cassirer, Chris Jones, Elena Buchatskaya, David Budden, Laurent Sifre, Simon Osindero, Oriol Vinyals, Jack Rae, Erich Elsen, Koray Kavukcuoglu, and Karen Simonyan. 2022. Unified scaling laws for routed language models. Preprint, arXiv:2202.01169.

Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. 2019. Boolq: Exploring the surprising difficulty of natural yes/no questions. Preprint, arXiv:1905.10044.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018. Think you have solved question answering? try arc, the ai2 reasoning challenge. Preprint, arXiv:1803.05457.

Marco Cuturi. 2013. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26.

Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. 2024. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. Preprint, arXiv:2401.06066.

DeepSeek-AI, Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Hanwei Xu, Hao Yang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, J. L. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li, Jin Chen, Jingyang Yuan, Junjie Qiu, Junxiao Song, Kai Dong, Kaige Gao, Kang Guan, Lecong Zhang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Liyue Zhang, Meng Li, Miaojun Wang, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Mingming Li, Ning Tian, Panpan Huang, Peiyi Wang, Peng Zhang, Qiancheng Wang, Qihao Zhu, Qinyu Chen, Qiushi Du, R. J. Chen, R. L. Jin, Ruqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, Runxin Xu, Ruoyu Zhang, Ruiyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuan Chen, Shaqoing Wu, Shengfeng Ye, Shiron Ma, Shiyu Wang, Shuang Zhou, Shuiqing Yu, Shunfeng Zhou, Size Zheng, T. Wang, Tian Pei, Tian Yuan, Tianyu Sun, W. L. Xiao, Wangding Zeng, Wei An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang, Xiao Bi, Xiaodong Liu, Xiaohan Wang, Xiaojin Shen, Xiaokang Chen, Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang, Xin Liu, Xin Xie, Xingkai Yu, Xinnan Song, Xinyi Zhou, Xinyu Yang, Xuan Lu, Xuecheng Su, Y. Wu, Y. K. Li, Y. X. Wei, Y. X. Zhu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yaohui Wang, Yi Zheng, Yichao Zhang, Yiliang Xiong, Yilong Zhao, Ying He, Ying Tang, Yishi Piao, Yixin Dong, Yixuan Tan, Yiyuan Liu, Yongji Wang, Yongqiang Guo, Yuchen Zhu, Yuduan Wang, Yuheng Zou, Yukun Sha, Yunxian Ma, Yuting Yan, Yuxiang You, Yuxuan Liu, Z. R. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhen Huang, Zhen Zhang, Zhenda Xie, Zhewen Hao, Zhihong Shao, Zhinui Wen, Zhipeng Xu, Zhongyu Zhang, Zhuoshu Li, Zihan Wang, Zhihu Gu, Zilin Li, and Ziwei Xie. 2024a. Deepseekv2: A strong, economical, and efficient mixture-of-experts language model. Preprint, arXiv:2405.04434.

DeepSeek-AI, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, J. L. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li, Jiawei Wang, Jin Chen, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, Junxiao Song, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lecong Wang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Litong Wang, Liyue Zhang, Meng Li, Miaojun Wang, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Mingming Li, Ning Tian, Panpan Huang, Peiyi Wang, Peng Zhang, Qiancheng Wang, Qihao Zhu, Qinyu Chen, Qiushi Du, R. J. Chen, R. L. Jin, Ruqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, Runxin Xu, Ruoyu Zhang, Ruiyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuan Chen, Shaqoing Wu, Shengfeng Ye, Shengfeng Ye, Shiron Ma, Shiyu Wang, Shuang Zhou, Shuiqing Yu, Shunfeng Zhou, Shuting Pan, T. Wang, Tao Yun, Tian Pei, Tianyu Sun, W. L. Xiao, Wangding Zeng, Wanjia Zhao, Wei
An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wenjin Yu, Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang, Xiao Bi, Xiaodong Liu, Xiaohan Wang, Xiaojin Shen, Xiaokang Chen, Xiaokang Zhang, Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xingxiang Wang, Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xingkai Yu, Xinnan Song, Xinxia Shan, Xinyi Zhou, Xinyu Yang, Xinyuan Li, Xuecheng Su, Xuheng Lin, Y. K. Li, Y. Q. Wang, Y. X. Wei, Y. X. Zhu, Yang Zhang, Yanhong Xu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yao-hui Wang, Yi Yu, Yi Zheng, Yichao Zhang, Yifan Shi, Yiliang Xiong, Ying He, Ying Tang, Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma, Yiyuan Liu, Yongjiang Guo, Yu Wu, Yuan Ou, Yuchen Zhu, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yukun Zha, Yunfan Xiong, Yunxin Ma, Yuting Yan, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Z. F. Wu, Z. R. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhean Xu, Zhen Huang, Zhen Zhang, Zhenda Xie, Zhengyan Zhang, Zhewen Hao, Zhibin Gou, Zhicheng Ma, Zhigang Yan, Zhihong Shao, Zhipeng Xu, Zhiyu Wu, Zhongyu Zhang, Zhuoshu Li, Zhiui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Ziyi Gao, and Zizheng Pan. 2024b. Deepseek-v3 technical report. Preprint, arXiv:2412.19437.

David Eigen, Marc’Aurelio Ranzato, and Ilya Sutskever. 2014. Learning factored representations in a deep mixture of experts. Preprint, arXiv:1312.4314.

William Fedus, Barret Zoph, and Noam Shazeer. 2022. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Preprint, arXiv:2101.03961.

Lester Randolph Ford. 1956. Network flow theory. Rand Corporation Paper, Santa Monica, 1956.

Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. 2022. Megablocks: Efficient sparse training with mixture-of-experts. Preprint, arXiv:2211.15841.

Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noach’h, Haonan Li, Kyle McDonell, Niklas Muenninghof, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. 2024. The language model evaluation harness.

Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhari, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schellen, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravan Kumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Cantón Ferrer, Cyrus Nikolaidis, Damien Alonsius, Daniel Song, Danielle Pintz, Danny Livshits, Danny Wyatt, David Esioub, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Govind Thattai, Graeme Nail, Gregoire Milanon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Jun-teng Jia, Kalyan Vasudean Alwala, Karthik Prasad, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Rantala-Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kam-badur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Ning Zhang, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng, Prajwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohan Maheswari, Rohit Giridhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhsale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vitor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyu Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang, Xi-aofang Wang, Xiaoqing Ellen Tan, Xide Xie, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmeen Babaei, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld,
Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Chingsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Duc Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric-Tuan Le, Erik Brinkman, Esteban Arcauta, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Filippo Kokkinos, Firat Ozenget, Francesco Caggioni, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Grant Herman, Grigory Sizov, Guangyi, Zhang, Gunasakthi Lakshminarayanan, Hakan Inan, Hamid Shojanazari, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Ilias Leontiadis, Irina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, Jean-Baptiste Gay, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reitzeinent, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matoshic, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kiran Jagadeesh, Kun Huang, Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Gard, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatkovsk, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Prithish Yuvaraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Rangaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shuqiang Zhang, Shuqiang Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumit Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Koehler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang, Xiaojiang Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhadouo Wen, Zhenyu Yang, Zhiwei Zhao, and Zhiyu Ma. 2024. The llama 3 herd of models. Preprint, arXiv:2407.21783.

Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, Xinrong Zhang, Zheng Leng Thai, Kaihuo Zhang, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ding Ding, Chao Jia, Guoyang Zeng, Dahai Li, Zhiyuan Liu, and Maosong Sun. 2024. Minicpm: Unveiling the potential of small language models with scalable training strategies. Preprint, arXiv:2404.06395.

Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, Joe Chau, Peng Cheng, Fan Yang, Mao Yang, and Yongqiang Xiong. 2023. Tutel: Adaptive mixture-of-experts at scale. Preprint, arXiv:2206.03382.

Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michał Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygódź, Piotr Sankowski, Marek Cygan, and Sebastian Jaszczur. 2024. Scaling laws for fine-grained mixture of experts. Preprint, arXiv:2402.07871.

Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. 2017. Race: Large-scale reading comprehension dataset from examinations. Preprint, arXiv:1704.04683.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2020. Gshard: Scaling giant models with conditional computation and automatic sharding. Preprint, arXiv:2006.16668.

Ilya Loshchilov and Frank Hutter. 2019. Decoupled weight decay regularization. Preprint, arXiv:1711.05101.

André F. T. Martins and Ramón Fernandez Astudillo. 2016. From softmax to sparsemax: A sparse model of attention and multi-label classification. Preprint, arXiv:1602.02068.

Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. 2018. Can a suit of armor conduct electricity? a new dataset for open book question answering. Preprint, arXiv:1809.02789.

Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wetting, Binyuan Hui, Tim Dettmers, Douwe Kiela, Ali Farhadi, Noah A. Smith, Pang Wei Koh, Amanpreet Singh, and Hannaneh Hajishirzi. 2024. Olmoe: Open mixture-of-experts language models. Preprint, arXiv:2409.02060.

Denis Paperno, Germán Kruszewski, Angeliki Lazari-dou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. 2016. The lambada dataset: Word prediction requiring a broad discourse context. Preprint, arXiv:1606.06031.

Ben Peters, Vlad Niculae, and André F. T. Martins. 2019. Sparse sequence-to-sequence models. Preprint, arXiv:1905.05702.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2019. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints.

Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. Zero: Memory optimizations toward training trillion parameter models. Preprint, arXiv:1910.02054.

Noam Shazeer. 2020. Glu variants improve transformer. Preprint, arXiv:2002.05202.

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. Preprint, arXiv:1701.06538.

Jianlin Su. 2024. After softmax: Finding a smooth approximation for top-k.

Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Muradtha, Bo Wen, and Yunfeng Liu. 2023. Roformer: Enhanced transformer with rotary position embedding. Preprint, arXiv:2104.09864.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023a. Llama: Open and efficient foundation language models. Preprint, arXiv:2302.13971.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashykhov, Soumya Batra, Prajwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esibou, Jude Fernandes, Jeremy Fu, Wenjin Wu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Harthorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molbyog, Yixin Nie, Andrew Poultren, Jeremy Reizenstein, Rashu Runeta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. 2023b. Llama 2: Open foundation and fine-tuned chat models. Preprint, arXiv:2307.09288.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. CoRR, abs/1706.03762.

Gary R Waissi. 1994. Network flows: Theory, algorithms, and applications.

Ziteng Wang, Jun Zhu, and Jianfei Chen. 2025. Remoe: Fully differentiable mixture-of-experts with relu routing. Preprint, arXiv:2412.14711.

Johannes Welbl, Nelson F. Liu, and Matt Gardner. 2017. Crowdsourcing multiple choice science questions. Preprint, arXiv:1707.06209.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierrette Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. 2020. Huggingface’s transformers: State-of-the-art natural language processing. Preprint, arXiv:1910.03771.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. Hellaswag: Can
a machine really finish your sentence?  Preprint, arXiv:1905.07830.

Biao Zhang and Rico Sennrich. 2019. Root mean square layer normalization. Preprint, arXiv:1910.07467.

Sheng Zhang, Xiaodong Liu, Jingjing Liu, Jianfeng Gao, Kevin Duh, and Benjamin Van Durme. 2018. Record: Bridging the gap between human and machine commonsense reading comprehension. Preprint, arXiv:1810.12885.

Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew Dai, Zhifeng Chen, Quoc Le, and James Laudon. 2022. Mixture-of-experts with expert choice routing. Preprint, arXiv:2202.09368.

Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. 2022. St-moe: Designing stable and transferable sparse expert models. Preprint, arXiv:2202.08906.

Reviewer paper summaries: This paper presents Maximum Score Routing (MaxScore), a novel routing mechanism for Mixture-of-Experts (MoE) models. The authors formulate MoE routing as a minimum-cost maximum flow problem and introduce a Soft-Topk operator to improve token-expert matching. MaxScore retains expert capacity constraints while achieving better load balancing than existing methods. Evaluations show that MaxScore exhibits lower training losses and higher evaluation scores compared to both capacity-constrained (GShard, ExpertChoice) and unconstrained (DropLess, DeepSeek) baselines under the same computational cost.

This paper introduces a new MoE routing method called Maximum Score Routing (MaxScore). The method formulates the routing problem as a minimum-cost maximum flow problem and implements a Soft-Topk operator. Evaluations are conducted using Llama, and the results are compared against various MoE baselines.
Summary of strengths: 1. Novel Formulation: This paper model MoE routing as a minimum-cost maximum flow problem is an innovative approach that addresses some limitations in existing methods (especially unbalanced routing for MoEs).

2. Theoretical Grounding: The paper provides solid theoretical justification for the proposed method, including detailed explanations of the graph structure, computational complexity analysis, and algorithmic details.

3. Performance Improvements: MaxScore consistently outperforms baselines across metrics, with particularly strong results when scaling model size and expert count.

4. Analysis of Load Balancing: The paper includes insightful analyses of token distribution and expert utilization, demonstrating how MaxScore achieves better load balancing.

1. Formulating the MoE routing problem as a minimum-cost maximum flow problem is both novel and reasonable, demonstrating that graph matching techniques can be effectively applied to MoE routing.

2. The experimental results support the claim that the proposed method achieves improved performance under certain capacity constraints.

3. The approach scales better as the number of experts in the MoE increases.
Summary of weaknesses: 1. Limited Model Scale: Experiments are conducted on relatively small models (162M-599M parameters), which might not fully demonstrate the benefits at scales where MoE models are typically deployed (billions of parameters). This raises questions about the method's effectiveness in large-scale deployments.

2. Insufficient Training for Convergence:  This paper's models are trained for approximately 57.7 billion tokens (163,840 steps), however in figure 4, where loss curves are still declining at the end of training. The scaling experiments in Figures 5 and 9 use even shorter training (only 7.2 billion tokens), maybe making their conclusions preliminary rather than definitive.

3. Limited Hyperparameter Sensitivity Analysis: While there is a brief experiment on the hyperparameter t in the Soft-Topk operator, a more comprehensive analysis of sensitivity to other routing parameters (capacity factor, load balancing loss weights, etc.) would strengthen the paper. The stability of the method across different configurations remains uncertain.

4. Incomplete Comparison with Related Work: The paper lacks comparison with methods that address load balancing during inference, such as "Turn Waste into Worth: Rectifying Top-k Router of MoE." Such comparisons would better position this work within the broader context of MoE routing optimizations and clarify the relative advantages of addressing routing efficiency during training versus inference.

5. Computational Overhead Concerns: While MaxScore-Iter is mentioned to have "almost no extra CPU computation compared to GShard," the computational complexity of the full MaxScore algorithm using the Hungarian algorithm is O(N³). Maybe a more detailed analysis of real-world computational overhead is needed.

1. Although the paper reports an improved average score, the method does not consistently perform well across all entries. It would be beneficial to explicitly highlight the best results for each task and provide a more detailed analysis of the model's behavior.

2. In Figure 6, the results show that MaxScore-iter consistently achieves lower loss than MaxScore at different scales, even though MaxScore is intended to be the superior variant. This discrepancy between the results and the conclusions appears contradictory and requires clarification.

3. Evaluation is only based on Llama and it would be better to include more models to validate the generalizability.
Comments, suggestions and typos: You can supplement and explain the issue above, (see weakness).

None
Overall assessment scores: 3

2
Metareview score: 3.5
Metareview: The paper proposes a new MoE routing paradigm developed by modeling MoE routing as the minimum cost maximum flow problem and integrating a Soft-Topk operator.
The reviews given are entirely solved during the rebuttal period, focusing on experimental comparison including hyper-param and explanations of the method.
Summary of reasons to publish: The paper proposes a method that is the trend of current LLM development while brings a sound method for MOE router.