Transformer Variants & Efficiency

BERT reads both ways, GPT only looks left, and T5 just converts everything into a text-to-text problem. Family dinners are awkward.

Norm, Architecturally Confused AI Agent

1. The Three Architectural Families

If someone asks you "which Transformer should I use?", your first question should be "what are you trying to do?" The original Transformer is an encoder-decoder model, but since 2017 three distinct families have emerged. Each family uses a different subset of the Transformer's components and targets different types of tasks. Understanding when to use which architecture is a fundamental skill in applied NLP.

Figure 4.3.2: The three Transformer architectural families. The decoder-only pattern dominates modern LLMs.

1.1 Encoder-Only (BERT Family)

Encoder-only models use bidirectional attention: every token can attend to every other token, including those to its right. This makes them excellent for understanding tasks (classification, token-level labeling, sentence similarity) but unsuitable for generation, since there is no causal structure.

BERT (Devlin et al., 2018) is trained with masked language modeling (MLM): 15% of input tokens are randomly masked, and the model must predict them. This pre-training objective forces the model to build rich bidirectional representations. Key descendants include RoBERTa (larger data, longer training), DeBERTa (disentangled attention with relative position), and ELECTRA (replaced-token detection, which is more sample-efficient than MLM).

1.2 Decoder-Only (GPT Family)

Decoder-only models use causal (left-to-right) attention with the auto-regressive language modeling objective. They are the dominant architecture for modern LLMs because: (a) the training objective is simple, scalable, and naturally aligns with text generation; (b) they can be prompted to perform virtually any task (classification, translation, reasoning) through in-context learning; and (c) they are straightforward to scale.

GPT-2 (Radford et al., 2019) demonstrated that a sufficiently large language model develops emergent abilities. GPT-3 (Brown et al., 2020) showed that these abilities improve reliably with scale. Today, virtually all frontier LLMs (GPT-4, Claude, LLaMA, Gemini, Mistral) are decoder-only.

1.3 Encoder-Decoder (T5 Family)

Encoder-decoder models process an input sequence with bidirectional attention (encoder) and generate an output sequence with causal attention (decoder), using cross-attention to condition the decoder on the encoder's output. This is the original Transformer architecture and remains the best choice for tasks with a clear input/output structure where the input benefits from bidirectional processing: translation, summarization, and speech recognition (Whisper).

T5 (Raffel et al., 2020) reframed all NLP tasks as text-to-text problems, demonstrating that a single encoder-decoder model could handle classification, translation, summarization, and question answering with the same architecture, just different input/output text formats.

2. Positional Encoding Variants

Transformers have no built-in notion of order. Without positional information, the model sees a sequence as an unordered set: "the cat sat on the mat" and "the mat sat on the cat" would be indistinguishable. Every positional encoding scheme solves this problem differently, and the choice has real consequences for how well a model generalizes to sequence lengths it has not seen during training. Four main approaches have emerged since 2017.

2.0 Sinusoidal Encoding (Vaswani et al., 2017)

The original "Attention Is All You Need" paper used a fixed, non-learned encoding: each position is represented by a vector of sine and cosine values at geometrically spaced frequencies. For position $pos$ and dimension index $i$:

The low-frequency dimensions change slowly across positions (giving a coarse sense of where in the sequence we are) while high-frequency dimensions oscillate rapidly (giving fine-grained local position information). Together they form a unique "fingerprint" for every position.

Key advantage: because the encoding is a fixed mathematical function, it can be evaluated at any position, including positions beyond the training sequence length. A model trained on sequences up to length 512 can in principle receive a sinusoidal encoding for position 1000. Whether it attends to that position correctly is a separate empirical question, but the encoding itself does not break.

Key limitation: sinusoidal encoding represents absolute positions. The model must implicitly learn that position 5 and position 15 are 10 steps apart; that relative relationship is not directly injected. This becomes a problem for tasks that depend heavily on relative structure rather than absolute location.

2.0b Learned Positional Embeddings (BERT, GPT-2)

A simpler alternative: treat positions exactly like tokens. Create an embedding table nn.Embedding(max_length, d_model) and look up a learned vector for each position index. This is precisely what our Section 4.2 implementation does with self.pos_emb.

