Faster, Smaller, and Smarter: Task-Aware Expert Merging for Online MoE Inference
Ziyi Han, Xutong Liu, Ruiting Zhou, Xiangxiang Dai, John C. S. Lui
Abstract
Sparse Mixture of Experts (SMoE) has become a preferred architecture for scaling Transformer capacity without increasing computational cost, as it activates only a small subset of experts for each input. However, deploying such an approach for online inference remains challenging due to the large size of a full SMoE model and the complexity of expert routing, especially in resource-constrained edge networks. Moreover, during the online inference, task information is often unavailable, making the task-level routing error-prone. In this work, we propose a novel tree-structured adaptive neural bandit router, Tanbr, to enable efficient and reliable online MoE inference. Instead of relying on explicit task tags, Tanbr estimates the task distribution over time from historical data and uses it to guide task-aware expert merging within a given pre-trained MoE. To handle the large continuous space of merging weights, Tanbr employs a binary tree to progressively partition the space and generate finer candidate weights. It then applies a neural bandit to learn the non-linear mapping from merging weights to model performance and decides optimal expert merging. We prove that Tanbr achieves a sublinear regret bound of over T rounds, despite operating over a continuous decision space, matching regret bounds compared to existing methods. Extensive experiments show that Tanbr reduces inference latency by at least 45% and memory usage by up to 25%, while maintaining a high accuracy compared to many state-of-the-art methods.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 377a9648-8f23-41b5-b1c8-ee19ea88dc93Builds on21
- Merging Models with Fisher-Weighted AveragingMichael Matena, Colin RaffelNeurIPS 2022 · 741 citations
- Model Fusion via Optimal TransportSidak Pal Singh, Martin JaggiNeurIPS 2020 · 330 citations
- Neural Contextual Bandits with UCB-based ExplorationDongruo Zhou, Lihong Li, Quanquan GuICML 2020 · 329 citations
- Unified Scaling Laws for Routed Language ModelsAidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch et al.ICML 2022 · 266 citations
- AdaMerging: Adaptive Model Merging for Multi-Task LearningEnneng Yang, Zhenyi Wang, Li Shen, Shiwei Liu et al.ICLR 2024 · 230 citations
Related papers
- Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing PolicyPingzhi Li, Zhenyu Zhang, Prateek Yadav, Yi-Lin Sung et al.ICLR 2024 · 97 citations
- TA-MoE: Topology-Aware Large Scale Mixture-of-Expert TrainingChang Chen, Min Li, Zhihua Wu, Dianhai Yu et al.NeurIPS 2022 · 31 citations
- TT-LoRA MoE: Using Parameter-Efficient Fine-Tuning and Sparse Mixture-Of-ExpertsPradip Kunwar, Minh N. Vu, Maanak Gupta, Mahmoud Abdelsalam et al.SC 2025 · 1 citation
- A Provably Effective Method for Pruning Experts in Fine-tuned Sparse Mixture-of-ExpertsMohammed Nowaz Rabbani Chowdhury, Meng Wang, Kaoutar El Maghraoui, Naigang Wang et al.ICML 2024 · 18 citations
- Tight Clusters Make Specialized ExpertsStefan K. Nielsen, Rachel S. Y. Teo, Laziz U. Abdullaev, Tan Minh NguyenICLR 2025
