11 — Architecture

Mixture of Experts#

Sparse · Router · Top-K✓ Mathematical
◆ The PatternMore parameters without proportionally more compute

Mixture of Experts (MoE) replaces the single FFN with multiple "expert" FFNs and a router that selects which experts process each token. Only the top-K experts activate per token — typically K=2 out of 8–64 experts.

MoE(x) = Σᵢ gᵢ(x) · Expertᵢ(x) where g(x) = TopK(softmax(W_router · x))
Router assigns weights to top-K experts. Each expert is a standard FFN. Inactive experts skip computation entirely.
Mixtral 8×7B: 8 experts, top-2 routing → 47B total, ~13B active per token
7B-quality performance at 13B-compute cost, with 47B parameters of capacity.

Key challenges: (1) Load balancing — prevent all tokens routing to the same expert, fixed with auxiliary loss. (2) Expert collapse — some experts never activate. (3) Communication — experts on different GPUs need token routing across devices.

MoE models need more RAM (all experts loaded) but less compute (only K active). This makes them memory-bound, not compute-bound — great for inference on high-memory hardware.
Interactive — expert routing animation

Python — simplified MoE layer#

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

class MoELayer(nn.Module):
    def __init__(self, d_model, n_experts=8, top_k=2):
        super().__init__()
        self.top_k = top_k
        self.router = nn.Linear(d_model, n_experts, bias=False)
        self.experts = nn.ModuleList([
            SwiGLU_FFN(d_model) for _ in range(n_experts)
        ])

    def forward(self, x):
        # x: (batch, seq, d_model)
        logits = self.router(x)                       # (B, T, E)
        weights, indices = logits.topk(self.top_k)    # top-K experts
        weights = F.softmax(weights, dim=-1)           # normalize
        out = torch.zeros_like(x)
        for k in range(self.top_k):
            for e in range(len(self.experts)):
                mask = indices[..., k] == e
                if mask.any():
                    out[mask] += weights[..., k:k+1][mask] * \
                                 self.experts[e](x[mask])
        return out
Pattern bridge: Routing inputs to specialized sub-networks — only a fraction active per token. In markets, sector rotation activates different expert sectors at different times. In statistics, mixture models combine multiple distributions.
← Previous
Context Windows
Open in the full reader, with the topic sidebar →