AI Engineering Primer | Mixture of Experts (MoE) Spelled Out

Chapter Overview: This chapter provides a complete architectural breakdown and step-by-step PyTorch implementation of the Mixture of Experts (MoE) mechanism. We explore how sparse activation transforms standard dense Transformer models by replacing static Feed-Forward Networks (FFNs) with dynamically routed specialist sub-networks, enabling scaling to massive parameter capacities while keeping compute overhead low.

1. Paradigm Shift: Dense vs. Sparse Architectures

In traditional dense Transformer models, every single input token triggers computations across 100% of the model's parameters during both forward pass training and inference. As model parameters scale into hundreds of billions, dense execution imposes prohibitive latency and hardware compute costs.

Mixture of Experts (MoE) breaks this bottleneck by transitioning models from dense activation to sparse activation:

• Total Parameters: Represents the full capacity of all experts combined plus shared layers (e.g., a 70 Billion or 47 Billion parameter model).
• Active Parameters: Represents the subset of parameters executed for a given token during a single forward pass (e.g., only 13 Billion or 20 Billion parameters).

By keeping the majority of expert networks idle at any given millisecond, MoE dramatically increases total parameter capacity (knowledge storage) without incurring the compute costs of dense models.

2. Core Mental Models & Real-World Analogies

To understand how MoE manages conditional computation, consider two key real-world mental models:

Real-Life Analogy 1: Hospital Receptionist & Specialized Departments

Imagine a patient entering a large hospital. The patient presents specific symptoms (the input token). Rather than forcing every patient through every department (Cardiology, Neurology, Orthopedics, Pediatrics), a receptionist (the Router) evaluates the symptoms and directs the patient only to the relevant specialist departments. The remaining departments remain focused on other patients or idle.

Real-Life Analogy 2: Subject-Matter Expert Panel

When faced with complex user prompts, the Router selects a custom team of domain specialists depending on the task demands:

3. Structural Anatomy: MoE Inside the Transformer Layer

In a standard Transformer block, the architecture alternates between Multi-Head Self-Attention (MHSA) and a Feed-Forward Network (FFN). MoE modifies this structure by replacing the static FFN block with an MoE Layer:

1. Self-Attention Block: Processes contextual dependencies across sequence tokens.
2. Router (Gating Network): Takes the hidden state output of the attention layer and computes selection probabilities for each available expert.
3. Expert Bank: A collection of $N$ independent, parallel FFN sub-networks.
4. Top-$K$ Routing Logic: Activates only the top $K$ experts with the highest routing scores (e.g., $K=2$ out of $N=8$).
5. Weighted Combination: Multiplies each selected expert's output by its normalized routing weight and sums the results to form the layer output.

4. Vertical Execution & Routing Data Flow

The vertical data flow for a sequence of tokens passing through an MoE layer proceeds as follows:

5. Mathematical Foundations of Top-K Routing

Let $x \in \mathbb{R}^{d_{\text{model}}}$ be the input token vector.

Step 1: Router Logit Generation

The routing layer calculates raw logits $h(x)$ by projecting $x$ through a trainable router matrix $W_g \in \mathbb{R}^{N \times d_{\text{model}}}$, where $N$ is the total number of experts:

Step 2: Gating Softmax Probability Distribution

Applying the softmax function yields a full probability distribution $P(x)$ across all $N$ experts:

Step 3: Top-$K$ Selection

We extract the $K$ highest probability values and their corresponding expert indices:

Step 4: Weight Renormalization

Because unselected experts are zeroed out, the remaining $K$ gating weights $S_i(x)$ are normalized so their sum equals $1$:

Step 5: Weighted Sum Output Formulation

The final output $y$ of the MoE layer is the weighted sum of outputs from the selected experts $E_i(x)$:

6. PyTorch Implementation: Building Expert Networks

An individual expert is implemented as a standard Feed-Forward Network (FFN) consisting of two linear transformations with an intermediate non-linear activation function.

import torch
import torch.nn as nn
import torch.nn.functional as F

class Expert(nn.Module):
    """
    Individual Expert Module representing a standard Feed-Forward Network (FFN).
    Projects d_model -> d_hidden -> d_model.
    """
    def __init__(self, d_model: int, d_hidden: int):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_hidden)
        self.w2 = nn.Linear(d_hidden, d_model)
        self.activation = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x shape: [num_assigned_tokens, d_model]
        out = self.w1(x)
        out = self.activation(out)
        out = self.w2(out)
        return out

7. PyTorch Implementation: The Sparse MoE Layer

Below is the complete, self-contained PyTorch implementation of the SparseMoE layer, featuring routing logic, probability normalization, token dispatch, expert execution, and output accumulation.

