A mixture-of-experts (MoE) model increases capacity without using every parameter for every token. Instead of sending each token through the same feed-forward network (FFN), it selects a small subset of specialized expert networks. This conditional computation keeps active computation lower than in a comparably large dense model. section 3.5 section 3.5.1
In a transformer block, an MoE layer usually replaces the FFN after self-attention. It contains multiple expert FFNs and a gating network, or router. The router scores experts for a token’s hidden representation and selects the top k, commonly one or two. Selected scores can then be normalized to weight the chosen experts. section 3.5.1
Listing 3.8 Top-k gating function for sparse MoE
class NativeGate(BaseGate): #1
def __init__(self, d_model, num_expert, world_size, top_k=2, gate_bias=True):
super().__init__(num_expert, world_size) #2
self.gate = nn.Linear(d_model, self.tot_expert, bias=gate_bias)
self.top_k = top_k
def forward(self, inp, return_all_scores=False):
gate = self.gate(inp) #3
gate_top_k_val, gate_top_k_idx = torch.topk( #4
gate, k=self.top_k, dim=-1, largest=True, sorted=False
)
gate_top_k_val = gate_top_k_val.view(-1, self.top_k) #5
gate_score = F.softmax(gate_top_k_val, dim=-1)
self.set_loss(torch.zeros(1, requires_grad=True).to(inp.device))
if return_all_scores: #6
return gate_top_k_idx, gate_score, gate
return gate_top_k_idx, gate_score
The PyTorch gate projects input representations to expert scores, selects expert indices with torch.topk, then applies softmax to the selected scores. It implements routing, not the expert FFNs themselves. section 3.5.1
The trade-off is between total capacity and computation per token: many experts can exist, but only a few run for each token. A GShard-style top-2 gate adds load balancing and random routing to the simpler top-k gate. section 3.5.1
