Phase 1.1 · OptionalOptionalModule 20 of 21

MoE

Scaling dense models (every parameter active per token) is compute-expensive. MoE activates only a subset of experts per token — e.g., 8 of 64 experts — getting 64× parameter capacity with ~8× compute. Powers GPT-4, Mixtral, DBRX.

A hospital with 64 specialists but each patient only sees 8 relevant ones. The receptionist (router) directs patients to the right specialists. The hospital has vast expertise but each visit is efficient.

Visual Workflows

Start here — study each diagram, then use + / − to zoom if needed.

100%

Loading diagram...

Scroll inside the frame · use + / − to zoom

Key Takeaways

1.Mixture of Experts (MoE) replaces the dense FFN in transformer layers with multiple 'expert' FFNs and a router that selects which experts process each token — dramatically increasing model capacity without proportional compute increase.
2.MoE layer: Router (linear gate) → softmax over experts → select top-k experts → each selected expert is a standard FFN → weighted sum of expert outputs.
3.Mixtral 8x7B: 8 experts per layer, top-2 routing, 47B total params but only ~13B active per token.
4.Challenges: load balancing (some experts overused), communication overhead in distributed training, memory (all experts must be in GPU memory even if only k active).
5.Switch Transformer uses top-1 routing for simplicity.

Real Example

Scenario

Mixtral 8x7B processing 'The cat sat on the mat': router sends 'cat' tokens to experts 2,5 (animal semantics), 'sat' to experts 1,7 (verb actions), 'mat' to experts 3,4 (object nouns). Each token uses only 2 of 8 experts.

What you would do

In Transformer & ML Foundations, apply MoE to this scenario: Mixtral 8x7B processing 'The cat sat on the mat': router sends 'cat' tokens to experts 2,5 (animal semantics), 'sat' to experts 1,7 (verb actions), 'mat' to experts 3,4 (object nouns). Identify the inputs, run the technique, validate the output, and note one thing you would monitor in production.

Practice Task

Open the Code Walkthrough below and run it locally. Change one parameter related to MoE (e.g. model, temperature, top_k, or tool name), observe the difference in output, and write 2–3 sentences explaining what changed.

Code Walkthrough

Highlighted lines show where MoE happens in the code.

MoE

1import torch  # import dependencies2import torch.nn as nn  # import dependencies3
4class MoELayer(nn.Module):  # define a data structure or component5    def __init__(self, d_model, d_ff, num_experts, top_k=2):  # define a reusable function6        super().__init__()7        self.router = nn.Linear(d_model, num_experts)8        self.experts = nn.ModuleList([9            nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))10            for _ in range(num_experts)11        ])12        self.top_k = top_k13
14    def forward(self, x):  # define a reusable function15        gate_scores = torch.softmax(self.router(x), dim=-1)16        top_weights, top_indices = gate_scores.topk(self.top_k, dim=-1)17        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True)18        output = torch.zeros_like(x)19        for i, expert in enumerate(self.experts):20            mask = (top_indices == i).any(dim=-1)21            if mask.any():22                output[mask] += expert(x[mask]) * top_weights[mask, (top_indices[mask] == i).int().argmax(dim=-1)].unsqueeze(-1)23        return output  # return the result

Cheat Sheet

Quick recap

quick ref

•MoE: router + N expert FFNs
•top-k experts per token
•Mixtral 8x7B: 8 experts, top-2
•Active params << total params
•Load balancing loss required
•All experts in GPU memory

Common Mistakes

✕Assuming MoE models use less memory — all experts must be loaded
✕Ignoring load balancing — some experts become useless
✕Comparing total MoE params to dense params directly — compare active params
✕Expecting linear speedup from MoE — routing overhead exists