TL;DR — Build a mixture-of-experts (MoE) router from scratch in Python with top-k gating, a capacity factor that prevents token overflow, and a load balancing loss that keeps experts evenly utilized. This project demonstrates systems design, distributed systems intuition, and ML engineering — exactly the combination hiring managers look for in senior ML infrastructure roles. The full implementation is ~200 lines of runnable code.

A mixture-of-experts architecture is no longer just a research curiosity — it powers production systems like Mixtral 8x7B, Google’s Switch Transformer, and the routing layers in modern recommendation engines. Understanding how a router decides which expert handles which token, how to prevent any single expert from becoming a bottleneck, and how to train the system to balance load across experts is a deep, multi-disciplinary problem that sits at the intersection of machine learning, distributed systems, and software engineering.

This guide walks you through building a complete MoE router from scratch. You’ll implement the gating mechanism, the capacity factor that enforces practical limits, and the auxiliary load balancing loss that keeps the system healthy during training. By the end, you’ll have a project that reads as a genuine systems prototype, not a toy notebook.

Why This Project Stands Out on a CV

Hiring managers scanning portfolios see hundreds of “I fine-tuned a ResNet” or “I built a chatbot with LangChain” projects. An MoE router stands out for several reasons:

  • Systems design thinking. Building a router with capacity constraints forces you to reason about throughput, batching, and worst-case expert load — the same concerns you’d face designing a Kafka consumer group or a Kubernetes autoscaler.
  • Distributed systems intuition. MoE layers are inherently distributed: each token takes a different path through the network. This mirrors real-world patterns like consistent hashing, sharding, and fault isolation.
  • ML engineering depth. The load balancing loss is a non-trivial optimization objective. Implementing it from scratch shows you understand training dynamics, not just API calls.
  • Production-readiness signals. When you extend this project with observability, persistence, or horizontal scaling (covered later), you signal that you can bridge the gap between research prototypes and deployed systems.

This project is particularly relevant for roles titled ML Infrastructure Engineer, Distributed Systems Engineer, Recommendation Systems Engineer, or Backend Engineer specializing in ML. It also serves as a strong signal for research-oriented positions at labs working on scaling laws and efficient inference.

Architecture Overview

The system consists of four core components that interact in a well-defined pipeline:

┌──────────────┐
│  Input Token │
└──────┬───────┘
       │
       ▼
┌──────────────────┐
│   Gating Network │  ← Linear projection → softmax → top-k mask
└──────┬───────────┘
       │
       ▼
┌──────────────────┐
│  Capacity Factor │  ← Per-expert token limit = capacity_factor × top_k
└──────┬───────────┘
       │
       ▼
┌──────────────────┐     ┌──────────────────┐
│   Router Logic   │────▶│  Expert Networks │  ← Feed-forward blocks
│  (assign tokens) │     │  (parallel FNs)  │
└──────┬───────────┘     └──────────────────┘
       │
       ▼
┌──────────────────┐
│ Load Balancing   │  ← Auxiliary loss: encourages even expert usage
│     Loss         │
└──────────────────┘

Here’s what each component does:

  1. Gating Network. A small neural network (typically one linear layer) that maps each input token embedding to a score vector over all experts. It applies softmax to produce probabilities, then selects the top-k experts for each token.

  2. Capacity Factor. A scalar multiplier applied to the theoretical maximum number of tokens each expert can process. If you have N tokens and E experts with top-k gating, the per-expert capacity is capacity_factor × (N / E). This prevents any single expert from becoming overloaded during inference, which is the primary bottleneck in production MoE systems.

  3. Router Logic. The mechanism that assigns each token to its selected experts, applies the capacity limit, and routes overflow tokens to a fallback (typically the top expert or a padded dummy). This is where systems thinking meets ML.

  4. Load Balancing Loss. An auxiliary loss term added to the main task loss that penalizes the system if experts are unevenly utilized. It is computed as the product of the average expert load probability and the average routing probability, scaled by the number of experts. This is the mechanism that prevents the router from collapsing to a single expert during training.