Key advantage: the model can learn whatever positional representation is most useful for the task. If the corpus has strong structural patterns (e.g., the first token is often a subject, the last token is often a period), the learned embeddings can capture those patterns directly. In practice, learned embeddings perform comparably to or slightly better than sinusoidal on standard benchmarks.

Key limitation: the embedding table has exactly max_length rows. If the model encounters a sequence longer than max_length at inference time, it has no embedding for those positions. BERT's maximum length of 512 tokens is a direct consequence of this fixed table. Workarounds (interpolating learned embeddings, fine-tuning on longer sequences) exist but add complexity.

2.1 Rotary Position Embedding (RoPE)

RoPE (Su et al., 2021) has become the dominant positional encoding in modern LLMs (used in LLaMA, Mistral, Qwen, Gemma). Instead of adding positional information to the input embeddings, RoPE applies a rotation to the query and key vectors in each attention head. The rotation angle depends on the position and the dimension index.

The key insight: after applying RoPE to queries and keys, their dot product depends only on their relative position, not their absolute positions. This is achieved by rotating pairs of dimensions by $pos \times \theta _{i}$: Code Fragment 4.3.1 below puts this into practice.

# Rotary Position Embedding (RoPE): rotate pairs of dimensions by
# position-dependent angles so relative distance is encoded in dot products.

def apply_rope(x, freqs_cos, freqs_sin):
 """Apply Rotary Position Embedding to queries or keys.
 x: (B, n_heads, T, d_k)
 freqs_cos, freqs_sin: (T, d_k//2) precomputed cos/sin of rotation angles
 """
 # Split into pairs and rotate
 x_r = x.float().reshape(*x.shape[:-1], -1, 2) # (..., d_k//2, 2)
 x0, x1 = x_r[..., 0], x_r[..., 1]

 cos = freqs_cos.unsqueeze(0).unsqueeze(0) # broadcast over B, n_heads
 sin = freqs_sin.unsqueeze(0).unsqueeze(0)

 # 2D rotation: [cos -sin; sin cos] @ [x0; x1]
 out0 = x0 * cos - x1 * sin
 out1 = x0 * sin + x1 * cos

 out = torch.stack([out0, out1], dim=-1).flatten(-2)
 return out.type_as(x)

Figure 4.3.3: RoPE rotates each pair of embedding dimensions by an angle proportional to the token's position. Because the dot product of two rotated vectors depends only on the angle between them (the relative position), RoPE naturally captures relative position without explicit position IDs.

RoPE advantages: (1) it naturally encodes relative positions, (2) it requires no additional parameters, (3) it can be extended to longer sequences through frequency scaling (NTK-aware scaling, YaRN), and (4) it has strong empirical performance.

RoPE is powerful but still fundamentally couples the model to a maximum trained sequence length. An alternative approach avoids learned or computed embeddings entirely, instead injecting position information directly into the attention computation itself.

2.2 ALiBi (Attention with Linear Biases)

ALiBi (Press et al., 2022) takes a minimalist approach: it adds a linear bias to the attention scores that penalizes distant positions. No positional encoding is added to the embeddings at all. For head $h$, a bias of $-m_{h} \cdot |i - j|$ is added to the attention score between positions $i$ and $j$, where $m_{h}$ is a head-specific slope (a fixed, non-learned geometric sequence: 1/2, 1/4, 1/8, …). The attention formula becomes:

The distance penalty grows linearly with separation. Each head has a different slope, so some heads are more "local" (high slope, sharp penalty) and others are more "global" (low slope, gentler penalty). The title of the original paper, "Train Short, Test Long," captures the key benefit: a model trained on sequences of length 1024 can extrapolate well to length 4096 because the bias function is just a linear ramp and never encounters an out-of-range index. ALiBi is used in BLOOM (176B) and several Falcon variants.

Key advantage: outstanding length extrapolation; extremely simple to implement (two lines of code); no additional parameters.

Key limitation: unlike RoPE, ALiBi biases are fixed and not learned, so the model cannot adapt its distance-sensitivity to the data. Empirically, RoPE tends to outperform ALiBi when the training and evaluation lengths are the same; ALiBi wins when lengths differ substantially.

2.3 Comparing the Four Approaches

The four positional encoding strategies occupy distinct points in a design space defined by whether position is absolute or relative, whether encodings are learned or fixed, and how well the approach generalizes to sequence lengths beyond training. The SVG diagram below shows each approach schematically, and the table beneath gives a quick reference.

