Capacity-Aware Inference: Mitigating the Straggler Effect in Mixture of Experts
Authors: Shwai He
Weilin Cai
Jiayi Huang
Ang Li
Venue: ICLR 2026 Poster (2026)
Keywords: Mixture of Experts
Load Balancing
Computation Efficiency
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose capacity-aware inference to mitigate imbalanced token assignments during inference, significantly improving efficiency without compromising performance.

Abstract: The Mixture of Experts (MoE) is an effective architecture for scaling large language models by leveraging sparse expert activation to balance performance and efficiency. However, under expert parallelism, MoE suffers from inference inefficiencies due to imbalanced token-to-expert assignment, where underloaded experts complete computations early but must wait for overloaded experts, leading to global delays. We define this phenomenon as the \textbf{\textit{Straggler Effect}}, as the most burdened experts dictate the overall inference latency. To address this, we first propose \textit{\textbf{Capacity-Aware Token Drop}}, which enforces expert capacity limits by discarding excess tokens from overloaded experts, effectively reducing load imbalance with minimal performance impact (e.g., $30\%$ speedup with only $0.9\%$ degradation on OLMoE).  
Next, given the presence of low-load experts remaining well below the capacity threshold, we introduce \textit{\textbf{Capacity-Aware Expanded Drop}}, which allows tokens to include additional local experts in their candidate set before enforcing strict local capacity constraints, thereby improving load balance and enhancing the utilization of underused experts. 
Extensive experiments on both language and multimodal MoE models demonstrate the effectiveness of our approach, yielding substantial gains in expert utilization, model performance, and inference efficiency, e.g., applying Expanded Drop to Mixtral-8$\times$7B-Instruct yields a {0.2\%} average performance improvement and a {1.85$\times$} inference speedup. The code is released at: https://github.com/CASE-Lab-UMD/Capacity-Aware-MoE.

