Sparse Mixture-of-Experts: Dynamic Expert Selection and Aggregation

Sparse Mixture-of-Experts (MoE) architectures represent one of the most transformative paradigms in modern deep learning, enabling models to achieve unprecedented scale while maintaining computational efficiency. Unlike dense neural networks where every parameter participates in every computation, MoE models employ a conditional computation strategy—activating only a subset of specialized "experts" for each input token.

The Architecture of Sparse Expertise

In a Sparse MoE layer, the architecture consists of two key components:

1. Router Network: A lightweight gating mechanism that examines each input token and produces a probability distribution over all available experts
2. Expert Networks: A collection of N independent feed-forward networks (experts), each specializing in different aspects of the input space

Modern flagship models leverage this architecture at massive scale:

• DeepSeek-V3 deploys 671B total parameters with 256 experts, activating only 37B parameters per token
• Kimi K2 utilizes 384 experts with only 8 active per token
• Mixtral 8x7B employs 8 experts with top-2 routing

The Top-K Routing Mechanism

The routing process operates as follows:

Step 1 - Score Computation: For each input token, the router computes raw logit scores for all N experts: $$\text{router_logits} = W_r \cdot x$$

Step 2 - Expert Selection: Select the top-K experts based on their logit scores. Let $\mathcal{T}_k$ denote the set of indices for the top-K experts.

Step 3 - Weight Normalization: Apply softmax only over the selected experts to compute normalized routing weights: $$w_i = \frac{\exp(\text{logit}i)}{\sum{j \in \mathcal{T}_k} \exp(\text{logit}_j)} \quad \text{for } i \in \mathcal{T}_k$$

Step 4 - Expert Output Aggregation: Combine the outputs of selected experts using the computed weights: $$\text{output} = \sum_{i \in \mathcal{T}_k} w_i \cdot E_i(x)$$

where $E_i(x)$ represents the output of expert $i$ for input $x$.

Your Implementation Task

Implement the complete top-K routing logic for a Sparse MoE layer. Given the router logits for each token, the pre-computed outputs from all experts, and the sparsity parameter K, your function should:

1. Identify the top-K experts for each token based on router logits
2. Compute softmax-normalized routing weights over only the selected experts
3. Return the weighted combination of the selected expert outputs

Key Insight: The softmax normalization occurs after expert selection, meaning we compute a probability distribution only over the K selected experts, not over all N experts. This ensures routing weights sum to 1.0 for each token.