Mixture of Experts

A Mixture of Experts (MoE) layer replaces a single, large, densely-activated feed-forward network with several smaller "expert" sub-networks and a learned router that selects only a few of them per input — dramatically increasing total model capacity without proportionally increasing compute cost per token.

The Core Idea — Sparse Activation

Recall the position-wise feed-forward network from Feed-Forward Network — every token passes through the exact same, single feed-forward network. An MoE layer instead maintains \(N\) separate expert networks (each structurally similar to a standard feed-forward network), and for each individual token, a lightweight router selects just a small number (commonly the top-1 or top-2) of experts to actually process that specific token — every other expert is skipped entirely for that token.

The Routing Formula

The router computes a softmax score over all \(N\) experts for the current token \(\mathbf{x}\); only the top-\(k\) highest-scoring experts (commonly \(k=1\) or \(k=2\)) are actually computed, and their outputs are combined, weighted by the router's own scores. Every other expert contributes exactly nothing for this specific token — no wasted computation on unused experts.

Why This Is a Genuinely Different Kind of Scaling

This decoupling — total capacity versus per-token compute cost — is exactly what makes MoE architectures able to reach enormous total parameter counts (sometimes trillions) while keeping the actual compute cost per token comparable to a much smaller dense model.

Diagram

Code

import torch
import torch.nn as nn

class MixtureOfExperts(nn.Module):
    def __init__(self, d_model, d_ff, num_experts, top_k=2):
        super().__init__()
        self.experts = nn.ModuleList([
            nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model))
            for _ in range(num_experts)
        ])
        self.router = nn.Linear(d_model, num_experts)
        self.top_k = top_k

    def forward(self, x):
        router_logits = self.router(x)
        weights, indices = torch.topk(router_logits.softmax(dim=-1), self.top_k, dim=-1)

        output = torch.zeros_like(x)
        for i in range(self.top_k):
            expert_idx = indices[..., i]
            expert_weight = weights[..., i].unsqueeze(-1)
            for e in range(len(self.experts)):
                mask = (expert_idx == e)
                if mask.any():
                    output[mask] += expert_weight[mask] * self.experts[e](x[mask])
        return output

Common Mistakes

• Ignoring load balancing between experts — without an additional balancing loss term encouraging roughly even usage across experts, the router can collapse to relying heavily on just a few "favorite" experts, wasting the extra capacity the other experts represent.
• Confusing MoE's total parameter count with its effective compute cost — a trillion-parameter MoE model can have a per-token compute cost comparable to a much smaller dense model, since only a small fraction of experts activate per token.

Interview Relevance

Q: "How does a Mixture of Experts layer let a model have far more total parameters without a proportional increase in compute cost per token?" Instead of one large, densely-activated feed-forward network processing every token, an MoE layer maintains many smaller expert networks and a lightweight router that selects only a small number (top-1 or top-2) to actually process each specific token. Total parameter count (summed across all experts) can be enormous, but compute cost per token depends only on the small number of experts actually activated, decoupling total capacity from per-token compute.

Practice Question

Why might a Mixture of Experts model risk under-utilizing most of its experts without an explicit load-balancing mechanism during training?

Related DL Notes

• Feed-Forward Network
• LLM Architecture
• Mixture of Experts (Modern LLMs)