Figure 4.3.4: The four major positional encoding strategies. Sinusoidal and learned embeddings inject absolute position before the first attention layer. RoPE rotates Q and K vectors in each attention layer so their dot product depends on relative distance. ALiBi bypasses embeddings entirely and subtracts a distance penalty from attention scores. RoPE has become the dominant choice in frontier LLMs.

For a survey of which frontier models use which positional encoding in their production configurations, see Section 7.1. The KV cache interaction with RoPE (specifically, how cached keys and values retain their rotated coordinates across steps) is covered in Section 9.2.

3. Efficient Attention Mechanisms

Standard attention has $O(T^{2})$ time and memory complexity, where $T$ is the sequence length. For $T = 128\text{K}$ (a common context window in 2024+ models), the naive attention matrix would be $128\text{K} \times 128\text{K} = 16$ billion entries per head. This section surveys the main approaches to making attention tractable at long sequences.

3.1 Sparse Attention

Instead of attending to all positions, sparse attention restricts each token to attend to a carefully chosen subset of positions. The challenge is choosing which positions to attend to while preserving the model's ability to capture long-range dependencies.

• Local/Sliding Window Attention: Each token attends only to its neighbors within a fixed window of size $w$. Complexity becomes $O(T \times w)$. Used in Longformer and Mistral (window size 4096).
• Strided/Dilated Attention: In addition to local attention, some heads attend to every $k$-th position (dilated pattern), covering longer ranges with fewer computations. Used in BigBird and Sparse Transformer.
• Global Tokens: A small set of tokens (typically [CLS] or special sentinel tokens) attend to and are attended by all other tokens, serving as information hubs. This restores the ability to propagate information across the full sequence.

3.2 Linear Attention

Linear attention replaces the softmax kernel with a decomposable kernel function, allowing the attention computation to be rewritten in O(T) time:

The trick is to compute $\phi (K)^{T}V$ first (a $d \times d$ matrix, independent of T), then multiply by $\phi (Q)$. The feature map $\phi$ can be the identity (giving a simple outer-product formulation), an exponential, or a random feature approximation. Linear attention has seen renewed interest through architectures like RWKV and RetNet (discussed below).

3.3 FlashAttention

FlashAttention (Dao et al., 2022) is not an approximation; it computes exact standard attention but with dramatically better hardware utilization. The key insight is that the standard attention implementation is memory-bound: it writes the full T × T attention matrix to GPU global memory (HBM), reads it back for the softmax, writes it again, and reads it for the value multiplication. FlashAttention fuses these operations into a single kernel that keeps the attention matrix in fast on-chip SRAM, never materializing the full T × T matrix in HBM.

The result: 2 to 4x wall-clock speedup and dramatically reduced memory usage (from $O(T^{2})$ to $O(T)$ HBM). FlashAttention-2 further optimizes the kernel with better work partitioning across GPU thread blocks. FlashAttention-3 (2024) leverages Hopper GPU features (warp specialization, FP8 tensor cores) for additional gains. We cover the algorithm in detail in Section 4.4.

3.4 Multi-Query Attention (MQA) and Grouped-Query Attention (GQA)

During auto-regressive inference, the keys and values from all previous positions are cached (the "KV cache"), a technique we examine in detail in Section 9.2: Memory Optimization. With standard multi-head attention and many heads, this cache becomes enormous. Code Fragment 4.3.2 below puts this into practice.

• Multi-Query Attention (MQA) (Shazeer, 2019): All heads share a single set of keys and values. Only the queries differ per head. This reduces the KV cache by a factor of $h$ (the number of heads).
• Grouped-Query Attention (GQA) (Ainslie et al., 2023): A compromise where heads are divided into groups, and each group shares one set of keys and values. With $h=32$ heads and $g=8$ groups, the KV cache is reduced by 4x. Used in LLaMA 2, Mistral, and Gemma. For a survey of which production models use GQA, see Section 7.2.

# Grouped-Query Attention (GQA): multiple query heads share fewer KV heads.
# Reduces KV cache size while retaining most of multi-head attention quality.