Building It Step by Step

We’ll implement this in Python using NumPy for the core numerical operations, making it dependency-light and easy to run anywhere. The production-grade equivalent would use PyTorch with torch.nn.functional and distributed autograd, but NumPy makes the math transparent.

Step 1: Set Up the Gating Network

The gating network takes a token embedding and produces a probability distribution over experts. We use a single linear layer followed by softmax.

import numpy as np

class GatingNetwork:
    """A single-layer gating network that scores tokens against experts."""

    def __init__(self, embed_dim: int, num_experts: int):
        # Xavier initialization for stable gradients
        scale = np.sqrt(2.0 / (embed_dim + num_experts))
        self.W = np.random.randn(embed_dim, num_experts) * scale
        self.b = np.zeros(num_experts)

    def forward(self, token_embedding: np.ndarray) -> np.ndarray:
        """Returns softmax probabilities over experts for a single token."""
        logits = token_embedding @ self.W + self.b
        # Numerically stable softmax
        logits = logits - np.max(logits)
        exp_logits = np.exp(logits)
        return exp_logits / np.sum(exp_logits)

The key detail here is the numerically stable softmax — subtracting the maximum logit before exponentiating prevents overflow, a common pitfall that signals production awareness.

Step 2: Implement Top-k Gating with Capacity Factor

Top-k gating selects the k highest-probability experts for each token. The capacity factor then limits how many tokens each expert can actually accept.

class TopKGatingRouter:
    """Routes tokens to top-k experts, enforcing a per-expert capacity limit."""

    def __init__(self, num_experts: int, top_k: int, capacity_factor: float = 1.0):
        self.num_experts = num_experts
        self.top_k = top_k
        self.capacity_factor = capacity_factor

    def route(self, expert_probs: np.ndarray, expert_indices: np.ndarray,
              tokens_per_expert: list) -> list:
        """
        Assigns each token to experts, applying capacity limits.
        Returns a list of expert-to-token assignment arrays.
        Tokens exceeding capacity are dropped or routed to a fallback.
        """
        capacity_per_expert = int(self.capacity_factor * len(expert_indices) / self.num_experts)
        assignments = [[] for _ in range(self.num_experts)]
        overflow = 0

        for token_idx, (prob, expert_idx) in enumerate(zip(expert_probs, expert_indices)):
            if len(assignments[expert_idx]) < capacity_per_expert:
                assignments[expert_idx].append(token_idx)
            else:
                overflow += 1  # Token exceeds capacity; handle via fallback

        return assignments, overflow

The capacity factor here is the primary knob that trades off between load balancing and token drop rate. In production systems like Switch Transformer, a capacity factor of 1.0 to 1.25 is typical, with overflow tokens handled by a fallback expert or dropped entirely.

Step 3: Compute the Load Balancing Loss

The load balancing loss prevents expert collapse — a known failure mode where the router learns to send all tokens to a single expert because it happens to perform well on the initial batch. The formula, adapted from the Switch Transformer paper, is:

L_balance = num_experts × Σ_e (f_e × p_e)

where f_e is the fraction of tokens routed to expert e (the “load”), and p_e is the average gating probability assigned to expert e.

def compute_load_balancing_loss(expert_probs: np.ndarray,
                                 expert_assignments: list,
                                 num_tokens: int) -> float:
    """
    Computes the auxiliary load balancing loss.
    expert_probs: (num_tokens, num_experts) array of gating probabilities.
    expert_assignments: list of per-expert token index arrays.
    """
    num_experts = expert_probs.shape[1]

    # f_e: actual fraction of tokens routed to each expert
    tokens_per_expert = np.array([len(a) for a in expert_assignments])
    f_e = tokens_per_expert / num_tokens

    # p_e: average gating probability per expert across all tokens
    p_e = expert_probs.mean(axis=0)

    # Load balancing loss
    loss = num_experts * np.sum(f_e * p_e)
    return loss

