Mixture of Experts: Sparse Activation, Smaller "Expert" Feed-Forward Networks

Mixture of Experts (MoE) is a popular architecture that uses different 'experts' to improve Transformer models, and the difference between Transformer and MoE architectures shows up specifically in the decoder block: where a standard Transformer uses one feed-forward network, MoE uses multiple experts — which are themselves feed-forward networks, but each individually SMALLER than the single FFN a dense Transformer would use at that position.

The key behavioral difference: during inference, only a SUBSET of experts is selected and actually run for each token, not all of them. This is exactly what makes inference faster in MoE relative to an equivalently-sized dense model — even though the total parameter count across all experts combined can be large (keeping overall model capacity high), any single token's forward pass only activates a small fraction of those parameters, not all of them. This is precisely the decoupling the previous subtopic identified as missing in dense architectures: MoE lets total parameter count (capacity) grow largely independently of per-token compute cost.

Since a Transformer has multiple decoder layers stacked, this expert-selection process compounds across the network in two ways: text passes through DIFFERENT experts across different layers (each layer makes its own independent selection), and the chosen experts also differ BETWEEN tokens — token A and token B in the same sequence, at the same layer, can each get routed to entirely different experts based on what each one individually needs.

This allows models to grow in capacity without proportional increases in compute — exactly the scaling lever dense architectures can't offer. Mixtral 8x7B by MistralAI is one famous real-world LLM built on this MoE architecture, demonstrating the approach at production scale rather than just as a research idea.

import torch
import torch.nn as nn

class Expert(nn.Module):
    """An individual MoE expert — a feed-forward network, but
    smaller than a dense Transformer's single FFN would be."""
    def __init__(self, d_model: int, d_ff_expert: int):
        super().__init__()
        self.up = nn.Linear(d_model, d_ff_expert)
        self.down = nn.Linear(d_ff_expert, d_model)
        self.activation = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down(self.activation(self.up(x)))

class MoELayer(nn.Module):
    """Only a SUBSET of experts runs per token — unlike the dense
    FFN, this layer's total parameter count is much larger, but
    per-token active compute stays small."""
    def __init__(self, d_model: int, d_ff_expert: int, num_experts: int = 8):
        super().__init__()
        self.experts = nn.ModuleList([Expert(d_model, d_ff_expert) for _ in range(num_experts)])

    def forward(self, x: torch.Tensor, chosen_expert_idx: int) -> torch.Tensor:
        # In practice this selection is per-token and learned (next
        # subtopic) — simplified here to illustrate sparse activation
        return self.experts[chosen_expert_idx](x)

d_model, d_ff_expert = 4096, 4096  # notably smaller than a dense FFN's d_ff
moe_layer = MoELayer(d_model, d_ff_expert, num_experts=8)
total_params = sum(p.numel() for p in moe_layer.parameters())
active_params_per_token = sum(p.numel() for p in moe_layer.experts[0].parameters())
print(f"Total MoE layer params: {total_params:,}")
print(f"Active params for ONE token: {active_params_per_token:,} (~1/8th)")