class GroupedQueryAttention(nn.Module):
 """GQA: groups of query heads share KV heads."""

 def __init__(self, d_model, n_heads, n_kv_heads):
 super().__init__()
 self.n_heads = n_heads
 self.n_kv_heads = n_kv_heads
 self.n_rep = n_heads // n_kv_heads # how many Q heads per KV head
 self.d_k = d_model // n_heads

 self.W_q = nn.Linear(d_model, n_heads * self.d_k, bias=False)
 self.W_k = nn.Linear(d_model, n_kv_heads * self.d_k, bias=False)
 self.W_v = nn.Linear(d_model, n_kv_heads * self.d_k, bias=False)
 self.W_o = nn.Linear(d_model, d_model, bias=False)

 # Forward pass: define computation graph
 def forward(self, x, mask=None):
 B, T, _ = x.shape

 q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
 k = self.W_k(x).view(B, T, self.n_kv_heads, self.d_k).transpose(1, 2)
 v = self.W_v(x).view(B, T, self.n_kv_heads, self.d_k).transpose(1, 2)

 # Repeat KV heads to match the number of query heads
 # (B, n_kv_heads, T, d_k) -> (B, n_heads, T, d_k)
 k = k.repeat_interleave(self.n_rep, dim=1)
 v = v.repeat_interleave(self.n_rep, dim=1)

 scores = (q @ k.transpose(-2, -1)) / (self.d_k ** 0.5)
 if mask is not None:
 # Mask future positions with -inf so softmax ignores them
 scores = scores.masked_fill(mask == 0, float('-inf'))
 # Convert scores to attention weights (probabilities summing to 1)
 attn = torch.softmax(scores, dim=-1)
 out = (attn @ v).transpose(1, 2).contiguous().view(B, T, -1)
 return self.W_o(out)

All the attention variants we have seen so far modify how attention scores are computed or how many heads share parameters. But they all share a common limitation: softmax attention always assigns some weight to every position, even irrelevant ones. The next approach tackles this "attention noise" problem directly.

3.5 Differential Attention (DIFF Transformer)

A more recent innovation in attention design is Differential Attention, introduced by Ye et al. (2024) in the DIFF Transformer paper. The core idea is that standard softmax attention allocates non-trivial weight to irrelevant context tokens (a phenomenon sometimes called "attention noise"). Even with a clear focus token, the long tail of softmax probabilities still assigns small but non-zero attention to unrelated positions, diluting the signal and contributing to hallucination.

Differential Attention addresses this by computing attention as the difference of two separate softmax attention maps. Each attention head is split into two sub-heads that compute independent attention distributions over the same keys, and the final attention weights are the element-wise subtraction of one from the other. Positions that receive similar attention from both sub-heads cancel out (analogous to noise cancellation), while positions that one sub-head strongly attends to but the other does not are amplified. The result is a sparser, more focused attention pattern that better distinguishes signal from noise. Code Fragment 4.3.3 below puts this into practice.

# Simplified Differential Attention (conceptual)
def diff_attention(Q1, K1, Q2, K2, V):
 """Two sub-attention maps; subtract to cancel noise."""
 attn_1 = F.softmax(Q1 @ K1.T / math.sqrt(d_k), dim=-1)
 attn_2 = F.softmax(Q2 @ K2.T / math.sqrt(d_k), dim=-1)
 # Noise cancellation: positions attended equally by both cancel out
 diff_weights = attn_1 - attn_2
 return diff_weights @ V

Empirically, DIFF Transformers show improved performance on long-context tasks, key-value retrieval, and hallucination benchmarks, with particular gains on tasks requiring precise information extraction from noisy contexts. The approach adds minimal overhead (two sets of query/key projections per head instead of one) and is compatible with FlashAttention and GQA. Microsoft has explored integrating Differential Attention into production models, and the technique represents a promising direction for making attention more robust without fundamentally changing the Transformer architecture.

3.6 Multi-Head Attention: Why Multiple Heads?

We established the mechanics of multi-head attention in Section 4.1 and implemented it in Section 4.2. Now let us dig into the why: what does having multiple heads actually buy us?

The short answer is that different heads learn to attend to fundamentally different kinds of relationships, simultaneously. A single attention head produces one attention distribution over the sequence for each query position: it is forced to aggregate all of its relational reasoning into a single scalar weight per (query, key) pair. Multi-head attention runs n_heads independent attention computations in parallel, each in a lower-dimensional subspace, and then concatenates the results. This lets the model represent multiple relational modes at once.