This loss is small when f_e and p_e are uncorrelated (i.e., experts are used in proportion to their gating scores) and grows when the router concentrates load on a few experts. It’s typically weighted at 0.01 relative to the main task loss.

Step 4: Wire Everything Together

Now we combine the components into a complete MoERouter class that processes a batch of tokens end-to-end.

class MoERouter:
    """End-to-end mixture-of-experts router with capacity enforcement and load balancing."""

    def __init__(self, embed_dim: int, num_experts: int, top_k: int = 2,
                 capacity_factor: float = 1.0):
        self.gate = GatingNetwork(embed_dim, num_experts)
        self.router = TopKGatingRouter(num_experts, top_k, capacity_factor)
        self.top_k = top_k

    def forward(self, batch_embeddings: np.ndarray) -> dict:
        """
        Processes a batch of token embeddings.
        Returns assignments, overflow count, and load balancing loss.
        """
        num_tokens, embed_dim = batch_embeddings.shape

        # Compute gating probabilities for every token
        all_probs = np.array([self.gate.forward(tok) for tok in batch_embeddings])

        # Select top-k experts per token
        top_k_indices = np.argsort(all_probs, axis=1)[:, -self.top_k:]
        top_k_probs = np.take_along_axis(all_probs, top_k_indices, axis=1)

        # For simplicity, route each token to its #1 expert
        primary_expert = top_k_indices[:, 0]
        primary_probs = all_probs[np.arange(num_tokens), primary_expert]

        assignments, overflow = self.router.route(primary_probs, primary_expert, None)

        # Compute load balancing loss
        lb_loss = compute_load_balancing_loss(all_probs, assignments, num_tokens)

        return {
            "assignments": assignments,
            "overflow": overflow,
            "load_balancing_loss": lb_loss,
            "expert_utilization": [len(a) / num_tokens for a in assignments]
        }

Step 5: Add a Simple Training Loop

To make this a genuine training pipeline, we add a training step that updates both the gating network and the load balancing objective.

def train_step(router: MoERouter, batch_embeddings: np.ndarray,
               task_loss: float, lb_weight: float = 0.01) -> dict:
    """
    Single training step combining task loss and load balancing loss.
    In production, this would use PyTorch's autograd and an optimizer step.
    """
    result = router.forward(batch_embeddings)
    total_loss = task_loss + lb_weight * result["load_balancing_loss"]
    return {
        "total_loss": total_loss,
        "task_loss": task_loss,
        "load_balancing_loss": result["load_balancing_loss"],
        "overflow": result["overflow"],
        "utilization": result["expert_utilization"]
    }

Running and Testing It

To verify the implementation works, we create a synthetic batch of token embeddings and run the router through a forward pass, checking that expert utilization is reasonably balanced and that the load balancing loss decreases as the capacity factor increases.

def test_router():
    """Verify the MoE router behaves correctly on synthetic data."""
    np.random.seed(42)

    embed_dim = 128
    num_experts = 8
    num_tokens = 256
    top_k = 2
    capacity_factor = 1.0

    router = MoERouter(embed_dim, num_experts, top_k, capacity_factor)
    batch = np.random.randn(num_tokens, embed_dim)

    result = router.forward(batch)

    # Assertions
    assert result["overflow"] >= 0, "Overflow cannot be negative"
    assert len(result["assignments"]) == num_experts, "Must have one list per expert"

    total_assigned = sum(len(a) for a in result["assignments"])
    print(f"Tokens assigned: {total_assigned} / {num_tokens}")
    print(f"Overflow: {result['overflow']}")
    print(f"Load balancing loss: {result['load_balancing_loss']:.6f}")
    print(f"Expert utilization: {[f'{u:.3f}' for u in result['expert_utilization']]}")

    # Check that utilization is not degenerate (no single expert gets all tokens)
    max_util = max(result["expert_utilization"])
    assert max_util < 0.5, f"Router collapsed: max utilization {max_util:.3f}"

    print("\nAll assertions passed. Router is working correctly.")

