Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Switch Transformer › The Switch Layer

Michael BrenndoerferUpdated January 5, 20268 min read

Language AI HandbookAll chapters

Part 33 of 67Mixture of Experts

Chapter 247Switch Transformer › The Switch Layer

Reading section 2 of 14 · The Switch Layer · 8 minSection 2 of 14 · 8 min

Section 2 of 14 · 8 min read

Switch Transformer › The Switch LayerOne section · Sections in this chapter2 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The core architectural innovation of Switch Transformer is the Switch layer, which replaces the standard feed-forward network (FFN) in each transformer block with a sparse, routed alternative. To understand why this substitution matters, we need to first recall what happens in a standard transformer and then see how the Switch layer turns this computation into something far more powerful.

Recall from our discussion of feed-forward networks in Part XV that a standard transformer block processes each token through a two-layer MLP. This feed-forward network applies the same transformation to every token in the sequence, using shared parameters regardless of what the token is or what context it appears in. The network has learned everything it knows about language compressed into a single set of weight matrices, and every token, whether it is a mathematical symbol, a proper noun, or a common preposition, passes through the exact same computation:

\text{FFN}(x) = W_2 \cdot \text{activation}(W_1 \cdot x + b_1) + b_2

where:

$x$ : the input vector to the layer, typically of dimension $d_{\text{model}}$
$W_1$ : learnable weight matrix projecting from $d_{\text{model}}$ to the larger hidden dimension $d_{ff}$
$W_2$ : learnable weight matrix projecting back from $d_{ff}$ to $d_{\text{model}}$
$b_1, b_2$ : learnable bias vectors
$\text{activation}$ : the non-linear activation function (typically ReLU or GeLU)

This standard FFN architecture has served transformers well, but it embodies a basic limitation: every token receives identical processing. A token representing a mathematical concept receives the same transformation as a token representing a cooking ingredient. The network must somehow compress all of its knowledge about syntax, semantics, world facts, and linguistic patterns into a single set of parameters. At small scales this is manageable; at large scales it creates a capacity bottleneck where the dense parameters struggle to simultaneously stand for all the heterogeneous knowledge the model needs.

In a Switch layer, we break free from this constraint. Instead of one FFN, we have $E$ expert networks, each with the same architecture as the original FFN but with independent parameters. Each expert can specialize in different types of tokens or different aspects of language understanding. A router network decides which single expert should process each token, letting the model to deploy specialized computation where it matters most. The router's decision is content-dependent: a token representing a mathematical formula might consistently route to expert 7, while tokens that are common function words might distribute across experts 2, 4, and 9.

This architectural choice creates a natural division of labor. One expert might become skilled at processing mathematical notation, another at handling named entities, and yet another at processing syntactic function words. The key insight is that this specialization emerges from training rather than being imposed by design. You do not have to tell the model what each expert should specialize in; the combination of routing decisions and gradient updates naturally drives experts toward complementary specializations. The router and the experts co-evolve during training, with the router learning to recognize which tokens should go where, and the experts learning to excel at the tokens they receive.

In[3]:

Code

import torch
import torch.nn as nn
import torch.nn.functional as F


class SwitchExpert(nn.Module):
    """A single expert network with standard FFN architecture."""

    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Standard FFN: up-project, activate, down-project
        return self.w2(self.dropout(F.gelu(self.w1(x))))

The router is the decision-making component of the Switch layer. It examines each incoming token representation and produces a probability distribution showing how suitable each expert is for processing that particular token. Despite the sophistication of the routing decisions it makes, the router itself is remarkably simple: just a linear layer that projects the token representation to a vector of expert scores, followed by a softmax to convert these scores into probabilities. This simplicity is intentional. A complex router would add computational overhead and introduce its own optimization challenges. The linear router's parameters are small relative to the expert networks, and the softmax ensures the output is interpretable as a probability distribution.

In[4]:

Code

class SwitchRouter(nn.Module):
    """Routes tokens to experts using learned linear projection."""

    def __init__(self, d_model: int, num_experts: int):
        super().__init__()
        self.num_experts = num_experts
        self.router = nn.Linear(d_model, num_experts, bias=False)

    def forward(
        self, x: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        # x shape: (batch_size, seq_len, d_model)
        router_logits = self.router(x)  # (batch, seq, num_experts)
        router_probs = F.softmax(router_logits, dim=-1)

        # Select the single best expert for each token
        expert_indices = torch.argmax(router_probs, dim=-1)  # (batch, seq)
        expert_weights = router_probs.gather(
            -1, expert_indices.unsqueeze(-1)
        ).squeeze(-1)

        return expert_indices, expert_weights, router_logits

The key difference from previous MoE architectures like those in GShard is the hard assignment to a single expert. Each token goes to exactly one expert, and the router's softmax probability for that expert becomes a multiplicative weight on the expert's output. This design choice, which might initially seem like a limitation, turns out to be the important insight that makes Switch Transformer so effective at scale. Previous architectures used the top-2 routing inherited from the original Shazeer MoE paper, justified by the argument that blending two experts produces richer representations. Switch's central empirical claim is that this reasoning, while intuitive, does not hold up when you account for the full training dynamics and the practical benefits of simplicity.

Out[5]:

Visualization

Diagram showing five colored token circles at the top all connecting down to a single gray FFN box, with arrows from the FFN back to output token circles below, illustrating that every token passes through the same network. — Dense feed-forward network architecture. All tokens are processed through a single shared FFN, applying identical transformations regardless of token content.

Diagram showing five colored token circles at the top connecting to a green Router box, which then routes tokens via colored arrows to three dark expert boxes labeled E1, E2, and E3, with output tokens below, illustrating sparse one-to-one routing. — Switch layer architecture with top-1 routing. Each token is routed to exactly one specialized expert based on content, letting parameter scaling without increasing computational cost per token.