Empirical evidence for this comes from interpretability work. Clark et al. (2019) analyzed BERT's attention patterns and found striking specialization across heads:

• Syntactic heads: some heads consistently attend from a verb to its direct object, or from a noun to its modifying adjectives, regardless of their absolute positions in the sentence.
• Coreference heads: certain heads track pronoun-antecedent relationships, attending from "she" back to the named entity it refers to, potentially dozens of tokens away.
• Local context heads: some heads attend almost exclusively to immediately adjacent tokens (a sliding-window pattern), providing a local smoothing effect.
• Global anchor heads: other heads attend heavily to special tokens like [CLS] or sentence-boundary markers, effectively routing global context signals.
• Positional heads: a subset of heads attend to fixed relative offsets (always one token back, always two tokens forward), encoding structured local order information.

Figure 4.3.5: Four types of attention head behavior observed in trained models (Clark et al., 2019). These specializations emerge from gradient descent alone; they are not explicitly programmed. Multiple heads allow the model to simultaneously represent syntactic structure, local context, global routing, and coreference, all within a single forward pass.

This specialization is an emergent property: no training objective explicitly requires a head to learn coreference. It arises because tracking coreference is useful for next-token prediction, and the architecture provides enough capacity (enough separate heads) for the model to allocate a head to this purpose without sacrificing others. Single-head attention cannot do this: with only one attention distribution, the model is forced to trade off between all these relational signals at every position.

For a discussion of how Grouped-Query Attention (GQA) reduces the number of key-value heads while retaining most of this multi-head representational power, see Section 9.2.

3.7 Pre-Norm vs. Post-Norm Layer Normalization

Layer normalization is present in every Transformer block, but where you apply it changes both training stability and model quality in ways that matter at scale. Two placement strategies have competed since 2018, and a simplified normalization variant has emerged as a third option.

3.7.1 Post-Norm (Original Vaswani 2017)

The original "Attention Is All You Need" paper placed LayerNorm after the residual add:

The intuition is clean: normalize the output of each sublayer before passing it forward. However, this placement creates a training instability problem at scale. During backpropagation, gradients flowing back to early layers must pass through the LayerNorm operations at every block. The norm statistics depend on the current activations and can produce high-variance gradients in early training before the model has settled. This forced practitioners to use very careful learning-rate warmup schedules (often thousands of steps at very low learning rate) to stabilize Post-LN training.

3.7.2 Pre-Norm (GPT-2 onward, all modern LLMs)

Pre-Norm applies LayerNorm before the sublayer:

The key difference: the residual path x + is not passed through LayerNorm. Gradients flowing back through the residual connection reach all earlier layers directly, without being modified by any normalization step. This produces much more stable gradient magnitudes throughout training and essentially eliminates the need for warmup warmup, or at least reduces the sensitivity to its length. Virtually every LLM trained since GPT-2 (Radford et al., 2019) uses Pre-LN. In our Section 4.2 implementation, the two lines x = x + self.attn(self.ln1(x)) and x = x + self.ffn(self.ln2(x)) are both Pre-LN.

Figure 4.3.6: Post-Norm (left) places LayerNorm after the residual add; gradients must pass through the norm on the backward pass, causing instability at scale. Pre-Norm (right) places LayerNorm before the sublayer; the raw residual signal bypasses the norm entirely, giving stable gradients throughout training.

3.7.3 RMSNorm (Zhang & Sennrich, 2019)

Standard LayerNorm normalizes by subtracting the mean and dividing by the standard deviation, then scales and shifts:

RMSNorm simplifies this by removing the mean-centering step entirely and using only the root-mean-square for scaling:

The motivation: the mean-centering in standard LayerNorm re-centers activations at zero, but empirically the centering provides little or no quality benefit while adding computation. RMSNorm drops both the mean subtraction and the shift parameter $\beta$, reducing the operation to a single division and element-wise scale. Zhang and Sennrich (2019) showed that RMSNorm matches LayerNorm quality with roughly 7-12% less computation on the normalization step.

Modern LLMs have adopted RMSNorm widely. LLaMA 1 through 3, Mistral, Gemma, and Qwen all use RMSNorm in the Pre-LN position. The combination of Pre-LN placement and RMSNorm is now the de facto standard for frontier decoder-only models.