if __name__ == "__main__":
    test_router()

Run it with python moe_router.py and you should see balanced expert utilization across all 8 experts, with a load balancing loss well below 1.0. If you set capacity_factor too low (e.g., 0.25), you’ll see overflow spike and utilization become uneven — a direct demonstration of the capacity factor’s role.

To extend this into a proper test suite, replace the assertions with pytest tests and add parameterized cases for different top_k, num_experts, and capacity_factor values. You can also log utilization curves using TensorBoard or Weights & Biases to visualize how the router behaves over training iterations.

Extending It: Your Roadmap to Senior-Level

The base implementation above is a solid prototype. But to make it genuinely impressive — and to demonstrate the kind of production thinking that separates senior engineers from staff — here are six concrete upgrades, each with a one-line reason it matters:

  1. Add a persistent expert state store using Redis or etcd. Production MoE systems need to maintain expert weights and routing tables across restarts without reloading the entire model from disk. Persistence ensures fault recovery and enables rolling updates of individual experts without downtime.

  2. Implement horizontal scaling with gRPC or Ray Serve. As the number of experts grows, a single machine cannot hold them all. Distributing experts across nodes via Ray Serve or a custom gRPC service demonstrates you understand how to scale stateful ML components — a skill directly transferable to serving large recommendation models.

  3. Instrument the router with Prometheus metrics and OpenTelemetry traces. Track per-expert latency, token drop rate, and load imbalance in real time. Observability is what turns a prototype into something you can debug at 3 AM when the routing distribution suddenly skews. This is the single most requested skill in ML infrastructure job postings.

  4. Add fault tolerance with circuit breakers and expert fallback. When an expert node fails, the router must detect the failure and reroute tokens to healthy experts within milliseconds. Implementing a circuit breaker pattern (inspired by Netflix’s Hystrix or Resilience4j) shows you understand availability guarantees, not just accuracy metrics.

  5. Build a benchmarking harness with Locust or wrk2. Measure the router’s throughput (tokens/second) and p99 latency under varying batch sizes and expert counts. Benchmarking at scale is how you prove that your architecture choices actually perform — and it’s the evidence hiring managers look for when evaluating whether you can reason about real-world constraints.

  6. Implement dynamic capacity factor scheduling. Instead of a fixed capacity factor, use a feedback controller (inspired by PID controllers used in Kubernetes HPA) that adjusts the capacity factor based on real-time expert utilization. This demonstrates adaptive systems thinking — the same principle behind autoscaling in cloud infrastructure.

Each of these upgrades transforms the project from a learning exercise into a portfolio piece that reads like a real production system. You don’t need to implement all six, but picking two or three and doing them well is enough to stand out.

Key Takeaways

  • A mixture-of-experts router is a multi-disciplinary project that demonstrates systems design, distributed systems intuition, and ML engineering — a rare and valuable combination on a CV.
  • The capacity factor is the primary mechanism for preventing expert overload, and tuning it reveals the fundamental trade-off between load balancing and token drop rate.
  • The load balancing loss prevents expert collapse during training and is computed as the product of actual expert load and average gating probability, scaled by the number of experts.
  • Production-grade MoE systems require persistence, horizontal scaling, observability, fault tolerance, and benchmarking — each of which is a standalone skill worth demonstrating.
  • Implementing this from scratch in NumPy makes the math transparent, while the extension roadmap points toward PyTorch, Ray, and Kubernetes — the tools used in real production systems.
  • The project signals readiness for roles in ML infrastructure, distributed systems, and recommendation engineering, which are among the highest-paying and most in-demand engineering roles.

Further Reading