SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention

ArXiv: 2312.07987

🎯 Pitch

Standard transformers waste compute storing an attention matrix for every head—SwitchHead cuts that number by up to 8× while matching perplexity, proving that value/output MoE with shared attention maps actually works where earlier attempts failed.

---

1. Executive Summary

This paper introduces SwitchHead, a novel Mixture-of-Experts (MoE) method for the self-attention layer of Transformers that reduces both compute and memory requirements while matching the language modeling performance of parameter-matched dense baselines. The authors evaluate SwitchHead on C4, Enwik8, peS2o, and Wikitext 103 using Transformer XL and RoPE-based models at 47M and 262M parameter scales, employing MoE projections for value and output — but notably not key or query — within each attention head (sharing a single attention matrix across multiple value/output experts). SwitchHead achieves up to 8× fewer attention matrices than standard Transformers, with their 262M parameter model on C4 matching baseline perplexity using only 44% of the compute (2.0G vs. 5.4G MACs) and 27% of the memory (2.9M vs. 21M floats), while also delivering wall-clock speedups of ~1.5× in training. The authors further combine SwitchHead with σ-MoE feedforward layers into fully-MoE "SwitchAll" Transformers, establishing that MoE-based attention can replace dense multi-head attention without regularization or loss of expressivity only when the value and output projections are parameterized as non-competitive (sigmoid-gated) expert mixtures while keys and queries remain head-specific.

2. Context and Motivation

The Core Problem: Attention's Untouched Computational Footprint

The fundamental problem this paper tackles is the asymmetric deployment of Mixture of Experts (MoE) in Transformers. Since Shazeer et al. (2017) introduced sparsely-gated MoE layers and Fedus et al. (2021) adapted them to Transformers, the standard recipe has been clear: replace the dense feedforward (MLP) blocks with MoE layers, but leave the self-attention mechanism entirely untouched. This paper asks a straightforward question that the field has largely bypassed: can we apply MoE to the attention layer itself and achieve comparable resource savings?

To understand why this matters, recall the computational anatomy of multi-head self-attention (MHA). For a sequence of length $T$ with $n_{heads}$ heads and head dimension $d_{head}$, the attention computation scales as $O(n_{heads} \cdot d_{head} \cdot T^2)$. The $T^2$ term—computing and storing the attention matrix—dominates when sequence length grows. Modern LLMs deploy tens of heads (GPT-3 uses 96 heads in its largest configuration; LLaMA-70B uses 64 heads). Each head requires its own $T \times T$ attention matrix, meaning the memory footprint during training stores $O(n_{heads} \cdot T^2)$ floats. For a sequence of 2048 tokens with 16 heads, that's roughly 67 million floats for attention matrices alone—per layer, per batch element.

The authors observe that prior MoE research has concentrated overwhelmingly on the MLP layers (Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2021; Clark et al., 2022; Csordás et al., 2023). This is a sensible first step because the MLP typically accounts for the majority of parameter count in a Transformer block. But it leaves a significant gap: attention consumes a disproportionate share of activation memory during training due to the $n_{heads} \cdot T^2$ storage requirement, and its quadratic compute scaling with sequence length makes it the bottleneck for long-context applications. The paper's explicit position (Section 1) is that "attention also accounts for a considerable amount of compute and memory usage in Transformers (especially for long context sizes)," and that using MoE for attention "has potential to further improve resource efficiency."

Why This Problem Is Important

The practical significance operates on multiple levels:

Training accessibility. Training large Transformers at the cutting edge costs millions of dollars—GPT-4's training run is estimated at over $100 million in compute. Even modestly sized models at the 1B+ parameter range remain out of reach for most academic labs and smaller companies. Resource-efficient attention mechanisms directly lower this barrier.

Inference deployment. Running LLMs in production requires substantial engineering effort (model sharding, quantization, specialized kernels like FlashAttention). While inference is typically less memory-intensive than training because intermediate activations need not be stored, the attention computation itself remains costly. A method that reduces the number of attention matrices computed at both training and inference time yields persistent savings across the model's entire lifecycle.

Scaling behavior with sequence length. The $O(T^2)$ quadratic scaling of attention is the fundamental computational bottleneck preventing Transformers from processing very long sequences efficiently. While approaches like sparse attention, linear attention, and state-space models address this from an algorithmic perspective, SwitchHead attacks it from a structural redundancy angle: if many attention heads are learning similar or redundant patterns, reducing the number of distinct matrices computes the same number of projections (via MoE) with fewer intermediate tensors.

Edge and on-device deployment. The paper notes (Appendix A.1) that "Flash Attention depends on GPU-specific memory bandwidth/compute trade-offs, which might not be available on all hardware, especially on edge devices." SwitchHead's savings come from reducing the computation that must happen, independent of hardware-specific memory hierarchies—making it complementary to and combinable with FlashAttention.

Prior Approaches and Their Shortcomings

The paper identifies three existing lines of work that attempt to make attention more efficient, each with specific limitations.

Mixture of Attention Heads (MoA) — Zhang et al. (2022). This is the most direct predecessor and the paper's primary comparative baseline. MoA defines each attention head as an expert but shares a single key and value projection (source-side projections) across all experts. At each token position, MoA selects $k$ active query and output projections from a pool of $E$ experts using a softmax-based competitive gating mechanism.

The paper identifies several critical failures of MoA (Section 3.2):

• High effective head count. Each selected expert requires computing its own attention matrix. If $k$ experts are active, MoA computes $k$ attention matrices per layer—which in practice must be nearly as large as the dense baseline's head count to match performance. Table 1 shows MoA with $n_{heads} = 8$ (262M model on Wikitext-103) achieves 9.50 perplexity vs. the Transformer baseline's 9.66, but uses 2.9G MACs and 9.9M floats of memory—compared to SwitchHead's 2.0G MACs and 2.9M floats at 9.55 perplexity with only 2 heads. MoA cannot achieve significant resource reductions while matching performance.

• Competitive activation causes training instability. MoA uses softmax-based gating, which tends toward expert collapse—a well-documented problem in MoE literature where the router converges to always selecting the same few experts, rendering others useless. To prevent this, MoA requires three separate regularizers (load balancing, importance, and variance penalties), adding hyperparameter complexity and making training brittle. SwitchHead sidesteps this entirely by using a non-competitive sigmoid activation (following σ-MoE, Csordás et al., 2023), which "does not require regularization or extra tricks for stable training" (Section 2).

• Architectural mismatch with resource reduction goals. Because MoA shares key/value projections but computes separate attention matrices per selected expert, the savings come only from projection computations—not from attention matrix computation or storage, which is the dominant cost for long sequences.

Token-Level Head Gating — Peng et al. (2020). This method, introduced under the title "A Mixture of $h-1$ Heads is Better than $h$ Heads," proposes to reweight the contribution of each head by a learned gating function, reducing the total number of attention heads by one (presumably to offset the parameters used by the selection logic). The paper identifies this as a fundamentally different goal: "Their goal was not to reduce resource usage but to have better predictive performance, which they achieve" (Section 5). Additionally, Peng et al. use softmax-based competitive selection and, to avoid collapse, train the gating function only in some steps—a workaround rather than a solution.

Low-Rank Attention via Global Matrices — Nguyen et al. (2022). This approach, motivated by the observation that attention matrices are often low-rank, constructs a small number (e.g., 2) of "global attention matrices" and computes each head-specific matrix as a weighted average of these globals. However, the averaging occurs in logit space before the softmax, meaning each head-specific attention matrix must still be individually computed before averaging. As the paper points out (Section 5), this means "in the best case, they can only save half of the computation associated with the attention matrix because the readout (Equation 3) is still needed. For the same reason, memory savings are also low."

Multi-Query Attention (MQA) — Shazeer (2019). MQA uses a single key and value projection shared across all heads while maintaining separate query projections. The paper references this in passing, noting that their own findings "show that such a configuration is suboptimal: using multiple output and value projections is the most important choice in our model design" (Section 5). The implication is that MQA throws away beneficial diversity on the wrong side—SwitchHead's ablation (Table 6) shows value and output projections need MoE, while key and query projections can remain head-specific.

The Unresolved Tension: Many Heads Are Necessary, but Redundant

The paper's motivation crystallizes around a specific empirical observation (Section 2.2):

"modern LLMs use tens of heads. Are so many of them all necessary? As we show later in Sec. 3, indeed, naively reducing the number of heads (while keeping the same number of parameters by increasing the head dimension) results in performance loss."

Table 2 provides this evidence concretely. On C4 with the 47M parameter model, reducing from 10 heads to 2 (with proportionally larger $d_{head}$ to maintain parameter count) degrades perplexity from 22.71 to 23.71. On the 262M model, reducing from 16 to 4 heads degrades from 16.28 to 17.09. Performance consistently deteriorates. Something about having many attention heads is structurally necessary.

The paper offers three hypotheses for why, without needing to resolve them experimentally (Section 2.2):

1. Multiple inputs for downstream operations: each head provides a distinct channel of information that subsequent layers process differently.
2. Specialization: different heads serve different linguistic or computational functions, and each downstream operation selectively attends to the outputs of specific heads. In this case, not all heads are needed simultaneously—a given token or operation might only require a subset.
3. Initialization diversity: heads with different random initializations converge to different useful functions; more heads simply provide more lottery tickets, some of which become more useful than others.

Hypotheses (2) and (3) are the key to SwitchHead's design. If heads are specialized and not all are needed at the same time for a given computation, then conditional computation—switching among heads based on context—can reduce resource usage without losing the expressive benefit of having many head-like projections available. The challenge is designing a switching mechanism that actually reduces the expensive part of attention (the $T \times T$ matrices) rather than just the projection overhead.

How SwitchHead Positions Itself

The paper frames SwitchHead as solving the structural problem that prevented prior MoE-for-attention methods from achieving meaningful resource savings. The core insight (Section 2.2) is that to reduce the number of computed attention matrices, expert selection must happen independently on the source and destination sides, before and after the attention computation:

• Source-side experts select which value projections to use, producing a weighted average of the selected expert value projections. This happens before the attention matrix is applied.
• Destination-side experts select which output projections to use, applied after the attention-weighted value aggregation.
• A single attention matrix per head is shared across all experts within that head.

This architecture means the number of attention matrices to compute is $n_{heads}$ (typically 2–4 in their experiments), not $n_{heads} \times k$ (as in MoA) or $n_{heads} \times E$ (as in dense). The expert multiplicity exists only in the linear projections—which are cheap relative to the $O(T^2)$ attention computation.

The paper explicitly contrasts its approach with MoA along this dimension (Section 3.2):

"MoA computes the attention map for each selected expert and computes their weighted average after the attention computation takes place. In contrast, SwitchHead calculates the weighted average of the K selected experts before and after attention computation. Because of this, in practice, the same perplexity is achieved with the required number of computed attention matrices ($n_{heads}$) which is much lower for SwitchHead compared to MoA."

The paper also positions itself within the broader MoE-for-transformers lineage but draws on σ-MoE's key innovation: non-competitive sigmoid gating (Csordás et al., 2023). Previous MoE methods (Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2021) use softmax-based competitive gating, where increasing one expert's gate value necessarily decreases another's. This creates a tendency toward expert collapse, requiring careful load-balancing regularization. SwitchHead inherits σ-MoE's sigmoid approach, where each expert's selection weight is computed independently via $\sigma(x W_S)$—experts don't compete. The top-$k$ operation then selects the $k$ highest-scoring experts. This means the model can use anywhere from 0 to $k$ experts for a given token without penalty, eliminating the need for regularization. The paper is explicit that this carries over from σ-MoE and applies it to attention for the first time.

A Taxonomy: Where the Resource Savings Come From

To understand the paper's positioning, it's helpful to decompose exactly where compute and memory are spent in attention and how different methods address different components. The resource analysis in Appendix A.2 formalizes this:

Standard Transformer XL with $n_{heads}$ heads:

• Projection compute: $4 T d_{head} d_{model}$ MACs per head (for K, Q, V, and output projections)
• Attention matrix compute: $2 C T^2 d_{head}$ MACs per head (where $C$ accounts for XL context size)
• Position encoding compute: $2 C T d_{head} d_{model}$ MACs per head
• Projection memory: $4 T d_{head}$ floats per head
• Attention matrix memory: $2 C T^2$ floats per head (stored before and after softmax)
• Position encoding memory: $2 C T d_{head}$ floats per head

The quadratic terms in $T$ dominate for long sequences. SwitchHead's key move is reducing $n_{heads}$ (the multiplier on all terms) from 10–16 down to 2–4, while using MoE to maintain the effective diversity of projections through $E$ experts per head with $k$ active. The increased $d_{head}$ partially offsets the reduction, but since $d_{head}$ appears linearly while $T^2$ appears quadratically, the net savings are substantial for realistic sequence lengths.

SwitchHead (value and output MoE only):

• Projection compute: $2 T d_{head} d_{model} + 2 T k d_{head} (d_{model} + 1)$ MACs per head (the fixed K, Q projections plus the MoE V and O projections with $k$ active experts)
• Attention matrix compute: unchanged per head, but now multiplied by far fewer $n_{heads}$
• Attention matrix memory: unchanged per head, same multiplication benefit

MoA, in contrast, shares K and V projections (reducing projection compute) but multiplies attention matrix costs by the number of active experts $k$ per head, failing to reduce the dominant $T^2$ term. This structural analysis underlies the paper's empirical finding that MoA cannot achieve resource reductions comparable to SwitchHead while maintaining performance.

The Parameter-Matched Evaluation Philosophy

A crucial aspect of how this paper positions its contribution is its evaluation methodology. Section 3 states:

"We conduct our experiments in a parameter-matched setting which better reflects the task of language modeling (than the FLOPS-matched setting often used to evaluate MoEs)."

This is a deliberate departure from the dominant MoE evaluation paradigm. Most MoE papers (Shazeer et al., 2017; Fedus et al., 2021) compare MoE models against dense baselines with equal compute budgets (FLOPS-matched). Because MoE models have more total parameters, they naturally outperform compute-matched dense models—the question is by how much. The parameter-matched setting asks a different question: given a fixed parameter budget, can we replace dense components with MoE components and maintain performance while reducing resource usage? This is a higher bar: the MoE gating mechanism consumes parameters that could otherwise go toward model capacity, so the experts must be genuinely more efficient per parameter than the dense equivalent.

This choice matters because it tests whether MoE attention can deliver resource efficiency (compute and memory savings) rather than just parameter scaling (more total parameters for the same compute). The paper's core claim—"matching baseline perplexity with only 44% compute and 27% memory"—is a parameter-matched claim and would not be meaningful under a FLOPS-matched comparison where the MoE model would have different total parameters.

3. Technical Approach

This is primarily a systems-design paper that introduces a specific architectural modification—SwitchHead—for replacing dense multi-head attention with a conditional-computation Mixture-of-Experts attention layer that reduces the number of attention matrices computed while maintaining expressivity.

3.1 Reader Orientation

The paper builds a resource-efficient attention layer that uses Mixture-of-Experts (MoE) gating on the value and output projections to serve the role that many independent attention heads serve in standard Transformers, but with far fewer attention matrices needing to be computed. The system solves the problem of reducing the $O(n_{heads} \cdot T^2)$ compute and memory cost of multi-head attention by replacing spatial multiplicity (many parallel heads) with expert multiplicity (a pool of expert projections that are conditionally activated per token), where the expensive attention matrix is computed only once per head and shared across multiple experts for the value readout and output projection.

