Expert Capacity and Sizing Considerations

Choosing the right number of experts, their individual size, and the processing capacity allocated to each is fundamental to designing effective Mixture of Experts (MoE) models. These choices directly influence the model's parameter count, computational requirements (FLOPs), communication overhead in distributed settings, and ultimately, its ability to learn specialized functions. Getting these hyperparameters right involves navigating a complex set of trade-offs between model performance, training stability, and computational efficiency.

Understanding Expert Capacity

A core concept in MoE design is expert capacity. Since computation is typically batched and parallelized, especially on accelerators like GPUs, we need to define a fixed buffer size for the number of tokens each expert can process within a given computation batch (e.g., per forward pass within a specific device group). This buffer size is the expert's capacity, denoted as $C$.

The capacity is usually determined using a capacity factor (CF), a hyperparameter typically greater than 1. The calculation ensures that the capacity is slightly larger than what would be needed if tokens were perfectly distributed among experts. For a group of $T$ tokens being processed and $N$ experts available for that group, the capacity $C$ for each expert is often calculated as:

$$C = \text{round}\left( \text{CF} \times \frac{T}{N} \right)$$

If the routing mechanism assigns more than $C$ tokens to a particular expert within a processing group, the excess tokens are considered "dropped". Dropped tokens typically bypass the expert computation for that layer, meaning their representations are passed directly through the residual connection. This constitutes an information loss.

Setting the capacity involves a direct trade-off:

• Low Capacity (Low CF, e.g., CF close to 1.0):
  • Pros: Reduces computational cost and memory footprint, as less padding is required for the expert computation tensors.
  • Cons: Increases the likelihood of dropped tokens, especially if the routing is imbalanced. This can hinder learning and degrade model quality. Requires effective load balancing mechanisms (discussed in Chapter 3) to work well.
• High Capacity (High CF, e.g., CF >= 2.0):
  • Pros: Minimizes dropped tokens, providing more stable training signals and potentially higher model quality. More stable to routing imbalances.
  • Cons: Increases computational cost and memory usage significantly. If capacity is much higher than the actual number of tokens routed to most experts, computation is wasted on padding tokens. In extreme cases, computation approaches that of a dense layer.

The optimal CF value is often determined empirically, balancing the acceptable percentage of dropped tokens against the computational budget. Values like 1.25 or 1.5 are common starting points.

Increasing the capacity factor reduces dropped tokens, especially with imperfect load balancing. However, higher CF increases computational overhead due to padding.

Determining the Number of Experts (N)

The total number of experts, $N$, is a primary architectural choice. Increasing $N$ allows for:

• Higher Parameter Count: MoE layers scale model size efficiently. Adding experts increases the total parameters substantially, but the computational cost (FLOPs per token) only increases moderately if $k$ (the number of experts each token is routed to) remains small.
• Finer Specialization: More experts potentially allow the model to learn more distinct, specialized functions for different types of inputs or contexts.
• Hardware Mapping: In distributed settings using Expert Parallelism (Chapter 4), $N$ is often chosen based on the number of available processing devices (e.g., GPUs), with one or more experts assigned to each device. Common values for $N$ in large models include 8, 16, 64, 128, or even more.

However, increasing $N$ also introduces challenges:

• Communication Overhead: Expert parallelism requires All-to-All communication (Chapter 4) to shuffle tokens between devices. More experts generally mean more communication overhead, which can become a bottleneck.
• Load Balancing Difficulty: Ensuring that a larger number of experts are all utilized effectively can be more challenging for the router and auxiliary loss functions (Chapter 3).
• Diminishing Returns: At some point, adding more experts may not yield significant improvements in model quality if the data doesn't naturally decompose into that many distinct specialties, or if the router struggles to assign tokens effectively.

Sizing Individual Experts

Complementary to the number of experts is their individual size, typically referring to the hidden dimensions within the expert network (e.g., the intermediate dimension in a standard Transformer Feed-Forward Network).

• Expert Size vs. Dense Equivalent: An MoE layer replaces a single dense FFN block. If the dense FFN has an intermediate dimension $d_{\text{ff}}$, and the MoE layer has $N$ experts, each with an intermediate dimension $d_{\text{expert}}$, the total parameter count is roughly proportional to $N \times d_{\text{expert}}$. You might choose $d_{\text{expert}}$ such that it's smaller than $d_{\text{ff}}$, but the total parameters ($N \times d_{\text{expert}} \times d_{\text{model}}$) might be much larger than the original dense layer's parameters ($d_{\text{ff}} \times d_{\text{model}}$).
• Computational Cost: The FLOPs per token depend on the size of the experts actually activated. For top-k routing with $k=1$, the FLOPs are related to $d_{\text{expert}}$, not $N \times d_{\text{expert}}$.
• Trade-off with N: For a fixed parameter budget or computational target, there's a trade-off:
  • Fewer, Larger Experts: Each expert has more capacity to model complex functions, potentially requiring less fine-grained routing. Might be easier to load balance.
  • More, Smaller Experts: Encourages higher specialization, but each expert is less powerful individually. Relies more heavily on the router identifying very specific contexts.

The choice often depends on the nature of the task and data, and practical considerations around implementation complexity and distributed training performance.

Interaction with Top-k Routing

Most modern MoE implementations use top-k routing, where the gating network selects the top $k$ experts for each token (typically $k=1$ or $k=2$). This choice interacts significantly with capacity considerations:

• k=1: Each token is routed to a single expert. The total number of token assignments across the batch/group is $T$. The total available capacity is $N \times C$. Load balancing aims to distribute the $T$ assignments roughly evenly among the $N$ experts.
• k=2: Each token is routed to two experts. The total number of token assignments is $2T$. The total available capacity is still $N \times C$. This places significantly more pressure on the system. To avoid excessive token dropping, you might need:
  • A higher capacity factor (CF).
  • More effective load balancing via auxiliary losses.
  • A larger number of experts ($N$) relative to the token batch size ($T$).

Using $k=2$ can sometimes improve model quality by allowing tokens to benefit from multiple specialized functions, but it comes at the cost of increased computation (activating two experts per token) and potentially higher communication and capacity requirements.

Practical Guidance and Monitoring

Important metrics to track during development and training include:

1. Percentage of Dropped Tokens: A primary indicator of whether capacity is sufficient. High values (>1-2%) often signal a problem.
2. Expert Utilization / Load Balance: Monitor metrics like the coefficient of variation (CV) of the number of tokens assigned to each expert, or visualizations of expert assignments. This helps diagnose router pathologies and tune load-balancing losses (Chapter 3).
3. Computational Cost (FLOPs) and Training Throughput: Measure the actual performance impact of different configurations.
4. Overall Model Performance: Track standard metrics like perplexity (for language models) or accuracy on downstream tasks. Ensure that architectural choices translate to better final performance.
5. Memory Usage: Especially relevant in distributed settings, ensure the chosen configuration fits within device memory constraints.

By systematically adjusting $N$, expert size, and CF while observing these metrics, you can converge on an MoE architecture that balances performance, specialization, and computational feasibility for your specific application and hardware environment.