Building a Pure-Python MoE Router with Noisy Top-K Gating and Load Balancing

TL;DR — This post walks you through building a pure-Python Mixture of Experts (MoE) router with noisy top-k gating and load balancing loss. You’ll write ~150 lines of runnable code, test it with synthetic data, and learn how to extend it into a production-grade component. It’s a portfolio project that demonstrates both ML algorithm knowledge and systems engineering discipline — exactly what senior roles look for.

Why This Project Stands Out on a CV

Hiring managers for senior ML engineer and ML systems roles often face a flood of applicants who can tune hyperparameters but few who can explain why a routing decision was made under load. A MoE router project directly signals four high-value competencies:

Transformer internals beyond attention. You’ve implemented the gating mechanism that modern models like Mixtral, Qwen-MoE, and GPT-4 rely on. This shows you understand the sparse-activation paradigm that’s replacing dense transformers at scale.
Load-aware systems thinking. The load balancing loss isn’t just a training trick — it’s a production concern. Unbalanced expert utilization leads to OOM errors on the expert hosting GPU while other experts sit idle. Demonstrating this awareness translates directly to infrastructure cost conversations.
Pure-Python discipline. Writing it without PyTorch or TensorFlow (just NumPy) forces you to confront matrix shapes, gradient flow, and numerical stability at the metal level. Recruiters recognize this as the difference between calling nn.Linear() and understanding what happens under the hood.
Portfolio narrative. You can pair this with a FastAPI wrapper that serves inference requests, a Prometheus exporter that tracks expert utilization, and a Dockerfile. That progression — from algorithm to serving to observability — is the exact story staff engineers tell when they want to showcase end-to-end ownership.

In terms of role targeting, this project speaks to: Staff ML Engineer, ML Infrastructure Engineer, Applied Research Engineer (at model-serving companies), and even Senior Backend Engineer roles at AI startups where “knowing transformers” is now table stakes.

Architecture Overview

Here’s the component breakdown of the system you’ll build:

Token Embeddings (B, S, D)
        │
        ▼
┌─────────────────────────┐
│   Linear Gate (D → E)   │  Produces routing logits per expert
│   + Gaussian Noise      │  Noisy gating for exploration & load spread
└───────────┬─────────────┘
            │
            ▼
┌─────────────────────────┐
│   Top-K Selection       │  Keep only k highest logits per token
│   (Softmax over selected)│  Normalize weights for the chosen experts
└───────────┬─────────────┘
            │
            ▼
┌─────────────────────────┐
│   Expert Dispatch       │  Route token to its selected experts
│   (Weighted aggregation) │  Each expert processes, results weighted-sum
└───────────┬─────────────┘
            │
            ▼
┌─────────────────────────┐
│   Load Balancing Loss   │  Auxiliary loss on expert assignment counts
│   (CV of token distribution) │  Penalizes imbalance, encourages uniform load
└─────────────────────────┘

Key components:

Gate network: A single linear layer mapping hidden dimension D to E experts. This is deliberately simple — in production you’d add a hidden layer, but the linear gate is the canonical starting point from the GShard paper.
Noise injection: Additive Gaussian noise N(0, σ) to logits before top-k selection. This is the “noisy” part of noisy top-k gating. It serves two purposes: (1) it encourages exploration during training so experts don’t collapse to a single specialist, and (2) it naturally spreads load because tokens with similar logits may route differently across forward passes.
Top-k selector: For each token, keep only the k largest logits after noise. Typically k=2 for inference efficiency. Apply softmax only over the selected experts to get routing weights.
Expert assignment: Each selected expert processes the token (in this simplified version, we simulate expert computation with a linear projection). The final output is the weighted sum of expert outputs.
Load balancing loss: Compute the fraction of tokens assigned to each expert, then calculate the coefficient of variation (CV = std/mean) across experts. Multiply by the mean routing weight to get a differentiable auxiliary loss. This is the simplified version from the Switch Transformer paper; the full Shazeer et al. 2017 formulation uses a more complex formulation.

Building It Step by Step

Let’s implement this in NumPy. Create a file called moe_router.py.

Step 1: Set up imports and the router class skeleton.

import numpy as np
from typing import Tuple

class MoERouter:
    """
    Noisy Top-K Gating MoE Router.
    
    Attributes:
        num_experts (int): Number of experts (E).
        top_k (int): Number of experts to route each token to.
        hidden_dim (int): Token embedding dimension (D).
        noise_std (float): Standard deviation of Gaussian noise.
        gate_weights (np.ndarray): Gate projection matrix (D, E).
    """
    
    def __init__(self, num_experts: int = 8, hidden_dim: int = 512, 
                 top_k: int = 2, noise_std: float = 1.0):
        self.num_experts = num_experts
        self.top_k = top_k
        self.hidden_dim = hidden_dim
        self.noise_std = noise_std
        
        # Initialize gate weights: (D, E)
        self.gate_weights = np.random.randn(hidden_dim, num_experts) * 0.02
        
        # Expert projection weights: (E, D, D) — each expert is a linear transform
        self.expert_weights = np.random.randn(num_experts, hidden_dim, hidden_dim) * 0.02

Step 2: Implement the gate forward pass with noisy top-k selection.

    def gate(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Compute routing logits, apply noise, and select top-k experts.
        
        Args:
            x: Token embeddings of shape (batch, seq_len, hidden_dim).
        
        Returns:
            routing_weights: (batch, seq_len, top_k) — softmax over selected experts.
            expert_indices: (batch, seq_len, top_k) — selected expert IDs.
            logits: (batch, seq_len, num_experts) — pre-noise logits (for debugging).
        """
        B, S, D = x.shape
        
        # Reshape for matrix multiply: (B*S, D) @ (D, E) → (B*S, E)
        flat_x = x.reshape(-1, D)
        logits = flat_x @ self.gate_weights  # (B*S, E)
        
        # Add Gaussian noise for exploration
        noise = np.random.normal(0, self.noise_std, logits.shape)
        noisy_logits = logits + noise
        
        # Top-k selection along expert axis
        # argsort descending: get indices of top-k logits
        top_k_indices = np.argsort(-noisy_logits, axis=1)[:, :self.top_k]  # (B*S, k)
        
        # Gather the noisy logits for selected experts
        # Using advanced indexing: for each row, pick the columns at top_k_indices
        batch_idx = np.arange(flat_x.shape[0])[:, None]
        selected_logits = noisy_logits[batch_idx, top_k_indices]  # (B*S, k)
        
        # Softmax over selected logits only
        selected_logits_max = np.max(selected_logits, axis=1, keepdims=True)
        selected_logits_exp = np.exp(selected_logits - selected_logits_max)
        routing_weights = selected_logits_exp / np.sum(selected_logits_exp, axis=1, keepdims=True)
        
        # Reshape back
        routing_weights = routing_weights.reshape(B, S, self.top_k)
        expert_indices = top_k_indices.reshape(B, S, self.top_k)
        logits = logits.reshape(B, S, self.num_experts)
        
        return routing_weights, expert_indices, logits

Step 3: Implement expert dispatch and output aggregation.

    def route(self, x: np.ndarray, expert_indices: np.ndarray, 
              routing_weights: np.ndarray) -> np.ndarray:
        """
        Route tokens to selected experts and aggregate weighted outputs.
        
        Args:
            x: Input embeddings (B, S, D).
            expert_indices: (B, S, k) selected expert IDs.
            routing_weights: (B, S, k) routing weights.
        
        Returns:
            output: (B, S, D) aggregated token representations.
        """
        B, S, D = x.shape
        output = np.zeros((B, S, D))
        
        for b in range(B):
            for s in range(S):
                token = x[b, s]  # (D,)
                for k_idx in range(self.top_k):
                    expert_id = expert_indices[b, s, k_idx]
                    weight = routing_weights[b, s, k_idx]
                    
                    # Expert computation: linear transformation
                    expert_out = token @ self.expert_weights[expert_id]  # (D,)
                    output[b, s] += weight * expert_out
        
        return output

Step 4: Add the load balancing loss.

    def load_balancing_loss(self, expert_indices: np.ndarray, 
                            routing_weights: np.ndarray) -> float:
        """
        Compute auxiliary load balancing loss.
        
        Uses the coefficient of variation of expert assignment counts,
        weighted by mean routing weight (simplified Switch Transformer formulation).
        
        Args:
            expert_indices: (B, S, k) selected expert IDs.
            routing_weights: (B, S, k) routing weights.
        
        Returns:
            loss: Scalar float.
        """
        B, S, k = expert_indices.shape
        flat_indices = expert_indices.reshape(-1)  # (B*S*k,)
        flat_weights = routing_weights.reshape(-1)   # (B*S*k,)
        
        # Count tokens per expert
        counts = np.bincount(flat_indices, minlength=self.num_experts).astype(np.float64)
        total_tokens = flat_indices.shape[0]
        
        # Fraction of tokens assigned to each expert
        f = counts / total_tokens  # (E,)
        
        # Mean routing weight per expert (average weight when expert is selected)
        # Sum weights per expert, divide by count (avoid div by zero)
        weight_sums = np.bincount(flat_indices, weights=flat_weights, 
                                  minlength=self.num_experts).astype(np.float64)
        with np.errstate(divide='ignore', invalid='ignore'):
            mean_weights = np.where(counts > 0, weight_sums / counts, 0.0)
        
        # Load balancing loss: E * sum(f * mean_weight)
        # This is the standard auxiliary loss from the Switch Transformer paper
        loss = self.num_experts * np.sum(f * mean_weights)
        
        return float(loss)

Step 5: Full forward pass with loss.

    def forward(self, x: np.ndarray) -> Tuple[np.ndarray, float, dict]:
        """
        Full forward pass: gate → route → aggregate, plus load balancing loss.
        
        Args:
            x: Input token embeddings (B, S, D).
        
        Returns:
            output: Routed token representations (B, S, D).
            lb_loss: Load balancing loss scalar.
            metrics: Dict with routing stats for logging.
        """
        routing_weights, expert_indices, logits = self.gate(x)
        output = self.route(x, expert_indices, routing_weights)
        lb_loss = self.load_balancing_loss(expert_indices, routing_weights)
        
        # Compute metrics for observability
        flat_indices = expert_indices.reshape(-1)
        counts = np.bincount(flat_indices, minlength=self.num_experts).astype(np.float64)
        expert_utilization = counts / flat_indices.shape[0]
        
        metrics = {
            'expert_utilization': expert_utilization.tolist(),
            'max_utilization': float(np.max(expert_utilization)),
            'min_utilization': float(np.min(expert_utilization)),
            'load_imbalance_ratio': float(np.max(expert_utilization) / (np.min(expert_utilization) + 1e-8)),
        }
        
        return output, lb_loss, metrics

Step 6: Training loop sketch (for portfolio demonstration).

def train_step(router: MoERouter, x: np.ndarray, lr: float = 0.01):
    """
    One training step: forward, compute loss, update gate weights via gradient descent.
    
    In a real setup, you'd use autodiff (PyTorch/JAX). Here we approximate
    with finite differences for the portfolio demo.
    """
    output, lb_loss, metrics = router.forward(x)
    
    # Dummy task loss (e.g., reconstruction loss)
    task_loss = np.mean((output - x) ** 2)
    total_loss = task_loss + lb_loss
    
    # Simple gradient update on gate weights (numerical approximation)
    eps = 1e-5
    grad_gate = np.zeros_like(router.gate_weights)
    for i in range(router.hidden_dim):
        for j in range(router.num_experts):
            router.gate_weights[i, j] += eps
            _, lb_plus, _ = router.forward(x)
            router.gate_weights[i, j] -= eps
            
            _, lb_minus, _ = router.forward(x)
            grad_gate[i, j] = (lb_plus - lb_minus) / (2 * eps)
    
    router.gate_weights -= lr * grad_gate
    
    return total_loss, metrics

Running and Testing It

Save the code above as moe_router.py and create a test script test_router.py:

import numpy as np
from moe_router import MoERouter

# Initialize router
router = MoERouter(num_experts=8, hidden_dim=512, top_k=2, noise_std=0.5)

# Synthetic input: batch=4, seq_len=16, hidden=512
np.random.seed(42)
x = np.random.randn(4, 16, 512).astype(np.float32)

# Forward pass
output, lb_loss, metrics = router.forward(x)

print(f"Output shape: {output.shape}")  # Expected: (4, 16, 512)
print(f"Load balancing loss: {lb_loss:.4f}")
print(f"Expert utilization: {metrics['expert_utilization']}")
print(f"Load imbalance ratio: {metrics['load_imbalance_ratio']:.2f}")

# Verify output shape matches input
assert output.shape == x.shape, "Output shape mismatch!"

# Verify all experts are used (with noise, they should be)
used_experts = sum(1 for u in metrics['expert_utilization'] if u > 0)
print(f"Active experts: {used_experts}/{router.num_experts}")

print("All tests passed!")

Run it:

python test_router.py

You should see all 8 experts active (thanks to noise injection) and an imbalance ratio close to 1.0. If the ratio is high (>3.0), increase noise_std or adjust the loss scaling.

Proving it works in a portfolio context:

Unit test for gate correctness: Write a test that verifies routing weights sum to 1.0 per token across selected experts.
Load balancing stress test: Feed 10,000 random tokens and assert that no expert receives more than 30% of assignments (tunable threshold). This demonstrates the loss is doing its job.
Gradient sanity check: Compare numerical gradients to analytical ones for the gate weights. Even a simple assert np.allclose(grad_num, grad_analytic, atol=1e-4) is a strong signal to recruiters.

Extending It: Your Roadmap to Senior-Level

A pure NumPy implementation is great for interviews, but senior roles want to see production thinking. Here are six concrete upgrades that transform this toy into a portfolio piece that speaks to systems engineering:

Add a FastAPI serving layer — Wrap the router in a /route POST endpoint that accepts token embeddings and returns routed outputs. Use Pydantic models for input validation and orjson for fast serialization. This demonstrates you can move an algorithm from a notebook into a service that handles concurrent requests.
Implement expert persistence with SQLite — Store expert weights in a SQLite database with a expert_id primary key. Add a load_expert(expert_id) function that streams weights on demand. This shows you understand memory constraints: you can’t always fit all experts in GPU memory, so you need a cold/warm storage strategy.
Add Prometheus metrics export — Expose moe_expert_utilization_total, moe_routing_latency_seconds, and moe_load_balancing_loss as Prometheus metrics. Use the prometheus_client library. This signals observability maturity — you’re not just building a model, you’re building a system that can be monitored in production.
Horizontal scaling with a task queue — Integrate with Celery or RQ to distribute expert computation across multiple worker processes. When an expert is selected, enqueue the token processing task. This demonstrates you understand distributed systems patterns and can handle scaling beyond a single machine.
Fault tolerance via expert fallback — Implement a try/except around expert dispatch that falls back to a “dummy expert” (identity transformation) if a specific expert fails to load or compute. Add circuit breaker logic that disables flaky experts after N consecutive failures. This is the difference between a research prototype and a production system that survives real-world failures.
Benchmarking harness with pytest-benchmark — Create a benchmarks/ directory with benchmarks for: (a) gate forward pass latency at different batch sizes, (b) memory usage as expert count scales, (c) load balancing convergence over training steps. Output results in JSON and commit them to the repo. This gives hiring managers concrete evidence of performance characteristics and scalability limits.

Key Takeaways

Noisy top-k gating is the mechanism that makes MoE practical: additive Gaussian noise prevents expert collapse and naturally balances load during training.
Load balancing loss is an auxiliary objective that penalizes uneven expert utilization. The coefficient of variation of token counts, weighted by mean routing weight, is the standard formulation from the Switch Transformer paper.
Pure-Python implementation forces you to confront matrix operations, softmax numerical stability, and gradient flow at a low level — a signal of foundational understanding to recruiters.
Portfolio narrative: Pair this with FastAPI, Prometheus, and a Dockerfile to show end-to-end ownership from algorithm to serving to monitoring.
Senior-level extensions: Persistence, horizontal scaling, fault tolerance, and benchmarking are the upgrades that transform a toy into a production-grade system.

Why This Project Stands Out on a CV#

Architecture Overview#

Building It Step by Step#

Running and Testing It#

Extending It: Your Roadmap to Senior-Level#

Key Takeaways#

Further Reading#