Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Expert Networks › Code Implementation

Michael BrenndoerferNovember 13, 20259 min read

Language AI HandbookAll chapters

Part 33 of 67Mixture of Experts

Chapter 240Expert Networks › Code Implementation

Reading section 8 of 12 · Code Implementation · 9 minSection 8 of 12 · 9 min

Section 8 of 12 · 9 min read

Expert Networks › Code ImplementationOne section · Sections in this chapter8 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's implement a simple expert network structure and explore capacity-constrained processing.

In[7]:

Code

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


# Expert as a standard FFN with SiLU activation
class Expert(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=True)
        self.w2 = nn.Linear(d_ff, d_model, bias=True)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)))


# Create experts with shared configuration
d_model = 256
d_ff = 512
num_experts = 8

experts = nn.ModuleList([Expert(d_model, d_ff) for _ in range(num_experts)])

Each expert is an independent FFN with its own weight matrices $W_1$ and $W_2$ . The nn.ModuleList container ensures that PyTorch registers all expert parameters for gradient tracking and optimizer updates. Note that creating experts inside a ModuleList is needed: if you used a plain Python list, the expert parameters would not appear in model.parameters() and would not receive gradient updates during training.

Let's examine the parameter counts to verify our earlier analysis:

In[8]:

Code

def count_parameters(module):
    return sum(p.numel() for p in module.parameters())


single_expert_params = count_parameters(experts[0])
total_expert_params = count_parameters(experts)

Out[9]:

Console

Parameters per expert: 262,912
Total expert parameters: 2,103,296
Expansion factor: 8x over single FFN

The total parameter count (approx. 2.1 million) reflects the sum of all 8 experts, but the active parameters per token remain small (approx. 260,000, since only 2 experts activate per token in a typical top-2 routing scheme). This is the MoE tradeoff made concrete: 8 times the stored parameters, but only 2 times the per-token compute.

Now let's implement capacity-constrained expert processing. This simulates how tokens are assigned to experts with overflow handling:

In[10]:

Code

def process_with_capacity(
    tokens,  # (batch_size * seq_len, d_model)
    expert_assignments,  # (num_tokens,) - which expert each token goes to
    experts,  # list of expert modules
    capacity_factor=1.25,
):
    """
    Process tokens through experts with capacity constraints.
    Returns processed tokens and drop statistics.
    """
    num_tokens = tokens.shape[0]
    num_experts = len(experts)
    d_model = tokens.shape[1]

    # Calculate capacity per expert
    base_capacity = num_tokens // num_experts
    capacity = int(capacity_factor * base_capacity)

    # Track assignments and outputs
    output = torch.zeros_like(tokens)
    expert_counts = torch.zeros(num_experts, dtype=torch.long)
    dropped_counts = torch.zeros(num_experts, dtype=torch.long)

    # Process each token
    for i in range(num_tokens):
        expert_idx = expert_assignments[i].item()

        if expert_counts[expert_idx] < capacity:
            # Expert has capacity: process the token
            with torch.no_grad():  # For demonstration only
                output[i] = experts[expert_idx](tokens[i : i + 1]).squeeze(0)
            expert_counts[expert_idx] += 1
        else:
            # Expert at capacity: token is dropped (uses residual)
            output[i] = tokens[i]  # Pass through unchanged
            dropped_counts[expert_idx] += 1

    return output, expert_counts, dropped_counts

This implementation processes tokens sequentially to clearly show the capacity-checking logic. In a production implementation, you would batch all tokens for each expert together and process them simultaneously, but the sequential version makes the capacity constraint logic transparent. Notice that dropped tokens receive tokens[i] as their output: this simulates the residual connection carrying the token representation unchanged past the MoE layer.

Let's simulate imbalanced routing to see capacity constraints in action:

In[11]:

Code

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

# Recreate experts for this cell (needed since earlier cell has output: false)
d_model = 256
d_ff = 512
num_experts = 8


class Expert(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=True)
        self.w2 = nn.Linear(d_ff, d_model, bias=True)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)))


experts = nn.ModuleList([Expert(d_model, d_ff) for _ in range(num_experts)])

# Generate random tokens
torch.manual_seed(42)
batch_tokens = 256
tokens = torch.randn(batch_tokens, d_model)

# Create imbalanced routing: some experts more "popular"
# Expert 0 and 1 get most tokens
probs = torch.tensor([0.3, 0.25, 0.1, 0.1, 0.08, 0.07, 0.05, 0.05])
expert_assignments = torch.multinomial(probs, batch_tokens, replacement=True)

# Define capacity factor
capacity_factor = 1.25

# Process with capacity constraint
output, counts, dropped = process_with_capacity(
    tokens, expert_assignments, experts, capacity_factor=capacity_factor
)

