A Replicate-and-Quantize Strategy for Plug-and-Play Load Balancing of Sparse Mixture-of-Experts LLMs

2026-02-23

AI Summary

This work addresses the severe expert load imbalance in sparse mixture-of-experts (SMoE) large language models during inference, which leads to inefficient resource utilization. The study presents the first systematic analysis of SMoE routing behavior and introduces a plug-and-play, training-free dynamic load rebalancing framework. By duplicating frequently accessed experts, quantizing low-importance ones, and leveraging a lightweight calibration set to evaluate expert importance and compute a Load Imbalance Score (LIS), the method enhances parallel efficiency while adhering to the original memory budget. Experiments across multiple SMoE architectures and benchmarks demonstrate up to a 1.4× improvement in load balance, with accuracy variations tightly controlled within ±0.6%.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Planning, Routing, and Scheduling: Planning with Language Models
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Sparse Mixture-of-Experts (SMoE) architectures are increasingly used to scale large language models efficiently, delivering strong accuracy under fixed compute budgets. However, SMoE models often suffer from severe load imbalance across experts, where a small subset of experts receives most tokens while others are underutilized. Prior work has focused mainly on training-time solutions such as routing regularization or auxiliary losses, leaving inference-time behavior, which is critical for deployment, less explored. We present a systematic analysis of expert routing during inference and identify three findings: (i) load imbalance persists and worsens with larger batch sizes, (ii) selection frequency does not reliably reflect expert importance, and (iii) overall expert workload and importance can be estimated using a small calibration set. These insights motivate inference-time mechanisms that rebalance workloads without retraining or router modification. We propose Replicate-and-Quantize (R&Q), a training-free and near-lossless framework for dynamic workload rebalancing. In each layer, heavy-hitter experts are replicated to increase parallel capacity, while less critical experts and replicas are quantized to remain within the original memory budget. We also introduce a Load-Imbalance Score (LIS) to measure routing skew by comparing heavy-hitter load to an equal allocation baseline. Experiments across representative SMoE models and benchmarks show up to 1.4x reduction in imbalance with accuracy maintained within +/-0.6%, enabling more predictable and efficient inference.

Problem

• Sparse Mixture-of-Experts
• load balancing
• inference-time
• expert routing
• workload imbalance

Innovation

• Replicate-and-Quantize
• Sparse Mixture-of-Experts
• load balancing
• inference-time optimization
• expert quantization