Gating Networks › Router Architecture

The gating network in modern MoE transformers is remarkably simple: a linear projection from the token representation to a score for each expert. Despite this simplicity, the router must learn to make complex decisions about which experts are best suited for each input. To understand why this seemingly minimal architecture succeeds, we need to examine both its mathematical structure and the reasoning behind key design choices.

At the heart of every MoE layer lies a deceptively simple question: given a token's representation, which expert should handle it? The standard router architecture answers this question with a single linear layer without bias. Before we write down the formula, consider what we need: a way to transform each token's high-dimensional representation into a set of scores, one for each expert. These scores should capture how well-suited each expert is for processing that particular token. A linear projection accomplishes exactly this:

$$G(x) = xW_g$$

where:

• $G(x)$: the vector of raw routing scores (logits) for the $n$ experts
• $x$: the input token representation vector of dimension $d$
• $W_g$: the learnable router weight matrix of dimension $d \times n$
• $n$: the number of experts

This formula performs a matrix multiplication that projects the $d$-dimensional token representation onto $n$ scores, one for each expert. The result is a vector of routing logits that will later be converted to probabilities.

This design choice, a simple linear projection, might seem limiting at first glance. After all, modern deep learning has shown the power of multi-layer networks for capturing complex patterns. Why not use a more expressive multi-layer network for routing? Several factors favor the linear design:

• Computational efficiency: The router runs for every token, so it must be fast. A linear layer adds minimal overhead compared to the expert computations themselves.
• Sufficient expressivity: The input representation $x$ already encodes rich semantic information from prior layers. By the time a token reaches the MoE layer, attention mechanisms have already captured contextual relationships. The router just needs to project this pre-processed information to expert preferences.
• Training stability: Simpler architectures are easier to train and less prone to the instabilities that plague MoE systems. Adding nonlinearities to the router can create optimization challenges.
• Interpretability: Linear routers allow clearer analysis of what input features influence routing decisions, since the mapping from input dimensions to expert scores is direct.

What exactly feeds into the router? Understanding the input is important because it determines what information the router can use to make its decisions. In a transformer's MoE layer, the router typically receives the same input that would go to a standard FFN block. This is the output of the attention sublayer (plus residual connection and normalization):

$$x = \text{LayerNorm}(\text{Attention}(h) + h)$$

• $x$: the input vector to the router
• $h$: the input state from the previous layer
• $\text{Attention}(h)$: the output of the self-attention mechanism
• $\text{LayerNorm}$: the normalization operation that stabilizes the representation

This formula shows that the router's input incorporates both the attention output and the original hidden state through the residual connection, all normalized for stable computation. This representation already captures contextual information about the token's role in the sequence. The attention mechanism has allowed the token to gather needed information from other positions, so the router is not making decisions based on the token in isolation. Instead, the router's job is to translate this contextual representation into expert preferences.

Some architectures experiment with alternative inputs. For instance, some routers use the pre-normalization representation or concatenate multiple representations to give richer information. However, the standard approach (using the same input as the FFN) remains dominant due to its simplicity and effectiveness. Using the same input also ensures consistency: the expert receives the same representation that was used to decide to route to it.

Understanding the router weight matrix $W_g$, the only learnable component, gives intuition about how routing decisions are made. Each column $w_i$ of the matrix is a "prototype" for expert $i$. You can think of these prototypes as characteristic patterns that expert $i$ is designed to handle. The routing score for expert $i$ is in effect the dot product between the input and that expert's prototype:

$$G(x)_i = x \cdot w_i$$

where:

• $G(x)_i$: the routing score (logit) for expert $i$
• $x$: the input token representation
• $w_i$: the column vector from the weight matrix representing expert $i$'s prototype

This framing as a dot product helps build intuition about how routing works. Recall that the dot product measures similarity between vectors: it is large when vectors point in similar directions and small (or negative) when they point in different directions. Therefore, experts with prototypes that align well with the input representation receive higher scores. Conversely, experts whose prototypes point away from the input in representation space receive lower scores.

During training, the router learns prototypes for token groups. One prototype might match scientific terms, while another matches conversational text. Training then sends each token group to a suitable expert.

The router parameters are typically initialized with standard techniques. As we covered in Part X, proper initialization prevents the router from starting with extreme biases toward particular experts. This keeps all experts have a chance to receive tokens during early training.