total_dropped = dropped.sum().item()
drop_rate = 100 * total_dropped / batch_tokens

# Calculate per-expert statistics
expert_stats = []
for i in range(num_experts):
    assigned = (expert_assignments == i).sum().item()
    processed = counts[i].item()
    dropped_count = dropped[i].item()
    expert_stats.append((i, assigned, processed, dropped_count))

The routing probabilities here are intentionally skewed: experts 0 and 1 receive over half of all tokens, while experts 6 and 7 receive only 5% each. This kind of imbalanced routing is what the capacity mechanism must handle without letting the popular experts to become bottlenecks.

Out[12]:

Console

Expert utilization with capacity_factor=1.25:
Base capacity: 32 tokens per expert
Actual capacity: 40 tokens per expert

Expert 0: 82 assigned, 40 processed, 42 dropped
Expert 1: 65 assigned, 40 processed, 25 dropped
Expert 2: 28 assigned, 28 processed, 0 dropped
Expert 3: 22 assigned, 22 processed, 0 dropped
Expert 4: 22 assigned, 22 processed, 0 dropped
Expert 5: 15 assigned, 15 processed, 0 dropped
Expert 6: 10 assigned, 10 processed, 0 dropped
Expert 7: 12 assigned, 12 processed, 0 dropped

Total drop rate: 26.2%

Out[13]:

Visualization

Bar chart of token counts routed to experts E0 through E7. Overloaded experts (E0, E1) have bars split into accepted tokens (blue) and dropped tokens (red) above the capacity threshold dashed line; underutilized experts (E5-E7) show only partial blue bars. — Token load distribution across 8 experts under imbalanced routing conditions. Popular experts (e.g., E0, E1) exceed the capacity threshold (dashed line) and must drop excess tokens (red), while underutilized experts leave allocated capacity unused.

The imbalanced routing causes popular experts to drop tokens. Let's visualize how different capacity factors affect drop rates:

In[14]:

Code

capacity_factors = [1.0, 1.1, 1.25, 1.5, 1.75, 2.0]
drop_rates = []

for cf in capacity_factors:
    _, _, dropped = process_with_capacity(
        tokens, expert_assignments, experts, cf
    )
    drop_rate = 100 * dropped.sum().item() / batch_tokens
    drop_rates.append(drop_rate)

Out[15]:

Visualization

Line chart showing token drop rate decreasing from about 32% at capacity factor 1.0 to about 7% at capacity factor 2.0. — Relationship between capacity factor and token drop rate under skewed routing. Increasing the capacity factor reduces drops substantially, but because experts 0 and 1 receive a highly disproportionate share of tokens, even C=2.0 does not eliminate dropping entirely, illustrating that capacity alone cannot compensate for severe routing imbalance.

The plot demonstrates that increasing the capacity factor reduces the drop rate. In this simulation, the heavy skew in routing (30% of tokens going to a single expert) means some dropping persists even at $C = 2.0$ . This is an important lesson: capacity factors are a management tool, not a cure. If the routing mechanism produces severely imbalanced assignments, capacity alone cannot compensate. You need both appropriate capacity and good routing, which is exactly why load-balancing auxiliary losses (covered in upcoming chapters) are so important for making MoE training work.

The diversity across experts also becomes apparent when we process the same token through all 8 experts. In an untrained model, the outputs differ simply because the weight matrices are randomly initialized. After training, this diversity becomes meaningful: each expert has learned to produce outputs optimized for its assigned token distribution.

In[16]:

Code

import torch

# Simulated hidden representations for different token types
# (In practice, these come from earlier transformer layers)
code_token = torch.randn(1, d_model)  # Represents "def"
english_token = torch.randn(1, d_model)  # Represents "the"
math_token = torch.randn(1, d_model)  # Represents a math symbol

# Each expert produces different outputs
# due to their independent weights
with torch.no_grad():
    code_outputs = [experts[i](code_token) for i in range(num_experts)]
    english_outputs = [experts[i](english_token) for i in range(num_experts)]

# Measure output diversity across experts
code_var = torch.stack(code_outputs).var(dim=0).mean().item()
english_var = torch.stack(english_outputs).var(dim=0).mean().item()

Out[17]:

Console

Code token output variance across experts: 0.0396
English token output variance across experts: 0.0416

Each expert transforms the same input differently. After training, a well-functioning MoE learns to route tokens to experts whose transformations are most appropriate for that token type. The router, which we will explore in the next chapter, learns to make these assignments based on the input representation.

Out[18]:

Visualization

Bar chart of L2 distance from the mean output for each of 8 experts processing the same code token. Bars vary in height across experts, showing that each expert produces a distinct transformation. — L2 distance between individual expert outputs and the mean output for the same code token. The substantial deviation of each expert from the mean demonstrates that experts learn distinct transformations, giving specialized processing for the same input.