Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Switch Transformer › Worked Example: Routing Visualization

Michael BrenndoerferUpdated January 5, 20263 min read

Language AI HandbookAll chapters

Part 33 of 67Mixture of Experts

Chapter 247Switch Transformer › Worked Example: Routing Visualization

Reading section 10 of 14 · Worked Example: Routing Visualization · 3 minSection 10 of 14 · 3 min

Section 10 of 14 · 3 min read

Switch Transformer › Worked Example: Routing VisualizationOne section · Sections in this chapter10 / 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

Let's visualize how tokens get routed to experts in a trained Switch layer. This visualization helps build intuition about what the router learns and how it distributes tokens across the available experts. We will trace through a full forward pass with concrete numbers to make the routing mechanism tangible.

In[23]:

Code

# Create a Switch layer and process some text
torch.manual_seed(7)

num_viz_experts = 8
switch_layer = SwitchLayer(
    d_model=128, d_ff=512, num_experts=num_viz_experts, capacity_factor=1.5
)

# Simulate token embeddings for a sentence
# Different semantic categories should route differently
batch_size, seq_len, d_model = 1, 20, 128
tokens = torch.randn(batch_size, seq_len, d_model) * 5.0

# Get routing decisions
with torch.no_grad():
    expert_indices, expert_weights, router_logits = switch_layer.router(tokens)
    router_probs = F.softmax(router_logits, dim=-1)

Out[24]:

Visualization

Heatmap with token position on the y-axis (0 to 19) and expert index on the x-axis (0 to 7). Most cells are pale, with a few dark-blue cells per row showing the single high-probability expert assigned to each token, illustrating sparse top-1 routing. — Router probability distribution heatmap for 20 tokens across 8 experts. Each row shows one token's probability distribution over all experts. Most tokens have a single dominant expert (dark cell), with probabilities near zero elsewhere, confirming the sparse routing behavior characteristic of Switch Transformer.

Out[25]:

Visualization

Bar chart showing the number of tokens (y-axis) assigned to each of 8 experts (x-axis, labeled 0 through 7) by a randomly initialized router. Steel-blue bars vary in height, and a horizontal red dashed line marks the perfect balance level at 2.5 tokens per expert. — Expert load distribution for a 20-token sequence with 8 experts. The bar heights show how many tokens each expert receives, with the red dashed line marking the ideally balanced allocation of 2.5 tokens per expert. A well-trained router would push this distribution toward the balance line through repeated exposure to the auxiliary load balancing loss.

The heatmap shows each token's probability distribution over experts. Notice how most tokens have one dominant expert (high probability in one column, near-zero in the others). This shows the router's confident routing decisions even with random initialization. This confidence emerges naturally from the softmax: when one logit is substantially higher than the others, the softmax output is a near-one-hot vector. The bar chart shows the actual token distribution, with the red line showing perfect balance. With random initialization, we see large variance in expert popularity. During actual training, the auxiliary load balancing loss would apply pressure to even out this distribution, driving the bar heights toward the red line while letting real departures driven by token content.