4. Beyond Attention: State Space Models

State Space Models (SSMs) represent a fundamentally different approach to sequence modeling. Instead of computing pairwise attention between all tokens, SSMs process sequences through a linear recurrence with continuous-time dynamics. The key innovation is that this recurrence can be computed in O(T log T) time during training (using a parallel scan or convolution) while maintaining O(1) per-step cost during inference.

4.1 The S4 Foundation

S4 (Gu et al., 2022) models the mapping from input $u(t)$ to output $y(t)$ through a linear state-space equation:

After discretization with step size $\Delta$, this becomes a linear recurrence: $h_{t} = \bar{A}h_{t-1} + \bar{B}u_{t}$. The discrete recurrence can also be unrolled as a convolution (over the full sequence), enabling parallel training. S4's breakthrough was showing how to parameterize the matrix $A$ (using the HiPPO initialization) to capture long-range dependencies that RNNs struggle with.

4.2 Mamba: Selective State Spaces

Mamba (Gu and Dao, 2023) introduced selective state spaces, where the parameters B, C, and $\Delta$ are input-dependent (functions of the current token). This breaks the linear time-invariance that allows convolution-based training, but the authors developed a hardware-aware parallel scan algorithm that achieves efficient training on GPUs.

Mamba's advantages over Transformers:

• Linear time complexity in sequence length (vs. quadratic for attention)
• Constant memory during inference: no KV cache that grows with context
• Fast inference: each new token requires only a fixed-size state update
• Strong performance on tasks requiring long-range reasoning (up to millions of tokens)

Mamba's disadvantages:

• Cannot perform arbitrary lookback into the full context the way attention can; information must be compressed into the finite-dimensional state
• Harder to parallelize across the sequence dimension during training compared to attention
• Fewer mature optimization tools (no equivalent of FlashAttention yet for all SSM variants)

The success of hybrid architectures raises a natural question: are SSMs and attention really so different? Recent theoretical work suggests they may be closer than anyone expected.

4.3 Mamba-2 and the Connection to Attention

Mamba-2 (Dao and Gu, 2024) revealed a deep connection between SSMs and attention. The structured state space duality (SSD) framework shows that selective SSMs are equivalent to a form of structured masked attention, where the mask has a specific semiseparable structure. This unification opens the door to transferring optimization techniques between the two paradigms.

5. RWKV: Linear Attention as an RNN

RWKV (Peng et al., 2023) combines the training parallelism of Transformers with the inference efficiency of RNNs. It uses a variant of linear attention with time-dependent decay, formulated so that it can be computed either as a parallel attention-like operation (for training) or as a sequential RNN (for inference).

The core idea: replace softmax attention with a weighted sum using exponentially decaying weights:

Here $w$ is a learned decay factor (how quickly the model "forgets" past tokens) and $u$ is a bonus for the current token. This can be computed in O(T) during both training and inference. RWKV-6 and later versions add more sophisticated mechanisms (data-dependent linear interpolation, multi-scale decay) while maintaining linear complexity.

6. Mixture-of-Experts (MoE)

Mixture-of-Experts is an orthogonal scaling strategy: instead of making every layer wider, you create multiple "expert" sub-networks and route each token to only a few of them. This dramatically increases the total parameter count (and thus model capacity) while keeping the computation per token roughly constant.

6.1 Architecture

In a typical MoE Transformer, the FFN in each block is replaced by a set of $E$ expert FFNs plus a router (gating network). For each token, the router selects the top-$k$ experts (typically k=1 or k=2), and the token is processed only by those selected experts. The output is a weighted combination of the expert outputs. Code Fragment 4.3.4 below puts this into practice.

# Mixture-of-Experts FFN: a learned router selects top-k experts per token.
# Only the selected experts run, keeping compute constant as total params grow.

