Back to Curriculum
Advanced•Modern AI
Mixture of Experts (MoE)
Sparse routing mechanisms, gating networks, Top-K dispatch, and auxiliary load balancing.
Interactive Playground
Initializing Interactive Playground...
Research-Level Deep Dive & Equations
Mixture of Experts (MoE) replaces standard dense Feed-Forward Network (FFN) layers in Transformers (e.g., Mixtral 8x7B, DeepSeek-V3) with parallel expert networks and a trainable router (gating network) .
•Gating Network Equations:
•Expert Output Combination:
•Sparse Computation Advantage: Even with experts (total 47B parameters), routing only experts per token executes only 13B active FLOPs per inference step!
Key Equations
PyTorch Sparse MoE Router & Top-K Dispatch Blockpython
import torch
import torch.nn as nn
import torch.nn.functional as F
class SparseMoERouter(nn.Module):
def __init__(self, d_model: int, num_experts: int = 8, top_k: int = 2):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
self.gate = nn.Linear(d_model, num_experts, bias=False)
# Parallel expert networks
self.experts = nn.ModuleList([
nn.Sequential(
nn.Linear(d_model, d_model * 4),
nn.SiLU(),
nn.Linear(d_model * 4, d_model)
) for _ in range(num_experts)
])
def forward(self, x: torch.Tensor):
# x: [Batch, Seq_len, d_model]
batch_size, seq_len, d_model = x.shape
x_flat = x.view(-1, d_model)
# Router logits & Top-K gating
logits = self.gate(x_flat)
weights, indices = torch.topk(F.softmax(logits, dim=-1), self.top_k, dim=-1)
out = torch.zeros_like(x_flat)
for i, expert in enumerate(self.experts):
# Find tokens dispatched to expert i
batch_idx, kth_idx = torch.where(indices == i)
if len(batch_idx) > 0:
expert_in = x_flat[batch_idx]
expert_out = expert(expert_in)
gating_weight = weights[batch_idx, kth_idx].unsqueeze(-1)
out[batch_idx] += expert_out * gating_weight
return out.view(batch_size, seq_len, d_model)Test Your Knowledge
Check whether you have mastered this concept with a quick quiz.
Was this lesson helpful?
Your feedback helps us continuously improve the curriculum and interactive visualizations.