3.2 Big-Picture Architecture (Diagram in Words)

The SwitchHead attention layer has five major components:

1. Fixed Key/Query Projections — one matrix each per head ($W^h_K, W^h_Q$), taking the input $x$ and producing standard keys $K^h$ and queries $Q^h$ for computing the attention matrix. These are NOT MoE-based.

2. Source-Side Expert Selection Network — a learned linear router $W^h_S$ that takes the input $x$ and produces, via sigmoid activation and top-$k$ selection, which $k$ value experts to activate for this head on this token.

3. Value Expert Pool — for each head $h$, a set of $E$ expert projection matrices $W^{h,e}_V$. The selected experts' projections are weighted by their gate scores and summed to produce a single values tensor $V^h$ before attention.

4. Shared Attention Matrix — exactly one $T \times T$ attention matrix $A^h$ per head, computed from $Q^h$ and $K^h$, used for the value readout exactly as in standard attention. The key architectural insight: multiple value experts all read through the same attention matrix.

5. Destination-Side Expert Selection Network and Output Expert Pool — a learned linear router $W^h_D$ selects $k$ output experts from a pool of $E$ matrices $W^{h,e}_O$. The attention-weighted value output (a single tensor $A^h V^h$) is then projected through the selected output experts and summed to produce the head's contribution to the final output.

Information flows as follows: input $x$ → source router selects value experts → value expert projections are computed and weighted-averaged → keys and queries are computed (fixed projections) → attention matrix $A^h$ is computed from keys and queries → values are projected through $A^h$ → destination router selects output experts → output expert projections are applied to the attention-weighted values and summed → the result is added to the layer output (along with contributions from other heads).

3.3 Roadmap for the Deep Dive

• First, the formal definition of standard multi-head attention (Equations 1–3), establishing the exact computational costs that SwitchHead targets. This grounds the resource analysis.
• Second, the naive head-gating approach (Equations 4–6) and why it fails to deliver resource savings, which motivates the independent source/destination expert design. This establishes what does NOT work.
• Third, the full SwitchHead mechanism: the expert selection networks, the value and output MoEs, the shared attention matrix, and the final output computation (Equations 7–10). This is the core technical contribution.
• Fourth, the ablation study that determines WHICH projections need to be MoEs (Table 6 and Section 3.1), establishing that value and output projections are necessary while key/query projections can remain fixed. This explains a critical design choice.
• Fifth, the resource accounting formulas (Equations 11–15 in Appendix A.2) that formally quantify where the compute and memory savings come from compared to standard attention and to MoA, showing precisely why the architecture is efficient.

3.4 Detailed, Sentence-Based Technical Breakdown

Standard Multi-Head Attention: The Baseline to Improve Upon

The paper begins by formalizing the standard multi-head self-attention layer (Section 2.1) to establish exactly what operations consume resources. Let $x \in \mathbb{R}^{T \times d_{model}}$ be the input sequence of length $T$ with hidden dimension $d_{model}$. For each head $h \in \{1, ..., n_{heads}\}$, there are three projection matrices: $W^h_K, W^h_Q, W^h_V \in \mathbb{R}^{d_{model} \times d_{head}}$, where $d_{head}$ is the per-head dimension. The keys, queries, and values are computed as:

$K^h = x W^h_K$, $Q^h = x W^h_Q$, $V^h = x W^h_V$

These all have shape $\mathbb{R}^{T \times d_{head}}$. The attention matrix for head $h$ is:

$A^h = \text{softmax}\left(\frac{1}{\sqrt{d_{head}}} Q^h K^{h^{\intercal}}\right)$

where $A^h \in \mathbb{R}^{T \times T}$, the softmax is applied over the last dimension (the source token dimension), and the scaling factor $1/\sqrt{d_{head}}$ prevents the dot products from growing too large as $d_{head}$ increases.

The output is typically written as a concatenation of all heads followed by an output projection $W_O \in \mathbb{R}^{n_{heads} d_{head} \times d_{model}}$:

$y = (A^1 V^1 | A^2 V^2 | ... | A^{n_{heads}} V^{n_{heads}}) W_O$

where $|$ denotes concatenation along the last dimension.

However, the authors decompose $W_O$ into per-head submatrices to reveal a critical structural property. Let $W^h_O \in \mathbb{R}^{d_{head} \times d_{model}}$ be the submatrix of $W_O$ corresponding to head $h$, such that $W_O = (W^1_O{}^\intercal | W^2_O{}^\intercal | ... | W^{n_{heads}}_O{}^\intercal)^\intercal$. Then the output can be rewritten as:

$y = \sum_{h} A^h V^h W^h_O$

What this decomposition reveals: Each head's contribution to the output is a self-contained computation: the attention matrix $A^h$ reads from the values $V^h$ specific to that head, and the result is projected through a head-specific output matrix $W^h_O$. No cross-head interactions occur until the summation.

Why highlighting this matters for the mechanism: This head-local formulation shows that the only operations coupling different tokens are (a) the computation of $A^h$ which involves products of $Q^h$ and $K^h$ across all token pairs, and (b) the application of $A^h$ to $V^h$ which mixes information across tokens. The linear projections ($W^h_{\{K,Q,V,O\}}$) operate independently per token. SwitchHead exploits this separation: it modifies the projections to be MoE-based (operating per-token, cheap) while keeping the token-mixing attention core shared within each head (reducing $n_{heads}$ dramatically, saving the expensive $T^2$ operations).

The paper quantifies the resource costs (Appendix A.2). For standard attention with Transformer XL's extended context of size $CT$ (where $C$ is the context multiplier, typically 2 for one additional chunk):

• Compute per head: $4 T d_{head} d_{model} + 2 C T^2 d_{head} + 2 C T d_{head} d_{model}$ MACs. The first term is projections; the middle term (quadratic in $T$) is the attention matrix computation and value readout; the last term is position encoding.
• Memory per head: $4 T d_{head} + 2 C T^2 + 2 C T d_{head}$ floats. The middle term—$2CT^2$—is the attention matrix stored before and after softmax, which dominates for long sequences.

The critical multiplier is $n_{heads}$. With $n_{heads} = 16$, this means 16 separate $T \times T$ matrices must be computed and stored. SwitchHead's goal is to reduce this multiplier to 2–4 while maintaining model quality through the MoE mechanism.

The Naive Head-Gating Approach and Why It Fails

Before presenting the full SwitchHead method, the paper describes a conceptually simpler approach (Section 2.2) and explains why it does not achieve resource savings. This serves as a pedagogical stepping stone: it motivates the more complex independent source/destination expert design.

The naive idea: use a learned gating function to select which attention heads are active for each token, and only compute the selected heads. Let $W_S \in \mathbb{R}^{d_{model} \times n_{heads}}$ be a learned selection matrix. For each token position $t$, compute a gate vector $s \in \mathbb{R}^{n_{heads}}$ as:

$s = \sigma(x W_S)$

where $\sigma$ is the sigmoid function applied elementwise (non-competitive selection, following $\sigma$-MoE). Then select the top-$k$ highest-scoring heads:

$E = \text{arg topk}(s, k), \quad E \subset \{1, ..., n_{heads}\}$

The output at position $t$, channel $c$ would be:

$y[t, c] = \sum_{h \in E} s[t, h] (A^h V^h W^h_O)[t, c]$

What this computes: For each destination token, the model selects which attention heads contribute to its output, weights each selected head's contribution by its gate score, and sums. Heads not in the top-$k$ contribute nothing.

Why this fails to deliver resource savings: The head selection $E$ is based solely on the destination side—the token receiving the attention output. But the attention matrix $A^h$ depends on pairs of tokens (source and destination). In the worst case, each destination selects a different subset of heads that collectively require all source-side projections (keys and values) to be computed for all heads, because you cannot know in advance which destination tokens will need which source projections. The paper states: "in the worst case, for each destination, a different source might be chosen, in which case all possible source projections have to be computed for the keys and values, which we would like to avoid."

Concretely: if destination token 1 needs heads {1,3} and destination token 2 needs heads {2,4}, you still need to compute $K^h$ and $V^h$ for all heads $h=1,2,3,4$ before attention can be applied. The expensive attention matrices $A^h$ still need to be computed for all heads that any destination might need. In the worst case, all heads are needed by SOME destination, so no savings occur. The paper's preliminary experiments confirm this method is feasible for language modeling (it learns something useful) but "it is difficult to achieve acceleration and memory savings with this method."

The key insight from this failure: To achieve genuine resource reduction, the expert/head selection must be separated into source-side selection (which value projections to compute) and destination-side selection (which output projections to use), and both must happen independently before the attention matrix couples tokens together. This way, the attention matrix itself is computed once per head and shared across all experts within that head.

SwitchHead: Independent Source and Destination MoEs with Shared Attention

The full SwitchHead architecture (Section 2.2) implements the insight above. The paper explicitly redefines the notion of a "head": in SwitchHead, a head is an instance of a computed attention matrix. The total number of such heads is $n_{heads}$—dramatically smaller than in dense models (2–4 instead of 10–16). Each head $h$ has a pool of $E$ experts, giving a total of $n_{heads} \cdot E$ expert projections across the layer. The projection matrices are now $W^{h,e}_K, W^{h,e}_Q, W^{h,e}_V, W^{h,e}_O \in \mathbb{R}^{d_{model} \times d_{head}}$ where $h$ indexes the head and $e \in \{1, ..., E\}$ indexes the expert.

Source-side expert selection. For each head $h$, a learned routing matrix $W^h_S \in \mathbb{R}^{d_{model} \times E}$ produces a gate vector for every token position:

$s^h_S = \sigma(x W^h_S)$

where $s^h_S \in \mathbb{R}^{T \times E}$ and $\sigma$ is the elementwise sigmoid, so each entry is independently in $(0, 1)$. The top-$k$ experts are selected:

$E^h_S = \text{arg topk}(s^h_S, k), \quad E^h_S \subset \{1, ..., E\}$

What this computes: For each token position and each head, the router independently scores each of the $E$ value experts. The $k$ experts with the highest scores are activated. Critically, because the sigmoid is non-competitive, the model can activate anywhere from 0 to $k$ experts per token without penalty—if no expert scores above threshold, the weighted sum simply has small weights.

The value projection for head $h$ is then a weighted sum of the selected experts' outputs:

$V^h = \sum_{e \in E^h_S} s^h_S[e] \; x W^{h,e}_V$

What this computes physically: For each active expert $e$, compute $x W^{h,e}_V$ (a $T \times d_{head}$ matrix), multiply by the scalar gate score $s^h_S[e]$ at each token position, and sum across the $k$ active experts. The result is a single values tensor $V^h \in \mathbb{R}^{T \times d_{head}}$—exactly the same shape as in standard attention, but now it is a mixture of multiple expert value projections rather than a single fixed projection. The same gate scores are used to compute the key projections:

$K^h = \sum_{e \in E^h_S} s^h_S[e] \; x W^{h,e}_K$

Why the source-side MoE works this way: By computing the weighted average BEFORE attention, the attention matrix $A^h$ operates on a single $V^h$ tensor. If the averaging happened after attention (as in MoA), you would need separate attention matrices for each expert, defeating the purpose. The source-side gate $s^h_S$ depends only on the source token's representation $x$—the token that provides the value. This is local information available before the attention matrix couples tokens.

Destination-side expert selection. Symmetrically, the destination side has its own routing matrix $W^h_D \in \mathbb{R}^{d_{model} \times E}$ producing gate scores:

$s^h_D = \sigma(x W^h_D)$ $E^h_D = \text{arg topk}(s^h_D, k)$

The query projections are computed as:

$Q^h = \sum_{e \in E^h_D} s^h_D[e] \; x W^{h,e}_Q$

Critically, the attention matrix $A^h$ is computed exactly once per head, using these single $Q^h$ and $K^h$ tensors (which themselves may be MoE-weighted mixtures but are single tensors of shape $T \times d_{head}$):

$A^h = \text{softmax}\left(\frac{1}{\sqrt{d_{head}}} Q^h K^{h^{\intercal}}\right)$

The attention-weighted output is $A^h V^h \in \mathbb{R}^{T \times d_{head}}$. Then the destination-side experts are applied to project this back to the model dimension:

$y = \sum_{h=0}^{n_{heads}-1} \sum_{e \in E^h_D} s^h_D[e] \; A^h V^h W^{h,e}_O$

What this computes: For each head, the attention-weighted values $A^h V^h$ (which already mix information across tokens) are projected through each active output expert and weighted by the destination-side gate scores. The outer sum over heads $h$ aggregates all heads' contributions.

Overall information flow, step by step:

1. Input $x$ arrives at the SwitchHead layer.
2. For each head $h$, compute source gate scores $s^h_S = \sigma(x W^h_S)$ and select top-$k$ experts $E^h_S$.
3. Compute destination gate scores $s^h_D = \sigma(x W^h_D)$ and select top-$k$ experts $E^h_D$.
4. Compute value tensor $V^h = \sum_{e \in E^h_S} s^h_S[e] (x W^{h,e}_V)$—a weighted mixture of expert value projections.
5. Compute key tensor $K^h = \sum_{e \in E^h_S} s^h_S[e] (x W^{h,e}_K)$—a weighted mixture of expert key projections.
6. Compute query tensor $Q^h = \sum_{e \in E^h_D} s^h_D[e] (x W^{h,e}_Q)$—a weighted mixture of expert query projections.
7. Compute attention matrix $A^h = \text{softmax}(Q^h K^{h^{\intercal}} / \sqrt{d_{head}})$—one matrix per head, shared across all experts.
8. Compute attention-weighted values $A^h V^h$.
9. Compute head output $\sum_{e \in E^h_D} s^h_D[e] (A^h V^h W^{h,e}_O)$.
10. Sum over all heads to produce the final layer output.

Why this architecture achieves resource savings: The number of attention matrices computed is $n_{heads}$ (typically 2–4), not $n_{heads} \cdot k$ (MoA) or $n_{heads} \cdot E$ (dense equivalent). The expert multiplicity ($E$ experts, $k$ active) manifests only in the linear projections, which scale as $O(T \cdot d_{head} \cdot d_{model})$—linear in sequence length, not quadratic. For realistic sequence lengths where the $T^2$ term dominates, the savings are substantial.

The paper quantifies the MACs for SwitchHead (Appendix A.2, Equation 13):

$N^{\text{SwitchHead}}_{\text{MAC}} = n_{heads} \left( 2 T d_{head} d_{model} + 2 T k d_{head} (d_{model} + 1) + 2 C T^2 d_{head} + 2 C T d_{head} d_{model} \right)$

where the terms are: (1) $2 T d_{head} d_{model}$ for the fixed key and query projections (not MoE), (2) $2 T k d_{head} (d_{model} + 1)$ for the MoE value and output projections with $k$ active experts (the $+1$ accounts for the weighted averaging operation), (3) $2 C T^2 d_{head}$ for the attention matrix compute and value readout (unchanged per head), and (4) $2 C T d_{head} d_{model}$ for position encodings.

