05. Mixture-of-Experts Architecture

• Primary References:

1. 1
   Jacobs, R. A., Jordan, M. I., Nowlan, S. J., & Hinton, G. E. (1991). Adaptive Mixtures of Local Experts. Neural Computation, 3(1), 79–87.
2. 2
   Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. In ICLR 2017.
3. 3
   Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Journal of Machine Learning Research (JMLR), 23(120), 1–39.
4. 4
   Jiang, A. Q., et al. (2024). Mixtral of Experts. arXiv:2401.04088.

• Why I Wrote This Note: While Early-Exit Neural Networks scale computation adaptively along depth, Mixture-of-Experts scales computation conditionally along width. Understanding how learned routers partition the input manifold—and why routers degenerate without explicit load-balancing constraints—is essential for both LLM systems engineering and multi-modal biological/clinical modeling.

• Decoupling Parameters from FLOPs: Given \(N\) FFN experts and top-\(k\) routing, total parameters scale as \(\mathcal{O}(N \cdot d_{\text{model}} d_{\text{ff}})\), while active compute per token scales as \(\mathcal{O}(k \cdot d_{\text{model}} d_{\text{ff}})\). For Mixtral (\(N=8, k=2\)), a \(46.7\text{B}\) parameter model executes the FLOPs of a \(\sim 12.9\text{B}\) dense model.
• Why Routers Collapse ("Rich-Get-Richer"): Early in training, whichever expert \(E_i\) happens to produce slightly lower loss receives larger gradients from the gating network, which increases its router logit \(h_i(x)\), sending it even more tokens until \(1\text{–}2\) experts monopolize \(100\%\) of traffic and the remaining \(N-2\) experts die.
• Differentiable Load Balancing: Because the indicator function \(\mathbf{1}(\text{token } x \text{ routed to } i)\) has zero gradient almost everywhere, Switch Transformer multiplies the non-differentiable dispatch fraction \(f_i\) by the differentiable mean router probability \(P_i\), yielding \(\mathcal{L}_{\text{aux}} = \alpha N \sum_{i=1}^{N} f_i \cdot P_i\), which attains its global minimum of \(\alpha\) when \(f_i = P_i = 1/N\).
• Hardware Capacity Factor (\(CF\)): Static tensor shapes on GPUs/TPUs require bounding each expert's buffer to \(\text{Capacity} = \lceil \frac{T \cdot k}{N} \times CF \rceil\); excess tokens overflow via residual pass-through.

4.1 Sparse Top-K Gating Formulation

Let \(x \in \mathbb{R}^d\) represent a token or patch embedding entering an MoE FFN block containing \(N\) independent experts \(\{E_1, \dots, E_N\}\), where each expert is typically a SwiGLU MLP:

A lightweight linear router parameterized by \(W_g \in \mathbb{R}^{N \times d}\) computes raw affinity logits \(\mathbf{h}(x) = W_g x \in \mathbb{R}^N\). To enforce sparsity, we select the indices \(\mathcal{T}_k(x) \subset \{1, \dots, N\}\) corresponding to the \(k\) largest logits:

Applying the softmax operator yields a sparse probability vector \(G(x) \in \Delta^{N-1}\) with exactly \(k\) non-zero entries:

Switch (\(k=1\)) vs. Mixtral (\(k=2\)) Order of Operations: In Mixtral (\(k=2\)), TopK is applied before Softmax, so the two selected experts' weights sum to \(1\). In Switch Transformer (\(k=1\)), applying Softmax before TopK is strictly required; otherwise \(\text{Softmax}(\text{Top1}(\mathbf{h}))\) would trivially equal \(1.0\), killing the gradient \(\partial y / \partial W_g\)!

flowchart LR
    X["Token Embedding x ∈ ℝᵈ"] --> R["Router Logits h(x) = W_g x"]
    R --> TK["Top-K Selection (k=2 of N)"]
    TK --> E2["Expert 2: E₂(x)"]
    TK --> E5["Expert 5: E₅(x)"]
    TK -. "Zero Dispatch (G_i = 0)" .-> E1["Inactive Experts {E₁, E₃, E₄, ...}"]
    E2 & E5 --> AGG["Weighted Sum: y = G₂(x)E₂(x) + G₅(x)E₅(x)"]
    R --> AUX["Auxiliary Loss: L_aux = α N Σ f_i · P_i"]

4.2 Preventing Expert Collapse: The Load-Balancing & Z-Loss Objectives

Given a batch of \(T\) tokens \(\mathcal{B} = \{x_1, \dots, x_T\}\) and \(N\) experts, let \(p(x) = \text{Softmax}(W_g x) \in \mathbb{R}^N\) denote the dense router probability distribution over all \(N\) experts. We define two batch-level vectors:

1. 1
   Actual Token Dispatch Fraction (\(f_i\)): The fraction of tokens in batch \(\mathcal{B}\) routed to expert \(i\):

1. 1
   Mean Router Probability Mass (\(P_i\)): The average probability assigned to expert \(i\) across batch \(\mathcal{B}\):

The Switch Auxiliary Load-Balancing Loss (Fedus et al., 2022) penalizes the dot product between \(\mathbf{f}\) and \(\mathbf{P}\):

