🤖 machine learning

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

This paper proposes Replicate-and-Quantize (R&Q), a training-free, inference-time framework that dynamically rebalances load in Sparse Mixture-of-Experts models by replicating frequently selected experts and quantizing others, thereby significantly reducing routing skew while maintaining model accuracy within a fixed memory budget.

Original authors: Zijie Liu, Jie Peng, Jinhao Duan, Zirui Liu, Kaixiong Zhou, Mingfu Liang, Luke Simon, Xi Liu, Zhaozhuo Xu, Tianlong Chen

Published 2026-07-15

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Zijie Liu, Jie Peng, Jinhao Duan, Zirui Liu, Kaixiong Zhou, Mingfu Liang, Luke Simon, Xi Liu, Zhaozhuo Xu, Tianlong Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine the world of artificial intelligence as a bustling, high-tech kitchen where giant robots are trying to cook the perfect meal for millions of hungry people at once. To handle this massive demand, these robots don't just have one giant brain; they have a team of specialized chefs, each an expert in a specific ingredient or recipe. This setup is called a "Sparse Mixture-of-Experts" (SMoE). Instead of every chef cooking every single dish, a smart manager (called a "router") looks at each order and sends it to the few chefs who are best at that specific task. This is incredibly efficient, allowing the robots to be huge and smart without needing a kitchen the size of a city.

However, there's a catch. Just like in a real kitchen, the manager sometimes gets a bit biased. If a customer orders a popular dish, the manager might keep sending all those orders to the same "star chef," while the other talented chefs stand around with empty hands, waiting for work. This is called "load imbalance." The star chef gets overwhelmed and slows down the whole line, while the others waste energy doing nothing. This paper tackles a very specific problem: how to fix this mess while the kitchen is already open for business, without firing anyone, hiring new staff, or rewriting the manager's rulebook.

The Problem: The Overworked Star Chef

The researchers noticed that even though these AI kitchens are designed to be fair, they often get stuck in a rut. When a large group of orders comes in at once (a "batch"), the manager tends to funnel almost everything to a tiny group of "heavy-hitter" experts. These experts become the bottleneck, slowing down the entire system, while the rest of the team sits idle.

The team first checked if simply retraining the manager to be fairer would fix the problem. They tried various training tricks, like adding penalties for unfairness, but found a frustrating trade-off: making the manager fairer often made the robot's answers less accurate. It was like trying to force a chef to stop cooking their signature dish so everyone else could cook; the food just didn't taste as good. They also discovered that the problem gets worse when you try to serve more people at once. As the batch size grew, the imbalance exploded, leaving the system inefficient and slow.

The Discovery: Popularity Isn't Everything

Before building a solution, the team made a surprising discovery. They looked closely at which experts were getting the most work and compared that to which experts were actually the most important for getting the right answer. They found that being popular doesn't mean being important.

Some experts were getting a massive amount of work (high load) but, if you removed them, the robot's performance barely dropped. Conversely, some experts were rarely called upon but were absolutely critical when they were used. The manager was confusing "frequent requests" with "high value." This meant the team didn't need to treat all experts the same; they could treat the overworked ones and the underused ones differently.

The Solution: The "Replicate-and-Quantize" Strategy

To fix the bottleneck without retraining the whole system, the authors proposed a clever, plug-and-play strategy called Replicate-and-Quantize (R&Q). Think of it as a dynamic kitchen rearrangement that happens instantly when the orders start coming in.

Replicate the Heavy-Hitters: When the system spots an expert who is drowning in work (the "heavy-hitter"), it doesn't fire them. Instead, it creates a "clone" of that expert. But here's the trick: the clone is a "lightweight" version. It's the same chef, but they are working with a slightly smaller, more efficient toolkit (using lower-precision math, or "quantization"). This allows the system to split the massive pile of orders between the original chef and the clone, doubling the speed without needing a bigger kitchen.
Quantize the Underused: To make room for these new clones without running out of memory, the system looks at the experts who are barely doing anything. These "less important" experts are also given the lightweight toolkit. Since they aren't handling many orders, shrinking their toolkit doesn't hurt the quality of the food.

By doing both at the same time—splitting the load of the busy chefs and shrinking the tools of the idle ones—the system stays within its original memory budget but runs much faster and more fairly.

The Results: A Balanced Kitchen

The team tested this strategy on several different AI models and a wide variety of tasks, from solving math problems to answering tricky questions about the world. They measured the imbalance using a new score they invented called the Load Imbalance Score (LIS), where a score of 1 is perfect balance and higher numbers mean more chaos.

The results were impressive. The R&Q strategy managed to reduce the load imbalance by up to 1.4 times compared to the original, unmodified models. For example, on a difficult math dataset called GSM8K, the imbalance score for one model dropped from 1.9709 down to 1.3937. On another dataset called PIQA, it dropped from 4.3504 to 3.2925.

Crucially, this massive improvement in efficiency came with almost no cost to the quality of the answers. The accuracy of the models stayed within ±0.6% of the original. In some cases, like on the MMLU dataset, the accuracy actually improved slightly (from 23.0% to 25.2% for one model). The researchers also tested this in a "streaming" scenario, where orders come in one after another like a continuous river, and found that the system remained stable and balanced over time, proving it works even when the workload is unpredictable.

What This Means

This paper suggests that we don't need to completely rebuild our AI models to make them efficient. By simply recognizing that some parts of the brain are overworked while others are underused, and then dynamically cloning the busy parts while shrinking the idle ones, we can create a much smoother, faster, and fairer system. It's a practical, "plug-and-play" fix that makes these giant AI brains run like a well-oiled machine, ready for the real world without needing a massive overhaul.

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

Problem Statement

Sparse Mixture-of-Experts (SMoE) architectures have become a dominant paradigm for scaling Large Language Models (LLMs) efficiently, allowing models to increase capacity without linearly increasing computation per token. However, SMoE models suffer from severe load imbalance during inference. A minority of "heavy-hitter" experts receive the majority of tokens, while others remain underutilized. This imbalance leads to:

Latency Bottlenecks: Overloaded experts act as stragglers, determining the batch completion time.
Resource Inefficiency: Underutilized experts leave compute resources idle, reducing effective parallelism.
System Overhead: Increased synchronization and buffer management costs.

Existing solutions primarily focus on training-time mitigation through auxiliary losses, entropy regularization, or reinforcement learning. These methods often fail to generalize to inference-time scenarios, particularly when facing non-stationary inputs, streaming data, or varying batch sizes. Furthermore, empirical analysis reveals that selection frequency does not reliably indicate expert importance; heavily used experts are not necessarily the most critical for task performance, and infrequently used experts can be highly valuable. Current inference-time balancing methods often require custom runtimes, architectural changes, or retraining, limiting their applicability to pre-trained models.

Methodology: Replicate-and-Quantize (R&Q)

The authors propose Replicate-and-Quantize (R&Q), a training-free, plug-and-play framework designed to dynamically rebalance expert workloads during inference without modifying the router or retraining the model. The approach relies on two complementary mechanisms:

1. Load Imbalance Score (LIS)

To quantify the severity of imbalance, the paper introduces the Load Imbalance Score (LIS). For a given MoE block $i$ with $m$ experts and $n$ total tokens, where $k$ experts are activated per token:
$l_i = \frac{m \cdot \max_{j \in [m]} n_{i,j}}{nk}$
where $n_{i,j}$ is the number of tokens routed to expert $j$ . A score of 1 indicates perfect balance; higher values indicate increasing imbalance.

2. The R&Q Strategy

The framework operates in two phases based on a small calibration set:

Identify Heavy-Hitters: The method analyzes routing traces to identify the expert in each layer that receives the maximum number of tokens (the heavy-hitter).
Identify Less Important Experts: Using a gradient-free, pruning-based metric inspired by Wanda, the method estimates expert importance by computing the product of weight magnitudes and input activation norms ( $|W| \times \|X\|_2$ ). Experts with low scores are deemed less important.
Replication: The identified heavy-hitter experts are replicated within the same layer. These replicas provide additional parallel capacity to handle the token load, effectively distributing the computational bottleneck.
Quantization: To maintain the original memory and compute budget, the replicas are deployed at a lower precision (e.g., 8-bit or 16-bit). Simultaneously, the identified less important experts are aggressively quantized. This dual action ensures that the added capacity of the replicas does not exceed resource constraints while preserving the model's representational fidelity.

The entire process is post-hoc, requiring no modification to the routing mechanism or model weights (other than quantization), making it applicable to any pre-trained SMoE model.

Key Contributions

Empirical Analysis of Inference Dynamics: The paper provides a systematic analysis revealing that load imbalance persists and intensifies at inference time (especially with larger batch sizes) and that selection frequency is a poor proxy for expert utility.
Load Imbalance Score (LIS): Introduction of a principled metric to rigorously evaluate and compare routing strategies and load balancing approaches.
Replicate-and-Quantize (R&Q): A lightweight, training-free method that dynamically balances workloads by replicating overloaded experts and quantizing underutilized ones, achieving near-lossless performance.
Comprehensive Evaluation: Validation across five model architectures (Switch Transformer, LLaMA-MoE, DeepSeek-MoE, DeepSeek-V2 Lite) and six diverse benchmarks (MMLU, GSM8K, TruthfulQA, etc.), covering both static and streaming inference scenarios.

Experimental Results

Extensive experiments demonstrate the efficacy of R&Q:

Load Balancing: R&Q achieves up to a 1.4 $\times$ reduction in load imbalance (LIS) across various models and tasks. For instance, on the Switch Transformer (8 experts) with GSM8K, LIS dropped from 1.97 to 1.39.
Accuracy Preservation: The method maintains model accuracy within $\pm$ 0.6% of the baseline. In several cases (e.g., MMLU on Switch Transformer 16-expert), accuracy slightly improved (from 23.0% to 25.0%).
Streaming Stability: Under streaming inference conditions with adversarial inputs, R&Q consistently maintains lower LIS compared to raw models, demonstrating robustness over time.
Ablation Studies: Results confirm that combining replication and selective quantization yields superior trade-offs compared to using either technique alone or quantizing all experts.

Significance and Claims

The paper claims to bridge a critical gap between training-time optimization and practical inference deployment. By shifting focus from retraining-based solutions to a post-hoc, inference-time mechanism, R&Q offers a practical path toward predictable and efficient SMoE inference. The authors emphasize that their approach:

Does not require retraining or architectural modifications.
Is compatible with existing pre-trained models and standard deployment pipelines.
Addresses the specific challenge of dynamic, real-world workloads where static routing policies often fail.
Provides a scalable solution for heterogeneous workloads by balancing efficiency (through replication) and resource constraints (through quantization).

The work concludes that R&Q enables more balanced computation and reliable inference without sacrificing accuracy, making it a viable strategy for deploying large-scale MoE systems in latency-sensitive and resource-constrained environments.