class MoELayer(nn.Module):
 """Mixture-of-Experts feed-forward layer."""

 def __init__(self, d_model, d_ff, n_experts, top_k=2):
 super().__init__()
 self.n_experts = n_experts
 self.top_k = top_k

 # Router: maps d_model -> n_experts (logits for each expert)
 self.router = nn.Linear(d_model, n_experts, bias=False)

 # Expert FFNs
 self.experts = nn.ModuleList([
 nn.Sequential(
 nn.Linear(d_model, d_ff, bias=False),
 nn.SiLU(),
 nn.Linear(d_ff, d_model, bias=False),
 )
 for _ in range(n_experts)
 ])

 def forward(self, x):
 B, T, C = x.shape
 x_flat = x.view(-1, C) # (B*T, C)

 # Compute routing weights
 router_logits = self.router(x_flat) # (B*T, n_experts)
 weights, indices = torch.topk(router_logits, self.top_k, dim=-1)
 weights = torch.softmax(weights, dim=-1) # normalize top-k

 # Dispatch tokens to experts and combine results
 output = torch.zeros_like(x_flat)
 for k in range(self.top_k):
 expert_idx = indices[:, k] # (B*T,)
 weight = weights[:, k].unsqueeze(-1) # (B*T, 1)
 for e in range(self.n_experts):
 mask = (expert_idx == e)
 if mask.any():
 expert_input = x_flat[mask]
 expert_output = self.experts[e](expert_input)
 output[mask] += weight[mask] * expert_output

 return output.view(B, T, C)

6.2 Notable MoE Models

The key insight: a 46.7B parameter MoE model (Mixtral 8x7B) can match or exceed a dense 70B model in quality while using only 12.9B parameters of computation per token. This means it runs at roughly the speed of a 13B dense model despite having the knowledge capacity of a much larger one.

The 2024 and 2025 releases confirm MoE as the dominant scaling paradigm. DeepSeek-V3 (December 2024) pushed the frontier with 671B total parameters and 256 fine-grained experts, using auxiliary-loss-free load balancing and multi-token prediction heads. Despite its enormous capacity, only 37B parameters activate per token, keeping inference costs manageable. Meta's Llama 4 family (2025) completed Meta's own transition from dense to MoE: Llama 4 Scout (109B total, 16 experts) and Llama 4 Maverick (400B total, 128 experts) both use top-1 expert routing, meaning each token activates a single expert plus shared parameters. Both Llama 4 variants are natively multimodal, processing text and images in a unified architecture. The trend is clear: leading labs have converged on MoE as the way to scale model knowledge without proportionally scaling inference cost.

7. Gated Attention and Gated Linear Units

The concept of gating (element-wise multiplication of two parallel pathways) has become ubiquitous in modern Transformers, appearing in both the FFN and the attention mechanism.

7.1 Gated FFN Variants

The standard ReLU FFN computes ReLU(xW1)W2. Gated variants split the first projection into two branches and multiply them:

• GLU (Dauphin et al., 2017): (xW1 ⊙ σ(xWgate))W2
• GEGLU: Replace σ with GELU activation
• SwiGLU (Shazeer, 2020): Replace σ with SiLU (Swish). This is the standard in LLaMA, PaLM, Mistral, Gemma

Why does gating help? The gate branch learns to selectively amplify or suppress features produced by the value branch. This provides a richer form of nonlinearity than a single activation function and consistently improves performance at a given compute budget.

7.2 Gated Attention Units

GAU (Hua et al., 2022) applies the gating principle to attention itself. Instead of the standard residual attention pattern, GAU computes:

where U is a gating signal derived from the input, and V is the value signal. This allows single-head attention to be competitive with multi-head attention, since the gate provides the diversity that multiple heads normally provide. GAU reduces both the number of attention heads needed and the overall parameter count.

8. Multi-Head Latent Attention (MLA)

Multi-Head Latent Attention (DeepSeek-V2, 2024) addresses the KV cache bottleneck through a different lens than GQA. Instead of sharing KV heads across groups, MLA compresses the keys and values into a low-dimensional latent space before caching.

8.1 How MLA Works

In standard attention, we cache the full K and V tensors of shape $(T, n_{heads}, d_k)$. MLA instead caches a compressed representation $c_{KV}$ of shape $(T, d_c)$ where $d_c << n_{heads} \times d_k$. The full K and V are reconstructed from the compressed representation when needed:

The compression ratio can be 4x to 16x, dramatically reducing the KV cache memory. The decompression matrices $W_{UK}$ and $W_{UV}$ are small and fast to apply. DeepSeek-V2 reports comparable quality to standard multi-head attention while reducing KV cache memory by 93%.

9. Putting It All Together: Modern LLM Recipes

No production LLM uses a single technique in isolation. Here is how several prominent models combine the building blocks discussed in this section: