Mixture-of-Experts Sparse Model Scaling
# Mixture-of-Experts & Sparse Model Scaling
## Introduction & Motivation
Mixture-of-Experts (MoE) architectures address a fundamental tension in scaling large neural networks: model quality tends to improve with parameter count, but the compute cost of a dense model grows proportionally with every additional parameter, since every parameter participates in every forward and backward pass. Sparse MoE models break this coupling by partitioning the model's capacity into many separate "expert" sub-networks and activating only a small subset of experts for any given input, so that total parameter count can grow far faster than the actual compute cost incurred per token.
The core idea traces back to work in the early 1990s on adaptive mixtures of local experts, but it became a central technique in large language model scaling following its integration into Transformer architectures, most visibly in Google's Switch Transformer and GShard, and later in production-scale systems such as Mixtral, DeepSeek-MoE, and the expert-routing layers reportedly used in some of the largest proprietary language models. In these architectures, the standard dense feedforward network found in each Transformer block is replaced with a bank of many feedforward "expert" networks, and a lightweight routing mechanism selects a small number of experts (commonly one or two) to process each token, leaving the vast majority of the model's parameters unused for any individual token.
The practical appeal of this approach is substantial: a sparse MoE model can have, for example, ten or more times the total parameter count of a comparably-priced dense model while incurring roughly the same per-token inference and training compute, because only a fraction of the parameters are activated per token. This lets practitioners trade a large, cheap-to-store parameter budget (memory and disk) for improved model quality without a proportional increase in the far more expensive floating-point compute budget, an especially attractive trade-off given that compute, not parameter storage, is typically the binding constraint on large-scale model training and serving.
Understanding MoE systems requires grappling with a set of concerns largely absent from dense model training: how to route tokens to experts in a way that is differentiable or at least trainable, how to prevent routing collapse (where the router learns to send all tokens to only a few experts, wasting the rest of the model's capacity), how to balance load across experts for efficient distributed execution, and how to manage the substantial engineering complexity of distributing experts across many accelerators while keeping communication overhead manageable.
## Core Concepts & Theory
A Mixture-of-Experts layer replaces a single dense sub-network (most commonly the feedforward network within a Transformer block, though attention layers can in principle be made sparse as well) with a collection of N expert networks, each typically sharing the same architecture as the dense layer it replaces but with independently learned parameters, alongside a routing network (also called a gating network) that examines each input token's representation and decides which expert or experts should process it.
The routing network is usually a small linear layer followed by a softmax, producing a probability distribution over the available experts for each token. In a "top-k" routing scheme, only the k experts with the highest routing scores are actually invoked for a given token (with k commonly set to 1 or 2 in practice), and the token's representation is only passed through those selected experts, with their outputs combined (typically via a weighted sum using the router's softmax scores as weights) to produce the layer's output for that token. All other experts are skipped entirely for that token, which is the source of the compute savings: only a small fraction of the total expert parameters participate in processing any individual token.
This sparse, discrete routing decision introduces a fundamental training difficulty: the choice of which expert to route a token to is a discrete, non-differentiable operation (selecting the top-k highest-scoring experts is not a smooth function of the router's parameters), which means gradients cannot flow through the routing decision itself in the same way they flow through a standard continuous neural network computation. In practice, gradients still flow through the router's softmax scores (which are continuous and differentiable) used to weight the selected experts' outputs, and the router learns primarily through this pathway, alongside auxiliary losses specifically designed to shape routing behavior, discussed below.
A related but architecturally distinct family of sparse models uses "soft" mixture-of-experts, where instead of hard top-k selection, tokens are routed via a fully differentiable, continuous mixing mechanism (such as computing a weighted combination of expert inputs before any single expert sees a token), which sacrifices some of the discrete-routing compute savings but avoids the training instabilities associated with hard discrete routing decisions, representing a design point trading off between sparsity-driven efficiency and routing training stability.
## Mathematical Formulation
For an input token representation x, a Mixture-of-Experts layer with N experts and a router producing logits g_1 through g_N computes a softmax distribution over experts:
$$ p_i(x) = \frac{\exp(g_i(x))}{\sum_{j=1}^{N} \exp(g_j(x))} $$
where g_i(x) is the routing logit computed for expert i given the token representation x, typically a simple linear transform of x.
In a top-k sparse MoE layer, only the k experts with the largest routing probabilities are activated, and the layer's output is the weighted sum of just those selected experts' outputs, with all other experts contributing nothing (their weight is effectively set to zero rather than merely being small):
$$ y = \sum_{i \in ext{TopK}(p(x), k)} p_i(x) \cdot E_i(x) $$
where E_i(x) denotes the output of expert i applied to the token representation x, and TopK selects the indices of the k experts with the highest routing probability for this particular token.
A critical addition to the basic routing formulation is an auxiliary load-balancing loss, since without explicit encouragement, the router tends to collapse onto favoring a small subset of experts (a self-reinforcing pattern, since experts that receive more training tokens early on tend to become better at handling common inputs, attracting even more routing weight over time). The Switch Transformer's load-balancing loss multiplies, for each expert, the fraction of tokens actually routed to that expert by the fraction of the router's total probability mass assigned to that expert, summed across experts:
$$ L_{aux} = N \cdot \sum_{i=1}^{N} f_i \cdot P_i $$
where f_i is the fraction of tokens in the batch actually dispatched to expert i (a hard, non-differentiable count), and P_i is the average routing probability assigned to expert i across the batch (a soft, differentiable quantity), and the leading factor of N normalizes the loss so that a perfectly uniform routing distribution (f_i equals P_i equals one over N for all experts) yields a loss value of exactly one. This auxiliary loss is added to the main training objective with a small weighting coefficient, penalizing routing distributions that are far from uniform across experts even though its gradient only flows cleanly through the P_i term.
## Advanced Theory & Extensions
Expert capacity is a crucial practical mechanism layered on top of the basic routing formulation: because tokens are routed dynamically and the number of tokens routed to any given expert within a batch is not known in advance, MoE implementations set a fixed "capacity" (a maximum number of tokens any single expert can process within a given batch, usually expressed as a capacity factor multiplied by the expected average load) to bound compute and memory usage for distributed execution. Tokens that would exceed an expert's capacity are "dropped" for that layer, meaning they bypass the expert entirely (often handled via a residual connection that lets the unprocessed token representation pass through unchanged), a design choice that trades a small amount of model quality for predictable, bounded computation, which is essential for efficient parallel execution across many accelerators.
Auxiliary z-loss, introduced alongside the Switch Transformer's load-balancing loss, penalizes the router's logits from growing too large in magnitude, since large logits push the softmax toward an extremely peaked, near-deterministic distribution that can destabilize training and produce numerical precision issues, particularly when training in reduced-precision formats such as bfloat16, which are standard for large-scale training.
Expert-choice routing, in contrast to token-choice (top-k) routing where each token selects its experts, inverts the assignment problem: each expert instead selects the top tokens it wants to process from the batch, up to its capacity. This naturally guarantees balanced expert utilization by construction (each expert always fills exactly its capacity, assuming enough tokens are available), eliminating the need for an auxiliary load-balancing loss, though it introduces a different complication: a given token may end up processed by zero experts (if it is not selected by any expert's top-tokens list) or multiple experts, altering the semantics of routing relative to the token-choice formulation.
Fine-grained expert architectures, as used in DeepSeek-MoE and related systems, split what would traditionally be a single expert into a larger number of smaller experts (increasing N while proportionally decreasing each expert's hidden dimension, to keep total expert parameter count roughly constant), combined with an increased top-k to compensate, on the theoretical and empirical basis that finer-grained experts allow more flexible, precise combinations of specialized sub-networks per token, alongside dedicated "shared experts" that process every token unconditionally to capture common, broadly useful computation that would otherwise need to be redundantly re-learned by many separate routed experts.
## Computational Considerations
The dominant systems challenge in training large MoE models is distributing experts efficiently across many accelerators while managing the all-to-all communication required to route tokens to their assigned experts, which typically live on different devices than the tokens that need to reach them. Expert parallelism, the standard distribution strategy, places different experts on different devices and uses all-to-all collective communication operations to shuffle token representations to the correct expert-hosting device and then shuffle the resulting outputs back, and the cost of this communication step can become a significant fraction of total training time if not carefully overlapped with computation or optimized for the underlying network topology.
Memory requirements for MoE models are dominated by the need to store all expert parameters (even though only a fraction are used per token, all experts' weights must reside somewhere accessible during training and, for low-latency serving, ideally in fast memory), meaning the memory footprint of a sparse MoE model scales with total parameters much like a dense model, even though its compute cost scales with only the active parameters, a distinction that has significant implications for hardware provisioning: MoE models are often more memory-bound than compute-bound relative to a dense model of equivalent active-parameter count.
Training stability issues specific to MoE models include routing collapse (discussed above), numerical instability from extreme routing logit magnitudes (addressed via the z-loss), and higher sensitivity to precision choices, since the discrete, data-dependent routing decisions mean that small numerical differences (for instance between training and inference precision, or across different hardware) can occasionally flip a routing decision for a token near a routing boundary, producing outputs that are not perfectly reproducible across hardware or precision settings, an issue essentially absent in dense architectures.
## Practical Implementation Strategies
Choosing the number of experts, the top-k value, and the capacity factor requires balancing several competing concerns: more experts with a small top-k maximizes the ratio of total to active parameters (and thus the potential quality gain per unit of compute), but increases communication overhead, memory footprint, and the risk of routing degeneracies, while a lower expert count with larger top-k moves the model closer to a dense architecture's behavior and infrastructure simplicity at the cost of reduced sparsity benefits.
Initializing and warming up MoE training carefully matters more than in dense models, since early in training the router has not yet learned useful specialization, and a poorly initialized router can quickly enter a self-reinforcing collapse toward a small subset of experts before the load-balancing loss has a chance to correct the imbalance; practitioners commonly use a higher initial weight on the auxiliary load-balancing loss early in training, sometimes annealing it down as training progresses and routing stabilizes.
For practitioners without access to large-scale distributed training infrastructure, implementing an MoE layer within a single device or a small cluster (rather than large-scale expert parallelism across hundreds of accelerators) is still valuable for understanding the mechanics, since the core routing logic, auxiliary losses, and capacity-based token dropping can all be implemented and tested at small scale before considering the additional complexity of distributed expert placement.
## Benchmark Datasets & Evaluation
Because MoE architectures are primarily an efficiency and scaling technique rather than a task-specific innovation, they are evaluated on the same broad language modeling and downstream task benchmarks used for dense large language models, including perplexity on held-out language modeling corpora, and downstream accuracy on suites such as MMLU (broad multi-subject knowledge and reasoning), GSM8K and MATH (mathematical reasoning), HumanEval and MBPP (code generation), and HellaSwag or ARC (commonsense and science reasoning), with the key comparison of interest being quality achieved per unit of training or inference compute (FLOPs) relative to a dense model, rather than quality achieved per parameter, since parameter count alone is a misleading measure of cost for sparse architectures.
Routing-specific diagnostics, distinct from downstream task performance, are also commonly reported in MoE research: expert utilization histograms (how evenly tokens are distributed across experts over the course of training), routing entropy (how confidently or diffusely the router assigns tokens to experts), and the fraction of tokens dropped due to exceeding expert capacity, all of which provide insight into whether the routing mechanism is functioning as intended, independent of downstream task accuracy.
## Key Challenges & Limitations
Load imbalance and routing collapse remain persistent challenges despite auxiliary losses and architectural mitigations like expert-choice routing, and severe imbalance directly translates into wasted compute (some accelerators sit idle waiting for overloaded experts to finish processing their capacity-limited token batches) and degraded model quality (underused experts receive too few training examples to specialize effectively).
The engineering complexity of efficient distributed MoE training and serving is substantially higher than for dense models, requiring careful implementation of all-to-all communication, dynamic capacity management, and load balancing across heterogeneous hardware and network topologies, which raises the barrier to entry for organizations without significant distributed systems engineering investment, in contrast to dense models, which can be scaled using comparatively simpler data and tensor parallelism strategies.
Fine-tuning and inference serving of sparse MoE models present their own difficulties: fine-tuning on narrow downstream tasks can exacerbate routing collapse (since a narrow task may naturally favor only a subset of experts, causing catastrophic forgetting or under-training in the unused experts), and efficient low-latency serving requires either keeping all experts resident in fast memory (an expensive proposition given the model's large total parameter count) or accepting the latency cost of loading experts on demand, both of which complicate deployment relative to dense models of comparable active-parameter count.
## Hyperparameter Tuning
The number of experts N and the top-k routing value are the most consequential architectural hyperparameters, and empirical scaling studies (following the pattern established by scaling laws research more broadly) suggest that, for a fixed active-parameter compute budget, increasing total parameters via more or larger experts continues to improve model quality up to a point, after which returns diminish and the added communication and memory overhead outweighs quality gains, with fine-grained expert designs (many small experts, larger top-k) generally demonstrating better quality-per-compute than coarse-grained designs (few large experts, small top-k) at comparable total active parameter counts.
The capacity factor controls the trade-off between token-dropping frequency and compute or memory overhead, with a capacity factor of exactly one meaning each expert can process only exactly its "fair share" of tokens on average (causing meaningful token dropping under any routing imbalance), while larger capacity factors (commonly 1.25 to 2.0 in practice) reduce dropping at the cost of reserving additional, sometimes unused, compute and memory headroom per expert.
The weighting coefficient on the auxiliary load-balancing loss requires tuning to strike a balance between enforcing sufficiently even expert utilization and not so strongly penalizing imbalance that it interferes with the router's ability to learn genuinely useful, non-uniform specialization patterns where certain experts are legitimately better suited to certain kinds of input, since forcing perfectly uniform routing regardless of input characteristics would defeat the purpose of learned, input-dependent specialization.
## Real-World Applications & Case Studies
Mixtral, an open-weight sparse MoE language model, demonstrated that a model with a relatively modest active-parameter count per token (achieved via top-2 routing among eight experts per layer) could match or exceed the quality of substantially larger dense models on standard benchmarks, while incurring inference compute costs much closer to those of the smaller active-parameter count, illustrating the practical efficiency argument for MoE architectures in a widely studied, publicly available system.
DeepSeek-MoE and its successors popularized fine-grained expert segmentation combined with dedicated shared experts, reporting that this design achieves better performance-per-compute than coarser-grained MoE designs at comparable total training cost, contributing to a broader industry trend toward fine-grained MoE architectures in subsequent large-scale model releases.
Google's Switch Transformer and GShard were among the earliest large-scale demonstrations that trillion-parameter-class sparse models could be trained with reasonable stability using auxiliary load-balancing losses and expert-capacity mechanisms, establishing much of the engineering and algorithmic groundwork (load-balancing losses, capacity factors, expert-parallel distributed training strategies) that subsequent production MoE systems have built upon.
## Integration with Other Methods
MoE layers are commonly combined with standard dense scaling techniques rather than used as a complete substitute for them: modern large MoE language models still use techniques such as rotary positional embeddings, grouped-query attention, and standard Transformer pretraining objectives in their non-MoE components, with sparsity applied specifically to the feedforward sub-layers where the technique has proven most effective and well-understood.
MoE architectures interact meaningfully with quantization and model compression techniques (covered in dedicated treatments of quantization and distillation), since the large total parameter count of MoE models makes them attractive candidates for aggressive weight quantization to reduce memory footprint, though the dynamic, input-dependent nature of expert activation can complicate certain compression techniques (such as pruning entire experts) relative to compressing a dense network, since pruning an expert that is rarely but importantly used for certain inputs can cause disproportionate quality degradation on that input subset.
Distillation from a large sparse MoE "teacher" model into a smaller dense "student" model is an increasingly common deployment strategy, allowing organizations to capture much of the quality benefit of large-scale sparse pretraining while deploying a simpler, more predictable dense model that avoids the routing and infrastructure complexity of serving a sparse MoE model directly in latency-sensitive production settings.
## Future Research Directions
Improving routing algorithms to more reliably achieve balanced expert utilization without relying on auxiliary losses that must be carefully weighted against the primary training objective remains an active research area, with expert-choice routing and other capacity-based assignment schemes representing early steps toward routing mechanisms with more inherent load-balancing guarantees.
Reducing the communication overhead of distributed MoE training and inference, through techniques such as improved expert placement strategies that co-locate frequently co-activated experts, communication-computation overlap, and hardware-aware routing that accounts for network topology, continues to be an important systems research direction as MoE models scale to ever-larger total parameter counts across ever-larger accelerator clusters.
Understanding what different experts actually specialize in, whether that specialization is interpretable or useful for tasks like model editing and targeted capability removal, and whether MoE routing patterns reveal anything about the underlying structure of the training data or task distribution, connects MoE research to the broader interpretability and mechanistic-understanding research agenda applied to large language models generally.
## Summary & Key Takeaways
Mixture-of-Experts architectures decouple total model parameter count from per-token compute cost by routing each token through only a small subset of available expert sub-networks, enabling substantially larger effective model capacity without a proportional increase in training or inference compute.
Effective MoE training requires carefully managing routing behavior through auxiliary load-balancing losses, expert capacity limits, and architectural choices such as expert-choice routing or fine-grained expert segmentation, all aimed at preventing routing collapse and ensuring balanced, efficient utilization of the model's full expert capacity.
The primary costs of MoE architectures relative to dense models are increased systems and distributed-training complexity (particularly all-to-all communication for expert parallelism), higher memory requirements relative to active-parameter count, and additional training instabilities tied to the discrete, data-dependent nature of routing decisions.
Production systems such as Mixtral, DeepSeek-MoE, Switch Transformer, and GShard have demonstrated that these challenges are surmountable at scale, establishing sparse MoE architectures as a standard, practically important technique for scaling large language models efficiently.
Keywords: mixture of experts, MoE, sparse model scaling, expert routing, gating network, top-k routing, Switch Transformer, GShard, Mixtral, DeepSeek-MoE, load balancing loss, expert capacity, token dropping, expert-choice routing, fine-grained experts, expert parallelism, all-to-all communication, routing collapse, z-loss, sparse activation
---
## Appendix: Practical Labs
### Lab 1: Top-K Router with Softmax Gating and Weighted Expert Combination
import numpy as np
np.random.seed(0)
def softmax(logits):
shifted = logits - np.max(logits, axis=-1, keepdims=True)
exp = np.exp(shifted)
return exp / np.sum(exp, axis=-1, keepdims=True)
class ToyExpert:
"""A simple linear expert: y = x @ W + b."""
def __init__(self, in_dim, out_dim, seed):
rng = np.random.RandomState(seed)
self.W = rng.randn(in_dim, out_dim) * 0.5
self.b = rng.randn(out_dim) * 0.1
def forward(self, x):
return x @ self.W + self.b
class TopKRouter:
def __init__(self, in_dim, n_experts, seed=42):
rng = np.random.RandomState(seed)
self.W_gate = rng.randn(in_dim, n_experts) * 0.3
self.n_experts = n_experts
def route(self, x, k):
logits = x @ self.W_gate # shape (batch, n_experts)
probs = softmax(logits)
top_k_idx = np.argsort(-probs, axis=-1)[:, :k]
return probs, top_k_idx
class SparseMoELayer:
def __init__(self, in_dim, out_dim, n_experts, k):
self.experts = [ToyExpert(in_dim, out_dim, seed=i) for i in range(n_experts)]
self.router = TopKRouter(in_dim, n_experts)
self.k = k
self.n_experts = n_experts
def forward(self, x):
batch_size = x.shape[0]
out_dim = self.experts[0].W.shape[1]
probs, top_k_idx = self.router.route(x, self.k)
output = np.zeros((batch_size, out_dim))
expert_token_counts = np.zeros(self.n_experts, dtype=int)
for b in range(batch_size):
selected = top_k_idx[b]
# Renormalize the selected experts' probabilities so weights sum to 1
selected_probs = probs[b, selected]
selected_probs = selected_probs / selected_probs.sum()
for weight, expert_idx in zip(selected_probs, selected):
expert_out = self.experts[expert_idx].forward(x[b:b+1])
output[b] += weight * expert_out.squeeze(0)
expert_token_counts[expert_idx] += 1
return output, expert_token_counts
def test_sparse_moe_layer():
in_dim, out_dim, n_experts, k = 6, 4, 8, 2
layer = SparseMoELayer(in_dim, out_dim, n_experts, k)
x = np.random.randn(20, in_dim)
output, expert_counts = layer.forward(x)
print(f"Output shape: {output.shape}")
print(f"Expert token counts: {expert_counts}")
print(f"Total token-expert assignments: {expert_counts.sum()} (expected {20 * k})")
assert output.shape == (20, out_dim), "Output shape mismatch"
assert expert_counts.sum() == 20 * k, "Each token should be routed to exactly k experts"
# With 8 experts and only 20 tokens routed to 2 each, not every single expert
# is guaranteed to be used, but multiple distinct experts should be active.
assert (expert_counts > 0).sum() >= 2, "Routing should engage more than one expert"
print("Sparse MoE layer test passed.")
if __name__ == "__main__":
test_sparse_moe_layer()### Lab 2: Switch-Transformer-Style Load-Balancing Auxiliary Loss
import numpy as np
np.random.seed(1)
def softmax(logits):
shifted = logits - np.max(logits, axis=-1, keepdims=True)
exp = np.exp(shifted)
return exp / np.sum(exp, axis=-1, keepdims=True)
def compute_load_balancing_loss(routing_probs, chosen_expert_indices, n_experts):
"""
routing_probs: (batch, n_experts) full softmax distribution over experts per token
chosen_expert_indices: (batch,) the single top-1 expert index chosen per token
n_experts: total number of experts N
Implements L_aux = N * sum_i f_i * P_i
where f_i = fraction of tokens routed to expert i (hard assignment count)
P_i = average routing probability mass on expert i (soft, differentiable)
"""
batch_size = routing_probs.shape[0]
f = np.zeros(n_experts)
for idx in chosen_expert_indices:
f[idx] += 1
f = f / batch_size
P = routing_probs.mean(axis=0)
aux_loss = n_experts * np.sum(f * P)
return aux_loss, f, P
def simulate_routing(batch_size, n_experts, collapse_strength=0.0, seed=0):
"""collapse_strength in [0, 1]: 0 = uniform random routing preference,
1 = heavily biased toward expert 0 (simulating routing collapse)."""
rng = np.random.RandomState(seed)
logits = rng.randn(batch_size, n_experts) * 0.5
logits[:, 0] += collapse_strength * 5.0 # bias toward expert 0
probs = softmax(logits)
chosen = np.argmax(probs, axis=-1)
return probs, chosen
def test_load_balancing_loss_detects_collapse():
n_experts = 6
batch_size = 200
balanced_probs, balanced_chosen = simulate_routing(
batch_size, n_experts, collapse_strength=0.0, seed=10
)
collapsed_probs, collapsed_chosen = simulate_routing(
batch_size, n_experts, collapse_strength=1.0, seed=10
)
balanced_loss, f_bal, P_bal = compute_load_balancing_loss(
balanced_probs, balanced_chosen, n_experts
)
collapsed_loss, f_col, P_col = compute_load_balancing_loss(
collapsed_probs, collapsed_chosen, n_experts
)
print(f"Balanced routing: aux_loss={balanced_loss:.3f}, token fractions={np.round(f_bal, 2)}")
print(f"Collapsed routing: aux_loss={collapsed_loss:.3f}, token fractions={np.round(f_col, 2)}")
# The theoretical minimum of N * sum(f_i * P_i) is 1.0, achieved at perfectly
# uniform f and P. A collapsed router should show a substantially higher loss.
assert balanced_loss < collapsed_loss, (
"Load-balancing loss should be higher for collapsed routing than balanced routing"
)
assert collapsed_loss > 1.5, "Collapsed routing should produce a clearly elevated aux loss"
assert balanced_loss < 1.5, "Balanced routing should produce an aux loss close to the minimum of 1.0"
print("Load-balancing loss collapse-detection test passed.")
if __name__ == "__main__":
test_load_balancing_loss_detects_collapse()### Lab 3: Expert Capacity Limits and Token Dropping Simulation
import numpy as np
np.random.seed(2)
def simulate_capacity_constrained_routing(routing_choices, n_experts, capacity):
"""
routing_choices: array of shape (n_tokens,) giving the chosen expert index per token,
in the order tokens arrive (arrival order matters for capacity dropping).
n_experts: number of experts
capacity: max tokens each expert can process
Returns a boolean array of shape (n_tokens,) indicating whether each token
was successfully processed (True) or dropped due to exceeding expert capacity (False).
"""
expert_load = np.zeros(n_experts, dtype=int)
processed = np.zeros(len(routing_choices), dtype=bool)
for token_idx, expert_idx in enumerate(routing_choices):
if expert_load[expert_idx] < capacity:
expert_load[expert_idx] += 1
processed[token_idx] = True
else:
processed[token_idx] = False # dropped: expert at full capacity
return processed, expert_load
def compute_drop_rate(processed):
return 1.0 - processed.mean()
def test_capacity_and_drop_rate():
n_tokens = 100
n_experts = 4
average_load = n_tokens / n_experts # 25 tokens per expert if perfectly balanced
# Scenario A: balanced routing, generous capacity factor of 1.5
balanced_choices = np.random.randint(0, n_experts, size=n_tokens)
capacity_generous = int(average_load * 1.5)
processed_a, load_a = simulate_capacity_constrained_routing(
balanced_choices, n_experts, capacity_generous
)
drop_rate_a = compute_drop_rate(processed_a)
# Scenario B: imbalanced routing (skewed toward expert 0), tight capacity factor of 1.0
skewed_choices = np.concatenate([
np.zeros(60, dtype=int), # 60 tokens want expert 0
np.random.randint(1, n_experts, size=40), # remaining spread over experts 1-3
])
np.random.shuffle(skewed_choices)
capacity_tight = int(average_load * 1.0)
processed_b, load_b = simulate_capacity_constrained_routing(
skewed_choices, n_experts, capacity_tight
)
drop_rate_b = compute_drop_rate(processed_b)
print(f"Scenario A (balanced, generous capacity): drop_rate={drop_rate_a:.2%}, loads={load_a}")
print(f"Scenario B (skewed, tight capacity): drop_rate={drop_rate_b:.2%}, loads={load_b}")
assert drop_rate_a < 0.05, "Balanced routing with generous capacity should drop very few tokens"
assert drop_rate_b > drop_rate_a, (
"Skewed routing with tight capacity should drop noticeably more tokens than the balanced case"
)
assert load_a.max() <= capacity_generous, "No expert should exceed its capacity limit"
assert load_b.max() <= capacity_tight, "No expert should exceed its capacity limit"
print("Expert capacity and token-dropping test passed.")
if __name__ == "__main__":
test_capacity_and_drop_rate()### Lab 4: Fine-Grained vs. Coarse-Grained Expert Configurations at Fixed Parameter Budget
import numpy as np
np.random.seed(3)
def total_expert_parameters(n_experts, hidden_dim, in_dim=64, out_dim=64):
"""Parameter count for a bank of feedforward experts, each with one hidden layer."""
per_expert_params = (in_dim * hidden_dim) + hidden_dim + (hidden_dim * out_dim) + out_dim
return n_experts * per_expert_params, per_expert_params
def active_parameters_per_token(per_expert_params, top_k):
return per_expert_params * top_k
def configure_fine_grained_from_coarse(coarse_n_experts, coarse_hidden_dim, coarse_top_k,
granularity_factor):
"""Splits each coarse expert into `granularity_factor` smaller experts with
proportionally reduced hidden dimension, and scales top_k up by the same factor,
so total expert parameter count and per-token active parameter count are
approximately preserved (the fine-grained DeepSeek-MoE style transformation)."""
fine_n_experts = coarse_n_experts * granularity_factor
fine_hidden_dim = max(1, coarse_hidden_dim // granularity_factor)
fine_top_k = coarse_top_k * granularity_factor
return fine_n_experts, fine_hidden_dim, fine_top_k
def test_fine_grained_preserves_budget():
in_dim, out_dim = 64, 64
coarse_n_experts, coarse_hidden_dim, coarse_top_k = 8, 256, 2
coarse_total, coarse_per_expert = total_expert_parameters(
coarse_n_experts, coarse_hidden_dim, in_dim, out_dim
)
coarse_active = active_parameters_per_token(coarse_per_expert, coarse_top_k)
granularity_factor = 4
fine_n_experts, fine_hidden_dim, fine_top_k = configure_fine_grained_from_coarse(
coarse_n_experts, coarse_hidden_dim, coarse_top_k, granularity_factor
)
fine_total, fine_per_expert = total_expert_parameters(
fine_n_experts, fine_hidden_dim, in_dim, out_dim
)
fine_active = active_parameters_per_token(fine_per_expert, fine_top_k)
print(f"Coarse config: {coarse_n_experts} experts x hidden={coarse_hidden_dim}, "
f"top_k={coarse_top_k}")
print(f" Total params: {coarse_total:,}, Active params/token: {coarse_active:,}")
print(f"Fine-grained config: {fine_n_experts} experts x hidden={fine_hidden_dim}, "
f"top_k={fine_top_k}")
print(f" Total params: {fine_total:,}, Active params/token: {fine_active:,}")
# Total parameters won't be perfectly identical due to integer rounding of hidden_dim,
# but should be within a reasonably close tolerance of the coarse configuration.
relative_diff = abs(fine_total - coarse_total) / coarse_total
print(f"Relative difference in total parameters: {relative_diff:.2%}")
assert relative_diff < 0.30, (
"Fine-grained reconfiguration should approximately preserve total parameter budget"
)
assert fine_n_experts > coarse_n_experts, "Fine-grained config should have more, smaller experts"
assert fine_hidden_dim < coarse_hidden_dim, "Fine-grained experts should be individually smaller"
print("Fine-grained vs. coarse-grained MoE budget test passed.")
if __name__ == "__main__":
test_fine_grained_preserves_budget()