class SparseMoELayer(nn.Module):
    """
    Sparse Mixture of Experts (MoE) Layer built from scratch in PyTorch.
    Routes flattened input tokens dynamically to Top-K active experts.
    """
    def __init__(self, d_model: int, d_hidden: int, num_experts: int, top_k: int):
        super().__init__()
        self.d_model = d_model
        self.num_experts = num_experts
        self.top_k = top_k

        # 1. Instantiate Expert Networks
        self.experts = nn.ModuleList([
            Expert(d_model, d_hidden) for _ in range(num_experts)
        ])

        # 2. Instantiate Router / Gating Network
        self.router = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Input tensor shape: [batch_size, seq_len, d_model]
        batch_size, seq_len, d_model = x.shape

        # Flatten input to a 2D matrix of individual tokens: [num_tokens, d_model]
        flat_x = x.view(-1, d_model)  # num_tokens = batch_size * seq_len
        num_tokens = flat_x.shape[0]

        # Step A: Calculate Router Logits and Softmax Probabilities
        router_logits = self.router(flat_x)  # Shape: [num_tokens, num_experts]
        router_probs = F.softmax(router_logits, dim=-1)  # Shape: [num_tokens, num_experts]

        # Step B: Select Top-K Experts per token
        topk_probs, topk_indices = torch.topk(router_probs, self.top_k, dim=-1)
        # topk_probs shape: [num_tokens, top_k]
        # topk_indices shape: [num_tokens, top_k]

        # Step C: Renormalize Top-K probabilities to sum to 1.0 per token
        topk_probs_norm = topk_probs / topk_probs.sum(dim=-1, keepdim=True)

        # Initialize output buffer with zeros
        final_output = torch.zeros_like(flat_x)  # Shape: [num_tokens, d_model]

        # Step D: Route tokens expert-by-expert
        for expert_idx in range(self.num_experts):
            # Find which tokens routed to this specific expert
            # expert_mask shape: [num_tokens, top_k] (boolean)
            expert_mask = (topk_indices == expert_idx)
            
            # token_indices: row indices in flat_x; k_indices: position in top-k
            token_indices, k_indices = torch.where(expert_mask)

            # If no tokens are routed to this expert, skip execution
            if token_indices.numel() == 0:
                continue

            # Extract assigned tokens
            expert_inputs = flat_x[token_indices]  # Shape: [assigned_tokens, d_model]

            # Execute expert forward pass
            expert_outputs = self.experts[expert_idx](expert_inputs)

            # Retrieve normalized routing weights for these specific assignments
            routing_weights = topk_probs_norm[token_indices, k_indices].unsqueeze(-1)
            # Shape: [assigned_tokens, 1]

            # Scale expert output by routing probability
            weighted_outputs = expert_outputs * routing_weights

            # Accumulate results into final output tensor
            final_output.index_add_(0, token_indices, weighted_outputs)

        # Reshape output back to original 3D input structure: [batch_size, seq_len, d_model]
        return final_output.view(batch_size, seq_len, d_model)

8. Line-by-Line Execution Walkthrough with Tensor Shapes

To illustrate the mathematical transformations during execution, consider a concrete tensor trace using standard hyperparameters:

Hyperparameter Setup

• batch_size = 2
• seq_len = 4
• d_model = 16
• num_experts = 8
• top_k = 2

Tensor Shape Transformations

1. Input Flattening

The sequence structure is flattened into 8 independent token vectors of dimension 16.

2. Router Softmax Probability Matrix

Each row represents a token; each column represents the routing score for one of the 8 experts.

Token 0 Raw Router Probabilities:
[0.05, 0.14, 0.11, 0.33, 0.02, 0.01, 0.28, 0.06]

3. Top-$K$ Selection ($K=2$)

Extracting the top 2 probabilities and indices for Token 0:

• Highest value at Index 3: 0.33 (33%)
• Second highest value at Index 6: 0.28 (28%)

4. Renormalization Calculation

5. Masking and Accumulation Example for Expert 0

When the loop reaches expert_idx = 0:

1. Check expert_mask: Identifies tokens where Expert 0 was selected in topk_indices (e.g., Token 1 and Token 4).
2. Extract input vectors for Token 1 and Token 4 ([2, 16]).
3. Pass vectors through Expert_0 FFN to get outputs ([2, 16]).
4. Multiply Token 1 output by 0.46 and Token 4 output by 0.49.
5. Add weighted outputs into the zero-initialized final_output buffer at indices 1 and 4.

9. Computational Efficiency & Production Realities

Operational Advantages

1. High Capacity, Low Latency: Models can store vast domain knowledge across hundreds of billions of total parameters while evaluating only a small fraction during inference, keeping FLOP requirements low.
2. Specialization: Experts naturally specialize in distinct syntactical structures, mathematical domains, coding languages, or stylistic forms during pre-training.

Production Edge Cases & Pitfalls

• Expert Imbalance (Routing Collapse): Without regularizing constraints, the router network can bias towards a small subset of experts, causing them to receive all tokens while remaining experts stay unassigned.

  Implementation Note: Production MoE architectures (such as Mixtral or DeepSeek) use an auxiliary load-balancing loss to enforce an even distribution of tokens across all available experts during training.

• Empty Masks: If an expert is assigned zero tokens in a batch, calling torch.where() on the mask returns empty tensors. The guard condition if token_indices.numel() == 0: continue avoids unnecessary compute overhead and potential zero-dimension runtime errors.

10. Summary and Key Takeaways

• Sparse Activation Mechanism: MoE layers replace dense FFN blocks with dynamic routing networks, executing only a small subset of available expert parameters per token.
• Total vs. Active Parameters: Total parameters represent the overall model knowledge capacity, whereas active parameters dictate the computational FLOPs required per token forward pass.
• The Router Module: A lightweight linear projection layer followed by a Softmax operation that computes assignment probability distributions over the expert pool.
• Top-$K$ Selection: Limits expert execution to the top $K$ scoring experts (e.g., $K=2$), zeroing out unselected experts to conserve compute.
• Weight Renormalization: Top-$K$ gating probabilities are divided by their sum to ensure the scaling coefficients sum to $1.0$, preserving standard signal scale across layers.
• In-Place Accumulation: Output vectors from selected experts are scaled by their normalized gating weights and combined using tensor index addition (index_add_).
• Real-World Analogy: Like a hospital receptionist routing patients strictly to specialized medical departments, the Router dispatches input tokens only to the most relevant expert sub-networks.