Compare this to standard attention (Equation 11): the coefficient on the critical $T^2$ term is $n_{heads} \cdot 2 C d_{head}$. SwitchHead reduces $n_{heads}$ from 10–16 to 2–4, saving directly on the quadratic term, while increasing $d_{head}$ to maintain parameter count. The linear terms increase due to multiple experts ($k$) and larger $d_{head}$, but the net effect for realistic sequence lengths is a large reduction because the quadratic term dominates.

Similarly, memory usage (Equation 12 vs. SwitchHead's memory formula): the $2 C T^2$ floats needed for the attention matrix before and after softmax are multiplied by $n_{heads}$. Reducing heads from 16 to 2 directly reduces this memory by a factor of 8.

The choice of non-competitive sigmoid gating. The paper explicitly states that SwitchHead follows σ-MoE (Csordás et al., 2023) in using sigmoid activation $\sigma$ rather than softmax. In the softmax formulation used by most MoE work (Shazeer et al., 2017; Fedus et al., 2021; Zhang et al., 2022 for MoA), the gate vector is:

$s = \text{softmax}(x W_S)$

where the entries sum to 1. This creates competition: increasing one expert's score necessarily decreases another's. In practice, this causes the router to collapse—assigning all probability mass to one or a few experts, rendering the rest useless. Preventing this requires complex load-balancing regularization (typically adding auxiliary loss terms that penalize uneven expert utilization).

Sigmoid $\sigma(x) = 1/(1 + e^{-x})$ produces independent scores in $(0,1)$ for each expert. The top-$k$ operation then selects the $k$ highest. Because scores are independent, the router can freely assign high scores to multiple experts or low scores to all experts without penalty. The model learns to route without external regularization. The paper emphasizes this stability property: "our method performs well without any regularization, while MoA requires three different regularizers" (Section 3.2).

Determining Which Projections Need MoE

The full SwitchHead formalization (Equations 7–10) allows all four projection types (key, query, value, output) to be MoE-based. However, Section 3.1 asks a crucial empirical question: which of these actually benefit from being MoEs in a parameter-matched setting?

Why this matters: In the parameter-matched regime, every parameter allocated to an unnecessary MoE mechanism (expert projections, routing matrices) is a parameter removed from other parts of the model (head dimension $d_{head}$, feedforward width $d_{ff}$, or number of layers). If a projection type doesn't benefit from being an MoE, making it one actively hurts performance.

The paper conducts an exhaustive ablation on a 47M parameter model trained on WikiText-103 with $n_{heads} = 2$ and $E = 5$ experts (Table 6 in Appendix A.3). All 16 possible combinations of which projections ({V, K, Q, O}) are MoE vs. fixed are tested.

Results of the ablation (Table 6, sorted by perplexity):

The dense baseline with $n_{heads}=10$ achieves 12.31 perplexity. The dense baseline with $n_{heads}=2$ (simple head reduction without MoE) achieves 12.74—this is the lower bound showing the performance cost of naive head reduction.

Key findings from this ablation:

1. Output projection MoE is critical. The top-performing configurations all have output MoE (O=Y). The model with ONLY output MoE (row 2, V=N, K=N, Q=N, O=Y) achieves 12.30—actually outperforming the 10-head baseline (12.31). The output projection is where the head's contribution is mapped back to the model dimension $d_{model}$—this is where diversity of projections matters most, because different contexts demand different output transformations to integrate with the residual stream.

2. Value projection MoE provides an additional small gain. The best model (V=Y, O=Y) achieves 12.27, slightly better than output-only (12.30). The value projection determines what information is read from each token—having multiple value experts allows the model to extract different features depending on the source-side context.

3. Key and query MoEs are unnecessary or harmful. Adding K or Q MoE to the V+O configuration degrades performance (12.37 with K MoE, 12.42 with Q MoE, 12.57 with all four). The paper explains this as a parameter budget effect: "This is possible because we perform all our experiments in a parameter-matched setting. Allocating parameters to these projections uses the budget that can be otherwise spent on other parts of the network." The keys and queries are used to compute the attention matrix itself—a single shared computation per head—and apparently a single projection per head provides sufficient quality for this purpose.

4. Without output MoE, performance collapses. Configurations without output MoE (O=N) all score above 12.45, with the worst at 12.90. The output projection is the information bottleneck: after attention mixes information across tokens, the output projection determines what gets communicated to the next layer. Having only one such projection per head is insufficient when $n_{heads}$ is small.

The paper's design decision based on this ablation: Use MoE for value and output projections only. Keys and queries remain as single, head-specific fixed projections. This is the configuration used in all subsequent experiments in Sections 3.2–3.7.

This decision also yields practical implementation benefits. By keeping keys and queries as simple linear projections, the attention matrix computation $Q^h K^{h^{\intercal}}$ remains a standard dense operation—no expert mixing occurs inside the attention core. Only the value projection $V^h$ (computed before attention) and the output readout (computed after attention) involve MoE-weighted averages.

SwitchHead Compared to MoA: Structural Differences in Resource Usage

The paper uses the formal MAC and memory formulas (Appendix A.2) to explain why SwitchHead achieves resource savings that MoA cannot. Understanding this comparison is essential to grasping why the architecture is designed as it is.

MoA's architecture (Equations 14–15): MoA shares a single key and value projection across all experts but selects $n_{heads}$ active query and output projections from a pool of $E$ experts. Each selected query head requires its own attention matrix. The resource formulas are:

$N^{\text{MoA}}_{\text{MAC}} = (2 n_{heads} + 2) T d_{head} d_{model} + 2 n_{heads} C T^2 d_{head} + 2 C T d_{head} d_{model}$

$N^{\text{MoA}}_{\text{mem}} = (2 n_{heads} + 2) T d_{head} + 2 n_{heads} C T^2 + 2 C T d_{head}$

The crucial structural difference: In MoA's MAC formula, the quadratic term $2 n_{heads} C T^2 d_{head}$ is multiplied by $n_{heads}$—the number of active heads/experts. In SwitchHead's MAC formula (Equation 13), the quadratic term $2 C T^2 d_{head}$ is multiplied by $n_{heads}$—but SwitchHead's $n_{heads}$ is much smaller (2–4) than MoA's (which must typically be 8–12 to match performance, as shown in Table 1).

Concretely, from Table 1: on the 262M parameter model, SwitchHead uses $n_{heads}=2$ and achieves 9.55 perplexity with 2.0G MACs and 2.9M floats of memory. MoA with $n_{heads}=8$ achieves 9.50 perplexity but uses 2.9G MACs (45% more) and 9.9M floats of memory (3.4× more). MoA with $n_{heads}=4$ uses 1.7G MACs and 5.1M floats—less resource than the 8-head MoA but still more memory than SwitchHead (5.1M vs. 2.9M)—and its perplexity degrades to 9.69, already worse than the baseline Transformer (9.66).

Why MoA cannot reduce $n_{heads}$ as aggressively: MoA selects query/output experts per token, but each selected expert requires computing a full attention matrix because the attention operation couples the query (destination-side) to all source tokens. Reducing $n_{heads}$ in MoA directly reduces the number of distinct attention patterns the model can simultaneously employ. SwitchHead separates the problem: the attention pattern (computed once per head from $Q^h$ and $K^h$) can be reused across multiple value and output expert combinations. A single attention pattern, combined with $E$ different value experts and $E$ different output experts, can express $E^2$ different token-mixing behaviors (one per combination of value expert and output expert) without computing additional attention matrices.

The memory story is even starker: The dominant memory term is $2 n_{heads} C T^2$. For a sequence length $T$ of 512 with XL context $C=2$, this is $2 \cdot n_{heads} \cdot 2 \cdot 512^2 \approx n_{heads} \cdot 1.05$ million floats. With $n_{heads}=16$ (dense baseline), that's ~16.8M floats just for attention matrices. With $n_{heads}=2$ (SwitchHead), it's ~2.1M floats. MoA with $n_{heads}=8$ requires ~8.4M floats—a 4× reduction over dense, but SwitchHead achieves 8×. Table 5 confirms this: the 262M SwitchHead uses 12.5GB per GPU vs. 20.5GB for the dense Transformer and 16.4GB for MoA.

The Full Formal Model of SwitchHead

Bringing together the design decisions described above, the paper defines SwitchHead as:

For each head $h \in \{0, ..., n_{heads}-1\}$:

1. Source router: $s^h_S = \sigma(x W^h_S)$, with $W^h_S \in \mathbb{R}^{d_{model} \times E}$. Select $E^h_S = \text{arg topk}(s^h_S, k)$.

2. Destination router: $s^h_D = \sigma(x W^h_D)$, with $W^h_D \in \mathbb{R}^{d_{model} \times E}$. Select $E^h_D = \text{arg topk}(s^h_D, k)$.

3. Key and query projections (fixed, not MoE): $K^h = x W^h_K$, $Q^h = x W^h_Q$, with $W^h_K, W^h_Q \in \mathbb{R}^{d_{model} \times d_{head}}$. These are standard linear projections—one per head, no experts.

4. Value projection (MoE): $V^h = \sum_{e \in E^h_S} s^h_S[e] \; x W^{h,e}_V$, where each expert projection is $W^{h,e}_V \in \mathbb{R}^{d_{model} \times d_{head}}$.

5. Attention matrix: $A^h = \text{softmax}(Q^h K^{h^{\intercal}} / \sqrt{d_{head}})$.

6. Output projection (MoE): The head's contribution is $\sum_{e \in E^h_D} s^h_D[e] \; A^h V^h W^{h,e}_O$, where each expert projection is $W^{h,e}_O \in \mathbb{R}^{d_{head} \times d_{model}}$.

7. Final output: $y = \sum_{h} \sum_{e \in E^h_D} s^h_D[e] \; A^h V^h W^{h,e}_O$.

What distinguishes this from a generic MoE formulation: The key and query projections remain simple, head-specific linear transformations. This is the empirical result of the ablation in Table 6—adding MoE to K and Q degrades performance in the parameter-matched regime, and even if it didn't, it would provide no additional resource savings because the attention matrix is already shared. By keeping K and Q fixed, the attention computation $Q^h K^{h^{\intercal}}$ remains a standard, highly optimized operation with no conditional computation inside it. The MoE complexity is pushed to the edges: before attention (value projections) and after attention (output projections).

The practical parameter counting: For a head $h$ with $E$ value experts and $E$ output experts:

• Fixed projections: $W^h_K, W^h_Q$ = $2 \cdot d_{model} \cdot d_{head}$ parameters
• Value experts: $E \cdot d_{model} \cdot d_{head}$ parameters
• Output experts: $E \cdot d_{model} \cdot d_{head}$ parameters
• Routers: $2 \cdot E \cdot d_{model}$ parameters (source and destination, $d_{model} \times E$ each, shared across all tokens in the sequence)

The total is roughly $(2 + 2E) \cdot d_{model} \cdot d_{head}$ plus $2 E d_{model}$ for the routers. Compared to dense: $4 \cdot d_{model} \cdot d_{head}$ per head (K, Q, V, O projections). SwitchHead replaces $2 \cdot d_{model} \cdot d_{head}$ worth of V and O parameters with $2E \cdot d_{model} \cdot d_{head}$ worth of expert parameters—an $E$-fold increase in the projection parameter count—but only $k$ of the $E$ experts are active per token, so the compute scales with $k$ not $E$. The total parameter count is matched to the dense baseline by adjusting $d_{head}$ and $d_{ff}$.

The systematic hyperparameter selection procedure (Section 3):

The paper describes a deliberate procedure for configuring SwitchHead to maintain parameter count while maximizing resource savings:

1. Set $n_{heads} \cdot E$ equal to $n_{heads}$ of the dense baseline. This means the total number of expert projections equals the baseline's head count.
2. Start with $n_{heads} = 2$ and $k = 2$ (most aggressive resource reduction). Train and evaluate.
3. If the model underperforms the baseline, increase $k$ to 4. If still underperforming, set $n_{heads} = 4$ and reset $k = 2$.
4. Adjust $d_{head}$ so that the total parameter count matches the baseline.

This procedure is described as "reasonably simple" and "ensures a good amount of resource savings, while avoiding doing an expensive hyperparameter search." For the 47M models in Table 2, this yields $n_{heads}=2$ with $E=5$ experts and varying $k$ (2–3 depending on dataset). For the 262M models in Table 2, the configurations vary: $n_{heads}=2$ with $E=8$ and $k=4$ for WikiText-103, $n_{heads}=4$ with $E=4$ and $k=2$ for C4 and peS2o.

---

SwitchAll: Combining SwitchHead with σ-MoE MLP Layers

Section 3.4 describes "SwitchAll"—a fully-MoE Transformer where both the attention (SwitchHead) and feedforward (σ-MoE) layers use mixture-of-experts. This is not a separate technical contribution but rather an integration test: "it remains unclear whether it can be efficiently combined with our SwitchHead, or can have some negative interaction effect if combined."

The setup is straightforward: the authors "take the baseline architecture of Csordás et al. [17] without any hyperparameter change and replace the attention layer with SwitchHead." The attention hyperparameters are directly taken from the SwitchHead configurations in Table 2. No additional tuning is performed.

The results (Table 3) show that SwitchAll matches or exceeds the dense baselines:

• 47M on WikiText-103: SwitchAll 12.17 perplexity vs. Transformer 12.32 (better with SwitchAll)
• 262M on WikiText-103: Transformer 9.80 vs. SwitchAll 9.81 (essentially tied)
• 47M on C4: SwitchAll 22.09 vs. Transformer 22.63 (better with SwitchAll)
• 262M on C4: Transformer 16.58 vs. SwitchAll 16.45 (slightly better with SwitchAll)

The resource usage for SwitchAll is identical to SwitchHead-only in the attention component—the σ-MoE MLP savings stack on top of the SwitchHead attention savings. The paper does not report separate MAC/memory numbers for SwitchAll's MLP component, but the attention savings are the same as in Table 2.

A notable detail: the parameter count for SwitchAll models is sometimes slightly less than the dense baseline (e.g., 259M vs. 262M for the large model configuration). The paper explains in Appendix A.6 that this is due to the granularity of σ-MoE parameter matching: "the size of all experts must be increased at once, and the CUDA kernel supports only sizes of multiple of 4. Therefore, increasing the size of the experts would add too many parameters and the model would outgrow the baseline." Rather than overshooting the parameter budget, they keep the σ-MoE hyperparameters from Csordás et al. (2023) unchanged and accept the small parameter deficit.

Two configurations from the dense baseline are shown in Table 2 for each dataset and model size: one with the same $n_{heads}$ as the SwitchHead model (e.g., 2 heads for 47M), and one with $n_{heads}$ equal to $n_{heads} \cdot E$ (e.g., 10 heads for 47M, 16 for 262M). The SwitchHead model consistently:

• Matches the performance of the many-head baseline (e.g., 22.53 vs. 22.71 for 47M on C4)
• Substantially outperforms the few-head baseline (e.g., 22.53 vs. 23.71 for 47M on C4)

This pattern—shown across all four datasets and both model sizes—demonstrates that the MoE mechanism successfully recovers the performance lost by reducing $n_{heads}$, while retaining the resource savings of having fewer attention matrices.

---

The Shared Selection Variant

Section 3.6 introduces a further optimization: sharing the expert selection between source and destination sides. Instead of having two separate routers $W^h_S$ and $W^h_D$ with separate top-$k$ sorting operations, a single router and selection is used for both value and output experts. The paper states this "results in a minor performance loss, which might be tolerated in some cases where the acceleration is more important."

Concretely, shared selection means $s^h_S = s^h_D = \sigma(x W^h)$ and $E^h_S = E^h_D = \text{arg topk}(s^h, k)$ for a single routing matrix $W^h \in \mathbb{R}^{d_{model} \times E}$. The acceleration comes from "reducing the number of sorting and top-k steps compared to the full SwitchHead"—there is one top-$k$ operation instead of two. The performance impact is shown in Table 4: for the 47M model on C4, shared selection degrades perplexity from 22.53 to 22.81 (a 0.28 increase, or roughly 1.2% relative). For the 262M model, the degradation is from 16.23 to 16.49 (0.26 increase, roughly 1.6% relative). Zero-shot task performance is similarly slightly degraded (BLiMP drops from 75.7% to 74.6% for the small model; 79.6% to 79.4% for the large one).

---

MAC-Matched Setup: A Complementary Evaluation

Section 3.5 answers a different question than the main experiments: instead of fixing parameters and reducing resources, "what is the performance of SwitchHead in a MAC-matched setup, where the compute requirements of our model are matched to those of the baseline?"

This is achieved by increasing $d_{head}$ and $n_{heads}$ until the SwitchHead model's MAC count equals the dense baseline's. Because SwitchHead is more efficient per MAC (using conditional computation), this results in a model with more total parameters than the baseline—the opposite of the parameter-matched philosophy.

The reported configurations (Section 3.5 and Table 4):

• Small Transformer XL: increase $d_{head}$ from 76 to 112, $n_{heads}$ from 2 to 3. Parameters increase from 47M to 63M. Perplexity drops from 22.53 to 21.18 (vs. dense baseline 22.71).
• Large Transformer XL: increase $n_{heads}$ from 4 to 6, $d_{head}$ from 112 to 168. Parameters increase from 262M to 376M. Perplexity drops from 16.23 to 15.43 (vs. dense baseline 16.28).
• Small RoPE: increase $n_{heads}$ from 2 to 3, $d_{model}$ from 64 to 84. Zero-shot BLiMP improves from 77.3% to 77.4%.

In all cases, MAC-matched SwitchHead outperforms both the parameter-matched SwitchHead and the dense baseline, demonstrating that the MoE mechanism can be "scaled up" to fill a compute budget more effectively than dense projections.

---

Wall-Clock Time and Memory Measurements

While the paper primarily reports theoretical MAC counts for hardware independence (Section 3.7), it also provides real-world training speed measurements using a Triton kernel adapted from σ-MoE (Csordás et al., 2023). The measurements include the entire training pipeline (MLP layers, optimizer, gradient synchronization for multi-GPU training), not just the attention layer.

47M parameter models (Table 5):

• Dense Transformer: 473ms per iteration, 20.5GB GPU memory (RTX 3090)
• SwitchHead: 342ms per iteration (0.72×, ~1.38× speedup), 13.5GB memory (0.65×)
• MoA (best matching model, $H=4$): 412ms per iteration (0.87×), 15.3GB memory (0.75×)

262M parameter models, 8× V100 GPUs:

• Dense Transformer: 670ms per iteration, 20.5GB per GPU
• SwitchHead: 442ms per iteration (0.65×, ~1.52× speedup), 12.5GB per GPU (0.61×)
• MoA (best matching model, $H=8$): 851ms per iteration (1.27×, actually slower than dense!), 16.4GB per GPU (0.80×)

The SwitchHead speedup of ~1.5× is achieved despite the Triton kernel being suboptimal: the paper notes that the kernel "is currently around 60% of the speed of a single dense matrix multiplication of the size of a single expert with cuBLAS" and estimates that "80-90% should be achievable with a more optimal kernel." This means the reported speedups are lower bounds—a better kernel would increase the gap. The MoA result (slower than dense at the 262M scale) is particularly telling: MoA's computational overhead of computing multiple attention matrices negates any savings from shared key/value projections.

The memory reduction is even more dramatic: SwitchHead uses 61–65% as much memory as the dense baseline. This is the direct consequence of reducing $n_{heads}$ from 16 to 2–4: the $O(n_{heads} \cdot T^2)$ activation storage dominates training memory, and SwitchHead's reduction in head count directly reduces this term. The paper's MoA comparison on memory (80% of baseline for the large model vs. SwitchHead's 61%) reflects MoA's inability to reduce $n_{heads}$ sufficiently (MoA needs 8 heads to match performance, while SwitchHead uses 2–4).

---

Summary of Design Choices and Their Justifications

• MoE on value and output only (not key/query): determined by the exhaustive ablation in Table 6. Value and output projections benefit from conditional computation; key/query projections do not in the parameter-matched regime. This also minimizes MoE complexity inside the attention computation itself.

• Non-competitive sigmoid gating over softmax: inherited from σ-MoE. Prevents expert collapse without regularization. Each expert's activation is independent—the model can use 0 to $k$ experts per token without penalty.

• Separate source and destination routers over a shared router: enables different expert selections for what information to extract from source tokens (value projection) versus how to integrate the attention output back into the residual stream (output projection). The shared-selection variant exists as a speed-accuracy tradeoff.

• Weighted average of expert projections before attention (for values) and after attention (for output) over selecting one expert: allows a single attention matrix per head to serve multiple experts simultaneously, maintaining expressivity while minimizing the number of $T \times T$ matrices.

• Parameter-matched evaluation over FLOPS-matched: tests whether MoE attention can deliver genuine resource efficiency (same performance, less compute/memory) rather than just more parameters for the same compute. This is the harder test.

• Systematic hyperparameter selection ($n_{heads}=2, k=2$ first, increase $k$, then $n_{heads}$): a practical procedure to find configurations that maximize resource savings subject to a perplexity constraint, avoiding expensive search over the full hyperparameter space.

4. Key Insights and Innovations

Innovation 1: Separating the Multiplicity of Projections from the Multiplicity of Attention Matrices

Before SwitchHead, the dominant assumption in multi-head attention—implicit in both standard Transformers and prior MoE-for-attention attempts—was that a diverse attention function requires computing a distinct attention matrix for each distinct "head" or expert. Standard multi-head attention allocates projection diversity and attention-matrix diversity together: one set of projections (K, Q, V, O) maps to exactly one attention matrix. MoA (Zhang et al., 2022) preserves this one-to-one mapping: each selected query expert computes its own attention matrix. Multi-query attention (Shazeer, 2019) breaks the coupling on the source side (sharing K and V) but still computes one attention matrix per query head. In all these designs, the number of attention matrices scales with the number of distinct output projections.

SwitchHead severs this coupling entirely. By computing a weighted average of expert value projections before the attention matrix is applied, and a weighted average of expert output projections after, a single attention matrix per head serves an entire pool of E value experts and E output experts. The attention matrix—the expensive T × T operation—is computed once and reused across all E² possible (value expert, output expert) combinations within that head. The diversity of the attention function (which information is extracted from source tokens and how it is integrated into the destination representation) is achieved through the MoE projections, not through additional attention matrices.

This is a fundamental architectural insight rather than an incremental optimization. It redefines what a "head" is: no longer a self-contained attention unit with its own matrix and projections, but a shared attention pattern that multiple expert projections can read from and write to. The paper makes this explicit: "The concepts of 'heads' are no longer well defined in the conventional sense: we redefine a head as an instance of a computed attention matrix" (Section 2.2).

The evidence supporting this insight's validity comes from three places. First, the ablation in Table 6 shows that the model with MoE value and output projections but fixed key/query projections (the configuration that shares one attention matrix across multiple experts) is the best-performing variant—better than adding MoE to keys and queries, which would complicate the attention matrix computation. Second, the comparison with MoA in Table 1 demonstrates the practical consequence: MoA with 8 heads achieves similar perplexity to SwitchHead with 2 heads precisely because MoA cannot decouple expert count from attention-matrix count. Third, the analysis of attention maps on ListOps (Figures 2–5) shows that SwitchHead's per-head attention maps are qualitatively similar to dense attention maps—the single shared matrix per head is genuinely expressive enough to serve multiple expert projections, not a bottleneck.

Why this matters beyond SwitchHead: This decoupling principle suggests a taxonomy for future attention methods: operations can be classified as (1) per-token linear projections (cheap, can be made conditional via MoE), and (2) token-mixing operations (expensive T × T matrices, should be minimized and shared). Any method that conflates these two categories is leaving resource savings on the table.

---

Innovation 2: Non-Competitive Sigmoid Gating as an Enabling Technology for MoE Attention

Competitive softmax gating—where the router outputs a probability distribution over experts—has been the default in MoE research since Shazeer et al. (2017). It has an appealing theoretical motivation: the router learns to partition the input space among experts. But in practice, it causes expert collapse: the softmax concentrates probability mass on a few experts, others receive vanishing gradients, and the effective capacity of the MoE layer collapses. The standard remedy is auxiliary load-balancing loss terms that penalize uneven expert utilization, but these introduce additional hyperparameters and do not fully solve the problem—they merely trade off routing optimality against utilization fairness.

The σ-MoE method (Csordás et al., 2023) introduced non-competitive sigmoid gating for MoE feedforward layers, showing it eliminates the need for load-balancing regularization entirely. SwitchHead's contribution is not inventing sigmoid gating, but demonstrating that it is equally critical for MoE attention—and, importantly, that MoA's competitive gating is a primary reason it requires three regularizers and fails to scale down its head count.

The evidence is stark. The paper states (Section 3.2): "our method performs well without any regularization, while MoA requires three different regularizers." MoA's softmax gating forces competition among query/output experts, which means the router cannot independently assess each expert's relevance. This creates a structural pressure toward using fewer experts (collapse), which in turn means MoA must keep its effective head count high to maintain expressivity—defeating the resource-saving purpose. SwitchHead's sigmoid gating allows each expert's selection score to be computed independently, so the model can use many experts simultaneously without competition, or use none if no expert is relevant to a given token.

This is a diagnostic contribution: it identifies competitive gating as the specific mechanism that caused prior MoE attention methods to underperform on resource reduction, not any fundamental limitation of MoE attention itself. The field had attributed MoE attention's limited success to the difficulty of the problem; SwitchHead shows it was a gating design issue.

The two-fold cross-validation strategy selection and the shared-selection variant (Section 3.6) provide further evidence of stability. The shared router, which uses a single sigmoid-gated selection for both source and destination sides, still performs competently (Table 4, minor perplexity increase of 0.28 for the 47M model)—something that would likely fail catastrophically with a competitive softmax router forced to serve two different expert pools simultaneously.

---

Innovation 3: Empirically Establishing That Value and Output Projections Are the Bottleneck, Not Keys and Queries

The paper's ablation study (Table 6) answers a question that had never been systematically asked: which attention projections benefit from conditional computation, and which do not? The result is clean and surprising in its asymmetry: the output projection is by far the most important to make conditional; the value projection provides a small additional gain; and making keys or queries conditional is actively harmful in the parameter-matched regime.

This finding runs counter to what one might expect from the attention mechanism's symmetry. Keys and queries determine the attention pattern—the routing of information between tokens—which seems like it should be highly context-dependent and benefit from conditional computation. But the data say otherwise: the best model uses a single, fixed key and query projection per head, and MoE is reserved for the projections that determine what information is read (values) and how it is integrated back into the residual stream (output).

The paper offers a parameter-budget explanation: "Allocating parameters to these projections uses the budget that can be otherwise spent on other parts of the network" (Section 3.1). This is plausible but incomplete—it explains why key/query MoEs are unnecessary (a fixed projection suffices), but not why the output projection is so disproportionately important. A deeper interpretation, consistent with the attention-map analysis (Section 4), is that the attention pattern itself is largely determined by positional relationships and learned syntactic structures that are head-specific but not token-context-dependent—a single W_K and W_Q per head can learn these patterns. What varies strongly with token context is what features to extract from the source (value projection) and how to transform the mixed information before passing it forward (output projection). Multi-head attention's well-known functional specialization (different heads attending to different syntactic relations, as documented in the interpretability literature) manifests primarily through different attention patterns (key/query), but within a pattern type, the value extraction and output transformation need to adapt to the specific tokens involved—and this is where conditional computation helps.

This is an empirical discovery rather than a theoretical advance, but it has significant practical consequences. It means SwitchHead can keep the attention core (K, Q, attention matrix computation) as a standard, highly optimized dense operation, with MoE complexity pushed to the edges where it doesn't interact with the T² computation. It also tells future designers of efficient attention that investing in conditional key/query projections is likely a poor use of parameters—a finding that could have saved substantial effort had it been known earlier.

The downstream consequence—that SwitchHead works with standard attention implementations, FlashAttention, and various positional encoding schemes (Transformer XL and RoPE, Section 3 and Appendix A.4)—is a direct result of keeping keys and queries standard. If SwitchHead had required MoE keys and queries, integrating with optimized attention kernels would have been substantially harder.

---

Innovation 4: The Parameter-Matched Evaluation Philosophy for MoE Systems

Most MoE research since Shazeer et al. (2017) has evaluated MoE models under a FLOPS-matched or compute-matched paradigm: compare an MoE model against a dense model that uses the same amount of training compute, where the MoE model naturally has more total parameters. This tests whether sparse conditional computation can scale parameters more efficiently than dense computation—and the answer has consistently been yes. But it does not test whether MoE can deliver resource efficiency: same performance with less compute and memory.

SwitchHead explicitly adopts a parameter-matched evaluation (Section 3): fix the total parameter count and ask whether replacing dense components with MoE components maintains performance while reducing resource usage. This is a deliberately higher bar. In the parameter-matched setting, every parameter consumed by the MoE gating mechanism (routers, extra expert projections) is a parameter removed from somewhere else—the MoE must be genuinely more parameter-efficient than the dense equivalent to break even.

The paper demonstrates that SwitchHead clears this bar: the 262M model on C4 matches the dense baseline's 16.28 perplexity with 16.23, while using 44% of the compute and 27% of the memory (Table 2). This is the parameter-matched claim—the same total parameter count, the same performance, dramatically less resource usage.

This is a methodological contribution that reframes what success means for MoE systems. It shifts the goal from "more parameters for the same compute" (which primarily benefits organizations with large compute budgets) to "same performance with less compute" (which benefits everyone, especially resource-constrained researchers). The paper's MAC-matched experiments (Section 3.5) complement this by showing the other direction: what happens when you fill the compute budget with a SwitchHead model (more parameters, better performance). Both perspectives are valid, but the parameter-matched framing is the one that demonstrates resource efficiency as a primary claim rather than a secondary observation.

The choice of reporting both theoretical MACs (hardware-independent, verifiable from the formulas in Appendix A.2) and wall-clock measurements (Table 5, implementation-dependent but practically meaningful) is part of this methodological rigor. The paper acknowledges the gap between the two: the Triton kernel achieves ~60% of cuBLAS efficiency, so the wall-clock speedup (~1.5×) is a lower bound on what better engineering could achieve. This transparency about the gap between theoretical and realized efficiency is itself a methodological contribution—many systems papers report only one or the other, obscuring where the bottleneck actually lies.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper uses four language modeling datasets: C4 (Raffel et al., 2020) — a large cleaned Common Crawl corpus; Enwik8 (Hutter, 2006) — a character-level Wikipedia compression benchmark; peS2o (Soldaini & Lo, 2023) — a scientific papers dataset for pretraining; and Wikitext-103 (Merity et al., 2017) — a standard word-level language modeling corpus. For sub-word-tokenized datasets (all except Enwik8), a SentencePiece tokenizer with an 8k vocabulary is used. For Enwik8, character-level modeling is used as is standard for that benchmark. All models are trained for 100k batches; for large datasets (C4, peS2o), training uses the first 10⁵ × T × N_{batch} tokens only.

• Base model(s). The paper uses Transformer XL (Dai et al., 2019) with the context size set to twice the active chunk size as the primary architecture, since the authors found it "significantly more resource-efficient than the standard setup" (Section 3). To demonstrate generality, additional experiments use a standard Transformer with Rotary Position Embeddings (RoPE; Su et al., 2021) without the XL cache (Appendix A.4). Two model scales are evaluated: ~47M parameters (47M for Transformer XL, 45M for RoPE, 41M for Enwik8) and ~262M parameters (262M for Transformer XL, 244M for RoPE). The models train with the Adam optimizer (Kingma & Ba, 2015), batch size 64, learning rate 0.00025 (with 4k-step warmup for models >200K parameters), and gradient clipping with a configurable maximum norm κ. Dropout is applied to MLP layers (0.1 for small models, 0.2 for large models, except SwitchAll which uses no MLP dropout).

• Metrics. The primary metric is perplexity for sub-word-tokenized datasets and bits per character (bpc) for Enwik8. For downstream evaluation, zero-shot accuracy is reported on Lambada (Paperno et al., 2016), BLiMP (Warstadt et al., 2020), and Children's Book Test (CBT; Hill et al., 2016). Resource usage is measured in multiply-accumulate operations (MACs) for compute and number of floats for memory, both computed per attention layer per sequence (the formulas in Appendix A.2 provide the exact accounting). Wall-clock time and GPU memory are reported separately (Table 5) for a real-world training-speed comparison including MLP layers, optimizer, and gradient synchronization.

• Baselines. The primary baselines are: (1) Standard dense Transformer with n_heads equal to the total number of expert projections in the SwitchHead model (n_heads × E), serving as the parameter-matched upper bound; (2) Standard dense Transformer with n_heads equal to the number of attention heads in the SwitchHead model (e.g., 2 heads for a SwitchHead with n_heads=2), serving as a lower bound showing the cost of naive head reduction; (3) Mixture of Attention Heads (MoA) (Zhang et al., 2022), the closest prior MoE-for-attention method, evaluated at varying numbers of active heads; (4) Baseline σ-MoE Transformer (Csordás et al., 2023) for the SwitchAll experiments, where the attention layer is replaced with SwitchHead while keeping σ-MoE MLP layers. The MAC-matched baselines (Section 3.5) use SwitchHead models scaled up (increased d_head and n_heads) to equal the dense baseline's MAC count.

• Generation budget / compute accounting. All comparisons are conducted in a parameter-matched setting: the total number of model parameters is held constant between SwitchHead and the dense baseline. The systematic procedure for configuring SwitchHead (Section 3) is: set n_heads × E equal to the dense baseline's n_heads; start with n_heads=2 and k=2; if the model underperforms, increase k to 4; if still underperforming, increase n_heads to 4 and reset k=2; adjust d_head to match the parameter count. This procedure maximizes resource savings subject to matching baseline perplexity. Compute is measured in MACs per attention layer using the formulas in Appendix A.2 (Equations 11–15), which account for projection MACs, attention-matrix computation, value readout, and position encoding—all scaled by n_heads. Memory is measured in number of floats stored, including the 2CT² per-head attention matrix storage that dominates for long sequences.

• Cross-validation / statistical protocol. No formal cross-validation or statistical significance testing is reported. The experiments train each configuration once for 100k batches. For the main results (Tables 1–4), the systematic configuration procedure described above is the only reported model selection protocol—it chooses the most resource-efficient configuration that matches baseline perplexity, not the best-performing configuration overall. The two-fold cross-validation mentioned in the prior sections' template refers to the compute-optimal policy paper (the reference example) and is not part of this SwitchHead paper.

Main Quantitative Results

SwitchHead vs. Dense Baselines Across Datasets

The paper's core empirical claim is that SwitchHead matches the language modeling performance of parameter-matched dense Transformers while using substantially less compute and memory. Table 2 provides the evidence across four datasets and two model scales.

47M parameter models (Table 2):

On C4, SwitchHead (n_heads=2, E=5, k=3) achieves 22.53 perplexity versus the 10-head dense baseline's 22.71 — actually a slight improvement (0.18 lower perplexity). Meanwhile, the 2-head dense baseline (naively reducing heads without MoE) degrades to 23.71 — a full 1.00 perplexity worse than the 10-head baseline. SwitchHead uses 203M MACs and 0.8M floats of memory versus 453M MACs and 3.5M floats for the 10-head dense model — a 55% reduction in compute and a 77% reduction in memory. The 2-head dense baseline theoretically uses the same 453M MACs as the 10-head baseline (the formulas in Equation 11 scale with n_heads, but the paper appears to report either total model MACs including MLP or there is a reporting discrepancy — the table shows 453M for both dense Transformer configurations, which is inconsistent with Equation 11's linear scaling with n_heads if reporting attention only). Regardless, the resource reduction from 10-head dense to 2-head SwitchHead is substantial.

On Wikitext-103, SwitchHead (n_heads=2, E=5, k=2) achieves 12.31 perplexity versus 12.32 for the 10-head baseline — essentially identical. The 2-head dense baseline degrades to 12.73. SwitchHead uses 170M MACs and 0.8M floats of memory, a 62% compute reduction and 77% memory reduction.

On peS2o, SwitchHead (n_heads=2, E=5, k=3) achieves 12.84 perplexity versus 12.83 for the 10-head baseline — again essentially identical. The 2-head baseline degrades to 13.37. Resource usage is the same as C4 (203M MACs, 0.8M memory).

On Enwik8 (41M parameters, character-level), SwitchHead (n_heads=2, E=4, k=2) achieves 1.10 bits per character — exactly matching the 8-head dense baseline (both 1.10 bpc). The 2-head baseline degrades to 1.13. SwitchHead uses 709M MACs and 2.8M floats versus 1.6G MACs and 10M floats for the 8-head baseline.

The consistent pattern across all datasets and scales is: (a) SwitchHead matches or slightly exceeds the many-head dense baseline's perplexity, (b) substantially outperforms the few-head dense baseline (by 0.3–0.6 perplexity for word-level models, 0.03 bpc for Enwik8), and (c) uses dramatically less compute and memory than the many-head baseline (44–62% of the compute, 23–28% of the memory). The second finding — that SwitchHead outperforms the same-n_heads dense baseline — demonstrates that the MoE mechanism is not merely recovering lost capacity but actively providing useful conditional computation: if n_heads=2 with fixed projections degrades perplexity but n_heads=2 with MoE value/output projections matches the n_heads=10 baseline, then the MoE is successfully serving the functional role that 8 additional attention heads served in the dense model without computing 8 additional attention matrices.

262M parameter models (Table 2):

On C4, SwitchHead (n_heads=4, E=4, k=2) achieves 16.23 perplexity versus the 16-head dense baseline's 16.28 — again slightly better. The 4-head baseline degrades to 17.09. SwitchHead uses 2.4G MACs and 5.6M floats versus 5.4G MACs and 21M floats — a 56% compute reduction and 73% memory reduction.

On Wikitext-103, SwitchHead (n_heads=2, E=8, k=4) achieves 9.77 perplexity versus the 16-head baseline's 9.80 — essentially tied. SwitchHead uses 2.0G MACs and 2.9M floats versus 5.4G MACs and 21M floats — a 63% compute reduction and 86% memory reduction.

On peS2o, SwitchHead (n_heads=4, E=4, k=2) achieves 9.86 perplexity versus 9.78 for the 16-head baseline — a slight degradation (0.08 perplexity), but still dramatically outperforming the 4-head baseline at 10.11. Resource usage is the same as C4 for the large model.

Comparison with MoA (Table 1)

The paper positions SwitchHead as a direct improvement over MoA (Zhang et al., 2022), and Table 1 provides the empirical basis for this claim on Wikitext-103.

47M parameter models (Table 1):

SwitchHead (n_heads=2, E=5, k=2) achieves 12.27 perplexity with 170.4M MACs and 0.8M floats memory. The 10-head dense Transformer achieves 12.31 with 453.4M MACs and 3.5M memory. MoA with n_heads=4 achieves 12.60 with 223.5M MACs and 1.3M memory — worse perplexity than SwitchHead despite using 31% more compute. MoA increases head count to n_heads=6 (12.64 perplexity, 306.8M MACs) and n_heads=8 (12.77 perplexity, 390.2M MACs) and perplexity actually degrades — more heads produce worse results, a clear sign of training instability or overparameterization issues. MoA with n_heads=2 achieves 12.84 perplexity, the worst in the table other than the 2-head dense baseline (which isn't shown in this table but is 12.74 from Table 2).

The takeaway is unambiguous: SwitchHead with 2 heads outperforms all tested MoA configurations in both perplexity and resource usage. MoA needs n_heads ≥ 4 to approach SwitchHead's perplexity, and even the best MoA configuration (12.60 at 4 heads) underperforms while using more compute (223.5M vs. 170.4M MACs) and more memory (1.3M vs. 0.8M floats).

262M parameter models (Table 1):

SwitchHead (n_heads=2, E=8, k=4) achieves 9.55 perplexity with 2.0G MACs and 2.9M floats. The dense baseline achieves 9.66 with 5.4G MACs and 21M memory. MoA with n_heads=8 achieves the best perplexity of all methods at 9.50 — slightly better than SwitchHead's 9.55 — but uses 2.9G MACs (45% more) and 9.9M memory (3.4× more). MoA with n_heads=4 drops to 9.69 perplexity (worse than the dense baseline) despite using 1.7G MACs and 5.1M memory (still more memory than SwitchHead). MoA with n_heads=2 degrades further to 9.87.

The 262M results reveal the fundamental trade-off: MoA can match or slightly exceed SwitchHead's perplexity, but only by using many more attention heads (8 vs. 2), which multiplies the O(n_heads × T²) attention-matrix cost. SwitchHead achieves 9.55 perplexity with 2 heads; MoA needs 8 heads to achieve 9.50. This 4× difference in head count is the direct consequence of the architectural difference — MoA's selected experts each need their own attention matrix, while SwitchHead's experts share one matrix per head. The paper summarizes this finding:

"Given a similar computation and memory budget, our method consistently outperforms MoA" (Section 3.2)

The wall-clock measurements in Table 5 reinforce this. At the 47M scale, MoA with H=4 achieves 12.60 perplexity (vs. SwitchHead's 12.27) while training at 412ms/iteration (vs. SwitchHead's 342ms) and using 15.3GB memory (vs. 13.5GB). At the 262M scale, MoA with H=8 achieves 9.50 perplexity (vs. 9.55) but trains at 851ms/iteration — 27% slower than the dense baseline (670ms) and 93% slower than SwitchHead (442ms) — while using 16.4GB memory (vs. 12.5GB for SwitchHead). MoA at 262M is not just less efficient than SwitchHead; it is actually slower than the dense Transformer it was designed to accelerate, while SwitchHead delivers a 1.52× wall-clock speedup.

SwitchAll: Combining SwitchHead with σ-MoE MLP

Table 3 evaluates the "SwitchAll" configuration — SwitchHead attention plus σ-MoE feedforward layers (Csordás et al., 2023) — against the dense baseline on three datasets and two scales.

47M parameter models:

SwitchAll with n_heads=2 outperforms the dense baseline on all three datasets:

• Wikitext-103: 12.17 vs. 12.32 (Transformer with 10 heads)
• C4: 22.09 vs. 22.63
• peS2o: 12.56 vs. 12.83

The resource savings are identical to the SwitchHead-only configuration for the attention component, with additional savings from the σ-MoE MLP layers (not separately quantified in the paper, but see Csordás et al., 2023 for those savings). SwitchAll with 2 heads uses 170M MACs and 0.8M floats of attention memory (Wikitext-103 configuration) vs. 453M MACs and 3.5M floats for the dense baseline.

262M parameter models:

On Wikitext-103, SwitchAll achieves 9.81 perplexity vs. the dense baseline's 9.80 — essentially tied. On C4, SwitchAll achieves 16.45 vs. 16.58 — slightly better. On peS2o, the dense baseline (9.78) slightly edges out SwitchAll (9.86). The resource savings are as reported for the SwitchHead-only 262M models: 2.0–2.4G MACs vs. 5.4G for dense, 2.9–5.6M floats of attention memory vs. 21M.

The key finding from Table 3 is the absence of negative interactions: combining two independently-developed MoE methods (σ-MoE for MLP, SwitchHead for attention) does not cause degradation beyond what each method achieves individually. The paper explicitly frames this as answering a question about potential interaction effects: "it remains unclear whether it can be efficiently combined with our SwitchHead, or can have some negative interaction effect if combined in a 'SwitchAll', where every layer is MoE-based" (Section 3.4). The results show no such negative interaction — in fact, SwitchAll outperforms the dense baseline in 4 of 6 configurations and ties in one.

A minor caveat noted in Appendix A.6: the SwitchAll parameter counts are sometimes slightly below the dense baselines (e.g., 259M vs. 262M for the large model) due to the granularity of σ-MoE expert sizing — expert dimensions must be multiples of 4 for CUDA kernel compatibility, and increasing to the next multiple would overshoot the parameter budget. This means the performance comparisons slightly favor the dense baseline (SwitchAll has a small parameter deficit), making the SwitchAll results a conservative estimate.

RoPE-Based Transformer Results

All main experiments use Transformer XL. Appendix A.4 verifies that SwitchHead is not specific to this architecture by replicating key results with standard RoPE-based Transformers (the architecture used by modern LLaMA-style models).

Table 7 (perplexity):

On Wikitext-103 at the 45M scale, SwitchHead (n_heads=2, E=5, k=3) achieves 12.75 perplexity vs. the 10-head dense baseline's 12.78 — essentially tied, with the 2-head baseline at 12.96. SwitchHead uses 285.6M MACs and 1.3M floats of memory vs. 560.9M MACs and 6.1M floats for the 10-head baseline — a 49% compute reduction and 79% memory reduction. At 244M, SwitchHead (n_heads=4, E=4, k=2) achieves 10.00 vs. 10.17 for the 16-head baseline — actually better — with 4.2G MACs and 18.4M memory vs. 6.4G MACs and 37.7M memory.

On C4 at 45M, SwitchHead achieves 23.69 vs. 23.79 for the 10-head baseline. At 244M, SwitchHead achieves 16.41 vs. 16.35 for the 16-head baseline — a negligible 0.06 degradation. The resource savings are comparable to the Transformer XL experiments.

Table 8 (zero-shot downstream tasks on RoPE models):

At the 45M scale, SwitchHead achieves 20.9% on Lambada vs. 20.3% for the dense baseline, and 77.3% on BLiMP vs. 73.8% — a substantial 3.5% absolute improvement on BLiMP. CBT is not reported for the 45M models. At 243M, SwitchHead achieves 30.5% Lambada vs. 29.8% for dense, 79.9% BLiMP vs. 76.1% — another large 3.8% absolute improvement — and 83.8% CBT vs. 83.9% (tied).

The 3.5%+ absolute BLiMP improvements with SwitchHead appear across all configurations (RoPE and Transformer XL, small and large models). This is notable because BLiMP is a benchmark of linguistic acceptability judgments — the fact that SwitchHead, which uses drastically fewer attention heads, achieves better linguistic judgment scores than the dense baseline suggests that the MoE mechanism is not just preserving but actively improving the quality of the learned representations for linguistic phenomena. The paper does not analyze which linguistic phenomena benefit most, but the magnitude and consistency of the improvement rule out random variation.

Zero-Shot Performance of C4-Trained Models (Table 4)

Table 4 reports zero-shot performance of the Transformer XL-based SwitchHead models trained on C4, complementing the RoPE results in Table 8.

At the 47M parameter scale, SwitchHead and the dense baseline achieve identical Lambada accuracy (20.4%), but SwitchHead substantially outperforms on BLiMP (75.7% vs. 73.6%) — a 2.1% absolute improvement. CBT is not reported for the small models. At 262M, SwitchHead achieves 29.4% Lambada (vs. 28.2%), 79.6% BLiMP (vs. 76.1%, a 3.5% absolute improvement matching the RoPE result), and 83.3% CBT (vs. 83.6% — essentially tied).

These zero-shot results are important because they address a potential concern: that SwitchHead might match or exceed perplexity on the training distribution while learning representations that are worse for transfer. The zero-shot results show the opposite — SwitchHead's representations transfer at least as well, and for BLiMP consistently transfer better, than dense models with matched perplexity.

The MAC-matched SwitchHead models (Section 3.5, also in Table 4) show the expected improvement from scaling up to fill the compute budget. The 63M MAC-matched model (vs. 47M parameter-matched) improves perplexity from 22.53 to 21.18, Lambada from 20.4% to 23.5%, and BLiMP from 75.7% to 77.1%. The 376M MAC-matched model improves perplexity from 16.23 to 15.43, Lambada from 29.4% to 30.2%, and CBT from 83.3% to 84.2%. All MAC-matched models substantially outperform their parameter-matched counterparts, confirming that SwitchHead can productively use additional compute when available — making it viable in both resource-constrained (parameter-matched) and performance-maximizing (MAC-matched) scenarios.

The shared selection variant (Section 3.6, Table 4) shows a consistent but modest performance cost. At 47M, shared selection increases perplexity from 22.53 to 22.81 and reduces BLiMP from 75.7% to 74.6% (Lambada drops from 20.4% to 20.0%). At 262M, perplexity increases from 16.23 to 16.49, Lambada drops from 29.4% to 28.6%, and CBT drops from 83.3% to 82.7%. The BLiMP score for the large model with shared selection (79.4%) is essentially unchanged (79.6% for full SwitchHead), but this single-data-point exception is likely noise. The paper frames shared selection as an option "where the acceleration is more important" (Section 3.6) — the cost is ~0.3–0.5 perplexity points and 1–2% on downstream tasks, which may be acceptable when the additional top-k savings are significant relative to the total computation.

Ablation Studies and Robustness Checks

Which projections benefit from MoE (Table 6): The exhaustive 16-configuration ablation on the 47M SwitchHead model (n_heads=2, E=5) on Wikitext-103 reveals that the output projection MoE is the single most important factor. The configuration with only output MoE (V=N, K=N, Q=N, O=Y) achieves 12.30 perplexity, already outperforming the 10-head dense baseline (12.31). Adding value MoE (V=Y, O=Y) provides a further 0.03 improvement to 12.27. Conversely, any configuration without output MoE performs worse than the 10-head baseline, with perplexity ranging from 12.45 to 12.90. Key and query MoEs are unnecessary or harmful — the full 4-projection MoE (all Y) achieves only 12.57, substantially worse than the V+O-only configuration. This ablation is the empirical foundation for the paper's central design decision to use MoE only on value and output projections.

Naive head count reduction (Table 2, dense 2-head vs. dense 10/16-head): Across all datasets and scales, reducing head count while proportionally increasing d_head to maintain parameters consistently degrades perplexity: 23.71 vs. 22.71 on C4 (47M), 12.73 vs. 12.32 on Wikitext-103 (47M), 13.37 vs. 12.83 on peS2o (47M), 1.13 vs. 1.10 bpc on Enwik8 (41M), 17.09 vs. 16.28 on C4 (262M), and 10.09 vs. 9.80 on Wikitext-103 (262M). This degradation demonstrates that many heads are structurally necessary — they cannot be trivially replaced by fewer, wider heads — and establishes the baseline that SwitchHead must overcome.

Transformer XL vs. RoPE (Table 2 vs. Table 7): SwitchHead is evaluated under two different attention mechanisms and positional encoding schemes. The results are consistent across both: on Wikitext-103 at comparable scales, SwitchHead matches or beats dense baselines while using 49–63% of the compute and 21–27% of the memory. This robustness check confirms that SwitchHead's savings are not tied to a specific attention implementation.

σ-MoE MLP integration (SwitchAll, Table 3): The combination of independently-designed MoE methods for attention and MLP shows no negative interaction effects. SwitchAll outperforms the dense baseline in 4 of 6 perplexity comparisons and ties in one — the combination is at least additive and possibly complementary. The only case where dense outperforms SwitchAll is peS2o at 262M (9.78 vs. 9.86), a 0.08 perplexity gap within reasonable experimental noise.

MoA as a comparative ablation (Table 1): By systematically varying MoA's head count and comparing to SwitchHead, the paper effectively ablates the architectural difference between the two methods. The fact that MoA needs 4× more heads than SwitchHead (8 vs. 2) to achieve comparable perplexity while using more resources isolates the shared-attention-matrix design as the critical efficiency mechanism. The fact that MoA at 8 heads (262M) is actually slower than the dense baseline (Table 5, 851ms vs. 670ms) demonstrates that conditional computation in the attention core (computing separate attention matrices per selected expert) can be counterproductive when the overhead exceeds the savings.

Shared selection vs. independent routers (Table 4): The shared-selection variant reduces to a single router and top-k operation. The performance cost — approximately 0.3 perplexity points and 1–2% on downstream tasks — establishes the value of independent source and destination routing. This is a meaningful but not catastrophic degradation, validating that the two sides benefit from different gating decisions.

MAC-matched scaling (Tables 4 and 8, Section 3.5): When SwitchHead models are scaled up (increased n_heads and d_head) to match the dense baseline's MAC count, they substantially outperform both the parameter-matched SwitchHead and the dense baseline. This demonstrates that SwitchHead's efficiency is not a zero-sum trade between resources and performance — the mechanism can use additional compute productively. For the Transformer XL models on C4, MAC-matched SwitchHead achieves 21.18 vs. 22.71 (dense) at 63M parameters. For the RoPE models, MAC-matched SwitchHead achieves 22.18 vs. 23.76 (dense) at 54M parameters.

Critical Assessment

Central Claim: SwitchHead matches baseline performance with substantially less compute and memory.

The evidence broadly supports this claim, but with important nuance about which configurations achieve the matching and how close the matching is.

Strengths of the evidence:

• The claim is tested across four datasets (C4, Wikitext-103, peS2o, Enwik8), two model scales (47M and 262M), and two attention architectures (Transformer XL and RoPE). The consistency across these dimensions is the strongest argument for generality.
• The perplexity matching is genuinely close: in 11 of the 12 main comparisons in Table 2 and Table 7, SwitchHead is within 0.08 perplexity of the dense baseline or better. The largest gap is 0.08 perplexity (peS2o at 262M, 9.86 vs. 9.78), which is well within typical experimental variance for language model training.
• The resource reductions are substantial and consistent: compute reductions of 44–63%, memory reductions of 73–86% in the attention component.
• The wall-clock measurements (Table 5) confirm that theoretical MAC reductions translate to real speedups (~1.5×) despite a suboptimal Triton kernel, and that memory reductions (61–65% of baseline) are real.

Weaknesses and limitations of the evidence:

The most significant limitation is that difficulty of matching performance varies by dataset and scale, and the paper's systematic configuration procedure does not always find exactly matching perplexity. On some configurations, the baseline is matched or exceeded; on others, SwitchHead slightly underperforms. The paper handles this by selecting the most resource-efficient configuration that approximately matches, but "approximately matches" is not precisely defined — there is no stated threshold for acceptable perplexity degradation. A more rigorous approach would report confidence intervals or run multiple seeds to establish whether the gaps (e.g., 0.08 on peS2o at 262M, 0.05 on Wikitext-103 at 47M for SwitchHead vs. 12.31 for dense vs. 12.27 for SwitchHead in Table 1) are statistically significant or within noise. At 100k training steps with these model sizes, the between-run variance on perplexity is likely around 0.05–0.15 based on typical language modeling experiments, meaning many of the reported differences are likely within noise.

Model scale is modest by current standards. The 262M parameter models are roughly 1/30th the size of GPT-3 (175B) and 1/250th the size of current frontier models. The paper acknowledges this limitation (Section 6): "Our models are modest in size compared to the current state-of-art LLMs. However, training such models is estimated to cost millions of dollars, which we cannot afford." This is a legitimate resource constraint, but it leaves open the question of whether the findings scale. Several potential failure modes at scale: (1) expert utilization patterns might change — the sigmoid gating might saturate differently with larger models, (2) the optimal n_heads and E configuration might differ at billion-parameter scales, (3) the relative cost of attention vs. MLP shifts with model size, potentially making attention savings less impactful (or more impactful, depending on sequence length). The paper's consistent results across a 6× parameter range (47M to 262M) provide some evidence of scalability, but there is a substantial extrapolation gap to production-scale models.

Single training run per configuration. All results are from single training runs. Without multiple seeds, it's impossible to distinguish genuine architectural advantages from random initialization or training-order effects. This is particularly relevant for the SwitchAll results where the performance gaps are small (e.g., 16.45 vs. 16.58 on C4 at 262M). Reporting mean and standard deviation across 3–5 seeds for the main comparisons would significantly strengthen the evidence.

Limited sequence length analysis. The paper's resource savings depend critically on the T² term in attention dominating over the linear projection terms. All experiments use context sizes of 256–1024 tokens (Table 9). For these lengths, the attention matrix is a meaningful but not dominant fraction of total compute — the linear projection terms are still significant. The paper does not systematically vary sequence length to show how savings scale with T, which is the regime where SwitchHead's advantage should be largest. A sweep of sequence lengths (e.g., 128, 512, 2048, 4096) showing growing relative speedup with T would validate the core architectural motivation. The paper's resource formulas in Appendix A.2 allow readers to extrapolate, but empirical confirmation would be more convincing.

No comparison to FlashAttention. The paper mentions (Appendix A.1) that FlashAttention can be combined with SwitchHead and that the RoPE experiments already do so, but no systematic FlashAttention baseline or combination is reported. Since FlashAttention is now standard in most Transformer implementations, the practical question for most practitioners is "does SwitchHead provide additional savings on top of FlashAttention?" The paper's MAC-based accounting is orthogonal to FlashAttention's memory-bandwidth optimizations, but demonstrating that the combination yields wall-clock speedup beyond FlashAttention alone would strengthen the practical case.

Central Claim: SwitchHead achieves up to 8× fewer attention matrices.

This claim is definitionally true given the architecture: SwitchHead uses 2–4 heads vs. 10–16, so the reduction factor is 4–8×. The evidence supporting this claim is the Table 2 resource numbers showing that this reduction in attention matrices translates to the reported MAC and memory savings. However, the claim is about the number of attention matrices, not about end-to-end speedup — the 8× figure should not be misinterpreted as an 8× wall-clock speedup (which the paper does not claim; the measured speedup is ~1.5×).

Central Claim: SwitchHead outperforms MoA.

This claim is strongly supported by Table 1 and Table 5. MoA cannot match SwitchHead's perplexity at comparable resource usage, and at the 262M scale with 8 heads, MoA is actually slower than the dense baseline. The evidence is clear and the architectural explanation (separate attention matrices per expert in MoA vs. shared attention matrix in SwitchHead) is well-motivated. A minor critique: the paper compares against MoA with the same number of total experts but does not explore whether MoA with different hyperparameters than those reported might perform better. Given that MoA uses three regularizers with tunable coefficients, the hyperparameter space is large, and the reported MoA numbers may not represent the best possible MoA performance. However, the structural argument — that MoA cannot reduce attention-matrix count without reducing expert count — means that even optimally-tuned MoA would face the same fundamental trade-off.

Central Claim: SwitchAll successfully combines MoE attention and MoE MLP.

This claim is supported (Table 3), but with the weakest evidence in the paper. The performance gaps are small and based on single training runs. The paper's explicit statement that the σ-MoE MLP hyperparameters were taken "without any hyperparameter change" (Section 3.4) means no attempt was made to optimize the combination — this is a strength (no cherry-picking) but also means the reported performance may understate what a jointly-optimized system could achieve. The small parameter deficits in SwitchAll models (Appendix A.6) further complicate interpretation.

Missing experiments that would strengthen the paper:

• Scaling with sequence length. A sweep of T from 256 to 4096 showing that the relative speedup increases with sequence length would validate the core architectural motivation (that reducing n_heads reduces the T² term).
• Multi-seed results for main comparisons. 3–5 seeds for the main Table 2 configurations would allow distinguishing signal from noise.
• Inference-only measurements. All timing results are for training. Inference, where activation memory is less of a concern but the attention computation is still the bottleneck for long sequences, might show different ratios between SwitchHead and baselines.
• Expert utilization analysis. The paper does not report how evenly the E experts are used. With sigmoid gating, it's possible that some experts are rarely selected, which would effectively reduce the model's capacity. Analyzing expert utilization would reveal whether the E experts are genuinely all contributing.
• Ablation on number of experts and k. The paper's systematic procedure varies n_heads and k but not E independently. The relationship between E (total experts per head) and k (active experts) is unexplored — would E=8, k=1 work as well as E=4, k=2? This has implications for implementation efficiency since the number of active experts determines the MoE compute overhead.
• Direct comparison to multi-query attention (MQA). The paper mentions MQA in Related Work and argues it is "suboptimal" based on the Table 6 ablation showing output MoE is the most critical, but never directly compares SwitchHead to MQA. This comparison would directly test the claim that value/output MoE is better than key/query sharing.

6. Limitations and Trade-offs

1. All Results Are from Modest-Scale Models and a Single Dataset Domain

The assumption or constraint. The paper evaluates SwitchHead exclusively on models with 41M–262M parameters and across a single task domain: autoregressive language modeling. The datasets (C4, Enwik8, peS2o, Wikitext-103) all require predicting the next token in text sequences. The authors explicitly acknowledge the scale limitation in Section 6:

"Our models are modest in size compared to the current state-of-art LLMs. However, training such models is estimated to cost millions of dollars, which we cannot afford."

and justify the choice by arguing:

"We believe that the evidence we provided is enough for a research group with a larger amount of resources at their disposal to verify our findings in a state-of-the-art model."

The consequence. Several plausible failure modes at scale cannot be ruled out. First, sigmoid gating behavior may change with model capacity — larger models might saturate the sigmoid differently, leading to either expert collapse (most tokens routing to the same few experts) or excessive expert activation (too many experts with non-trivial gate scores, diluting the MoE benefit). Second, the optimal n_heads, E, and k configuration is unknown beyond the 47M–262M range. The paper's systematic procedure (start with n_heads=2, k=2, increase k, then increase n_heads) was validated only up to 262M and could fail to find good configurations at billion-parameter scales. Third, the relative importance of attention versus feedforward compute changes with model size — in very large models where the MLP dominates parameter count and compute, attention savings may be proportionally less impactful even if the relative reduction remains large. Fourth, the paper provides no evidence about whether SwitchHead's efficiency extends to non-language tasks (vision, multimodal, protein modeling) or non-autoregressive settings (masked language modeling, encoder-decoder architectures, bidirectional attention). The exclusive focus on autoregressive language modeling leaves the generality of the method untested.

What evidence exists in the paper. The paper shows consistent results across a ~6× parameter range (47M to 262M) and across four language modeling datasets with two different attention architectures (Transformer XL and RoPE). This provides some evidence of robustness within the tested regime but does not constitute a scaling law that can be extrapolated. The paper does not report experiments at larger scales, on non-language tasks, or with non-autoregressive training objectives. The downstream zero-shot evaluations (Lambada, BLiMP, CBT) provide some evidence that the learned representations generalize, but these are all linguistic benchmarks derived from the same LM training.

Mitigation status. The authors explicitly flag this as a limitation and call for external verification by groups with larger compute budgets. No theoretical analysis or scaling law is provided that would allow practitioners to predict SwitchHead's behavior at larger scales. The paper does NOT address the domain-specificity question at all — there is no discussion of whether SwitchHead would transfer to vision Transformers, protein language models, or other domains where attention is used.

---

2. The Wall-Clock Speedup Is Substantially Less Than the Theoretical MAC Reduction, With No Inference Measurements

The assumption or constraint. The paper reports theoretical MAC reductions of 44–63% in the attention layer, corresponding to headline figures like "only 44% compute." However, the measured wall-clock training speedup in Table 5 is approximately 1.38–1.54× — roughly a 28–35% reduction in iteration time — far smaller than the theoretical MAC reduction. The paper acknowledges this gap in Section 3.7 and Appendix A.1:

"The Triton kernel that we used is currently around 60% of the speed of a single dense matrix multiplication of the size of a single expert with cuBLAS. We estimate that 80-90% should be achievable with a more optimal kernel."

Additionally, all wall-clock measurements are for training only. Inference speed is not measured at all.

The consequence. The practical speedup that a practitioner can expect today is ~1.5× for training and unknown for inference. This matters because the use cases that benefit most from SwitchHead — long-sequence processing where the T² attention term dominates — are typically inference-heavy (e.g., document summarization, code completion, conversational agents). During autoregressive inference, the attention computation differs from training: only the new token's query interacts with the cached keys and values, meaning the attention pattern is O(T) per step rather than O(T²) for the full matrix. The relative savings from reducing n_heads may be different in this regime — potentially larger (because the attention memory for the KV cache scales with n_heads) or smaller (because the T²-dominated training forward/backward pass is no longer the bottleneck). Without inference measurements, the most deployment-relevant performance metric is missing.

Furthermore, the gap between theoretical MAC reduction (down to 44% of baseline) and measured speedup (~65–72% of baseline iteration time) reveals that the current implementation has substantial overhead. The MoE routing (two sigmoid + top-k operations per head), the weighted averaging of expert projections, and the kernel launch overhead all consume time that the theoretical MAC count does not capture. The paper's estimate that "80-90% [of cuBLAS efficiency] should be achievable with a more optimal kernel" is speculative and not backed by any prototype.

What evidence exists in the paper. Table 5 provides the only wall-clock measurements. These cover training for two model scales (47M on RTX 3090, 262M on 8× V100) and include the full pipeline (MLP, optimizer, gradient sync). No inference measurements, no kernel-level profiling to identify the overhead sources, and no ablation showing how much of the gap is due to the Triton kernel vs. inherent overhead of the MoE mechanism (sigmoid + top-k + weighted averaging operations that don't exist in dense attention).

Mitigation status. The paper acknowledges the kernel efficiency gap explicitly and suggests it can be closed with engineering effort. However, the inference-speed question is not mentioned at all. The MoA comparison in Table 5 provides useful context — MoA is actually slower than the dense baseline at the 262M scale, while SwitchHead is faster — but this does not address the inference gap.

---

3. Expert Configuration Is Tuned by a Heuristic Procedure, Not Optimized

The assumption or constraint. The paper uses a simple heuristic to configure SwitchHead: set n_heads × E equal to the dense baseline's n_heads, start with n_heads=2 and k=2, increase k to 4 if performance is insufficient, then increase n_heads to 4 if still insufficient. The paper describes this procedure as "reasonably simple" and notes it "ensures a good amount of resource savings, while avoiding doing an expensive hyperparameter search" (Section 3).

No systematic exploration of the interaction between n_heads, E, and k is performed. Key questions are unaddressed: Is E=4, k=2 equivalent to E=8, k=1? Does the number of total experts E matter independently of the number of active experts k? What happens when E is very large (e.g., 16 or 32) with the same k? The systematic procedure also does not consider E as a free variable — it is determined by E = baseline_n_heads / n_heads, coupling expert count to head reduction.

The consequence. The reported resource savings may not be Pareto-optimal. There could exist configurations with the same perplexity and even lower resource usage that the heuristic misses. For example, a configuration with n_heads=3 (intermediate between the tested 2 and 4) might provide better resource-perplexity tradeoffs on some datasets. A configuration with E=8, n_heads=2, k=1 might match the performance of E=4, n_heads=2, k=2 while requiring half the MoE compute (one active expert instead of two). Conversely, the heuristic might select configurations that are unnecessarily resource-intensive because it stops at the first configuration that "works" rather than searching for the minimum-resource configuration.

For practitioners wanting to apply SwitchHead, the lack of a principled configuration method means they must either replicate the paper's heuristic (which may not transfer to different model scales, architectures, or domains) or conduct their own hyperparameter search — which is time-consuming and undermines the method's practical value.

What evidence exists in the paper. The paper's experimental section (Tables 2, 7, 9) shows the configurations selected by the heuristic for each dataset and model size: n_heads=2 with E=5 and varying k ∈ {2,3} for 47M models, n_heads=2, E=8, k=4 for the 262M WikiText-103 model, and n_heads=4, E=4, k=2 for the 262M C4 and peS2o models. The variation in configurations (different n_heads, E, k across datasets at the same parameter count) suggests that optimal choices are dataset-dependent and a heuristic is needed, but no ablation studies explore the sensitivity to these choices. The paper does NOT report the performance of configurations that the heuristic rejected (e.g., what perplexity did n_heads=2, k=2 achieve on C4 before it was rejected in favor of k=3?).

Mitigation status. The paper does not address this limitation. There is no suggestion for future work on automated configuration search, no scaling analysis of how E, k, and n_heads interact, and no recommendation for how practitioners should choose these hyperparameters for new settings. The heuristic is presented as a practical procedure, not validated as a near-optimal one.

---

4. No Analysis of Expert Utilization or Potential for Collapse with Sigmoid Gating at Larger Scales

The assumption or constraint. The paper inherits σ-MoE's sigmoid gating mechanism (Csordás et al., 2023) and claims it eliminates the need for load-balancing regularization because sigmoid is non-competitive: "using a non-competitive selection function (sigmoid in Eq. 4)... our method performs well without any regularization, while MoA requires three different regularizers" (Section 3.2). The paper does not report any metrics on how evenly the E experts within each head are utilized — no expert selection frequency histograms, no entropy of the gate distributions, and no analysis of whether some experts are "dead" (never selected).

The consequence. Sigmoid gating removes the competitive pressure that causes softmax-based routers to collapse to always selecting the same few experts, but it introduces a different potential failure mode: the model may never learn to use some experts at all. In a sigmoid-gated MoE, each expert's gate is computed independently as σ(x w_e). If the initialization or learning dynamics cause some experts' gates to be consistently low across all inputs, those experts receive small gradients and may never become useful — effectively dead weight consuming parameters without contributing to model capacity. Unlike softmax collapse, which produces visible symptoms (most probability mass on one expert) that load-balancing losses can detect and penalize, sigmoid collapse could be silent — the dead experts simply have small gate values that the top-k never selects, and the model performs fine using the remaining experts.

The consequence for practitioners is uncertainty about capacity utilization: if SwitchHead with E=5, k=2 only ever uses 3 of the 5 experts per head, the effective capacity is lower than expected, and resources are wasted on dead expert projections. At the modest scales tested (47M–262M), this might not matter because there is sufficient capacity in the remaining experts. At larger scales where pushing the parameter efficiency frontier matters more, dead experts would represent wasted memory and potential capacity.

What evidence exists in the paper. No evidence. The paper provides no expert utilization analysis whatsoever — no histograms, no selection frequency statistics, no per-expert perplexity or gradient analysis. The paper does not even report whether all E experts receive non-zero top-k selections during a typical training run. The ListOps visualization (Figures 5a–5l) shows the expert selection maps for a synthetic task, revealing that on that particular task, experts do specialize (output experts for different operations, value experts for numbers vs. parentheses). However, this is a small, interpretable, 6-layer model on a synthetic task — it provides no evidence about expert utilization in the 262M language models where the main claims are made.

Mitigation status. The paper does not acknowledge this as a potential issue. The claim that sigmoid gating "does not require regularization" is a statement about training stability (no collapse during training), not about efficient utilization of all experts after training. The shared selection variant (Section 3.6), which forces the same experts to be selected for both value and output, implicitly assumes that the selected experts are useful on both sides — but without utilization analysis, there is no evidence that this assumption holds. Future work on monitoring and ensuring uniform expert utilization under sigmoid gating is not suggested.

---

5. No Systematic Analysis of Sequence Length Scaling — Where the Method Should Shine

The assumption or constraint. The core architectural motivation for SwitchHead is that reducing n_heads reduces the O(n_heads × T²) attention-matrix cost, which dominates for long sequences. The resource formulas in Appendix A.2 explicitly show that the quadratic term 2CT²d_head multiplies n_heads, so reducing n_heads from 16 to 2 should provide an 8× reduction in attention-matrix compute and memory. However, all experiments use fixed, relatively modest sequence lengths: 256 tokens for the 47M models, 512 tokens for the 262M Transformer XL models, and 512–1024 tokens for the RoPE models (Table 9). The paper does not vary sequence length to demonstrate that savings increase with T.

The consequence. Without a sequence-length sweep, the paper cannot empirically demonstrate its central scaling claim: that the relative advantage of SwitchHead over dense attention grows with sequence length. At short sequence lengths (T=256), the linear projection terms in the resource formulas (4T d_head d_model per head for projections vs. 2CT² d_head for attention) are comparable in magnitude, meaning the savings from reducing n_heads are partially offset by the larger d_head (needed to maintain parameter count) and the MoE overhead. At long sequence lengths (T=2048 or T=4096), the T² term dominates, and the 4–8× reduction in n_heads should provide proportionally larger savings.

For practitioners considering SwitchHead for long-context applications — which is precisely where attention cost is most problematic — the absence of this data means they must extrapolate from the resource formulas without empirical validation. The formulas assume ideal kernel efficiency; real implementations may have different scaling properties, especially with the MoE operations (which scale with T × k × d_head × d_model, also linear but with a larger constant than dense projections).

What evidence exists in the paper. None. All experiments use the context sizes in Table 9, which span 256–1024 tokens. There is no experiment that fixes all hyperparameters and varies T from, say, 128 to 4096 tokens. The resource formulas in Appendix A.2 provide a theoretical basis for extrapolation, but no empirical confirmation.

Mitigation status. The paper does not acknowledge this gap. The resource formulas are provided as analytical tools, and the narrative around the T² term's importance appears in the introduction and method sections, but the experimental design does not test the implication. This is particularly striking because varying sequence length is one of the cheapest ablations to run — it requires no new hyperparameter search, just changing the data preprocessing or chunk size.

---

6. The Parameter-Matched Setting Constrains the Design Choices and May Not Transfer to FLOPS-Matched Production Deployments

The assumption or constraint. The paper's entire experimental philosophy is parameter-matched: both SwitchHead and the dense baseline have the same total parameter count, and the goal is to maintain perplexity while reducing compute and memory. This drives the critical design decision to NOT use MoE for key and query projections. As the paper explains (Section 3.1):

"Since experts use a significant part of the parameter budget, they can reduce the number of parameters available for the more useful parts of the model."

In a parameter-matched setting, allocating parameters to K and Q MoE removes them from V, O, d_head, or d_ff — and the ablation shows this tradeoff is negative. But this conclusion is conditional on the parameter-matched constraint. In a FLOPS-matched setting, where compute is fixed but parameters can increase, the tradeoff is different: adding K and Q MoE would increase total parameters and compute per token, but the additional capacity might improve perplexity more than spending those FLOPs on wider fixed projections.

The consequence. The paper's central design finding — "value and output MoE are sufficient; key and query MoE are unnecessary" — is widely cited and likely to be interpreted as a general architectural principle. But it is not established as such. It is a finding about parameter efficiency, not about intrinsic utility. It could be that K and Q MoEs genuinely provide useful conditional computation, and the negative result in Table 6 simply reflects the parameter-matched tradeoff where the cost (diverted parameters) outweighs the benefit. In a FLOPS-matched production setting where the goal is maximizing quality for a given compute budget rather than minimizing compute for a given parameter budget, K and Q MoEs might be beneficial.

The SwitchAll experiments (Table 3) provide indirect evidence for this concern: the combined MoE-MLP and MoE-attention model sometimes slightly underperforms the dense baseline (peS2o at 262M: 9.86 vs. 9.78), and the model has a small parameter deficit (259M vs. 262M) because of the granularity of MoE sizing. This suggests that parameter-matched constraints can produce configurations that are slightly resource-starved — in a production deployment where slightly exceeding the parameter budget is acceptable, the results might shift.

For practitioners operating in the more common FLOPS-matched or quality-maximizing regime, the paper provides no guidance on whether adding K/Q MoE (or increasing E or k beyond what the heuristic selects) would help. The MAC-matched experiments (Section 3.5) partially address this by showing that SwitchHead can productively use additional compute when n_heads and d_head are scaled up, but they do not test expanding the MoE dimension (increasing E, adding K/Q experts) rather than the head dimension.

What evidence exists in the paper. Table 6 provides the ablation that establishes the V+O-only design, but all configurations in this table are parameter-matched at 47M. The paper notes in Appendix A.3: "In our preliminary experiments, we found that, allowing the parameter budget to increase, more experts always help." This sentence is the only acknowledgment that the design choices are budget-dependent, but it is buried in the appendix and provides no quantitative data. No experiment directly compares parameter-matched and FLOPS-matched configurations with K/Q MoE to test whether the negative result is budget-artifact or a genuine architectural principle.

Mitigation status. The paper acknowledges the "parameter-matched setting" as a deliberate choice and argues for its relevance, but does not discuss the conditional nature of the K/Q MoE finding. The MAC-matched experiments (Section 3.5) provide an alternative evaluation perspective but don't test K/Q MoE. Readers are left to infer that the V+O-only design is the recommended configuration without clear guidance on whether this recommendation depends on the evaluation regime.

7. Implications and Future Directions

How This Work Changes the Landscape

SwitchHead does not propose a new learning algorithm or a new theoretical framework for attention. It proposes a structural reorganization of the multi-head attention mechanism — one that separates the multiplicity of linear projections from the multiplicity of attention matrices — and demonstrates that this separation delivers genuine resource savings without performance degradation. The magnitude of this contribution is best characterized as a strong architectural refinement with diagnostic implications: it identifies a specific, previously unarticulated design degree of freedom (how many attention matrices you compute versus how many expert projections you maintain) and shows that prior MoE-for-attention methods failed because they did not exploit this degree of freedom.

What shifts conceptually. Before SwitchHead, the assumption embedded in virtually all attention variants — standard MHA, multi-query attention, MoA — was that the number of distinct attention matrices must scale with the number of distinct output transformations. If you wanted the model to produce 16 different output representations per token, you needed 16 attention matrices. SwitchHead demonstrates that this coupling is unnecessary: a single attention matrix per head can serve many expert output projections, because the attention-weighted values (a single tensor after the A^h V^h computation) contain sufficient information for the output experts to differentiate themselves. This is not an obvious claim — one might reasonably expect that different output experts need to attend to different source tokens, requiring distinct attention patterns. The paper's empirical results (Table 2, the attention-map analysis in Section 4) show that this is not the case: the same attention pattern, combined with different value extractions and output transformations, provides sufficient expressivity to match or exceed dense multi-head attention.

This conceptual shift changes the optimization landscape for efficient attention research. Before SwitchHead, the primary axes for reducing attention cost were: (1) sparsifying or approximating the attention matrix itself (sparse attention, linear attention, low-rank approximations), (2) reducing precision (quantization, FlashAttention's IO-aware tiling), or (3) sharing keys and values across heads (multi-query attention). SwitchHead introduces a fourth axis: reduce the number of attention matrices by making the projections conditional rather than the attention patterns. This axis is complementary to the others — SwitchHead works with dense attention, FlashAttention, and various positional encodings — meaning it stacks with existing efficiency improvements rather than competing with them.

Reconciling contradictions. The paper implicitly resolves a tension in the MoE-for-attention literature. MoA (Zhang et al., 2022) showed that MoE-style attention was possible — you could conditionally select query/output heads — but the practical resource savings were modest because each selected head still required its own attention matrix. Subsequent work might have concluded that MoE attention is fundamentally limited in its ability to reduce the O(n_heads × T²) cost, since the attention matrix computation seems inherently tied to each distinct head. SwitchHead shows that this limitation was an architectural choice (selecting experts after computing separate attention matrices) rather than a fundamental constraint (that MoE attention must compute separate attention matrices). The negative results on MoA and the naive head-gating approach (Section 2.2) are reframed not as evidence against MoE attention, but as evidence that the MoE must be placed around the attention core (in the projections) rather than inside it (multiplying attention matrices).

Research directions that become more attractive. The paper makes three lines of investigation newly compelling:

1. The design space of projection-level conditional computation in attention is underexplored. SwitchHead uses sigmoid-gated top-k MoE for value and output projections, but this is one point in a larger space. Other conditional computation mechanisms (continuous gating without hard top-k, soft mixtures with learned temperature, low-rank adaptation-style expert projections) might provide different trade-offs. The paper provides a strong baseline and evaluation methodology against which such variants can be measured.

2. The parameter-matched evaluation philosophy should be adopted more broadly in MoE research. The paper demonstrates that parameter-matched experiments reveal different design trade-offs than FLOPS-matched ones — specifically, that key and query MoEs are unnecessary in the parameter-matched regime. This finding would be invisible in a FLOPS-matched study where adding K/Q experts simply consumes more parameters. The parameter-matched framework provides a cleaner signal about intrinsic parameter efficiency, which matters for resource-constrained deployment.

3. The interplay between attention and feedforward MoE requires joint optimization. The SwitchAll results (Table 3) show that independently-developed MoE methods for attention and MLP can be combined without negative interaction, but they also reveal that the combination is not optimized — the attention and MLP MoE configurations were chosen independently by separate heuristics. Jointly optimizing both MoE mechanisms under a unified resource constraint is an open problem.

Research directions that become less attractive. The paper's strong negative results on MoA — particularly that MoA at 8 heads is actually slower than the dense baseline (Table 5, 851ms vs. 670ms) — suggest that MoE designs that multiply attention matrices rather than sharing them are unlikely to yield practical speedups. The naive head-gating approach (Equations 4–6), which the paper shows is feasible but cannot deliver resource savings, also appears unlikely to be salvageable without the independent source/destination expert decomposition that SwitchHead introduces. Research effort is better directed at mechanisms that reduce the number of attention matrices (SwitchHead-like sharing) or that improve the efficiency of each matrix computation (FlashAttention-like kernel optimization), rather than at mechanisms that trade projection compute for attention-matrix compute (as MoA does).

---

Follow-Up Research This Work Enables

Scaling SwitchHead to billion-parameter models with controlled variations of n_heads, E, and k. The paper validates SwitchHead at 47M and 262M parameters. The most urgent follow-up is a scaling study at 1B, 7B, and potentially 13B parameters, systematically varying the expert configuration. Specific questions: Does the optimal n_heads grow with model size (as it does in dense Transformers, where GPT-3 uses 96 heads), or does the MoE mechanism allow n_heads to remain small even at large scales? Does the optimal number of active experts k increase with model capacity, or does k=2–4 remain sufficient? Does the finding that key and query MoEs are unnecessary persist at scale, or does the much larger parameter budget make them beneficial? A strong follow-up would train models at 1B and 7B parameters on a standard corpus (C4 or The Pile), sweep n_heads ∈ {2,4,8,16} and k ∈ {1,2,4,8} at each scale, and produce scaling-law-style plots showing how the optimal configuration and the resource savings evolve with model size. The paper's consistent results across a 6× parameter range suggest the findings may transfer, but extrapolation without data is speculation.

Measuring inference-time speedup and KV-cache memory reduction for long-sequence autoregressive generation. SwitchHead's resource formulas and wall-clock measurements are training-only. During autoregressive inference, the attention pattern is incremental (one new query attends to all cached keys and values), meaning the T² cost becomes O(T) per step. However, the KV cache — which stores keys and values for all previous tokens — scales as O(n_heads × T × d_head). SwitchHead reduces n_heads by 4–8×, which should directly reduce the KV-cache memory footprint by a similar factor. For long-context applications (e.g., 32k-token document QA), this memory reduction could be the difference between fitting in GPU memory and requiring offloading. A strong follow-up would: (1) implement autoregressive generation with SwitchHead and FlashAttention, (2) measure per-token latency and total memory usage for sequence lengths from 1k to 32k tokens, (3) compare against dense attention and multi-query attention at matched perplexity, (4) report both the KV-cache size in GB and the time-to-first-token and per-token latency. The paper's theoretical memory formulas (Appendix A.2) predict an 8× KV-cache reduction when going from 16 heads to 2, but empirical validation — especially with the interaction of FlashAttention's memory optimizations — is missing.

Combining SwitchHead with linear or sparse attention mechanisms. SwitchHead reduces the multiplier on the attention cost (n_heads) but does not change the O(T²) scaling of each attention matrix. Linear attention methods (Katharopoulos et al., 2020; Choromanski et al., 2021) replace the O(T²) softmax attention with O(T) kernelized approximations. Sparse attention methods (Child et al., 2019; Beltagy et al., 2020) compute only a subset of the T × T entries. SwitchHead is orthogonal to both — it reduces the number of attention matrices, whatever their computational form. A strong follow-up would: (1) replace the standard softmax attention in SwitchHead with a linear attention variant (e.g., Performer-style random features or Linear Transformer-style kernelization), (2) train matched models with and without the linear attention, (3) measure whether SwitchHead's expert-value/output mechanism is compatible with the approximations introduced by linear attention (i.e., whether the weighted-average of expert value projections interacts poorly with kernelized attention), (4) characterize the combined savings at sequence lengths of 4k–32k where both the n_heads reduction and the O(T) scaling matter. The paper's design — MoE only in the projections, standard attention core — is ideally suited for this combination since the attention computation remains a drop-in module.

Training a lightweight difficulty/context predictor to dynamically select k during inference. SwitchHead uses a fixed k (2–4) for all tokens and all sequences. But the paper's attention-map analysis on ListOps (Figure 5) shows that expert selections are context-dependent and interpretable — different token types (numbers vs. parentheses vs. operators) activate different experts. This suggests that some tokens may need more active experts than others. A strong follow-up would: (1) train a small predictor network (e.g., a 2-layer MLP taking the router's pre-topk logits as input) that predicts whether a token needs k=1, k=2, or k=3 experts to achieve low reconstruction error, (2) train this predictor using the gate scores from a pretrained SwitchHead model as supervision (tokens where only one expert has a high score get k=1, tokens where scores are diffuse get higher k), (3) measure the average k used across a test corpus and the corresponding MAC savings versus fixed-k SwitchHead, (4) verify that dynamic k does not degrade perplexity. If tokens vary substantially in how many experts they need — which the ListOps visualizations suggest but the language modeling experiments don't quantify — dynamic k could provide additional resource savings on top of the fixed-k baseline.

Probing what the value and output experts learn — a mechanistic interpretability study. The paper's attention analysis on ListOps (Section 4, Figures 3–5) provides suggestive evidence that output experts specialize by operation type and value experts by token category. For language models, the paper only shows that SwitchHead attention maps look qualitatively similar to dense attention maps (Figure 6) and that induction heads emerge in both. A strong follow-up would conduct a systematic mechanistic interpretability study: (1) for each head and expert in a trained SwitchHead language model, characterize which linguistic phenomena (syntactic relations, semantic roles, factual associations, copying behavior) each expert responds to, using techniques from the transformer circuits literature (activation patching, attention knockout, logit attribution), (2) measure whether different value experts within the same head extract genuinely different features from the same source tokens (testing the "diversity through MoE" hypothesis), (3) measure whether different output experts within the same head route information to different downstream computations (e.g., one output expert feeds noun phrase information to the next layer's subject-verb agreement head, another feeds semantic information to the entity-tracking head). This would transform SwitchHead from a "black-box acceleration" method to an interpretable architecture where resource allocation can be understood and potentially manually improved. The paper provides the foundation — the ListOps results show expert specialization is real — but the language modeling analysis is superficial.

Adversarial evaluation of sigmoid gating stability at scale: does expert collapse silently emerge? The paper claims that sigmoid gating "does not require regularization or extra tricks for stable training" (Section 2), but provides no quantitative evidence about expert utilization. The ListOps visualizations are encouraging but from a 6-layer model. At billion-parameter scales, training dynamics change, and silent expert collapse (where some experts' gate values are consistently below the top-k threshold across all inputs) could waste a substantial fraction of the parameter budget. A strong follow-up would: (1) train SwitchHead models at 1B and 7B parameters, (2) track per-expert selection frequency, average gate score, and gradient norm throughout training, (3) report the effective number of used experts (e.g., the inverse Herfindahl index of the selection distribution) as a function of training steps, (4) test whether initialization scale, learning rate, or router temperature affects collapse propensity, (5) attempt to induce collapse adversarially (e.g., by initializing all expert projections identically and checking whether they differentiate) to establish a stress test for the gating mechanism. If sigmoid gating remains stable at scale, this is strong evidence for its adoption in large-scale MoE systems. If it silently degrades, the community needs to know the failure conditions.

---

Practical Applications and Downstream Use Cases

Long-context batch processing on consumer GPUs. A common pain point for practitioners is running inference on long documents (legal contracts, scientific papers, code repositories) with open-source models on consumer hardware. The attention memory O(n_heads × T²) during prefill and the KV cache O(n_heads × T × d_head) during generation often cause out-of-memory errors at sequence lengths beyond 2k–4k tokens on GPUs with 8–12GB VRAM. SwitchHead addresses this directly: the 262M SwitchHead model uses 2.9M floats of attention memory versus 21M for the dense baseline (Table 2) — a 7.3× reduction — and the wall-clock training memory drops from 20.5GB to 12.5GB (Table 5, 39% reduction). For batched inference on long sequences, where attention matrices dominate memory, a similar factor of memory reduction would allow roughly 3× longer sequences or 3× larger batch sizes on the same hardware. The shared selection variant (Section 3.6) trades a small perplexity penalty (~0.3 points) for even faster routing, making it suitable for latency-sensitive deployments. The schematic in Figure 1 showing 2 heads with multiple experts makes the architectural simplicity clear: implementors need only replace the value and output projections with MoE versions while keeping the attention core unchanged.

Training budget-constrained academic language models. Academic groups training language models often operate under fixed compute budgets (e.g., a 10-GPU cluster for 2 weeks). Under a parameter-matched philosophy, SwitchHead allows these groups to either: (1) train a model with the same parameter count as planned but complete training in ~65% of the wall-clock time (the 1.5× speedup from Table 5), or (2) train a model with more parameters for the same wall-clock time (the MAC-matched results in Section 3.5 show significant accuracy gains from adding parameters to SwitchHead rather than to a dense baseline). The 262M SwitchHead on C4 achieves 16.23 perplexity at 2.4G MACs; the MAC-matched 376M SwitchHead achieves 15.43 perplexity — a substantial improvement for the same compute budget. The paper's systematic configuration procedure (n_heads=2, k=2 first, increase if needed) provides a concrete, low-risk recipe that does not require extensive hyperparameter tuning. The public code release (footnote 1) and the portability across Transformer XL and RoPE architectures (Appendix A.4) lower the adoption barrier further.

Efficient fine-tuning of pretrained dense models via structural conversion. A less obvious but potentially impactful use case: converting a pretrained dense Transformer to a SwitchHead architecture for efficient fine-tuning or deployment. Since SwitchHead only modifies the value and output projections (keeping keys and queries as standard dense projections), it is possible to: (1) take a pretrained dense checkpoint, (2) for each head, initialize E value expert projections as copies of the dense value projection plus small random noise (and similarly for output experts), (3) initialize the routers to small random values (producing roughly uniform gate scores initially), (4) fine-tune the converted model with a small amount of data to learn routing. This would amortize the cost of pretraining a SwitchHead model from scratch while reaping the inference efficiency benefits. The paper's finding that SwitchHead matches dense performance with fewer attention matrices suggests the converted model could retain most of the pretrained knowledge while becoming substantially cheaper to serve. This is speculative — no conversion experiment is performed in the paper — but the architectural compatibility (standard attention core, standard key/query projections) makes it technically straightforward to attempt.

---

When to Prefer This Method

The paper does not explicitly frame SwitchHead against a named set of alternative efficient-attention methods (e.g., multi-query attention, grouped-query attention, linear attention, sparse attention) with a head-to-head comparison or decision rubric. The only direct alternative evaluated is MoA, which SwitchHead strictly dominates in resource efficiency. The paper notes that SwitchHead is compatible with FlashAttention and various positional encodings (Appendix A.1, Appendix A.4), but does not provide quantitative comparisons against FlashAttention-only attention, multi-query attention, or other alternatives at matched perplexity.

Because the paper does not articulate a specific trade-off matrix against these named alternatives, I do not provide a conditional "Prefer SwitchHead when X; prefer FlashAttention when Y" decision rule. Such a rule would be synthesized from general knowledge rather than grounded in the paper's own experimental comparisons, which is explicitly disallowed by the constraints. Practitioners should note that the paper establishes SwitchHead's advantage over dense multi-head attention (the standard baseline) and over MoA (the prior MoE-for-attention method), but does not empirically position it against the broader landscape of efficient attention mechanisms.