Mathematical Formulation of Basic MoE Layers

Mixture of Experts (MoE) layers offer a sparse alternative to dense neural network architectures. The computations involved in a basic MoE layer are mathematically formulated. This mathematical description clarifies precisely how input data flows through the system, how routing decisions are made, and how expert outputs are combined. Understanding this process is important before exploring more complex architectural variants and training optimizations.

Consider an MoE layer integrated within a larger network, such as a Transformer block. This layer receives an input representation for each token in a sequence. Let $x \in \mathbb{R}^d$ represent the input embedding for a single token, where $d$ is the model's hidden dimension. The MoE layer consists of $N$ expert networks, denoted as $E_1, E_2, \dots, E_N$, and a gating network $G$.

Gating Network Operation

The gating network's primary function is to determine which expert(s) should process the input token $x$. A common implementation uses a simple linear transformation followed by a softmax function.

1. Linear Transformation: The input token $x$ is projected using a learnable weight matrix $W_g \in \mathbb{R}^{d \times N}$. This produces logits $h(x) \in \mathbb{R}^N$, one for each expert: $h(x) = x W_g$
2. Softmax Activation: The logits are then passed through a softmax function to generate normalized scores or probabilities, indicating the preference for each expert: $G(x) = \text{softmax}(h(x))$ The resulting vector $G(x) = [G(x)_1, G(x)_2, \dots, G(x)_N]$ contains non-negative values that sum to 1, where $G(x)_i$ represents the gating network's score for assigning token $x$ to expert $E_i$.

Data flow through the gating network, transforming the input token representation $x$ into expert assignment scores $G(x)$.

Expert Networks

Each expert $E_i$ is typically a neural network module, often a standard Feed-Forward Network (FFN) similar to those found in Transformer blocks. Critically, each expert $E_i$ possesses its own distinct set of parameters. For an input token $x$, the output of the $i$-th expert is denoted as $E_i(x)$. While structurally similar, the parameter differentiation allows each expert to potentially specialize in processing different types of input patterns or performing different sub-tasks.

Sparse Routing and Output Computation

While the gating network produces a score $G(x)_i$ for every expert, sparse MoE models leverage conditional computation by activating only a subset of these experts for each token. This is commonly achieved using Top-k routing, where only the $k$ experts with the highest gating scores are selected (typically $k=1$ or $k=2$).

1. Expert Selection: Identify the indices $I$ of the top $k$ experts based on the scores in $G(x)$: $I = \text{TopKIndices}(G(x), k)$ For instance, if $k=2$ and experts $j$ and $l$ have the highest scores, then $I = \{j, l\}$.
2. Gating Score Renormalization (Optional but Common): The gating scores for the selected experts ($i \in I$) are often retained and potentially renormalized. For simplicity here, we'll denote the score used for weighting expert $i$'s output as $g_i = G(x)_i$. Some implementations might re-apply softmax only over the top-k scores.
3. Expert Computation: The input token $x$ is processed only by the selected experts: $O_i(x) = E_i(x) \quad \text{for } i \in I$
4. Weighted Output Combination: The final output $y$ of the MoE layer for token $x$ is computed as the weighted sum of the outputs from the selected experts, using their corresponding gating scores as weights: $y = \sum_{i \in I} g_i O_i(x) = \sum_{i \in I} G(x)_i E_i(x)$

This formulation ensures that only a small fraction ($k/N$) of the experts' parameters are engaged for any given input token, leading to significant computational savings during the forward pass compared to activating all $N$ experts or using a single large expert.

Forward pass for a single token $x$ in an MoE layer with $k=2$ routing. The gating network produces scores, Top-k selects experts (E1 and EN shown here), the input is processed only by selected experts, and their outputs are combined using the gating scores.

Batch Processing Techniques

In practice, MoE layers operate on batches of tokens. For a batch of $T$ tokens, the input is typically represented as a matrix $X \in \mathbb{R}^{T \times d}$. The gating network computes scores $G(X) \in \mathbb{R}^{T \times N}$ for all tokens simultaneously. Crucially, the Top-k selection is performed independently for each token (row) in the batch. This means different tokens within the same batch can be routed to different sets of experts.

While efficient, this independent routing per token introduces challenges, particularly regarding load balancing: ensuring that experts receive a roughly equal amount of computation across the batch. If the gating network consistently routes most tokens to a small subset of experts, others remain underutilized, diminishing the benefits of specialization and potentially destabilizing training. We will address these training dynamics and associated optimization techniques in Chapter 3.

This mathematical framework provides the essential building blocks for understanding how basic MoE layers perform conditional computation. The interaction between the gating network, sparse routing mechanism, and expert networks forms the core of the MoE approach.