Sparse Mixture of Experts Forward Pass

The Mixture-of-Experts (MoE) architecture is a powerful paradigm in deep learning that enables efficient scaling of neural networks by conditionally activating only a subset of model parameters for each input. Instead of processing every input through the entire network, MoE employs a gating mechanism to dynamically route inputs to the most relevant "expert" sub-networks, dramatically improving computational efficiency while maintaining or even increasing model capacity.

Conceptual Foundation

In a Sparse MoE layer, we have:

• E experts: Each expert is a separate neural network (typically a feed-forward network) with its own parameters
• A gating network: Computes routing probabilities that determine which experts process each input token
• Top-k selection: Only the top k experts (out of E) are activated for each input, introducing sparsity

This sparsity is the key innovation—while the model has massive total capacity (all experts combined), the computation cost per input scales with k, not E.

Mathematical Formulation

Given an input tensor x of shape (batch_size, sequence_length, d_input), the Sparse MoE forward pass proceeds as follows:

Step 1: Compute Gating Logits

For each token, compute the raw gating scores by multiplying with the gating weight matrix: $$\text{logits} = x \cdot W_g$$ where W_g has shape (d_input, num_experts), producing logits of shape (batch_size, seq_len, num_experts).

Step 2: Apply Softmax for Probabilities

Convert logits to probabilities using the softmax function: $$g_i = \frac{e^{\text{logits}i}}{\sum{j=1}^{E} e^{\text{logits}_j}}$$

Step 3: Top-k Expert Selection

Select the indices of the top-k experts with the highest gating probabilities for each token. This ensures sparse activation—only k experts are utilized per token.

Step 4: Renormalize Selected Probabilities

Renormalize the gating probabilities for only the selected experts so they sum to 1: $$\hat{g}i = \frac{g_i}{\sum{j \in \text{top-k}} g_j}$$

Step 5: Expert Computation

For each selected expert, compute its output by applying its weight matrix to the input token: $$\text{expert_output}i = x \cdot W{e_i}$$ where W_e has shape (num_experts, d_input, d_output).

Step 6: Weighted Aggregation

Combine the expert outputs using the renormalized gating weights: $$y = \sum_{i \in \text{top-k}} \hat{g}_i \cdot \text{expert_output}_i$$

Your Task

Implement the sparse_moe_forward function that computes the forward pass of a Sparse Mixture-of-Experts layer. Your function should:

1. Compute gating logits and apply softmax to get routing probabilities
2. Select the top-k experts for each token based on gating scores
3. Renormalize the selected expert probabilities
4. Apply each selected expert's transformation to the input
5. Aggregate expert outputs weighted by the renormalized probabilities
6. Return the final output tensor