Published as a conference paper at ICLR 2026
CAPACITY-AWARE INFERENCE: MITIGATING THE
STRAGGLER EFFECT IN MIXTURE OF EXPERTS
Shwai He Weilin Cai Jiayi Huang" Ang Li'
'University of Maryland, College Park
"he Hong Kong University of Science and Technology (Guangzhou)
(shwaihe, angliece yumd, edu
wcai738@connect. wwush-hkteee gz. edu. cn, sayyhhkustgzzeduu
ABSTRACT
The Mixture of Experts (MoE) is effective architecture for scaling large language models by leveraging sparse expert activation to halance performance and efficiency. However, under expert parallelism, MoE suffers from inference inefficiencies due to imbalanced token-to-expert assignment, where underloaded experts complete computations early but must wait for overloaded experts, leading global delays. We define this phenomenon as the Straggler Effect, as the most burdened experts dictate the overall inference latency. To address this, we first propose Capacity- Aware Token Drop, which enforces expert capacity limits by discarding excess tokens from overloaded experts, effectively reducing load imbalance with minimal performance impact (e.g.. 30% speedup with only 0.9% degradation OLMoE). Next, given the presence of low-load experts remaining well below the capacity threshold, we introduce Capacity-Awar Expanded Drop, which allows tokens to include additional local experts their candidate set before enforcing strict local capacity constraints, thereby improving load balance and enhancing the utilization of underused experts. Extensive experiments on both language and multimodal MoE models demonstrate the effectiveness of our approach, yielding substantial gains in expert utilization. model performance, and inference efficiency, e.g., applying Expanded Drop Mixtral-8x 1-BBBBBBBstnuee yields a 0.2% average performance improvement and 1.85x inference speedup. The code is released at: Moo
https://gi thub.
INTRODUCTION
In recent years, the rapid evolution of Large Lan- guage Models (LLMs) (OpenAT, 2024 Team 2024a al. 2024b) has driven vations, continuously expanding the frontiers AI research and applications. Among the model architectural innovations, the Mixture of Experts (MoE) framework has emerged pivotal tech- nique for optimizing the cost-performance trade- off in LLMs. Specifically, MoE (Shazeer et al. 2017b) enhances scalability by integrating multi- ple experts while activating only subset per input. This selective activation substantially improves model performance without a corresponding crease in computational cost, effectively balancing efficiency and performance.
Higstoad
load
of inno-
Nur burnukadd Loud
Ex
Figure 1: Illustration of the Straggler Effect in MoE. Inference. The normalized load is com- puted each expert's load divided by the mean load across all experts. Example shown with OLMoE (Muennighoff' et al.) 2024) on Open- BookQA (Mihaylov et al. 20188).
Despite the success of MoE, key efficiency challenge lies in the imbalanced token-to-exper distribution, which results in some experts being overloaded while others remain underutilized (Lepikhin et al., 2021 Zoph et al. (022). In distributed GPU settings, experts typically sharded across multiple devices, with each GPU responsible for a subset of the experts. Under expert parallelism, low-load experts complete their computations earlier but must wait for overloaded experts

---

Published as a conference paper at ICLR 2026
to finish, as synchronization barriers are required before proceeding to the next stage. This expert- level straggler effect further propagates to device-level delays, where GPUs hosting lighter expert workloads are stalled by GPUs hosting heavier workloads, leading to inefficient resource utilization and increased end-to-end latency during inference. As illustrated in Figure [ this phenomenon is referred as the Straggler Effect, where the heavily loaded experts determine the overall latency of imbalanced MoE inference.
While auxiliary balance losses have been incorporated into the training process to alleviate imbalance (Shazeer et al., 20176] Fedus et al., 2022 et al. 2024b), these techniques remain ineffective in mitigating imbalance during inference. Specifically, as shown in Figure I our findings reveal a highly uneven token distribution among experts, with the highest-load expert handling more than seven times the expected average load. Moreover, managing such imbalance during inference often incurs additional resource overhead. For example, DeepSeek-VS V3 mitigates this issue by duplicating high-load experts and deploying them redundantly across devices (et al. (024b). This motivates us to explore efficient token-to-exper assignment by addressing the key question: How can we prevent extreme overloading of heavily utilized experts?
We propose Aapacity-Aware Inference to address this challenge. Specifically, for high-load experts, we introduce Capacity-Aware Token Drop, which imposes a maximum capacity constraint and discards excess tokens from overloaded experts. This approach alleviates severe load imbalance and significantly improves efficiency, while maintaining model performance since the dropped tokens represent only small fraction of the total workload, e.g.. OLMoE achieves 30% speedup in MoE layers with just a 0.9% performance degradation. After removing excess tokens from overloaded experts, we observe that some low-load experts remain significantly underutilized relative to the predefined capacity constraints, yet must still wait for other experts to complete their computations. This leads us second key question: How can we effectively leverage the available capacity of underutilized experts?
For low-load experts, we extend Token Drop with Capacity-Aware Expanded Drop. which further utilizes the available capacity of underutilized experts to handle overflow tokens from high-load experts. Specifically, under expert parallelism across multiple GPUs, Expanded Drop allows each token to consider additional candidate experts on the same device while still enforcing strict local capacity constraints. This expanded selection improves the utilization of low-load experts and enhances the representational capacity of MoE models under sapacity-constrain
scenarios.
Extensive experimental results validate the effectiveness of our proposed techniques, demonstrating significant improvements in both efficiency and performance, e.g., applying Expanded Drop to Mixtral-8x Bextral-877B--nannn 7B-Instruct yields a 0.2% average performance improvement and a 1.85x inference speedup. Moreover, in multimodal models, we identify redundancy among image tokens and show that applying aggressive capacity constraints (e.g., setting the maximum to half of the average expert load) can still maintain performance. In short, our contributions are in four folds:
We explicitly clarify and analyze the Straggler Effect caused by test-time token imbalance in the Mixture of Experts, highlighting the optimization potential for reducing latency. Toward the high-load experts, we propose Capacity-Aware Token Drop, which enforces capacity constraints by discarding excess tokens assigned to overloaded experts, thereby mitigating extreme load imbalance.
To better utilize underloaded experts, we introduce Capacity-Aware Expanded Drop. which expands the candidate expert set to include additional local experts, further improving load balance and model performance.
Extensive experiments on both language and multimodal models validate the effectiveness our approach, demonstrating substantial improvements in inference efficiency with comparable performance.
2
RELATED WORKS
Mixture of Experts Models The Mixture of Experts (MoE) is a kind of neural network architecture with extended set of parameters (referred to "experts") controlled by router, which is first
introduced in the context conditional computation (Jacobs et al. 1991 & Jacobs 1994). The potential of sparse activation in MoR is subsequently exploited by (Shazeer et al. (2017a) for efficient.
2

---

Published as a conference paper at ICLR 2026
training and inference pretrained models with special designs, opening the door for MoE in various
vision (Riquelme et al., 2021) and language (Lepikhin et al. 2020; Du et al. 2022 Fedus et al.
2022) scenarios. Attributed to its exceptional efficiency, MoE has been adopted
foundational
framework in the designs of large language models (Jiang et al. 2024 Dai et al. 2024 Xue et al.
2024a Zhu al. [024 Team 2024b), achieving superior scaling laws low computational costs.
Despite these advancements, MoE still faces efficiency challenges in both training and inference (Cai
al. 2024), and our work specifically focuses on enhancing inference-time efficiency.
Imbalance in Mixture of Experts The imbalance in token- to-expert assignments (Zhou et al. Chen et al., 2022) poses a significant challenge to the deployment of MoE. This imbalance leads to inefficiencies in computation, communication, and memory (He et al. 2023 Song et al. 2023 Xue al. 2024b), making it a critical bottleneck for MoE scalability and deployment. To mitigate this issue, an auxiliary balance loss (Shazeer et al.) 2017b) is incorporated into the training process to encourage more uniform token distribution across experts. Additionally, various training strategies have been introduced to further balance token assignments: Switch-Transfonee et al., 2022) and DeepSeek- V2 (et al. 2024a) implement Token Drop to alleviate expert overload, while DeepSeek- V3 (et al. 2024b) introduces additional sequence-lever auxiliary loss to prevent severe token imbalance.
Chen
(Fedus)
However, these techniques primarily focus on training and fail to ensure balanced token assignments during inference. Instead, addressing token imbalance at inference often incurs additional resource costs. For example, DeepSeek-V3. (et al.) 2024b) mitigates this issue by duplicating high-load experts and deploying them redundantly. In contrast, our approach effectively balances token assignments without introducing additional computational overhead.
BACKGROUND AND MOTIVATION
3.1 EXTREMELY IMBALANCED EXPERT UTILIZATION
Mixture of Experts (MoE) layer consists of a collection of n experts, {E1. Eg...., E,) and a router G that dynamically selects the most relevant experts for given input The router computes selection scores Gæx), for all experts and selects the top & experts, resulting in a sparse activation:
K TopK (Softmax (Softmax(G(())) ((z)),
(1)
The input æ is processed by the selected experts, and their outputs are combined into a weighted sum based on the router's scores. This process is mathematically expressed
y Σ REK G(x); (x))
where K. denotes the indices of selected experts, G(æ:); represents the selection score for the i-th expert, and E, (x) is the output from the i-th expert. In transformer models, the MoE layer usually replaces the feed-forward network (FFN) and only activates a subset of experts for each input.
While experts in MoE models are often deployed in parallel across distributed GPUs, imbalanced ooken-to-expedt assignments lead varying levels of expert utilization and potential latency. Despite the incorporation of balancing techniques during training, the load imbalance persists during inference. To further investigate this issue, we conduct preliminary experiments to analyze expert-specific utilization patterns and assess the impact of imbalance on practical latency.
To quantify expert utilization, we measure the load across different experts. Given an input batch x Rxx with batch size and sequence length the total number of tokens is ( bs. Since each token selects & out n experts, the expected token count per expert is:
Hommlized
tk
(3)
Export
However, due to imbalanced token assignments, some experts may receive more or fewer tokens than the expected value.
Figure 2: Test-time expert load of OLMoE across different datasets, where each load value is normalized by the mean load N for clarity.
3

---

Published as a conference paper at ICLR 2026
Figure 2 illustrates the normalized peak tokens load for each expert to accommodate all tokens within a single layer of OLMoE, where some experts receive an excessively large number of tokens (e.g., more than seven times the average number of tokens), leading to severe load imbalance and, consequently, significant latency. A detailed layer-by-layer analysis is provided Appendix FF
3.2 MOTIVATION - THE STRAGGLER EFFECT
Under the expert parallelism scenario, where the number of assigned tokens dictates the processing time of each expert, high-load experts become the bottleneck for overall latency within MoE layer. Specifically, low-load experts remain idle while waiting for high-load experts to complete, leading to synchronization delays. Therefore, the latency of an MoE layer is given by:
L x max ((N
(4)
where N, represents the number of tokens assigned to the i-th expert, with the total token allocation satisfying Dis1 N; lk. According to Eq. 4 the latency follows the Straggler Effect: the most burdened experts dictate the overall latency of the MoE layer. In the worst casc, all tokens are assigned to the same group of experts, underutilizing the parallel processing capability of MoE. Conversely, distributing tokens evenly across experts maximizes computational efficiency and fully leverages the parallelism of multiple experts. With the bounds of the ideal and worst cases, the range of the highest load is given by:
max (a)(( E IN, nN
(5)
k
However, existing MoE models often adopt dropless strategy during inference, which fails to address token imbalance and can lead to significantly increased latency.
Given that the imbalance stems from excessively high- and low-load experts, we address this issue by exploring the following questions: (1) For high-load experts, are there redundant tokens that can be dropped without causing significant performance degradation? (2) For low-load experts that must wait for high-load experts to complete forward passes, is there an opportunity to enhance their utilization and improve performance without incurring substantial additional cost?
4 METHODOLOGY
Token Drop Regulates the Latency of High-Load Experts address the question about over- loaded experts, we first regulate their maximum utilization. Specifically, we introduce expert capacity to control token allocation. Given a capacity factor y, the maximum number of tokens assigned to each expert (i.e., expert capacity) is defined as: C' =NN.
higher allows more tokens to be retained, but experts handling excessive tokens may introduce latency. Conversely, lower y enforces stricter capacity limits, reducing latency by discarding more tokens, but at the risk of performance degradation. With the involvement of expert capacity we constrain the upper bound of latency follows:
L within [N, N
max (NN.)E
7
(6)
7 2
where y typically much smaller than -. This constraint ensures that no expert exceeds the specified capacity limit, effectively mitigating severe load imbalances and reducing latency. Note that tokens are distributed across devices under expert and data parallelism. To avoid additional communication overhead, apply capacity constraints to tokens within each local device, similar to the constraints used during training (Fedus et al. 2022). This ensures that all experts respect the limits, maintaining strict control over token flow the experts.
Specifically, when a capacity constraint imposed on each expert, experts must evaluate the volume of assigned tokens before execution. For experts with a load below the predefined capacity, there is no difference between load exceeds the capacity, experts must discard excess tokens to adhere the constraint. To address
inference and traditional inference. However, when the
4

---

Published as a conference paper at ICLR 2026
this, introduce a scoring function S to evaluate each token:
S$1
$12
81n
$21
$22
82n
S(x)
811
$12
Stn
where sij denotes the importance score of the mapping from the i-th token the j-th expert. With this score, each overflowed expert selectively discards those with lower scores. Let J be the set of overflowed experts and Sy the corresponding columns of S, with the Token Drop threshold set
TJ Kth CS)
(8)
where T. represents the thresholds, i.e., C-th highest value in Sy, serving as a threshold filter out excess tokens:
Ty + [(t,t) | te [I....,, j J. Stt. 2 b 5 S3 @ My, where M, + 1 [S, 2 T
(9)
(10)
where Ty denotes the token indices retained by the experts indexed in J. The scores of rejected tokens are masked to prevent them from being routed to their corresponding overflowed experts.
Regarding the specific scoring function, we explore multiple metrics and summarize them as follows:
Discarding later tokens once earlier tokens have filled the expert capacity. This strategy was first introduced in Switch-Transfone (Fedus al. 2022) during training, and we extend the inference phase.
Order:
Reverse Order: Instead of discarding later tokens, this approach removes earlier tokens to comply with the expert capacity constraint.
Random: Dropping Excess tokens randomly to meet the predefined expert capacity constraints.
Using gating scores after the softmax and top-k operations as an importance indicator and discarding tokens.
Score:
Among these metrics, "Order" and "Reverse Order" are unstable, as shuffling sequences within batch may result in different tokens being dropped (Hayes et al. 2024). "Random" assumes all tokens have an equal probability of being dropped. In contrast, "Score" " is stable, unaffected by sequence order within batch. Notably, there virtually no additional computational overhead associated with calculating these metrics, and the dropping operation incurs minimal cost compared to the intensive computations performed by the experts.
(a) Token Drop
(b) Expanded Drop
Figure 3: Illustration of Capacity-Aware Token Drop (a) and Expanded Drop (b). Both methods first select experts based on gating scores. In Token Drop, tokens exceeding the local device capacity are discarded prior All-to-All communication. Expanded Drop enhances expert utilization by allowing each token to consider additional m candidate experts on the same device while still enforcing strict local capacity constraints.
Expanded Drop Enhances the Utilization of Low-load Experts Token Drop exclusively targets overloaded experts by discarding overflowed tokens that exceed expert capacity but does not address the underutilization of low-load experts. Next, we introduce Expanded Drop to ensure a more balanced token-to-expee to-exeert allocation.
naive approach to rerouting under- selected tokens is to mask the mapping scores of overflowed experts and then reselect experts for these tokens. However, the reselection may still result in
5

---

Published as a conference paper at ICLR 2026
overflows, necessitating multiple rounds of selection and dropping, which increases latency. Moreover, the repeated selection and dropping substantially raise the cost of token-to-exxe to-eererrt mapping.
Expanded Drop adopts simple yet effective strategy: for cach token, it selects additional candidate experts. Given m experts deployed on single GPU, a token not only selects the top-k experts based on gating scores, but also includes m local experts (e.g., experts per device under 8- B-way expert parallelism across GPUs for a total of 64 experts) for substitution if the initially selected experts overflowed. As result, each token may select up refined as experts drop tokens as needed to satisfy capacity constraints. This makes no change in the token assignments in experts that are overflowed by the top-k experts. Meanwhile, for under-utilized experts, the expanded top-k + m candidate pool increases the likelihood of receiving additional tokens. After top-k + m selection and dropping, there might exists tokens that select more than experts. Through empirical analysis (Appendix [D), we choose not to enforce constraint that limits each token selecting most k experts, thereby removing the need to explicitly retain the top-k experts at the end.
+ k experts. The final selection is then
Notably, the extra cost token routing is minimal; the only difference lies the negligible cost concatenation of the gating scores from either the top-k or experts on the local device. Moreover, processing expanded tokens within local device eliminates inter-device communication.
EXPERIMENTS
In this section, we conduct experiments under capacity-aware inference for MoE, with deployment details provided in Appendix A
5.1 TOKEN DROP FOR HIGH-LOAD EXPERTS
Table Performance comparison across different capacity factors and selection metrics (i.e., Order, Reverse Order, Random, and Score). The baseline operates without capacity constraints, represented as +o.. We report the average performance over multiple random seeds.
Method
OBQA
PIQA
RTE
WinoGrande
BoolQ
ARC-C
HellaSwag
MMLU
Avg.
Baseline
+:0
45.6
53.7
74.7
$4.5
79.4
$2.5
64.0
Order
42.0
71.5
53.1
71.2
74.2
49.5
76.6
48.4
60.8
Reverse Order
41.8
71.8
52.7
71.0
73.9
49.4
76.4
49.2
60.8
2.0
Random
41.2
75.2
52.7
71.0
74.1
50.1
76.8
49.4
61.3
Score
45.0
54.5
74.6
54.9
51.8
Order
38.8
67.1
48.7
68.5
73.3
46.3
54.0
43.7
55.1
Reverse Order
40.2
67.3
$2.7
70.1
72.7
45.5
$4.4
45.2
$6.0
1.5
Random
39.6
72.1
73.8
45.8
74.2
45.2
Score
44.8
77.5
55.2
70.8
74.3
53.4
78.6
50.0
63.1
Order
36.0
60.2
52.2
62.6
69.6
38.7
58.0
36.9
$1.8
Reverse Order
36.2
59.5
50.5
63.3
69.4
39.4
58.7
38.7
52.0
1.0
Random
34.0
63.1
53.2
70.2
40.5
66.9
35.7
Score
41.6
76.0
53.4
69.9
73.2
50.4
47.8
61.1
Investigation on Token Drop Metries To assess the effectiveness of different metrics in regulating token load to the target capacity, we compare various approaches OLMoE by discarding excess tokens and applying a range of capacity factors. As shown in Table varying the dropping metrics impacts performance different levels. With higher capacities, the model maintains comparable performance even when using naive selection methods like "Random". However, as the capacity factor decreases, performance degradation becomes more pronounced, particularly for "Order". "Reverse Order", and "Random". Notably, "Score" consistently outperforms other methods by large margin, demonstrating the effectiveness of leveraging gating scores as an importance measure. Consequently, adopt "Score" as the default metric.
Efficiency Gains from Capacity-Constrainct Inference We next explore the efficiency improve- ments achieved by imposing expert capacity. Specifically, employ distributed inference using

---

Published as a conference paper at ICLR 2026
Exparto (se/ GrU)
Eaperts Shared Expert Jes (or%)
Dyprt Shaned Dpprt (olt 5.58 (PV)
Experta (LE GTUI
V-220
2.0
1.4
GLAtE--esrrect
Qwen1 Swtr1--fat--httt
Matraa-ee Batra-BxTB--raane
Figure 4: Speedup of a single MoE layer compared to the baseline without capacity constraints, achieved through two capacity-awas inference methods: Token Drop and Expanded Drop.
Papert Compllileo
3.0
Papprtt
Save Tlme
Paliminthing Jownt
Figure 5: End-to-end speedup. "T.D." and "E.D." are abbreviations for Token Drop and Expanded Drop, respectively.
Figure 6: Breakdown Analysis of the inference latency on OLMoE with different capacity fac- tors (e.g., 1.0, 1.5, 2.0).
eight H20 GPUs, utilizing an 8-way Data Parallelism (DP) and 8-way Expert Parallelism (EP) strategy through the Megatron-LM framework (Shoeybi al. 2019]. The input batches configured with batch size of 8K and a sequence length of 512, simulating real-time serving scenarios with high query throughput, Notably, in Mixtral-8x> IBIIntruct model, each GPU typically hosts one or experts, whereas, in models like OLMoE- -Lsttu-nstruct, multiple experts must be deployed on a single GPU (e.g.. eight experts per GPU) due GPU resource constraints.
As illustrated in Figure 4 imposing constraints on expert capacity through Token Drop and Expanded Drop, considerably accelerates inference across the four tested MoE models, in comparison to the baseline model without capacity limitations. The enhanced efficiency of each MoE layer (Figure contributes faster end-to-end model inference (Figure 5). Moreover, as the capacity factor ? decreases, capacity-aware inference methods achieve significantly greater acceleration.
Notably, the efficacy of acceleration is influenced by the numerical relationship between the total experts and the engaged GPUs in Expert Parallelism. As illustrated in Figure 5 for Mixtral-8 7B- Instruct, deploying one or two experts per GPU maximizes the effectiveness of capacity-award inference. In this configuration, Token Drop and Expanded Drop achieve end-to-end model speedups of 1.87 x and 1.85x, respectively, with y = 1.5. Conversely, deploying a greater number of experts single GPU results in more modest acceleration gains, evidenced by the "8E/GPU" (OLMoE- Instruct and Deepseek-V2-Lieee the "10E/GPU" (Qwen1.a 1.seel..--o--haal in Figure 4 and Figure (5) This is because the aggregated load from multiple experts diminishes the proportion reduced load, which is achieved by limiting the straggler expert. Therefore, it is anticipated that allocating more GPUs for expert distribution, thereby reducing the number of experts per GPU, would enhance the acceleration effect of capacity-award inference.
The breakdown analysis presented Figure 6 demonstrates that our proposed capacity-aware inference methods substantially reduce the duration of expert computation, permutation and commu- nication, while preserving comparable cost for gate processing. Notably, the duration of permutation and communication increases when tokens are expanded across range of global experts. This is due to the increased communication workload required to transmit expanded global tokens across various GPU devices. Consequently, these results underscore the necessity of restricting the expanded tokens to be processed by local experts.
Mitigating the Straggler Effect with Minimal Token Discarding Given that expert capacity enforces MoE layers to discard overflowed tokens, we next establish the relationship between expert capacity and the corresponding number of dropped tokens. For capacity factor 7, the total proportion
7

---

Published as a conference paper at ICLR 2026
of dropped tokens is given by:
DT Ei=l ReLU(N; -NN
(II)
Σi=1 N₂
where ReLU(N, - yN represent the number of dropped tokens for the i-th expert.
Figure 7 visualizes the number of dropped tokens across different capacity factors for various test datasets, with more detailed illustration provided in Appendix
100
PIDA
RTE
G Although the most overloaded ex- pert receives much more tokens than the expected number of tokens N, regulating the maximum ca- pacity has limited impact on the overall number of accommodated tokens, thereby maintaining com- after discarding overflow
WinGGrande
Dropped Tokens
60
petitive performance tokens. Moreover, dropping a small proportion of overflowed tokens can significantly reduce the la- lency caused by overloaded experts (e.g., dropping 12% of overloaded tokens promotes the inference speed by 85% in Mixtral-8x x (B-x7BBIIsiii Instruct), highlight- ing the efficacy of capacity-awe aware inference in im- proving both performance and efficiency.
Capadiy Factor
Figure Analysis of dropped tokens with respect to capacity factors.
5.2 EXPANDED DROP TO LOW-LOAD EXPERTS
Some experts receive very few tokens, raising the question of whether they are redundant be leveraged for balanced allocation. We next examine their role and validate the effectiveness of Expanded Drop.
Table 2: Comparison of Expert Drop, Token Drop and Expanded Drop. The capacity factor is set to 2.0 for OLMoE and DeepSeek- -eepSeck-V2-Lite V2-Lite, and 1.5 for Swenl.5-MoE-Chae and Mixtral-8x Instruct. For Expert Drop, each forward pass skips one out of eight experts for Mixtral-8x 1an---xxBB--BccC and the bottom 10% of lowest load experts for other models.
Model
Method
WineGrande
BoolQ
ARC-C
HellaSwag
MMLU
GSM8K
Avg:
Baseline
47.6
80.2
67.9
69.9
80.7
57.0
80.6
52.8
35.1
63.5
Mo
Expert Drop
44.6
76.9
61.0
67.6
78.2
54.4
77.0
50.6
31.6
60.5
Drop
Expanded Drop
47.2
79.4
66.3
70.5
80.9
57.1
80.3
52.3
34.4
63.2
Baseline
42.4
79.9
73.9
70.0
81.3
54.1
80.4
59.8
52.0
61.9
Qwen Seenl.S-MoF-CChu
Expert Drop
Taken Drop
40.4
78.8
72.6
69.1
80.9
53.0
80.0
$9.3
51.9
65.1
Expanded Drop
43,4
79.1
72.6
69.6
81.1
53.4
80.3
59.3
52.1
65.6
DeepScekk-V V2-1 Ite- Chei
Expert Drop
41.8
77.6
71.9
72.5
81.6
57.1
75.5
53.3
56.0
65.3
Taken Drop
45.2
78.3
72.6
74.0
83.2
59.3
80.9
57.3
62.7
68.2
Expanded Drop
Baseline
47.4
84.8
71.8
82.5
88.5
71.7
87.5
70.2
64.2
24.3
Mixtral-8. 8x 7B Instract
Expert Drop
46.8
83.2
70.1
81.3
87.6
67.1
85.6
66.2
62.3
72.2
62.2
X8.3
64.7
73.8
Expanded Drop
47.8
85.0
71.8
83.0
88.6
71.5
87.6
70.2
64.6
74.5
The Critical Role of Low-Load Experts To explore the impact of low-load experts, we further compare dropping tokens (i.c., Token Drop) with skipping experts (i.c., Expert Drop). For Expert Drop, we adopt a conservative strategy that dynamically skips the 10% of experts with the lowest token loads. Notably, the proportion of tokens removed in Expert Drop is significantly lower than in Token Drop (2% in Expert Drop vs. 12% in Token Drop on OLMoE- Instrucct....
Despite this, shown in Table 2 Expert Drop experiences significant performance degradation and is outperformed by Token Drop by a large margin. Moreover, due to the small proportion of tokens assigned low-load experts, removing these experts provides only marginal improvements in
8

---

Published as a conference paper at ICLR 2026
inference speed (less than a 5% speedup). These findings indicate that retaining low-load experts better preserves the performance of MoE models.
To analyze expert selection and justify Expanded Drop, we sort, for each token, all experts by their gat- ing scores in descending order and record the ranked scores (top-1, 2, ..., top-N). Aggregating across tokens, we compute the average, maximum, and min- imum score at each rank (FigureBB) B The curves show that while top-ranked experts receive much higher scores, the decay across ranks is gradual rather than abrupt, yielding a relatively flat tail beyond the first few ranks. Consequently, experts just outside the top-k often have scores comparable to those near the bottom of the top-k, enabling rerouting or drop- ping without materially changing model behavior, ie., rerouted tokens still go to reasonably relevant experts with similar gating probabilities. This validates the design of Expanded Drop, as can exploit the flatness of the gating distribution to balance load without sacrificing accuracy.
Min sta
lapa
Moan
Topde
0.20
0.00
Expert Rank
Figure 8: Gating score distribution across ranked experts.
Effectiveness of Expanded Drop We examine the
Taken Drop
effectiveness of utilizing low-load experts by Ex-
M Mwm
M
panded Drop instead of simply discarding these to-
3
1.0
kens to meet the target capacity. Comparing Ex-
panded Drop with Token Drop, redistributing excess tokens low-load experts enhances performance, yielding a 0.9% improvement in the average perfor- mance of 2- sidering the performance degradation observed in Expert Drop, our findings highlight the crucial role of low-load experts in maintaining model effectiveness.
0.0
20 Expert ID
60
Expanded Drop
1.5-
Furthermore, con-
0.0
Expert
Figure 9: Normalized expert load after Token Drop and Expanded Drop.
Expanded Drop overselects experts for each token expand the selection scope prompt a more balanced token-exped voken-expert assignments. As shown in Figure D increasing the overselection ratio m allows tokens to consider more candidate experts after being dropped from overflowed ones, thereby improving low-load expert utilization and balancing the expert load.
Device-Level Capacity-Aware Inference In scenarios where multiple experts reside on the same device, the overall straggler effect determined by the aggregated load on each device. Expert-level capacity constraints impose strict limit. Specifically, the number of tokens assigned to a single expert must not exceed 7.N When extended to device with ny local experts, this implies that the total load is bounded by nu 7,N However, this strict constraint can lead to over-aggressive token dropping, as a single overloaded expert may force token removal even when sufficient spare capacity remains across other experts on the same device.
Advanced Variant:
To alleviate this issue, we introduce a device-level capacity-aware formulation, which applics the constraint at the device granularity rather than individual experts. For instance, enforce Ni + N2 + .... Nn, S n NN. Unlike expert-level constraints, which may drop tokens simply because single expert exceeds its local budget, the device-level constraint admits these tokens as long as the total device-level load remains within the allowable bound. This results in a smoother, less restrictive, and more utilization-ffficient alternative.
Table 3: Comparison of Device-Level and Expert-Leved capacity inference on Qwen3-MoE.
Method
Granularity
OBQA
PIQA
RTE
WinoGrande
Bool() ARC-C
HellaSwag
MMLU
GSM8K
Avg.
Baseline
45.2
80.6
82.3
69.5
88.6
69.5
77.6
77.9
89.5
75.6
Expanded Drop
Expert Device
1.5
43.6
79.4
80.3
68.4
86.5
68.8
74.4
76,8
87.0
73.9
1.0
42.4
79.7
81.2
69.3
88.0
69.4
77.1
77.1
88.9
74.8

---

Published as a conference paper at ICLR 2026
We evaluate this variant and observe that device-level constraints consistently outperform expert-level constraints in downstream performance, as shown in Table This more flexible constraint also leads to stronger practical speedups: on Qwen3-MoE (Yang et al. 2025 with y 1. the device-level formulation achieves a 1.31 x end-to-end speedup and a 1.51 x speedup on a single MoE layer, both surpassing the 1.23x and 1. x speedups obtained under expert-level constraints with y 1.5.
5.3 EXTENSION TO MULTI-MODAL MIXTURE OF EXPERTS
In addition to applying capacity- aware inference to MoE models for language tasks, we also explore its effectiveness in multi-modal MoE settings. Specifically, we evaluate the OLMoB based MolmoE (Detkee et al., 2024), across multi-modal henchmarks, including MME (Fu et al. 2023). MMBench
Liu et al. 2023). and SEED-Bench (Li et al. 2024a).
Given that the input sequence contains tokens from multiple modalities, we first investigate different to- ken dropping strategies. Specifically, we first treat all tokens equally and drops those with the lowest scores ("Uniform"). Beyond this, considering the redundancy often found in image tokens, we also experiment with strategy that prioritizes dropping image tokens before selectively removing text tokens ("Image First"). For comparison, also consider drop text tokens first ("Text First"). As shown in Table 4 on the MME benchmark, this image-first first strategy yields improved performance, highlighting the benefit of prioritizing dropping image tokens for the load balance in multi-modal MoE models.
Table 4: Capacity-aware inference for Multi-modality MoE models. "Percep." and "Cognit." " denote Perception and Cognition, respectively. is set to 1.0.
Method
Strategy
Percep.
Cognit.
Baseline
1358.1
269.6
Token Drop
1248.4
245.4
Uniform
Expanded Drop
1307.6
273.6
Token Drop Expanded Drop
1114.2
214.4
Text First
1163.6
241.3
Token Drop Expanded Drop
1346.5
288.9
Image First
1362.1
297.1
23
32
50
60
Expert ID
A6
CP
1P4
F75
Figure 10: Multi-modal token assignments across different experts.
Figure 11: Comparison on MMBench across
six multimodal capabilities in Appendix
Given the redundancy of image tokens and their large proportion in multi-modal tasks, we further investigate more aggressive capacity factors for Token Drop and Expanded Drop using the "Image First" strategy. Figure TT demonstrates the effectiveness of Capacity-Aware Inference under low capacity constraints (i.e., = 0.5). This is largely due to the high redundancy in image tokens (Chen al. 2024), which allows a higher dropping ratio without significantly affecting performance. The dominance of image tokens also enables the use of very low capacity factors without significantly affecting text token retention, as illustrated in Figure [T] Dropping image tokens higher ratios leads to more balanced token assignments and substantially improved inference efficiency.
6 CONCLUSION
In this paper, we identify the issue of imbalanced token- toeenttt--xpett assignment in Mixture of Ex- perts (MoE) models and introduce the Straggler Effect during inference, where high-load experts become efficiency bottlenecks and dictate overall latency. To address this problem, we propose Capacity- Aware Token Drop, which mitigates expert overload by enforcing strict capacity constraints. Additionally, to better utilize underloaded experts, we present Capacity- -Aare Expanded Drop, which allows tokens to select additional experts the same device while still respecting capacity limits, thereby improving expert utilization. Our findings and proposed methods offer valuable insights and effective strategies for improving MoE inference efficiency.
10

---

Published as a conference paper at ICLR 2026
7 ETHICS STATEMENT
Our work focuses on improving the efficiency of large-scale Mixture-of-Experts (MoE) models at inference time. The proposed methods are lightweight modifications to model execution and do not involve collecting or annotating new data. All experiments are conducted on publicly available pretrained models and standard open-source benchmarks. We emphasize that our contributions are intended solely for research and educational purposes in efficient machine learning. We will explicitly release code under a research license and include clear terms discouraging misuse in sensitive applications such as mass surveillance or privacy-intrusie deployments.
REPRODUCIBILITY STATEMENT
We are committed to ensuring full reproducibility of our results. Upon acceptance, we will release the complete implementation of Capacity-Awann Aware Inference, including inference- timee strategies (Token Drop and Expanded Drop), evaluation harness integration, and scripts for running all benchmarks. Detailed instructions will be provided in a reproducibility README covering environment setup, library versions, random seeds, and commands for end-to-end replication.
REFERENCES
Winogrande: An adversarial winograd schema challenge at scale. 2019.
Asma Ben Abacha, Wen-wai Yim, Yadan Fan, and Thomas Lin. An empirical study of clinical note generation from dloctor-patien encounters. In Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics, pp. 2291-2302, Dubrovnik, Croatia, May 2023. Association for Computational Linguistics. URL http::tts:/ssee hts_httss///aclnnhoolooo. wwtph///oooooo.o/22233333 macl-main.1oe 168
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019.
Weilin Cai, Juyong Jiang, Fan Wang, Jing Tang, Sunghun Kim, and Jiayi Huang. A survey mixture of experts, 2024. URL https: /tt:://aaxxvvvvoo http:::///axx..oor//b////////// 06204
Iñigo Casanueva, Tadas Temcinas, Daniela Gerz, Matthew Henderson, and Ivan Vulic. Efficient intent detection with dual sentence encoders. In Proceedings of the 2nd Workshop on NLP for ConvAl ACL 2020, mar 2020. URL https///iinbbboom/lyy---D-aaaaa
wwtpxxi.orrgabb//00000 04807 Data available saat.//ssssssscii..cdaassssssss
Liang Chen, Haozhe Zhao, Tianyu Liu, Shuai Bai, Junyang Lin, Chang Zhou, and Baobao Chang. An image is worth 1/2 tokens after layer 2: Plug-an aad--ndd-ply inference acceleration for large vision- language models, 2024.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss,, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojeiech Zaremba. Evaluating large language models trained code. 2021.
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding the mixture-of-experts layer in deep learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL http:://oppmrevieeeeeeeffoooonnp

---

Published as a conference paper at ICLR 2026
Christopher Clark, Kenton Lee, Ming-Was Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions, 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Rejichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture- of-experts language models. arXiv preprint arXiv:2401.0600 2024.
Matt Deitke, Christopher Clark, Sangho Lee, Rohun Tripathi, Yue Yang, Jae Sung Park, and et al. Molmo and pixmo: Open weights and open data for state-of-the-art multimodal models. arXiv preprint arXiv:2409. 141114499.. 2024.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference on Machine Learning, pp. 5547-5569. PMLR, 2022.
DeepSeek- AI et al. Deepseck-v2: A strong, economical, and efficient mixture-of-experts language model, 2024a. Wttphttss///axxx httpt//s///aaa..rr//aa////////// 2405. 04434
DeepSeek-AI AI et al. Deepseek-v3 technical report, 2024b. URL https://axxiv.orr///////1 19437
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39 2022.
Chaoyou Fu, Peixian Chen, Yunhang Shen, Yulei Qin, Mengdan Zhang, Xu Lin, Zhenyu Qiu, Wei Lin, Jinrui Yang, Xiawu Zheng, Ke Li, Xing Sun, and Rongrong Ji. Mme: comprehensive evaluation benchmark for multimodal large language models. ArXiv, abs/2306. 13394, 2023, URL
hhtpps:/app. semanticscholas semmannisssstoo.rrroog////rooo:: 259243922
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang. Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL httthttss://eeeoot wwtp://or///////ooe//oo/oosss/// 1025s66
Jamie Hayes, Ilia Shumailov, and Itay Yona. Buffer overflow in mixture of experts. In Neurips Safe Generative AI Workshop 2024, 2024. URL httt::::/openreevveee..etfoooooo
Shwai He, Liang Ding, Daize Dong, Boan Liu, Fugiang Yu, and Dacheng Tao. PAD-net: An efficient framework for dynamic networks. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Proceedings of the 6l Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 14354-14366 Toronto, Canada, July 2023, Association for Computational Linguistics. doi: 10. 1186633v1/20233 acn URL https::/aallatthooooo..o/////
scl-long- 803)
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding, 2021.
Robert A Jacobs, Michael Jordan, Steven Nowlan, and Geoffrey Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
12

---

Published as a conference paper at ICLR 2026
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXivv2401.04088, 2024.
Michael I Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2):181-2)) 1994.
Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of- experts from dense checkpoints. In The Eleventh International Conference Learning Represen- tations, 2023, URL htttp:://oppeeeeeeee
hett/forum?id=T5nuDDrrr.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming- Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association, for Computational Linguistics. 7:452-466 2019. doi: 10.1 00276. URL
httpss//aaclanthoooyy.or//////3 1026/
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.1666 2020.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. (GS) hard: Scaling giant models with condi- tional computation and automatic sharding. In International Conference on Learning Representa- tions, 2021. URL httpss:s///peeeevieew net (ttt//oorrmmmmire/ettttimmm
Bohao Li, Yuying Ge, Yixiao Ge, Guangzhi Wang, Rui Wang, Ruimao Zhang, and Ying Shan. Seed-bench: Benchmarking multimodal large language models. 2024 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 13299-13308, 2024a. URL https
//api.semanticschol.a org/ CttC//ooouS:::11155000 2711633388211
Tianle Li, Ge Zhang, Quy Duc Do, Xiang Yue, and Wenhu Chen. Long-context ilms struggle with long in-context learning, 2024b.
Yuanzhan Liu, Haodong Duan, Yuanhan Zhang, Bo Li, Songyang Zhang, Wangbo Zhao, Yike Yuan, Jiaqi Wang, Conghui He, Ziwei Liu, Kai Chen, and Dahua Lin. Mmbench: Is your multi- modal model all-around player? In European Conference Computer Vision, 2023, URL https://api. cishoolaa wwwp:ooppoooooii...... 259837088
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? new dataset for open book question answering, 2018a.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering, 2018b. URL htttt::://aaxvv..o/
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 Wettig, Binyuan Hui, Tim Dettmers, Douwe Kicla, Ali Farhadi, Noah A. Smith, Pang Wei Koh, Amanpreet Singh, and Hannaneh Hajishirzi. Olmoe: Open mixture-of-expprts language models, 2024. URL/https://al :/arxivv//axxv...o .rg. 2409. 02060
OpenAl. Gpt-4 technical report, 2024.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Newral Information Processing Systems, 0148583-8595, 2021.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of of-experts layer. arXiv preprint arXiv: /00l:/700.00000 065388 2017a.
13

---

Published as a conference paper at ICLR 2026
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The pparsely-gated mixture-ou ex-eressss layer, 2017b. URL https::/aaxii.orrg////a//////
1701. 06538
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro, Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint orXiv:190900805: 2019.
Yixin Song, Zeyu Mi, Haotong Xie, and Haibo Chen. Powerinfer: Fast large language model serving with consumer-grade gpu, 2023.
Gemini Team. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, 2024a.
Qwen Team, Swenl.5-moo Matching 7b model performance with 1/3 activated parameters", February 2024b. URL https://qwennm.. gi tet.t.i.obbla.......ooooiiii http:/looo//e/mmeo////
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. 2019. In the Proceedings of ICLR.
Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You. Openmoe: An early effort open mixture-of-expee language models. arXiv preprint arXiv:2402.01739, 2024a.
Leyang Xue, Yao Fu, Zhan Lu, Luo Mai, and Mahesh Marina. Moe-infinity: Activation-aware expert offloading for efficient moe serving, 2024b.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang. Jiaxi Yang. Jing Zhou, Jingren Zhou, Junyang Lin, Kai Dang, Keqin Bao, Kexin Yang, Le Yu, Lianghao Deng, Mei Li, Mingfeng Xue, Mingze Li, Pei Zhang, Peng Wang. Qin Zhu, Rui Men, Ruize Gao, Shixuan Liu, Shuang Luo, Tianhao Li, Tianyi Tang, Wenbiao Yin, Xingzhang Ren, Xinyu Wang, Xinyu Zhang, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yinger Zhang, Yu Wan, Yugiong Liu, Zekun Wang, Zeyu Cui, Zhenru Zhang, Zhipeng Zhou, and Zihan Qiu. Qwen3 technical report, 2025. URL pnttps://arxiit org. 2505. 09388
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence?, 2019.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du Yanping Huang, Vincent Y Zhao, Andrew M. Dai, Zhifeng Chen, Quoc V Le, and James Laudon. Mixture-of-expeee with expert choice routing. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL :ttp::///ppnreviee..eefoooo net/forum?id idJo1HIVinI
Tong Zhu, Xiaoye Qu, Daize Dong, Jiacheng Ruan, Jingqi Tong, Conghui He, and Yu Cheng. Llama-moe; Building mixture-of-experee from llama with continual pre-training. arXiv preprint arXiv:2406. 16554, 2024.
wrgpg/.////........ 16554
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2002.08906, 2022.
14

---

Published as a conference paper at ICLR 2026
A IMPLEMENTATION DETAILS
Models We mainly focus on lightweight MoE models (less than 20B parameter budget). We conduct experiments on OLMoE (Muennighoff et al. 2024), Qwen1.5-MoE (Team] 2024b), DeepSeek-V2- Lite fel al., 2024a). Mixtral (Jiang et al. 2024), MolmoE (Deitke et al., 2024). and Qwen3-MoE Yang et al. 2025). due their competitive performance and widespread applications.
Table 5: Experimental settings for eval- uated language tasks. "Norm" refers to the normalization performed with respect to the length of the input.
Datasets To evaluate model performance, we report normalized zero-shot or few-shot accuracy on the LM-
Hamess benchmark. The number of shots for each task is detailed in Table 5] which includes multiple tasks: ARC-C (Clark et al. 2018). BoolQ (Clark et al., 2019);
Task
Number few-shot
Metric
HellaSwag (Zellers et al.) 2019), MMLU (Hendrycks
BoolQ
Accuracy
Accuracy
et al. [021 1, OBQA (Mihaylov al., 2018a). PIQA
OBQA
Accuracy (Norm)
(Bisk et al.
2019). RTE (Wang et al. 2019). Wino-
PIQA
Accuracy (Norm)
MMLU
Accuracy
Grande (ai2
2019). and GSM8K (Cobbe al., (02I).
WineGrande
Accuracy
The evaluation code is based on EleutherAT's LM Har-
GSM8K
Exact Match
HellaSwag
Accurney (Norm)
ness framework (Gao et al. 2023).
ARC C
Accuracy (Norm)
B ADDITIONAL RESULTS ON LANGUAGE TASKS
Long-context understanding poses substantial challenges for language models. We therefore investi- gate whether Token Drop and Expanded Drop can maintain competitive performance in this setting. presents their results on long in-context learning benchmark. Following LongICL.encch (Li et al. 2024b), we evaluate on the BANKING77 dataset (Casanueva et al., 2020), with the only modification being that we sample 100 test examples. BANKING77 is an intent detection dataset in the banking domain with 77 classes. We consider 1-shot/label to 5-shot/label settings, corresponding to context lengths of 2K, 4K, 7K, 9K, and 14K tokens. Across different input sequence lengths, both Token Drop and Expanded Drop consistently match the performance of the baseline. This demonstrates the strong generalization ability and wide applicability of these techniques.
Table presents
Table 6: Performance of Token Drop and Expanded Drop on PrepSeeek-V2LLing in LongICLBench.
Method
2K
4K
7K
9K
11K
Avg.
Baseline
0.28
0.44
0.56
0.62
0.70
0.52
Token Drop
0.27
0.43
0.57
0.65
0.71
0.53
Expanded Drop
0.23
0.53
0.63
0.62
0.70
0.54
To further access the robustness of Token Drop and Expanded Drop in generation tasks, we test additional benchmarks such as HumanEval (Chen al. (021)), NQ-Open (Kwiatkowski et al. 2201))
and MTS- MTS-Dialog (Ben Abacha et al. 2023 in Table
Table 7: Performance of Token Drop and Expanded Drop across three MoE-based LLMs on Hu- EvnEval (HE), NQ-open (NQ), and MTS-Dialog Dialog (MTS). The capacity factor set as 1.0 here and we report pass @1 for HE, exact_matee motch for NQ and BERTSccre for MTS, respectively.
L..MoE-Instrur
Diven.55-MoE-Ce
DeepSeek V2-Lite-Chat
Method
HR
NQ
MTS
HE
NQ
MTS
HE
NQ
MTS
Baseline
0.311
0.175
0.827
0.280
0.226
0.874
0,506
0.265
0,830
Token Drop
0.299
0.159
0.823
0.281
0.214
0.869
0.494
0.252
0.831
Expanded Drop
0.303
0.170
0.829
0.291
0.224
0.872
0.508
0.268
0.831
The consistent performance on code generation, open-domain question answering, and dialog creation demonstrates the robustness of Token Drop and Expanded Drop for long-range generation tasks.
15

---

Published as a conference paper at ICLR 2026
C ADDITIONAL RESULTS ON MULTI-MODAL TASKS
In the scope of this paper, multimodal tasks refer to those involving both vision and language modali- ties. We evaluate model performance using three representative benchmarks: MME, MMBench, and SEED-Bench, each targeting different aspects of multimodal understanding and reasoning.
MME benchmark evaluates vision- Maioonlanguuae models along two dimensions: perception, which tests visual grounding and recognition, and cognition, which assesses reasoning abilities such as counting and relational understanding. It provides a lineggrined analysis of multimodal understanding.
MMBench is comprehensive benchmark designed to assess the general multimodal understanding ability vision-language models. It evaluates model performance across six core capabilities: Coarse Perception (CP), Fine-grained Perception-inctuddii single-instance (FP-S) and cross-instance (FP- C), Attribute Reasoning (AR), Logical Reasoning (L.R), and Relational Reasoning (RR). By covering both perception and reasoning-orienter tasks, MMBench provides detailed insights into the strengths and limitations of VLMs across diverse multimodal scenarios.
SEED-Bench a large-scale benchmark for evaluating the generative comprehension of Multimodal Large Language Models (MLLMs) across both image and video modalities. It includes 19K human- annotated multiple-choie questions spanning 12 evaluation dimensions, enabling objective and efficient assessment without human or GPT intervention. SEED-Bendh reveals model limitations and maintains public leaderboard to support fair comparison and future research.
Component
Content
Token Count
Image
576
Text Prompt Is artwork titled child with cathersne, cecilla, bartars, ursula? Flease
Table 8: An example multi-modal query in MME benchmark, showing the dominant proportion of image tokens compared to text tokens.
As shown Table these tasks typically introduce large number of image tokens. When faced with imbalanced token-to-expert assignments, dropping redundant image tokens significantly improves load balancing. Moreover, due to the high redundancy among image tokens, dropping a portion of them has minimal impact model performance.
As in MME and MMBench, the Image-First variants Token Drop and Expanded Drop also exhibit consistent effectiveness SEED-Bench (Li et al. 2024a),, maintaining strong performance even under low capacity factors such as 0.5. In addition to the redundancy in image tokens, Figure I0 shows that text tokens constitute only small portion of the total token assignments. This allows the regulation of to retain almost all text tokens under the Image- First strategy.
Method
Inst. Attr:
Inst. ID
Inst. Interact.
Inst. Loc.
Inst. Count
Scene
Spatial
Text
Reasoning
Overall
Baseline
00
74.2
71.4
58.8
62.8
$7.0
73.5
49.6
72.6
76.4
68.7
Token Drop
70.4
67.8
60.8
$7.8
51.8
71.5
43.5
58.3
71.3
64.9
0.5
Expanded Drop
71.2
67.3
57.7
58.8
53.6
71.0
45.2
64.3
71.9
65.5
Teken Drop
1.0
73.6
70.2
58.8
62.6
56.9
72.7
48.1
61.9
73.4
68.0
Expanded Drop
Drop
58.8
73.1
74.6
1.5
Expanded Drop
73.7
61.9
57.4
73.3
49.3
76.1
68.7
Table 9: Token Drop and Expanded Drop strategies for multi-modal MoE models evaluated on SEEDD-Bncch. Abbreviations: Inst. Attr. Instance Attributes; Inst. ID Instance Identification; Inst. Interact. = Instance Interaction; Inst. Loc. = Instance Localization; Inst. Count Instance Counting.
16

---

Published as a conference paper at ICLR 2026
D ABLATION STUDY
Model-Specific Imbalanced Property explore the imbalance property in various models, such as OLMoE, DeepSeek, Qwen and Mixtral, which differ in both architectures (e.g.. depth and width) and training strategies (c.g.. training from scratch (Muennighoff et al. 2024 ct al. 2024a) vs. training
after upcycling (Jiang al. 2024 Team 2024b)).
On the one hand, our findings in Appendix reveal different training strategies result in significantly varying levels of imbalance. Specifically. MoE models trained from scratch exhibit much higher degree of imbalance. For instance, OLMoE and DeepSeek-V2442 V2-Lite experience peak expert-wise token allocations exceeding 5.,, whereas Qwen Tiven1.5-Ma 1.5-MoF and Mixtral upeycled from dense language models, maintain a more balanced distribution, with peak expert-wise allocations staying below 3N. This is because upcyeling initializes all experts with identical parameters (Komatsuzaki et al. 2023 reducing divergence and promoting balanced training in the early stages.
On the other hand, despite the widespread of auxiliary balance loss in MoE training, it does guarantee balanced token assignments across experts, as token distribution still varies significantly during inference on test data. This necessitates integrating expert capacity into the inference process.
WinoSronde
A
HeloSwag
MULU
(a) OLMoE
(b) OLMMoE-Inttrct
(c) @went
(d)
Figure 12: Performance change capacity factors decrease from 3.0 0.0.
Capacity Factor Beyond the specific capacity values presented in Table we further investigate wide range of capacity factors Figure spanning from 0.0 3.0. We exclude values exceeding 3.0, their performance closely aligns with rapacity-agnostee changes when decreasing the capacity factor, we find that setting to 1.5 sufficient to maintain performance comparable to the original models. However, maintaining performance becomes challenging under excessively low capacity factors, as high-load experts are forced to drop a significant number tokens.
scenarios. By analyzing the performance
Speedup Measurement Under Different Workloads To assess the effectiveness of Token Drop in efficiency, measure speedup under a range of token workloads. Specifically, we vary both the sequence length and batch size, and present the speedup results on DeepSeek- V2 with a capacity factor of 2.0 in Table 10
Batch Size
8K
8K
4K
2K
2K
IK
Prompt Length
0.1K
0.2K
0.4K
IK
IK
2K
IK
Speedup
x.09
1.18x
1.24x
1.26x
1.27x
1.27
1.27x
1.24x
1.23x
Table 10: Speedup results across varying batch sizes and prompt lengths.
The straggler effect becomes more pronounced under heavier workloads, where GPUs operate at higher utilization with limited spare capacity, making the speedup more noticeable. In practical server-side MoE deployments, workloads are substantially higher, and our techniques effectively mitigate the resulting straggler effect.
Maximum Number Selected Experts Expanded Drop adopts an overselection--aa--drpiiiii strategy that not only maintains load balance but also improves the utilization of underloaded experts. Although this mechanism allows some tokens to select more than & experts, Table II demonstrates that such flexibility benefits downstream task performance, suggesting that enforcing strict maximum of experts is unnecessary. Allowing additional expert selections can enhance representational capacity, whereas rigid constraints may unnecessarily limit model performance.
17

---

Published as a conference paper at ICLR 2026
Table 11: Ablation study on limiting the maximum number of ks selected experts. "w/max" and
"wo max" indicate runs with and without this constraint, respectively. y is set to 1.0.
Method
OBQA
PIQA
RTE
WinoGrande
BoolQ
ARC-C
HellaSwag
MMLU
Avg.
w/ max
42.4
75.8
68.6
72.7
50.3
77.1
47.6
61.0
2/
max
42.2
75.9
53.4
69.7
73.1
50.3
77.0
47.8
61.2
3k
max
42.0
75.6
72.7
77.0
61.0
w/o max
42.0
75.6
53.8
69.8
72.9
50.3
77.1
47.6
61.1
E IMPLEMENTATION-LEVEE CODE FOR TOKEN DROP AND EXPANDED DROP
The detailed implementation-lewe code for Token Drop and Expanded Drop presented in Algo- rithms 1 and respectively.
Algorithm Token Drop
Algorithm 2 Expanded Drop
# x: input tokens;
x: input tokens;
# k: top-k experts;
k: top-k experts;
ganma: capacity factor;
ganna: capacity factor;
local_ids: local device expert ids;
def token_drop(self,
ganna):
logits self. gate(x)
def expanded_droppselff, x, ganma, local_ids):
scores torch. softsssstoxxxllitss dims-1)
logits self.gate(x) scores = torch. softmax dimF-1)
topk scores, topk_idx scores. topk(k, dim=1)
topk_scores, topk_idx topkekk dim=1)
topk_maae jask torch. (it.:eros__ikeesoorressssecc
topk_idx, 1)
N, scores.size()
masked_scores scores topk_mask
local_idx torch. tensor(local__de devicemx.
repica) 1) exp_idx torch.cat((topk.ida idx, local_idx], dim =1)
N, scores. size cap int (gamma
-. keep_idx = masked_scored seoisedd__oeeeeccee topk(cap, dim=0) cap_mask = zorch. likesoes. scatter keep_idx, 1)
local_scores scores(:, local_ids] exp_scores = torch.cat((topk_scom)
local_scores], dim=1)
final_map topk_mask cap_mask
exp_mask torch. zeros_ (ilkeeeeorrsss.aatttte..e
final_scores masked_scores final_map
exp_idx, 1)
masked_scores scores exp_nask
final_scores final_map
cap int (N k)
keep_idx masked_score topk dim=0)
cap_mask torch. zeros_ scatter
keep_idx, 1)
final_map exp_mask cap_mask final_scores masked_scores final_map
return final_scores, final_map
F WASE EXPERT LOAD
To analyze imbalanced token assignments, measure the expert load for each expert by tracking the peak expert load while running MoE models on various test datasets. Figure 13 15 and 16 present
the full results for the normalized layer-wise expert load for OLMoE, DeepSeek V2, QwenI.5-Mot. and Mixtral-8 - Instruct, respectively.
G CALCULATION DROPPED TOKENS
Based on Equation II we calculate the total number dropped tokens across experts in each layer under different capacity factors, illustrated in Figures 17
19
18
and 20
18

---

Published as a conference paper at ICLR 2026
Layer
Layer
Layer
ExpertID Layer
ExpertID
ExportI Layer
ExpertI0 Layer 0
ExportID Layer
Layer 10
10N0
ExpercIO
ExpertID
ExpertID Layer
Layer
ExportID Layer 15
Experti Layer 16
105
Enpor:ID
Exper:ID
BoolQ
Figure 13: Layer-wise expert load in Instructt
19

---

Published as a conference paper at ICLR 2026
Layer 2
Layer 3
Expert tD Layer 5
Layer 4
Lowd
Ex Layer? 7
Layer 6
Layer
LoNE
Expart Layer
ExowrL
Layer 11
Expert
Layer
Layer
Export 1D
Ex
Layer
Layer
Espeet In
Expert ID
Layer
Layer 17
Lo0d
Expert in Layer 18
Ex
Layer 19
Lood
Expart ID
Ex
Layer 20
Layer 21
Expert ID
Layer 22
Layer 23
Loud
Layer 24
Layer 25
Layer 26
Layer 27
15
Export 1D
Ex
BoolQ
Figure 14: Layer wise expert load in Deepseek- V2-Lite.
20

---

Published as a conference paper at ICLR 2026
Layer
Layer 2
Layer
Layer
ExportIO
1000
ExpertI
Expert
Layer
1002
EDperc
Exseet ID Layer
Layer
Lond
Exper:10
Expert Layer 12
Layer
ExportIO
ExpertID Layer 14
Layer
16
1000
ExportI
ExportID Layer
Layer
Lowd
ExpertID
Expert Layer
Layer
Load
ExperiID Layer 21
Expeett
Layer 22
Layer
Layer 24
15
Expar:I0
ExpertID
BoolQ
Figure 15: Layer-wise expert load in Qwen1.5-MoE. 15-MoEEChaa
21

---

Published as a conference paper at ICLR 2026
Layer
Layer 2
Expert
Expert 10
Layer
Layer
Layer
Export ID Layer
Export ID
Layer 8
Expert 1b
Expert 10 Layer
Layer
Layer 11
Layer 12
Expert ID Layer 13
Expert ID
Layer 14
Expert
Expert
Layer
Layer 17
Layer 18
Expert ID Layer 19
Expert ID
Layer 20
100
Expar 22
Expert Layer 23
Layer 24
Expert ID Layer
Expert 1D
Layer 26
Lood
Layer 28
Expert Layer 29
Layer 30
Expert ID Layer
Export 1D
Layer 32
Exsert
Expor
Figure 16: Layer-wise expert load in Mixtral-8x 7B-Instruct
22

---

Published as a conference paper at ICLR 2026
Layer
Layer
320
Oropec taharrt
20-
Layer
Layer
326
Droppi Mlate PM
Layer
Layer
Stopped Scarsr
Layer
Layer
320
Dropect Towee 1N6)
50-
Layer
Layer 10
320
Doposd Mlare
Layer
Layer
Dropped
Layer
Layer 14
320
Jeppos Towns 1h1
50
Layer 15
Layer 16
326
Dippred
OBQA
PIQA
WinoGrande
Baolq
Figure 17: Dropped tokens with respect to capacity factors in OLMoE-Instruet.
23

---

Published as a conference paper at ICLR 2026
Layer 2
Layer
07/
Y
Layer 4
Layer
07 1941
Y
Layer
Layer 7
DT 190
Layer
Layer
7541
Y
100
Y
Y
Layer
Layer
100
Y
Layer
Layer
100
DT (5)
Y
Y
Layer
Layer 17
Y
Y
Layer
Layer 19
D7 PNC
Y
Layer 20
Layer 21
OT 754
Y
Layer 22
Layer 23
DT 158
Layer
Layer 25
STEMN
Layer 27
100
2.0
1.0
1.5
2.0
3.0
35
0.0
0.5
2.0
25
OBQA
PIQA
WinoGrande
BoolQ
Figure 18: Dropped tokens with respect to capacity factors in -eepSeek-V2-Chat.
24

---

Published as a conference paper at ICLR 2026
Layer 1
Layer
(N4)
Y Layer
Layer
170
Layer
Layer
100
D7 152
Y
Y Layer
Layer
100
07 090
Y
Y
Layer
Layer 10
(54)
Y
Layer 11
Layer 12
Y
Layer 13
Layer 14
100
Y
Layer 16
100
Y
Y Layer
Layer
100
Y
Y
Layer
Layer 20
Y
Layer 21
Layer 22
DT
Layer 23
Layer 24
100
3.0
1.5
2.0
3.0
35
0.0
0.5
2.0
25
OBQA
PIQA
WinoGrande
BoolQ
Figure 19: Dropped tokens with respect to capacity factors in Qwen-I.5-MoE-CChat
25

---

Published as a conference paper at ICLR 2026
Layer 2
Layer
07/
Y
Layer 4
Layer
07 1941
Y
Layer
Layer 7
DT 190
Layer
Layer
7541
Y
100
Y
Y
Layer
Layer
100
Y
Layer
Layer
100
DT (5)
Y
Y
Layer
Layer 17
DF 1M4
Y
Y
Layer
Layer 19
D7 PNC
Y
Layer 20
Layer 21
OT 754
Y
Layer 22
Layer 23
DT 152
Layer
Layer 25
STEMN
Layer 27
100
2.0
1.0
1.5
2.0
3.0
35
0.0
0.5
2.0
25
OBQA
PIQA
WinoGrande
BoolQ
Figure 20: Dropped tokens with respect to capacity factors in Mixtral-8x7B- Insttrtttt....
26


Reviews:
[{"review_id": "8WduDumeMt", "reviewer": "Reviewer_yZTe", "summary": "The authors propose a method for reducing the inference time of mixture-of-experts (MoE) models by load balancing during inference. In MoE models, tokens are typically routed to top-k most relevant experts and uneven loads across experts significantly slows down inference since sync operations are performed after the slowest expert has completed (termed ‘straggler effect’). The work proposes dropping tokens associated with overloaded experts and potentially re-routing them to underloaded ones to reduce the peak load on any single expert. Experiments on several MoE models and benchmark datasets suggest that the proposed method results in huge speedups with minor drop in performance.", "questions": "Please refer to the weakness section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is mostly well written and easy to understand. The authors identify a specific issue with MoE models and propose a solution to address it. The idea is simple, interesting and well motivated. The experiments are thorough, with results on multiple models, benchmark datasets and ablations and analysis to understand the different components of the proposed method. The authors show that existing models suffer from the ‘straggler effect’ and that the proposed solution can alleviate it, achieving 7-18% speedup with up to 0.4% points drop in average performance. The solution is also extended to multi-modal models, where huge gains can be achieved by dropping the redundant tokens in the image domain.", "weaknesses": "The work does not have any major weaknesses. \n\n1. There is not much discussion/comparison with other approaches to address the ‘straggler effect’. For instance, the authors mention the use of additional GPUs for overloaded experts. For commercial use-cases where there are thousands/millions of concurrent requests to a service, would just having extra servers for overloaded experts be efficient or would the proposed solution still be necessary given that it comes with a performance drop? The proposed method is also shown to provide the biggest improvements when the number of experts per GPU is reduced by increasing the number of GPUs. \n2. Discussion on the effect of different sequence lengths within a batch and dataset is missing. Since sequences in a batch can be of different lengths and the proposed method prunes tokens at the batch level, how does the variance in length affect efficiency and performance? Also, how would a high variance in sequence lengths within a dataset affect the speed and performance?\n3. The expanded token drop method (in the Methodology section) is a bit unclear. How exactly are the additional ‘m’ experts chosen and how are they ensured to be ‘local’ experts? \n4. Since some experts specialize towards particular domains, the proposed method might have a bigger drop in performance on specific tasks/datasets. For instance, in table 2, the maximum drop in performance over the 8 datasets on DeepSeek-V2-Lite-Chat model is 2.4% points while the average drop is 0.4% points (although this seems to be a one-off scenario in the table). How can this be predicted / handled?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "K9xcIKiiDL", "reviewer": "Reviewer_TfSK", "summary": "This paper identifies that Mixture-of-Experts (MoE) models exhibit highly imbalanced token-to-expert assignments during inference, leading to inefficient inference computation. To address this, the authors propose Capacity-Aware Token Drop, where overloaded experts drop excess tokens beyond the compputation budget. To further compensate for potential performance loss and utilize idle capacity on underloaded experts, they introduce Capacity-Aware Expanded Drop, allowing tokens to leverage additional underutilized experts on the same device. Experiments across multiple models demonstrate that the proposed methods achieve notable inference speedups with minimal or no performance degradation.", "questions": "## Writing clarify\n\n* Figure 1: The definition of \"normalized load\" should be explained so that the figure is self-contained; it’s not clear until Section 3.\n\n* Which dataset and model are used in Figure 1?\n\n* Line 204 mentions that the capacity constraint is applied per device rather than per expert. However, later discussion (e.g., Line 217: \"each overflowed expert selectively discards those with lower scores\") seems to describe per-expert constraints. This discrepancy makes the algorithm unclear.\n\n* Token Drop: If I understand correctly, a token could potentially end up with no assigned expert after dropping, if its scores assigned to each expert are ranked below the threshold. How does the algorithm ensure each token still has a valid output?\n\n* Expanded Drop: It seems that each token additionally selects local experts (beyond the normal top-k experts), and these experts then drop tokens following the same rule as in \"Token Drop Regulates the Latency of High-Load Experts\" section. In that case, could an expert retain some tokens assigned as local experts while dropping others from the top-k assignment? Please clarify.\n\n* Does Table 1 include the Expanded Drop technique?\n\n* In Figure 5, how many experts per GPU are used for models other than Mixtral?\n\n* In Figure 6, do the x-axis labels \"1.0,\" \"1.5,\" and \"2.0\" correspond to the \\gamma values?\n\n* In Figure 9, which model and dataset are used?\n\n* Line 448 states that \"we also experiment with a strategy that prioritizes dropping image tokens before selectively removing text tokens (“Image First”).\" Could the authors explain the implementation details of this strategy?\n\n## Experiments\n\n* Eq. 4 resents the latency of an MoE layer. The intuition is sound, but latency in practice depends on many factors. It would strengthen the paper to:\n    * (1) Measure actual layer latency and compare it with this theoretical bound; and \n    * (2) Measure end-to-end model latency to verify whether max({N_i}i=1^n) remains the dominant factor.\n\n* Line 340: The claim that \"As illustrated in Figure 5, for Mixtral-8×7BInstruct, deploying a single expert per GPU maximizes the effectiveness of capacity-aware inference.\" requires experiments with multiple experts per GPU for the Mixtral model. Currently, Figure 5 only reports results for a single expert per GPU for the Mixtral model.\n\n* Line 352 states that \"Notably, the duration of permutation and communication increases when tokens are expanded across a range of global experts.\" Could the authors clarify what it means and how to see it from Figure 6.\n\n* Line 373: The paper claims \"regulating the maximum capacity has a limited impact on the overall number of accommodated tokens,\" yet Figure 7 shows that varying \\gamma dramatically changes the percentage of dropped tokens from nearly 100% to 0%. Please clarify.\n\n* Figure 3: Why does “Expanded Drop + Image First” significantly outperform the baseline? Additional explanation would be helpful.\n\n## Other minior issues\n\n* The citation format is incorrect in many places. For example, \\citep should be used in:\n    * \"In recent years, the rapid evolution of Large Language Models (LLMs) OpenAI (2024); Team (2024a); DeepSeek-AI et al. (2024b) ...\"\n    * \"Specifically, MoE Shazeer et al. (2017b); Fedus et al. (2022) enhances scalability ...\"\n    * \"We conduct experiments on OLMoE Muennighoff et al. (2024), Qwen1.5-MoE Team (2024b), DeepSeek-V2-Lite DeepSeek-AI et al. (2024a) and Mixtral Jiang et al. (2024)\"\n    * \"The number of shots for each task is detailed in Table 4, which includes multiple tasks: ARC-C Clark et al. (2018), BoolQ Clark et al. (2019), HellaSwag Zellers et al. (2019), MMLU Hendrycks et al. (2021), OBQA Mihaylov et al. (2018), PIQA Bisk et al. (2019), RTE Wang et al. (2019), WinoGrande ai2 (2019) and GSM8K Cobbe et al. (2021). The evaluation code is based on EleutherAI’s LM Harness framework Gao et al. (2023).\"\n\n* Section 1 claims the identification of \"the Straggler Effect caused by token imbalance at inference time in Mixture of Experts\" as a contribution. However, as discussed in the paper, this phenomenon is already mentioned in prior work (e.g., DeepSeek-V3). It should be framed as a motivation rather than a contribution.\n\n* Line 706: Add a comma before \"and\"\n\n* Line 713: Add a comma before \"and\"", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* The experimental section is rich and goes beyond reporting results—it includes several insightful analyses that provide independent value.\n\n* The proposed approach is simple  and easy to implement.", "weaknesses": "* When the number of experts per device exceeds one, the observed speedup becomes limited.\n\n* The writing could be clearer in several places (see “Questions” below).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SMa9wppGO0", "reviewer": "Reviewer_vtDH", "summary": "This paper identifies the **Straggler Effect** in Mixture-of-Experts (MoE) models during inference: overloaded experts cause global latency bottlenecks due to imbalanced token-to-expert assignments. To mitigate this, the authors propose **Capacity-Aware Inference**, introducing two lightweight strategies:\n\n1. **Token Drop**: enforces capacity limits on overloaded experts by discarding low-importance tokens, reducing latency with minimal performance loss.\n2. **Expanded Drop**: allows tokens to select additional local experts, improving utilization of underloaded experts and enhancing performance.\n\nExperiments on both language and multimodal MoE models show significant improvements in inference efficiency (up to **1.85× speedup**) with comparable or even slightly better accuracy. The methods are training-free, easy to implement, and broadly applicable.", "questions": "1. What is the specific model you used in Figures 1 and 2? You should provide the detailed information in the caption texts.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper is well written and organized, with clearly stated motivation and method, together with relatively comprehensive experiments.\n2. This paper investigated both LLM and MLLM and demonstrated efficiency gains, demonstrating its effectiveness and universality.", "weaknesses": "1. The novelty of this paper is insufficient, since there has been comprehensive research on load balancing in MoE model training, such as Switch Transformer (Google) and AuxFree load balancing (DeepSeek). This paper does not have fundamental distinctions from these well-known papers, where the only difference is taming the load balancing issue during the inference stage in a training-free manner.\n2. In Lines 322-323, you should clarify that DP is applied to attention modules while EP is applied to MoE modules. \n3. You should specify the prefilling stage or decoding stage in your efficiency-related experiments with the listed workload, such as batch size and context length, which include Figures 4, 5, and 6.\n4. Also, you should provide more comprehensive experimental results on efficiency benefits under different workloads. For example, you should scale the context length with a fixed batch size, and scale the batch size with fixed context length, to compare and visualize the speedup or latency, and therefore provide more insights with a scaled workload.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XBrVsTBd5T", "reviewer": "Reviewer_yGAq", "summary": "The core problem addressed in this paper is the Straggler Effect in the inference phase of MoE systems: due to the imbalanced allocation between tokens and experts, overloaded experts become the bottleneck for inference latency, causing underutilized experts to remain idle while waiting, thereby reducing overall efficiency.\nThe key solution proposed is a Capacity-Aware Inference framework, consisting of Capacity-Aware Token Dropping and Capacity-Aware Expanded Dropping. The core concept involves discarding low-importance tokens for overloaded experts while increasing token allocation for underutilized experts. It achieves an average 1.85x inference speedup when applied to Mixtral-8×7B-Instruct.", "questions": "1. Regarding the multiple experts per GPU scenario, the authors mention that aggregated workload diminishes the method's effectiveness but provide no specific improvement strategies. Have you considered dynamically adjusting the capacity limits for experts within a single GPU?\n2. In the token dropping strategy, the gating score G(x) is used as the importance metric, but its reliability during inference has not been sufficiently validated. What about other important metrics (such as token's contextual contribution or semantic relevance), and how robust is the gating score across different tasks (e.g., logical reasoning, sentiment analysis)?\n3. The experimental results show model inference speed improvements rely on reduced expert computation time, but fail to analyze changes in communication overhead. Does Expanded Drop's expansion of candidate expert sets increase token scheduling overhead within local devices?\n4. Are there no related works on MoE inference acceleration? For example, https://github.com/LINs-lab/DynMoE", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The significance of the addressed problem lies in being the first to explicitly identify the \"Straggler Effect\" during MoE inference, and through quantitative analysis (e.g., revealing how some experts bear 7× the expected load N), it demonstrates its critical impact on inference latency.\n2. The methodology is designed to be efficient: both strategies involve lightweight modifications during the inference phase, requiring neither adjustments to the model training process nor additional computational resources (such as the expert redundancy deployment in DeepSeek-V3). Token Dropping rapidly alleviates overload by discarding low-importance tokens, while Expanded Drop utilizes local expert expansion of candidate sets to avoid cross-device communication overhead.", "weaknesses": "1. The long-term impact of token dropping remains unvalidated: existing experiments only evaluate performance changes on static tasks (e.g., MMLU, OBQA), without investigating its effects on long-context generation (such as dialogue and text creation). For instance, whether persistent dropping of low-importance tokens may lead to degraded contextual coherence and generation quality remains unexamined, lacking verification in dynamic scenarios.\n2. The method's generalizability has limitations: experiments show the approach achieves optimal results with models deploying \"single expert per GPU\" (e.g., Mixtral-8×7B-Instruct), but exhibits diminished acceleration effects for models with \"multiple experts per GPU\" (e.g., OLMoE-Instruct with 8 experts/GPU). The aggregated workload across multiple experts reduces the benefits of constraining straggler experts. The authors suggest \"increasing GPU count to reduce experts per GPU\" without proposing adapted solutions for multi-expert single-GPU scenarios, thus limiting universal applicability.\n3. The study lacks comparison with existing inference optimization methods: it fails to benchmark the proposed approach against other MoE inference optimization solutions (such as PowerInfer's expert offloading or MoE-Infinity's activation-aware scheduling).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)