Under uniform routing (\(f_i = 1/N\) and \(P_i = 1/N\)), \(\sum_{i=1}^{N} f_i P_i = N \cdot (1/N^2) = 1/N\), so \(\mathcal{L}_{\text{aux}} = \alpha\) (typically \(\alpha \in [0.01, 0.02]\)). Additionally, ST-MoE (Zoph et al., 2022) introduces the Router Z-Loss to prevent numerical instability from large router logits:

4.3 PyTorch Implementation of a Sparse Top-K MoE Layer with Load-Balancing Loss

import torch
import torch.nn as nn
import torch.nn.functional as F

class SparseMoELayer(nn.Module):
    # Sparse Top-K Mixture-of-Experts FFN with Switch load-balancing and router Z-loss
    def __init__(self, d_model: int, d_ff: int, num_experts: int = 8, top_k: int = 2, aux_alpha: float = 0.01):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.aux_alpha = aux_alpha
        self.router = nn.Linear(d_model, num_experts, bias=False)
        self.experts = nn.ModuleList([
            nn.Sequential(nn.Linear(d_model, d_ff), nn.SiLU(), nn.Linear(d_ff, d_model))
            for _ in range(num_experts)
        ])

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        # x shape: (batch_size, seq_len, d_model)
        orig_shape = x.shape
        x_flat = x.view(-1, orig_shape[-1])  # (T, d_model)
        num_tokens = x_flat.size(0)

        router_logits = self.router(x_flat)  # (T, N)
        dense_probs = F.softmax(router_logits, dim=-1)  # P_i uses dense softmax

        topk_logits, topk_indices = torch.topk(router_logits, k=self.top_k, dim=-1)  # (T, k)
        topk_gates = F.softmax(topk_logits, dim=-1)  # Normalized over selected k experts

        # Compute Switch Auxiliary Load-Balancing Loss: alpha * N * sum(f_i * P_i)
        dispatch_mask = F.one_hot(topk_indices, num_classes=self.num_experts).float()  # (T, k, N)
        f_i = dispatch_mask.mean(dim=(0, 1))      # (N,) fraction of dispatches per expert
        P_i = dense_probs.mean(dim=0)             # (N,) mean router probability per expert
        aux_loss = self.aux_alpha * self.num_experts * torch.sum(f_i * P_i)

        # Dispatch tokens to active experts and accumulate weighted outputs
        out_flat = torch.zeros_like(x_flat)
        for expert_idx, expert in enumerate(self.experts):
            token_idx, k_pos = torch.where(topk_indices == expert_idx)
            if token_idx.numel() > 0:
                expert_out = expert(x_flat[token_idx])
                gates = topk_gates[token_idx, k_pos].unsqueeze(-1)
                out_flat.index_add_(0, token_idx, gates * expert_out)

        return out_flat.view(orig_shape), aux_loss

1. 1
   Multi-Modal Protein Function Prediction (IEEE ICEEICT 2024): In biological and clinical learning, samples vary wildly in modality completeness and structural complexity. An MoE gating network offers a principled mechanism to route sequence motifs, structure graphs, and physicochemical features to specialized sub-networks rather than forcing a single monolithic MLP to compress heterogeneous signals.
2. 2
   Uncertainty-Aware Routing: A fascinating open question in my research is Router Entropy: when \(\mathcal{H}[p(x)]\) across experts is high, should the token dynamically expand \(k\) (e.g., from \(k=1 \to k=3\), Dynamic-\(k\) MoE) to hedge epistemic uncertainty, while confident tokens use \(k=1\)?

• Do Experts Specialize by Semantic Topic or by Token Syntax? Empirical inspections of Switch Transformer and Mixtral routers (Jiang et al., 2024) reveal a surprising truth: MoE experts rarely specialize by macroscopic "domain" (e.g., an "all-biology expert" vs. an "all-history expert"). Instead, routing is overwhelmingly syntactic and structural—specific experts specialize in punctuation, whitespace in Python code, mathematical operators, or prepositions.
• Shared Expert + Routed Experts (DeepSeekMoE Pattern): Isolating \(1\) or \(2\) always-active shared experts alongside \(N\) finely segmented routed experts prevents redundant knowledge duplication (every routed expert no longer needs to re-learn basic grammar or edge filters).

• Implement SparseMoELayer with vectorized index_add_ dispatch and Switch \(\mathcal{L}_{\text{aux}}\) loss.Verified
• Implement Dynamic-\(k\) Uncertainty Routing, where the number of active experts \(k(x) \in \{1, \dots, K_{\max}\}\) is chosen until cumulative router probability exceeds \(1 - \epsilon\).Pending
• Compare router entropy \(\mathcal{H}[p(x)]\) against predictive class entropy \(\mathcal{H}[p(y\mid x)]\) on out-of-distribution inputs.Pending

"MoE is FLOP-sparse for a single token, but memory-dense for a batch of tokens: if you send 32 diverse tokens through an 8-expert layer, you will almost certainly wake up all 8 experts."

Switch Auxiliary Load-Balancing & Z-Loss Objective:

• Connected Research Notes:
• Early-Exit Neural Networks — Width-wise conditional compute (MoE) vs. depth-wise adaptive compute (Early-Exit).
• Deep Ensembles vs MC Dropout — MoE combines \(k\) specialized sub-networks via input-dependent gates \(G_i(x)\), whereas Deep Ensembles combine \(M\) full models via uniform averaging \(1/M\).
• LoRA: Mathematical Foundations — How low-rank adapters can parameterize lightweight experts (MoLoRA) without \(N\times\) VRAM blowup.