The implementation below is a minimal but faithful sparse MoE feed-forward layer: a router, top-k gating with renormalized weights, dispatch to the chosen experts, and the standard auxiliary load-balancing loss (the formulation used in the Switch Transformer and largely unchanged since).
python
import torch
import torch.nn as nn
import torch.nn.functional as F
class Expert(nn.Module):
"""A single expert: an ordinary two-layer MLP."""
def __init__(self, d_model, d_ff):
super().__init__()
self.fc1 = nn.Linear(d_model, d_ff)
self.fc2 = nn.Linear(d_ff, d_model)
def forward(self, x):
return self.fc2(F.gelu(self.fc1(x)))
class TopKMoE(nn.Module):
def __init__(self, d_model, d_ff, num_experts=8, top_k=2, aux_loss_coef=0.01):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
self.aux_loss_coef = aux_loss_coef
self.router = nn.Linear(d_model, num_experts, bias=False)
self.experts = nn.ModuleList([Expert(d_model, d_ff) for _ in range(num_experts)])
def forward(self, x):
# x: (batch * seq, d_model) -- flatten tokens before calling this
num_tokens, d_model = x.shape
logits = self.router(x) # (T, E)
probs = F.softmax(logits, dim=-1) # (T, E) -- soft routing distribution
top_vals, top_idx = probs.topk(self.top_k, dim=-1) # (T, k) each
top_vals = top_vals / top_vals.sum(dim=-1, keepdim=True) # renormalize the k chosen weights
out = torch.zeros_like(x)
# Dispatch: for each of the k slots, gather the tokens assigned to
# each expert and run only those through that expert.
for slot in range(self.top_k):
expert_ids = top_idx[:, slot] # (T,) which expert each token wants in this slot
weights = top_vals[:, slot].unsqueeze(-1) # (T, 1)
for e in range(self.num_experts):
mask = expert_ids == e
if mask.any():
out[mask] += weights[mask] * self.experts[e](x[mask])
aux_loss = self._load_balancing_loss(probs, top_idx)
return out, self.aux_loss_coef * aux_loss
def _load_balancing_loss(self, probs, top_idx):
"""Switch-Transformer-style auxiliary loss:
loss = E * sum_e (fraction_of_tokens_to_e * mean_router_prob_for_e)
Minimized when both distributions are uniform over experts."""
num_tokens = probs.shape[0]
# f_e: fraction of tokens for which e was ANY of the top-k choices
one_hot = F.one_hot(top_idx, num_classes=self.num_experts).float() # (T, k, E)
chosen = one_hot.sum(dim=1).clamp(max=1.0) # (T, E), 1 if e in top-k
f = chosen.mean(dim=0) # (E,)
# P_e: average router probability mass assigned to expert e
P = probs.mean(dim=0) # (E,)
return self.num_experts * (f * P).sum()
# --- Demonstrate the effect of the aux loss on utilization ---
torch.manual_seed(0)
d_model, d_ff, num_experts, top_k = 64, 256, 8, 2
moe = TopKMoE(d_model, d_ff, num_experts, top_k, aux_loss_coef=0.0) # start with aux loss OFF
opt = torch.optim.Adam(moe.parameters(), lr=1e-3)
x = torch.randn(512, d_model)
target = torch.randn(512, d_model)
def expert_utilization(moe, x):
with torch.no_grad():
logits = moe.router(x)
top_idx = logits.topk(moe.top_k, dim=-1).indices
counts = torch.bincount(top_idx.flatten(), minlength=moe.num_experts).float()
return counts / counts.sum()
for step in range(300):
out, aux = moe(x)
task_loss = F.mse_loss(out, target)
loss = task_loss + aux
opt.zero_grad(); loss.backward(); opt.step()
print("Utilization with aux_loss_coef=0.0:", expert_utilization(moe, x).round(decimals=2))