Mixture-of-Experts Layer Architecture

Introduction

In our last lesson, we established the "what" and "why" of MoE. We learned that by replacing the dense, computationally expensive Feed-Forward Network (FFN) layers in a Transformer with a sparse MoE layer, we can dramatically increase a model's parameter count (its knowledge) without a proportional increase in computational cost. This is achieved through conditional computation, where only a fraction of the network is activated for any given input token.

We identified the two main components of this architecture at a high level:

1. A set of specialized FFNs called experts.
2. A router or gating network that decides which expert(s) to use.

Today, we will move from the "what" to the "how." Our goal is to dissect the MoE layer and understand its internal architecture. By the end of this lesson, you will be able to describe the architecture of a Mixture-of-Experts layer, including its router network and parallel expert FFNs. We will look at how these components are structured and how they interact to achieve sparse activation.

1. Replacing the Dense FFN Layer

The core architectural change introduced by MoE occurs within the Transformer block. As a reminder, a standard Transformer block consists of a self-attention layer followed by a feed-forward network (FFN), with residual connections and normalization. In an MoE model, the FFN layer is swapped out for an MoE layer.

The animation below provides a clear visual comparison of a standard (dense) decoder block and an MoE decoder block.

As you can see, the fundamental structure of the block remains the same, but the FFN component becomes more complex. Let's break down the two key parts of this new MoE layer: the experts and the router.

2. The Experts: A Bank of Parallel FFNs

The "experts" in an MoE Transformer are, quite simply, Feed-Forward Networks. They have the same internal structure as the FFNs we've studied previously (e.g., using SwiGLU activations).

The key difference is that instead of having one single, large FFN, the MoE layer contains a "bank" of multiple, parallel FFNs. Each of these FFNs is an "expert."

You can think of this using our previous analogy of a large corporation. Instead of a single, monolithic "General Processing Department," the company now has several smaller, specialized departments (the experts), such as "Accounting," "Legal," and "Engineering." Each department is a fully functional unit capable of handling requests. The power of this model comes from not sending every request to every department.

3. The Router: The Decision-Making Component

How does the model decide which of the parallel experts to use for a given token? This is the job of the router, also known as the gating network.

The router is itself a small, learnable neural network. Its sole purpose is to look at an incoming token's hidden state (the vector representing that token after the attention layer) and assign a score to each expert, indicating how suitable that expert is for processing this specific token.

The following diagram illustrates the data flow within a complete MoE layer.

Let's formalize this process, as it's a cornerstone of MoE architecture.

Mathematical Formulation of the Router

For an input token represented by the hidden state vector $x$, the router performs the following steps:

1. Compute Logits: The router has a trainable weight matrix, $W_g$. It computes an "affinity score" or "logit" for each expert by taking the dot product of the input vector $x$ and the router's weight matrix. $$\text{logits}=x\cdot W_g$$
2. Compute Gate Probabilities: These logits are then passed through a softmax function. This converts the raw scores into a probability distribution, where each probability $G(x{)}_i$ represents the router's "confidence" that the i-th expert should be used. $$G(x)=\text{softmax}(\text{logits})$$

The video segment below provides a clear, animated walkthrough of how the router uses these probabilities to select experts and combine their outputs.

The Full MoE Forward Pass

As you saw in the video, the final output of the MoE layer, $y$, is a weighted sum of the outputs from all the individual experts $E_i(x)$.

where $N$ is the total number of experts.

However, computing the output of every expert for every token would defeat the purpose of sparsity. This is where Top-K routing comes in. In practice, the router doesn't use the full probability distribution. Instead, it selects only the top K experts with the highest probabilities (where K is a small number, typically 1 or 2). All other expert weights are set to zero.

The output calculation then becomes a weighted sum of only the top-K activated experts. This ensures that the computational cost remains low, as only a small subset of the total experts are ever used.

4. Architectural Variants: Shared and Routed Experts

The basic architecture of a router plus a set of parallel experts forms the foundation of MoE. However, state-of-the-art models like DeepSeek-V3, which this course focuses on, introduce a clever refinement to this structure.

The following video from a Stanford course explains this innovation.

As the video described, this advanced architecture splits the experts into two types:

1. Routed Experts: This is the bank of parallel FFNs we have been discussing, which are activated conditionally by the router (e.g., top-k selection).
2. Shared Experts: These are one or more additional FFNs that every single token passes through, regardless of the router's decision. They are thought to handle common, generic computations that are useful for all tokens, while the routed experts handle more specialized processing.

This hybrid approach allows the model to balance specialized knowledge (in the many routed experts) with general processing efficiency (in the few shared experts).

Conclusion

In this lesson, we have opened up the "black box" of the MoE layer and examined its internal architecture.

Key Takeaways:

• An MoE layer replaces the standard FFN layer within a Transformer block.
• It consists of two main components: a set of parallel FFNs called experts and a learnable router network.
• The router is typically a simple linear layer whose output logits are passed through a softmax function to generate probabilities, or "gates," for each expert.
• To maintain sparsity, only the Top-K experts with the highest gate values are selected to process a given token.
• The final output is a weighted sum of the outputs from these selected experts.
• Advanced architectures, like that of DeepSeek-V3, may also include shared experts that process all tokens, complementing the specialized routed experts.

Preview of the Next Lesson:

We now have a solid picture of the static components of an MoE layer. However, the dynamic process of routing tokens raises important questions. How does the model ensure that it uses a diverse set of experts, rather than always picking the same one or two "favorites"? This is known as the load balancing problem. In our next lesson, we will explore the gating/routing mechanism in more detail and discuss the auxiliary losses and other techniques used to ensure efficient and balanced training.