Chapter 3 Transformer Architecture Interview Essentials

3.1 The Baseline Decoder-Only Transformer

3.1.1 Residual Stream, Attention, and MLP Blocks

The opening question in virtually every LLM architecture deep-dive is some form of “walk me through the forward pass of a decoder-only transformer.” At mid-level, interviewers want to hear you trace the data flow; at senior level, they want you to estimate the parameter budget per block at a realistic model size, explain why the two sub-blocks are sequenced the way they are, and articulate what goes wrong if you change the structure. Getting this right is the baseline for every architecture discussion that follows.

A decoder-only transformer processes $n$ input tokens, each represented as a $d$-dimensional vector. These vectors form the residual stream - a horizontal channel of activations that flows from the bottom of the network to the top. The defining property of the residual stream is that neither attention nor the MLP writes directly over the previous layer’s output: both read from the stream and add their outputs back into it. The stream carries the cumulative representation of each token from embedding to logit head.

Formally, a single layer $\ell$ executes two sequential updates:

The attention sub-block computes multi-head scaled dot-product attention. With $h$ heads, each of dimension $d_{h} = d / h$, the combined projection matrices $W_{Q} , W_{K} , W_{V} \in \mathbb{R}^{d \times d}$ and output projection $W_{O} \in \mathbb{R}^{d \times d}$

define:

where $Q_{i} = X W_{Q_{i}}$, $K_{i} = X W_{K_{i}}$, $V_{i} = X W_{V_{i}}$ are per-head slices of the combined projections.

The MLP sub-block, in its modern SwiGLU form, uses three weight matrices instead of two:

where $W_{1} , W_{3} \in \mathbb{R}^{d \times d_{f f}}$ are up-projections, $\odot$ denotes element-wise multiplication (the gate), and $W_{2} \in \mathbb{R}^{d_{f f} \times d}$ projects back. The gate - $x W_{3}$ - modulates which hidden units are active, giving the MLP an adaptive selection mechanism absent in ReLU and GeLU variants.

The residual stream is the skeleton of the network: every sub-block reads from it and adds an increment, so a model with zero-initialized sub-block outputs is a pure identity - this is what makes stacks of 32 or 96 layers trainable from random initialization.

Residual connections solve gradient vanishing. Without an additive skip, the gradient of the loss with respect to the activations at layer $\ell$ requires a chain of $L - \ell$ Jacobians, which generically collapses to near-zero at depth. With the skip, the gradient sees a direct path of ones alongside each Jacobian chain, guaranteeing a non-vanishing lower bound. This is why transformers can sustain depths of 80 or 96 layers without the careful cell-state engineering that LSTMs relied on.

The functional division between attention and MLP reflects a genuine computational separation. Attention aggregates across positions: it lets each token construct a weighted mixture of all other tokens’ representations, gathering context from wherever it is most informative. The MLP processes per-position: once a token has gathered its context, the MLP mixes information across the feature dimension of that single vector. Stacking the two repeatedly allows the network to alternate between “gather context” and “refine representation.” This intuition is supported by mechanistic analyses showing that factual associations tend to live in MLP weights, while syntactic and positional relationships live in attention patterns.

The overwhelming consensus is to run attention before MLP in series. A minority of models (GPT-J, PaLM) explored parallel layers where both sub-blocks read the same normalized input simultaneously and their outputs are summed. Parallel layout can yield a roughly 15% wall-clock speedup at large scale because both sets of projections can be fused into a single large matrix multiply against a shared normalized input. However, the parallel layout means the MLP never sees the attention output within the same layer - a mild expressiveness loss. Most frontier models released after 2023 have reverted to the serial default, which I read as evidence that the expressiveness loss is real enough to matter at the model sizes where these labs operate.

A Worked Example

Take a 7B-parameter model: $d = 4096$, $h = 32$ heads ($d_{h} = 128$), $L = 32$ layers, SwiGLU MLP with $d_{f f} = 11008$ (the canonical $\lfloor 8 d / 3 \rfloor$ sizing for GLU variants, rounded to the nearest 128), vocabulary $V = 32000$.

Parameter budget per layer:

• Attention ($W_{Q} , W_{K} , W_{V} , W_{O}$, no biases): $4 \times d^{2} = 4 \times 4096^{2} \approx 67$M

• MLP ($W_{1} , W_{3} , W_{2}$, no biases): $3 \times d \times d_{f f} = 3 \times 4096 \times 11008 \approx 135$M

• Layer total: $\approx 202$M

Across 32 layers: $32 \times 202$M $\approx 6 . 46$B. Adding the embedding table ($32000 \times 4096 \approx 131$M) lands at roughly 6.6B - consistent with LLaMA-1 7B’s actual count of 6.7B. The MLP holds approximately two-thirds of each layer’s parameters, a ratio worth memorizing: when an interviewer asks you to split the parameter budget between attention and FFN, the answer for SwiGLU models is roughly 1:2.

Memory at bf16: $6 . 7$B $\times 2$ bytes $= 13 . 4$ GB for weights alone. On an 80 GB H100, that leaves $\sim$67 GB for activations, KV cache, and optimizer state - comfortable for inference, tight for training under AdamW, which needs at least $4 \times$ the weight size for gradients and first/second moment buffers.

What You Actually Decide in Practice

• Serial vs. parallel layers. Default to serial. Parallel layers require custom fused kernels to realize their speedup; without them, the benefit vanishes. Only consider parallel if you have ablated quality impact and verified kernel throughput on your specific hardware. Most teams find it is not worth the engineering overhead.

• MLP sizing with SwiGLU. Set $d_{f f} = \lfloor 8 d / 3 \rfloor$, rounded up to the nearest multiple of 64 for hardware alignment. The vanilla ReLU convention of $d_{f f} = 4 d$ over-counts your parameter budget by roughly 50% relative to a non-gated MLP, because SwiGLU introduces a third matrix. Getting this wrong means your “7B” model is actually closer to 10B.

• Head dimension. Keep $d_{h} = d / h = 128$, the near-universal default. Scaling law evidence shows that varying head count at fixed $d$ and $d_{h}$ has negligible effect on loss; what matters is that $d_{h}$ stays large enough to avoid rank bottlenecks in each head’s key-value projection.

• Bias terms. Drop them everywhere in attention and MLP. Biases add negligible expressiveness but have been linked empirically to training instability in large models. Virtually every frontier model since LLaMA-1 omits them, and the rule of thumb is: if you cannot point to a specific ablation that shows bias terms help, do not include them.

• Residual stream accumulation dtype. Keep the residual stream accumulation in bf16 or fp32 depending on your mixed-precision setup. Do not quantize the skip-connection additions - precision erosion in the accumulation compounds across 32+ layers and is very hard to diagnose after the fact.

How This Shows Up in Interviews

• “Walk me through the forward pass of a 7B decoder-only model. How many parameters are in the attention block versus the MLP?” A strong answer names all four projection matrices in attention, all three in SwiGLU, derives the $\sim$67M and $\sim$135M figures, and explains why the MLP is roughly $2 \times$ larger. A weak answer stops at “it has multi-head attention and a feed-forward network” without numbers.

• “Why is the residual connection important? What breaks without it?” Strong answer: gradient propagation. Without the additive skip, the gradient at layer $\ell$ involves a product of $L - \ell$ Jacobians that collapses to near-zero at depth; with the skip, there is always a direct gradient path of approximately ones from the top to every earlier layer.

• Senior follow-up: “Some models run attention and MLP in parallel instead of in series. What do you trade off, and when would you choose parallel?” Senior answer: parallel layout fuses both sets of projections against the same normalized input in a single large matrix multiply, saving memory roundtrips and yielding roughly 15% wall-clock speedup when kernels are properly written (as demonstrated in PaLM). The cost is that the MLP sees the pre-attention representation rather than the post-attention one - the MLP and attention cannot interact within the same layer. In practice, serial is the safer default unless you have custom kernel support and have ablated quality impact at your target model size; most post-2023 frontier models have reverted to serial for exactly this reason.

Key Takeaways

• The residual stream is the identity highway: every sub-block reads from it and adds an increment, enabling clean gradient flow through networks of 32-128 layers.

• Each layer contains two sub-blocks in series: multi-head attention (mixes across positions) followed by an MLP (mixes across features at each position).

• For a 7B model with $d = 4096$, $d_{f f} \approx 11008$: each layer holds $\approx$67M attention parameters and $\approx$135M MLP parameters - the MLP is roughly $2 \times$ the attention block by parameter count.

• SwiGLU MLPs use three weight matrices and require $d_{f f} = \lfloor 8 d / 3 \rfloor$ to stay parameter-matched with a two-matrix ReLU MLP at $4 d$; getting the sizing wrong inflates the model by $\sim$50%.

• Use serial layers and drop all bias terms by default; parallel layers require verified kernel-level speedups and quality ablations before committing.

3.1.2 Pre-Norm vs. Post-Norm and Why It Matters

The question “where does the LayerNorm go in your transformer block?” reads like trivia until you have watched a training run crash after a week of GPU-hours because gradient norms spiked and refused to come back down. Normalization placement directly controls training stability at scale, and every frontier lab I am aware of has converged on the same answer. Interviewers at these labs use this question to filter candidates who have only read the original 2017 paper from those who have actually thought about what happens at depth.

The original transformer applied normalization after each sub-block’s output was added back to the residual stream - the post-norm convention:

Pre-norm moves the normalization before the sub-block computation, operating on the residual stream input prior to branching off to attention or MLP:

The normalization now sits inside the non-residual branch. The residual stream itself - the $x$ that accumulates across layers - is never touched by a normalization layer. This single repositioning is the difference between a network that requires fragile learning-rate warmup to avoid divergence and one that trains stably from step one.

Pre-norm is the universal choice in production models: LLaMA, PaLM, Chinchilla, Mistral, Qwen, DeepSeek, Gemma, and virtually every frontier model since GPT-2 uses it. The only notable exception is OPT-350M, which used post-norm - an artifact widely attributed to an engineering oversight that was never corrected before the run completed.

Place the normalization inside the non-residual branch: pre-norm keeps the gradient highway clean, enables stable training without warmup, and is the choice every serious practitioner makes.

The core argument is gradient behavior at depth. With post-norm, the loss gradient traveling back to layer $\ell$ must pass through LayerNorm at every layer between $\ell$ and the top. LayerNorm rescales activations by the inverse of their empirical standard deviation - a data-dependent quantity that is volatile in early training when activations have not yet settled. This volatility causes the effective step size at early layers to oscillate by orders of magnitude from batch to batch, which is why post-norm training famously required careful warmup: you were buying time for activations to stabilize before the gradient became reliable.

Pre-norm sidesteps this entirely. The residual addition $x \leftarrow x + f ( \text{Norm} ( x ) )$ means the gradient of the loss at layer $\ell$ always sees an additive term of $1$ from the skip connection, regardless of what the sub-block gradient looks like. The normalization sits only on the branch and does not obstruct the residual gradient path. At depth $L = 32$ or $L = 96$, this difference is the margin between a gradient norm that stays in $[ 0 . 1 , 2 ]$ throughout training and one that spikes to $10^{3}$ and triggers a loss divergence.

Empirical evidence is consistent and old enough to be well-established: training curves from pre-norm models show lower gradient norm variance, fewer loss spikes, and equivalent or slightly better final loss compared with post-norm models that used careful warmup. Several research teams demonstrated that pre-norm models can remove the warmup phase entirely and still converge, which simplifies resuming from checkpoints mid-run.

A newer variant, which I will call the double-norm, adds a second normalization after each sub-block (before the residual addition) as well as the pre-block norm. Grok-1 and Gemma 2 use this:

The logic is that bounding both the input to the sub-block and its output gives two scale control checkpoints instead of one, providing additional protection against activations exploding inside very large sub-blocks. Early reports suggest marginal stability gains at very large model sizes ($>$30B parameters). It adds memory access overhead (LayerNorm is bandwidth-bound, not compute-bound) but is otherwise a small cost. Double-norm has not yet displaced plain pre-norm as the community default.

A Worked Example

Consider a 32-layer, 1.3B-parameter model trained with AdamW, peak learning rate $3 \times 10^{- 4}$, no warmup schedule, on 100B tokens. Compare the gradient dynamics across three configurations:

Configuration                   Grad norm at step 500    Stability
-----------------------------------------------------------------
Post-norm, no warmup            ~450  (highly variable)  Diverges
Post-norm, 2k-step warmup       ~12   (settling)         Converges
Pre-norm,  no warmup            ~1.8  (stable from step 1) Converges

Post-norm without warmup diverges reliably: by layer 10, the LayerNorm denominators are fluctuating because the activations are still in their initial transient. The effective step size at early layers swings by $5$-$10 \times$ per step, causing gradient norms to spike past $10^{2}$-$10^{3}$ within the first thousand steps.

Post-norm with 2,000-step warmup eventually converges because the warmup period allows activations to settle before the learning rate reaches its peak. But this coupling between normalization placement and the warmup schedule is a fragile dependency: if you need to restart from a checkpoint (e.g., after a hardware failure), you must re-warm carefully or risk a second divergence.

Pre-norm without warmup produces gradient norms in the range $0 . 1$-$2 . 0$ from the very first step. The residual stream accumulates small increments from each sub-block, and the gradient back-propagates through the skip connections without interference from normalization. You can restart from any checkpoint with a cold optimizer and the training will resume cleanly.

The practical upshot at scale: at 7B parameters trained on 1T tokens, a single gradient divergence that requires a rollback and rerun from the last stable checkpoint wastes days of compute on a cluster of 256 H100s. Pre-norm eliminates an entire class of these failures.

What You Actually Decide in Practice

• Always use pre-norm for new models. There is no empirical justification for post-norm in a fresh training run. If you inherit a post-norm codebase, do not switch normalization placement mid-training - it changes the loss landscape and will almost certainly cause a divergence.

• Consider double-norm only above $\sim$30B parameters. The extra normalization costs additional memory accesses (LayerNorm is bandwidth-bound; the extra norm adds $\sim$0.3% wall time per step) and a small number of additional parameters. Below 30B, plain pre-norm is stable enough that double-norm provides no measurable benefit. Above 30B, if you are seeing persistent gradient norm spikes despite pre-norm, double-norm is a reasonable next step before reaching for gradient clipping or modified optimizers.

• Never put a normalization layer on the residual stream itself. Any norm that sits on the accumulation path (i.e., on $x$ before the skip addition rather than on the branch input) re-introduces the post-norm pathology. The identity gradient path is the entire point.

• Monitor gradient norms from step 1. With pre-norm and a sane learning rate, gradient norms should be below 5.0 throughout the first 1,000 steps. If they are not, suspect a normalization placement bug, a bad initialization scale, or a learning rate that is too high - not a fundamental architecture problem.

• Treat pre-norm and RMSNorm as a package deal. The decision of where to place the norm and which norm to use are distinct but nearly always made together in modern models. Pre-norm with RMSNorm is the combination that appears in LLaMA, PaLM, Chinchilla, Mistral, and most other frontier open models.

How This Shows Up in Interviews

• “Where does the LayerNorm go in a modern transformer block, and why?” A strong answer places the norm before the sub-block (pre-norm), explains that it keeps the residual stream as a clean gradient highway, and notes this removes the need for warmup. A weak answer just says “before the block” without explaining the gradient argument.

• “What does the original 2017 transformer do differently, and what problem does that cause?” Strong answer: post-norm places the normalization after the residual addition, which means the gradient back-propagating to early layers must pass through LayerNorm at every layer. Because LayerNorm rescales by the inverse standard deviation - a volatile quantity in early training - this corrupts the effective step size at early layers and requires careful warmup to avoid divergence.

• Senior follow-up: “Gemma 2 applies a LayerNorm both before and after each sub-block (doublenorm). What problem does that solve, and what does it cost?” Senior answer: double-norm bounds both the input to the sub-block and its output before re-adding to the stream. At very large scale, sub-blocks can still produce large-magnitude outputs even when their inputs are normalized; bounding the output provides a second safety checkpoint. The cost is extra memory bandwidth (LayerNorm is bandwidth-bound, not flop-bound - each extra norm adds one additional read/write pass over a $\mathbb{R}^{n \times d}$ activation tensor per layer), plus the additional learned scale parameters $\gamma$. At 7B parameters this cost is negligible; at 70B or 100B, you should ablate whether the stability gain justifies it rather than adopting it by default.

Key Takeaways

• Pre-norm applies normalization before each sub-block, keeping the residual stream as a clean identity path for gradients; post-norm applies it after the residual addition, forcing gradients through a normalization at every layer.

• Post-norm requires careful learning-rate warmup to avoid divergence; pre-norm trains stably from step one and allows cold restarts from any checkpoint.

• Pre-norm is universal in modern frontier models: LLaMA, PaLM, Chinchilla, Mistral, DeepSeek, Gemma, and essentially every open model released since GPT-2.

• Double-norm (pre- and post-block norms, used in Gemma 2 and Grok-1) adds a second stability checkpoint for very large models; it is not yet the default and costs extra memory bandwidth.

• The gradient norm should be below 5.0 from step 1 with pre-norm and a sane learning rate; values above that signal a normalization placement bug or initialization issue, not a fundamental architecture problem.

3.2 Normalization Choices

3.2.1 LayerNorm vs. RMSNorm

In a system design round at a frontier lab, you will almost certainly be asked why your training stack uses RMSNorm instead of LayerNorm. The question sounds trivial - one formula versus another - but the real interview signal is whether you can connect normalization choices to memory bandwidth, training stability, and the arithmetic intensity of GPU kernels. This comes up at the mid-to-senior level whenever hardware efficiency enters the conversation, and getting the answer right separates engineers who have profiled real training runs from those who have only read the paper abstracts.

LayerNorm normalizes each activation vector by its empirical mean and standard deviation, then applies a learned affine transform:

where $\mu = \frac{1}{d} \sum_{j = 1}^{d} x_{j}$ and $\sigma^{2} = \frac{1}{d} \sum_{j = 1}^{d} ( x_{j} - \mu )^{2}$. Two statistics must be computed, and two parameter vectors $\gamma , \beta \in \mathbb{R}^{d}$ must be loaded from memory on every forward pass.

RMSNorm drops mean-centering entirely and eliminates the shift parameter:

Only one statistic (the root-mean-square) is computed, and only $\gamma \in \mathbb{R}^{d}$ is loaded. The $\beta$ parameter vanishes.

On paper, the FLOP difference is negligible - normalization operations are less than 0.2% of the total FLOPs in a transformer. The argument for RMSNorm is not arithmetic; it is memory. Careful profiling (Ivanov et al., 2023) shows that normalization kernels account for roughly 25% of wall-clock runtime despite their vanishing FLOP share, because they are memory-bandwidth-bound: every scalar in $\gamma$ and $\beta$ must be staged through HBM on every call. Halving the parameter count directly halves that memory traffic.

The case for RMSNorm is a memory-bandwidth argument dressed up as a simplification argument: fewer parameters to load per call, fewer operations per element, and empirically indistinguishable model quality.

GPU computation is dominated by two bottlenecks: arithmetic throughput (FLOPs) and memory bandwidth (bytes/s). Matrix multiplies are compute-rooflined at large batch sizes - the GPU is limited by how fast it can multiply, not by how fast it can read weights. Normalization kernels are the opposite: they touch each activation element a small, fixed number of times, so the runtime is bounded by how fast the GPU can stream data through HBM. On an H100, peak compute is around 989 TFLOP/s (bf16) while peak memory bandwidth is 3.35 TB/s; the ratio is roughly $300 \times$. Any kernel that is not doing at least 300 FLOPs per byte loaded will hit the memory wall, not the compute ceiling. Normalization is firmly in that regime.

LayerNorm’s $\beta$ parameter is small in count but adds a full load-and-multiply pass over $d$ scalars per invocation. For a model with $d = 4096$ and 64 normalization calls per forward pass (32 layers $\times$ 2 norms), that is $64 \times 4096 \times 2$ bytes $\approx 524$ KB of parameter tensor loads per sample per forward pass in bf16, entirely eliminated by RMSNorm.

The quality story is equally clean. Narang et al. (2020) ran controlled ablations on a transformer of comparable scale and found that RMSNorm reaches the same final cross-entropy as LayerNorm while running at 3.68 steps/second versus 3.50 - a free 5% throughput improvement. The LLaMA family, PaLM, Chinchilla, Mistral, and DeepSeek have all converged on RMSNorm. The near-consensus explanation: when pre-norm placement already keeps the residual stream well-conditioned before each sub-block, the mean of the activation vector entering the norm is already close to zero, so mean-centering provides negligible correction.

A Worked Example

Consider a 7B decoder-only model: $d = 4096$, 32 layers, batch size $B = 4$, sequence length $T = 2048$, bf16 precision. Each transformer layer has two normalization calls (pre-attention and pre-MLP). Per call:

1. LayerNorm: compute mean ($d$ additions), compute variance ($2 d$ multiply-adds), normalize ($2 d$ ops), apply $\gamma$ and $\beta$ ($2 d$ multiply-adds). Total: $\approx 7 d$ scalar ops. Memory loads: $2 d$ bytes for $\gamma$, $2 d$ bytes for $\beta$ (bf16).

2. RMSNorm: compute RMS ($d$ squares, one sum, one sqrt), normalize ($d$ divisions), apply $\gamma$ ($d$ multiplyadds). Total: $\approx 3 d$ scalar ops. Memory loads: $2 d$ bytes for $\gamma$ only.

Across the entire forward pass: $32 \times 2 = 64$ normalization calls. The $\beta$ removal saves $64 \times 4096 \times 2$ bytes $= 524$ KB per sample per forward pass. At $B = 4$, that is roughly $2$ MB per training step eliminated from HBM traffic. At 10,000 steps/hour, this is $\approx 20$ GB/hour of avoided memory movement - not the dominant cost, but meaningful when normalization is otherwise consuming 25% of wall time.

The throughput difference of 5% reported in ablations maps directly onto this: if normalization is 25% of runtime and RMSNorm cuts its overhead in half, the net gain is $0 . 25 \times 0 . 5 = 12 . 5 \%$ reduction in normalization time, which translates to approximately 3-5% overall throughput improvement accounting for the compute-rooflined portions of the forward pass.

What You Actually Decide in Practice

• Default to RMSNorm. Every major model family released since 2022 that publishes architecture details uses it. Cohere Command A and Command R+ are the notable exception - they retain LayerNorm for undisclosed reasons. Unless you have a specific compatibility constraint, there is no reason to start a new run with LayerNorm.

• Do not re-introduce $\beta$. Some codebases add a bias back to RMSNorm for “completeness.” This negates the memory saving and provides no quality benefit - the downstream linear layers already supply the affine freedom that $\beta$ would add.

• Watch $\varepsilon$ at bf16. The default $\varepsilon = 10^{- 5}$ can underflow in bf16 for activation vectors with very small magnitudes, producing a near-zero denominator. Use $\varepsilon = 10^{- 6}$ and monitor for NaN spikes on the first 1B tokens; if they appear, increase $\varepsilon$ to $10^{- 8}$ in fp32 accumulation mode before blaming the optimizer.

• Fused kernels are mandatory at scale. A naive PyTorch RMSNorm implementation still issues separate HBM transactions for each element-wise op. Production runs use Triton or cuDNN-fused kernels; unfused implementations can cost $2$-$3 \times$ the wall time of the fused version, erasing the RMSNorm advantage entirely.

• LoRA adapters on RMSNorm backbones need explicit bias handling. When fine-tuning with LoRA or prefix tuning, adapter layers sometimes implicitly assume a $\beta$ term exists. If your adapter introduces a bias, verify it is absorbing the role of $\beta$ intentionally; accidental double-counting or missing shift can cause subtle drift in generation quality that is hard to diagnose.

How This Shows Up in Interviews

• “Why do modern LLMs use RMSNorm instead of LayerNorm?” A strong answer names both the memorybandwidth argument (normalization is $\sim$25% of runtime despite $<$0.2% of FLOPs) and the quality-parity finding, and cites at least two concrete model families that made the switch. Answering only “it’s simpler” scores mid-level at best.

• “If RMSNorm drops mean-centering, doesn’t that hurt models where activation distributions are not zero-mean?” The correct answer is that pre-norm placement ensures the residual stream entering each normalization is already regulated by upstream norms, so the empirical mean is close to zero at initialization and remains close throughout training. Mean-centering provides negligible correction in that regime, which is exactly what the ablations confirm.

• Senior follow-up: “Walk me through the arithmetic intensity of a LayerNorm kernel versus a matrix multiply. Why does the roofline model determine how much you can save by switching to RMSNorm?” A senior answer articulates the roofline model: LayerNorm is bandwidth-rooflined (low arithmetic intensity), so reducing memory accesses directly reduces wall time proportionally. A matrix multiply at large batch size is compute-rooflined, so an equivalent FLOP reduction there yields a much smaller real-world speedup. The implication is that optimizing memory-bound kernels gives super-linear returns relative to their FLOP share.

Key Takeaways

• RMSNorm drops mean-centering and the shift parameter $\beta$; it computes only $\text{RMS} ( \mathbf{x} )$ and applies a single learned scale $\gamma$.

• Normalization kernels are memory-bandwidth-bound: they are $\sim$0.2% of FLOPs but $\sim$25% of runtime. The win from RMSNorm is in memory traffic, not arithmetic.

• Published ablations show quality parity with LayerNorm and a roughly 5% throughput improvement; the LLaMA family, PaLM, Chinchilla, Mistral, and DeepSeek all converged on RMSNorm.

• Do not add $\beta$ back, and always use fused kernels - naive implementations negate both the memory and arithmetic savings.

• The deeper lesson: once a model is architecture-stable, the bottleneck often moves from FLOPs to memory movement. Normalization is the canonical example of an operation that is cheap in FLOPs and expensive in bandwidth.

3.2.2 QK-Norm and Stability Interventions

At some point in every large training run, you will open your monitoring dashboard and see a gradient norm plot that looks like a seismograph during an earthquake - sudden spikes of $10 \times$ or $100 \times$ the baseline, occasionally followed by NaN propagation that ends the run. Training instability is one of the most expensive failure modes in LLM development, capable of wasting weeks of GPU time and hundreds of thousands of dollars on a single diverged job. Knowing which architectural interventions actually prevent these spikes, and being able to explain why they work at a mechanistic level, is a senior-level topic that surfaces in both system design rounds and deep technical interviews at labs actively training at scale.

The two locations in a transformer where numeric instability most commonly originates are the attention softmax and the output softmax. Both involve exponentiating logits that can grow without bound during training, producing either overflow, near-zero denominators, or pathologically peaked distributions that starve most positions of gradient signal.

QK-norm targets the attention softmax. Before computing $\mathbf{Q} \mathbf{K}^{\top} / \sqrt{d_{k}}$, apply a normalization layer - typically RMSNorm - to the query and key tensors independently:

This bounds the magnitude of pre-softmax attention logits regardless of how the weight matrices $\mathbf{W}_{Q}$ and $\mathbf{W}_{K}$ evolve during training.

Z-loss targets the output softmax. Let $z ( \mathbf{x} ) = \log \sum_{v =}^{V} \exp_{1} ( u_{v} ( \mathbf{x} ) )$ be the log-normalizer of the output distribution, where $u_{v}$ are the pre-softmax logits. Z-loss adds an auxiliary penalty to the training objective:

This encourages $z \rightarrow 1$ (equivalently, $\log z \rightarrow 0$), keeping the output logit scale in a numerically clean regime. When $z = 1$, the softmax reduces to the raw exponentials, which are bounded and smooth.

Logit soft-capping is a complementary intervention on the attention softmax. After computing the raw dot-product attention scores $l_{i j}$, apply:

where $c$ is a soft-cap hyperparameter (e.g. $c = 50$). The tanh clamps extreme logits softly while preserving gradients near zero.

In every case the goal is identical: keep inputs to any softmax bounded in magnitude so that neither the exponential overflows nor the distribution degenerates to a Kronecker delta.

Softmax is the only non-linear operation in the standard transformer that can be driven to numerical failure by the surrounding linear algebra. A weight matrix with bounded spectral norm produces bounded outputs if its inputs are bounded - but a softmax with growing logit magnitudes produces exponentially larger activations regardless of what came before it. QK-norm breaks this dependency: regardless of how $\mathbf{W}_{Q}$ and $\mathbf{W}_{K}$ evolve during optimization, the vectors fed to the dot product are normalized to unit scale, bounding pre-softmax logits to $O ( 1 / \sqrt{d_{k}} )$.

The contrast with gradient clipping is instructive. Clipping intervenes in the backward pass after a spike has already occurred; the forward-pass activations that triggered the spike are already resident in the residual stream and will propagate to downstream layers before the next update corrects them. QK-norm and z-loss act in the forward pass, preventing the spike from occurring in the first place. For runs longer than $\sim$100B tokens, this distinction matters: each uncaught spike degrades optimizer state in ways that take many subsequent steps to wash out.

Z-loss and QK-norm address orthogonal failure modes. Z-loss controls the global scale of output logits (across $V \sim 32 , 000$-$128 , 000$ vocabulary items), where failure is typically a slow drift. QK-norm controls attention logit magnitude (across $T$ positions), where failure is typically a sudden spike. Running both provides defense in depth.

A Worked Example

Consider a 7B model with head dimension $d_{k} = 128$, sequence length $T = 2048$. Suppose after 80B tokens of training the key vectors in one attention head have drifted to RMS magnitude $\| \mathbf{k} \|_{2} \approx 12$. The pre-softmax attention logit between positions $i$ and $j$ is:

If position $j^{\ast}$ has logit $12 . 7$ and all other 2047 positions have logit $- 12 . 7$, the attention weight on $j^{\ast}$ is:

The gradient of the loss with respect to any position $j \neq j^{\ast}$ through this softmax is essentially zero. That head has effectively stopped learning - and if the logit continues to grow, the next forward pass may overflow to inf, producing a NaN that propagates through the entire residual stream.

With QK-norm, both $\mathbf{q}_{i}$ and $\mathbf{k}_{j}$ are normalized to unit scale before the dot product. The post-normalization logit is:

The softmax operates on logits near zero; the resulting attention distribution is diffuse, every position receives meaningful gradient, and neither overflow nor gradient vanishing occurs. The cost is two RMSNorm calls per attention layer per forward pass - a kernel that is $\ll 1 \%$ of the compute of the surrounding QKV projections.

What You Actually Decide in Practice

• Use QK-norm and z-loss together, not as alternatives. They target different softmaxes and different failure modes. A model with QK-norm but no z-loss can still diverge through the output logit path as training progresses and the output weight matrix scale drifts upward.

• Set $\alpha = 10^{- 4}$ for z-loss as your starting point. This is the PaLM coefficient and has been reproduced by DCLM and OLMo 2. Raising it above $10^{- 3}$ will hurt token-level cross-entropy because the model is penalized for concentrating probability mass, which it must do to make high-confidence predictions. Lowering below $10^{- 6}$ provides no measurable stabilization.

• Do not add logit soft-capping without ablating. NVIDIA’s ablation found that soft-capping at $c = 50$ worsened perplexity by 0.1-0.3 points on their setup, while QK-norm improved it by a similar margin. Gemma 2 reports benefit from soft-capping. The inconsistency likely stems from interactions with the Adam optimizer’s second-moment estimates: tanh distorts the gradient landscape in a way that depends on the learning rate schedule and model scale. Treat it as a model-specific knob.

• Consider the double-norm pattern for very large runs. Grok and Gemma 2 place normalization both before and after each sub-block (attention and MLP), rather than only before (standard pre-norm). The post-block norm bounds the output magnitude entering the residual add, which is complementary to QK-norm inside the attention. OLMo 2 uses only the post-block norms, omitting the pre-block norms. Evidence so far suggests this is strictly more stable than pure pre-norm for models trained beyond 300B tokens.

• Gradient clipping at 1.0 is a floor, not a ceiling. Keep gradient clipping as a last-resort backstop, but do not rely on it as your primary stability mechanism. At the scale where instability becomes dangerous ($\geq$30B parameters, $\geq$100B tokens), the frequency of clipping events correlates with final model quality degradation. If you clip more than 2-3% of steps, investigate the architectural root cause rather than raising the clip threshold.

How This Shows Up in Interviews

• “You are training a 70B model and gradient norm spikes appear after 50B tokens. What do you check first and what architectural changes do you consider?” A strong answer mentions per-layer gradient monitoring to localize whether spikes originate in the attention blocks or the output projection, then proposes QK-norm for the attention path and z-loss for the output path. It also checks whether the current learning rate is still in the cosine schedule’s stable descent region.

• “What does QK-norm actually do to the attention distribution, and what is the cost?” Strong answer: it bounds pre-softmax attention logits to $O ( 1 / \sqrt{d_{k}} )$, preventing attention from collapsing to a one-hot distribution over long training; cost is two extra RMSNorm calls per attention layer, negligible compared to the surrounding projections.

• Senior follow-up: “Z-loss controls the output softmax normalizer. But cross-entropy already backpropagates through the softmax. How can z-loss provide information that the cross-entropy gradient does not?” The senior answer: cross-entropy gradient pushes the correct logit up relative to the others, but does not constrain the absolute scale of the logit vector. Z-loss penalizes the scale directly - even a model achieving perfect cross-entropy can have large-magnitude logits that create numerical instability. The two losses operate on orthogonal degrees of freedom in logit space (direction versus magnitude), which is why they are complementary and not redundant.

Key Takeaways

• Training instability in transformers concentrates at the two softmaxes: the attention softmax over $T$ sequence positions and the output softmax over the $V$ -token vocabulary.

• QK-norm applies RMSNorm to query and key tensors before the attention dot product, bounding pre-softmax logit magnitudes and preventing attention collapse or gradient vanishing. Gemma 2, OLMo 2, and DCLM all use it.

• Z-loss adds $\alpha \cdot [ z ( \mathbf{x} ) ]^{2}$ to the training objective, keeping the output softmax normalizer near 1 and preventing unbounded logit growth; PaLM set the standard with $\alpha = 10^{- 4}$.

• Use QK-norm and z-loss together - they target different failure modes and are complementary. Logit soft-capping is a model-specific choice; ablate before adding it.

• Gradient clipping is a backstop, not a primary stability tool; frequent clipping ($>$2% of steps) signals an architectural problem that needs to be fixed at the source.

3.3 Activation Functions in Modern LLMs

3.3.1 ReLU, GeLU, and SwiGLU Families

Every major open-weight model released since 2022-LLaMA, PaLM, Gemma, Mistral, DeepSeek, Qwenuses a gated activation in its MLP blocks. If you walk into an architecture deep-dive and cannot explain what SwiGLU is, why it replaced GeLU, and how it changes the hidden dimension sizing, you will stall on a warmup question before the real design challenge begins. This topic appears most often at mid-to-senior level interviews at frontier labs, where it quickly signals whether a candidate has read recent model cards or is still reasoning from undergraduate deep learning.

The MLP block in a transformer layer projects from the residual stream dimension $d$ up to a wider hidden dimension $d_{f f}$, applies a nonlinearity, then projects back down. Three activation families partition the design space.

ReLU. $\text{ReLU} ( x ) = \max ( 0 , x )$. The two-layer FFN is:

Fast, hardware-friendly, but a large fraction of hidden units saturate at zero during training, wasting capacity. The hard kink at zero also creates gradient cliffs that require careful initialization.

GeLU. The Gaussian Error Linear Unit smooths the hinge:

where $\Phi$ is the CDF of the standard normal. In code it is approximated as $x \cdot \sigma ( 1 . 702 x )$. Differentiable everywhere, no dead-neuron regime, and empirically better than ReLU on every benchmark published for it. GPT-1 through GPT-3 all use GeLU, and it held as the dominant choice through 2021.

Gated Linear Units (GLU) and SwiGLU. The core idea: multiply the activated hidden representation element-wise by a second learned linear projection that acts as a soft gate:

SwiGLU replaces the sigmoid gate with Swish, where $\text{Swish} ( x ) = x \cdot \sigma ( x )$. The full SwiGLU MLP, without bias terms as in all modern LLMs, is:

where $W_{1} , V \in \mathbb{R}^{d \times d_{f f}}$ and $W_{2} \in \mathbb{R}^{d_{f f} \times d}$.

GeGLU substitutes GeLU for Swish and is used by T5v1.1, Gemma 2, and Gemma 3:

Gated activations let the network learn to suppress or amplify each hidden dimension independently at each token, which is strictly more expressive than a fixed pointwise nonlinearity at the same parameter budget.

The $V$ matrix acts as a context-dependent gain: for every hidden dimension at every token, the network independently decides how much of the activated signal passes forward. A plain ReLU or GeLU MLP can only gate in a binary sense-either the value is positive and passes, or it does not. A GLU variant attenuates values softly, in a learned, input-dependent way. The empirical result is a consistent 0.3-0.8 nat improvement in pre-training perplexity at equal parameter count.

The smooth nonlinearity in SwiGLU and GeGLU matters separately. The hard kink at zero in ReLU means neurons near the boundary receive nearly zero gradients and stop participating in learning. Swish and GeLU both maintain a small non-zero gradient slightly below zero, which keeps more neurons active throughout training. The Swish curve dips very slightly below zero for large negative inputs-a property that alarmed some early reviewers-but in practice, optimizer momentum carries activations far away from that tiny regime and it has never been measured as a problem in production training runs.

Before GLU variants, the choice was ReLU (fast, rough) versus GeLU (smooth, slightly better). GLU variants deliver GeLU-level expressiveness plus a multiplicative gate, at the cost of one additional weight matrix $V$ . The question is always whether that matrix earns its keep-the evidence consistently says yes.

A Worked Example

Take LLaMA-2 7B: $d = 4096$, SwiGLU. The convention is that the total MLP parameter count should match a standard $4 \times$ ReLU FFN to keep the comparison fair.

Standard ReLU FFN with $d_{f f} = 4 d = 16 , 384$: two matrices, each $d \times d_{f f}$:

SwiGLU FFN has three matrices: $W_{1}$, $V$ , $W_{2}$. To match 134 M with three equal-shaped matrices:

Round up to the nearest multiple of 128 for hardware tile alignment: LLaMA-2 7B uses $d_{f f} = 11 , 008$.

The forward pass through one SwiGLU MLP layer at batch size $B$, sequence length $T$:

1. Gate projection: $g = \text{Swish} ( x W_{1} ) \in \mathbb{R}^{B \times T \times 11008}$

2. Value projection: $v = x V \in \mathbb{R}^{B \times T \times 11008}$

3. Element-wise product: $h = g \odot v \in \mathbb{R}^{B \times T \times 11008}$

4. Down projection: out $= h W_{2} \in \mathbb{R}^{B \times T \times 4096}$

Total FLOPs for the three matrix multiplications: $6 \times B \times T \times 4096 \times 11 , 008 \approx 270$B $\cdot B \cdot T$ FLOPs. The element-wise product in step 3 is memory-bound and negligible in total FLOP count. The two extra weight matrices compared to a GeLU FFN cost nothing in FLOPs beyond what a parameter-matched GeLU FFN would cost.

What You Actually Decide in Practice

• Round $d_{f f}$ to a multiple of 64, preferably 128. Hardware tiling requirements mean oddly sized dimensions can cost 10-20% throughput on A100 and H100. LLaMA-2 uses 11008 rather than the exact 10922.7 for this reason. Some labs prefer multiples of 256 for very large models.

• Default to SwiGLU. GeGLU and SwiGLU are empirically indistinguishable in pre-training loss; pick whichever your framework has a fused kernel for. Both beat GeLU, which beats ReLU. There is no reason to use ReLU in a new architecture today.

• Drop bias terms everywhere. Modern LLMs omit biases in $W_{1}$, $V$ , and $W_{2}$. Bias parameters add optimizer state, consume memory bandwidth on every forward pass, and have been repeatedly observed to destabilize large-scale training without any measurable quality benefit. Dropping them is free money.

• The $8 / 3$ factor only applies to the $4 \times$ baseline. If your starting $d_{f f}$ uses a different ratio-Gemma 2 uses $8 \times$, the original T5-11B used $64 \times$-the matching factor changes. Always derive from parameter matching, not from memorizing ‘$8 / 3$’.

• Profile fused kernel availability before committing. SwiGLU requires launching matmuls for both $W_{1}$ and $V$ before the element-wise product. Without kernel fusion, this is two separate GEMM launches. Unfused SwiGLU can be slower in wall-clock time than fused GeLU on some inference runtimes. Verify your serving stack has a fused GLU kernel before assuming the theoretical win materializes.

How This Shows Up in Interviews

• “Why did LLaMA switch from GeLU to SwiGLU?” Strong answer: gating adds a learned per-dimension gate that selectively suppresses hidden activations; ablations across multiple studies show consistent perplexity gains for the same parameter and FLOP budget; Shazeer’s original GLU paper and subsequent PaLM/T5 work provided the production validation. Mention the $8 / 3$ factor to demonstrate you have internalized the parameter consequences.

• “If $d_{\text{model}} = 8192$ and I want SwiGLU, what is $d_{f f}$?” Strong answer: $\frac{8}{3} \times 8192 = 21 , 845 . 3$, round to 21,888 (nearest multiple of 128). State your rounding rationale-hardware tile alignment-not just the number.

• Senior follow-up: “Roughly how much memory do the MLP weights consume in a 7B model with 32 layers, $d = 4096$, $d_{f f} = 11008$, stored in BF16?” Sketch the math on a whiteboard: one matrix is $4096 \times 11008 \times 2$ bytes $\approx 90$ MB; three matrices per layer gives $\approx 270$ MB; across 32 layers, $\approx 8 . 6$ GB. That is roughly half of the model’s total 14 GB (BF16 7B weights), consistent with the rule of thumb that MLPs and attention each consume about half of model parameters.

Key Takeaways

• SwiGLU and GeGLU add a learned gate ($x V$ ) that suppresses or amplifies each hidden dimension per token; this consistently improves pre-training perplexity at equal parameter count.

• Use $d_{f f} = \frac{8}{3} d$, rounded to the nearest hardware-friendly multiple (128 or 256), to match the parameter count of a standard $4 \times$ ReLU FFN.

• LLaMA-2 7B ($d = 4096$, $d_{f f} = 11 , 008$) is the canonical reference; verify any architecture number against it.

• Drop bias terms in all MLP linear layers-they add memory bandwidth and optimizer state with no measurable quality gain.

• Verify your runtime has a fused SwiGLU kernel; two unfused GEMM launches for $W_{1}$ and $V$ can erase the theoretical throughput advantage.

3.3.2 Why GLU Variants Dominate Production Models

If you survey the 19 dense LLM releases from 2024 alone-Command A, OLMo 2, SmolLM, Phi-4, Gemma 3, Qwen 2.5, LLaMA 3-almost every one uses SwiGLU or GeGLU. This is not groupthink or cargo-culting. It is the result of independently run ablations at multiple labs converging on the same answer. Understanding why GLU variants win tells you something more durable than a recipe: it gives you a mental model for reasoning about future architecture changes that interviewers at senior levels will probe for.

The empirical case for GLU variants rests on two independent lines of evidence. Noam Shazeer’s 2020 paper evaluated every GLU variant-ReGLU (ReLU gate), GeGLU (GeLU gate), SwiGLU (Swish gate), and others-against a standard GeLU baseline on CoLA and SST-2 fine-tuning tasks. The results were unambiguous: every GLU variant outperformed the GeLU baseline, and the gains were statistically significant. GeGLU scored 84.67 on CoLA against 84.20 for the GeLU baseline, a difference that might look small until you consider these are pre-training architecture choices where 0.5% gains are hard-won.

Narang et al. (2020) extended this to pre-training loss in a T5-style study covering a wide range of architecture decisions simultaneously. The bolded top-performers in their loss comparison tables are, without exception, the GLU variants. The non-gated activations (ReLU, GeLU, ELU) cluster together below them at every scale tested.

The convergent evolution after 2022 seals the argument. PaLM (Google, 2022) used SwiGLU in a 540B-parameter model-the most expensive training run of its era-after internal ablations. T5v1.1 (the improved version of T5) switched from the original T5’s ReLU to GeGLU and achieved better downstream performance. The entire LLaMA family, Mistral, Qwen, DeepSeek, Yi, and dozens of open-weight models followed. When multiple independent labs running trillion-token training runs independently converge on the same choice, the evidence is about as strong as you ever get in empirical deep learning.

GLU variants consistently win because gating is a universally useful architectural primitive: it appears in LSTMs (forget gate), GRUs (update gate), attention (the $V$ projection scaled by softmax weights), and now in MLP blocks.

The underlying mechanism is expressiveness for free. A standard FFN uses one matrix multiply to mix features and one pointwise function to introduce nonlinearity. A GLU FFN uses two matrix multiplies to create an activated branch and a linear gate, then combines them element-wise. The gate is not a fixed function of the input magnitude; it is learned, input-dependent, and can vary across the $d_{f f}$ hidden dimensions independently. This effectively gives the MLP a form of conditional computation at the feature level.

The memory bandwidth argument reinforces this. As discussed in the context of RMSNorm, the transformer’s bottleneck is not FLOPs-tensor contractions are 99.8% of FLOPs but not 99.8% of wall-clock time. What matters is how often you move data between memory hierarchies. A GLU variant with $d_{f f} = \frac{8}{3} d$ moves the same amount of data as a standard FFN with $d_{f f} = 4 d$-three smaller matrices versus two larger ones at the same total parameter count. The memory bandwidth cost is identical. The quality gain is free.

The contrast with what came before clarifies the win. The pre-GLU choice was: ReLU (fast but rough, dead neurons, gradient cliffs) versus GeLU (smooth, differentiable, better quality, slightly slower to compute). GLU variants deliver GeLU-level gradient behavior plus conditional gating at the same memory cost. The iteration from the Shazeer paper to PaLM to LLaMA 1 to LLaMA 3 is the clearest example of a clean architectural improvement compounding across the field.

A Worked Example

To see why production models chose SwiGLU over a scaled-up GeLU FFN, consider two configurations for a 70B model with $d = 8192$ and 80 transformer layers.

Option A: GeLU FFN, $d_{f f} = 4 d = 32 , 768$:

Option B: SwiGLU FFN, $d_{f f} = \frac{8}{3} \times 8192 = 21 , 845 \rightarrow 21 , 888$ (rounded to 128):

The parameter counts are within 0.2% of each other. The memory bandwidth cost per forward pass is essentially identical. Yet Option B (SwiGLU) consistently achieves 0.3-0.8 nat lower perplexity. This is why every lab that ran this ablation chose SwiGLU: it is a free lunch. LLaMA-2 70B uses exactly this configuration ($d = 8192$, $d_{f f} = 28 , 672$, which reflects a slightly larger ratio of roughly 3.5$\times$ rather than 2.67$\times$, a choice that trades slight parameter increase for wider matrices).

What You Actually Decide in Practice

• GLU is the default; exceptions require justification. GPT-3 uses GeLU and is still a strong model, Nemotron-4 340B uses squared ReLU ($\text{ReLU}^{2} ( x ) = \max ( 0 , x )^{2}$), and Falcon 2 11B uses plain ReLU. These are competitive models. But if someone asks why you chose GeLU over SwiGLU for a new pretraining run, “it’s simpler” is not a satisfying answer-you need an ablation or a systems constraint (e.g., your inference runtime has no fused GLU kernel).

• SwiGLU and GeGLU are interchangeable in practice. No production paper shows a measurable difference between them at the same parameter count. Pick the one with better kernel support in your stack.

• The T5 lesson: extreme $d_{f f}$ ratios are legal but risky. The original T5-11B used $d_{f f} = 65 , 536$ against $d = 1 , 024$, a ratio of 64$\times$. The model worked, but T5v1.1 (the improved version) dialed this back to a standard 2.5$\times$ ratio with GeGLU and got better results. Very large ratios spend parameters inefficiently on width rather than depth; the Kaplan et al. scaling experiments show a wide basin (1$\times$ to 10$\times$) but the optimum is around 4$\times$ (or 2.67$\times$ with GLU).

• Watch for the squared-ReLU surprise. Nemotron and some mixture-of-experts models use $\text{ReLU}^{2}$, which is sparsity-inducing: most hidden units fire at exactly zero. This creates naturally sparse MLP activations that can be exploited for inference acceleration. If your deployment stack supports sparse kernel optimizations, squared ReLU is worth benchmarking.

• Parameter matching is a convention, not a law. Some labs (Mistral, PaLM) use slightly larger $d_{f f}$ ratios than $8 / 3$ for their GLU models-effectively choosing to spend more parameters in the MLP. This is a deliberate capacity allocation decision, not an error. Always check the actual model card rather than assuming every SwiGLU model uses exactly $8 / 3$.

How This Shows Up in Interviews

• “Would you ever use ReLU in a new architecture? When?” Strong answer: only with a specific justification-squared ReLU for sparsity-based inference acceleration, or a constrained runtime with no GLU kernel and GeLU not available either. For standard pre-training with a modern stack, there is no reason to choose ReLU over SwiGLU.

• “What does it mean that LLaMA and PaLM both independently picked SwiGLU?” Strong answer: it validates the ablation-two groups with different data, hardware, and model scales ran independent experiments and converged. When multiple frontier labs agree on an architecture choice without coordinating, that is the strongest signal available that the choice is correct. The interviewer is testing whether you reason from evidence, not just pattern-match on popular choices.

• Senior follow-up: “You are designing a 100B-parameter model. Your inference stack runs on edge hardware where kernel fusion is limited to a single GEMM per layer. How does that change your activation function choice?” Sketch the answer: SwiGLU requires two back-to-back GEMMs ($W_{1}$ and $V$ ) before the element-wise product; without fusion, this is two memory round-trips to HBM. On memory-constrained edge hardware, this can be a throughput bottleneck. Options: (1) fall back to GeLU (one GEMM, almost as good empirically); (2) rewrite the kernel to fuse both GEMMs into a single wider matmul; (3) use squared ReLU (one GEMM, exploits sparsity for further gains if the hardware has sparse execution units). The answer demonstrates that activation function choice is not isolated from systems constraints.

Key Takeaways

• Multiple independent labs-Google, Meta, Alibaba, Mistral-running ablations at scale all converged on SwiGLU or GeGLU; that convergence is stronger evidence than any single paper.

• GLU variants win because gating is a universal architectural primitive: the same element-wise multiplicative gate that works in LSTMs and attention works in MLP blocks.

• SwiGLU and GeGLU are empirically indistinguishable; choose based on which your runtime can fuse efficiently.

• GPT-3 (GeLU), Nemotron 340B (squared ReLU), and Falcon (ReLU) are legitimate exceptions-strong models exist without GLU, but each exception needs a reason.

• Very wide MLPs ($d_{f f} > 10 \times d$) spend parameters inefficiently; the T5-11B to T5v1.1 regression from $64 \times$ to $2 . 5 \times$ ratio is the canonical cautionary example.

3.4 Position Encodings

3.4.1 Absolute, Relative, and Sinusoidal Encodings

Transformers are permutation-invariant by design: self-attention produces identical outputs regardless of the order in which tokens appear. Position encoding is therefore not an optional add-on - it is a load-bearing architectural choice, and “how does the model know token order?” is a near-universal architecture deep-dive question. The more revealing follow-up is “what breaks when you run the model on sequences longer than its training context?” That question separates engineers who memorized the formula from those who understand why a decade of iteration eventually converged on RoPE.

Three families of position encodings cover virtually every production model built before 2022.

Absolute learned embeddings. Add a dedicated vector $\mathbf{p}_{i} \in \mathbb{R}^{d}$ to the token embedding at position $i$. The $\mathbf{p}_{i}$ are free parameters trained end-to-end alongside everything else. This is exactly what GPT-1, GPT-2, GPT-3, OPT, and BERT use. The only constraint is that the vocabulary of positions $\{ 0 , 1 , . . . , L - 1 \}$ must be declared at model initialization, and any position beyond $L - 1$ is simply undefined.

Sinusoidal (fixed) embeddings. The original Attention Is All You Need (2017) proposed hard-coded vectors built from paired sine and cosine functions at geometrically spaced frequencies:

where $i$ is the token position and $k = 0 , 1 , . . . , d / 2 - 1$ indexes dimension pairs. The denominators form a geometric series of wavelengths from $2 \pi$ tokens (fastest oscillating pair) up to $10000 \cdot 2 \pi$ tokens (slowest). Because a fixed rotation matrix $M_{\delta}$ maps PE$( i )$ to PE$( i + \delta )$ for any offset $\delta$, the model can in principle represent relative distances as linear combinations of absolute embeddings.

Relative encodings. Rather than encoding absolute position into token representations, these schemes inject a position-dependent term directly into the attention score. Shaw et al. (2018) added a learned bias $a_{i j}^{K}$

to the query-key dot product for each pair $( i , j )$, clipped at some maximum offset. T5 generalized this with a scalar learned bias $b ( i - j )$ shared across all positions at the same distance. Because the bias depends only on $i - j$, these schemes are inherently relative and do not need to define a maximum absolute position.

The axis that matters: absolute schemes bake position into token representations before attention; relative schemes fold position into the attention score computation itself.

Sinusoidal embeddings have no learnable parameters, which means they add nothing to the model’s memory footprint. Their multi-scale structure - fast oscillations for nearby tokens, slow oscillations for distant ones - lets the model distinguish positional relationships across a wide range of distances with a single $d$-dimensional vector. The linear-rotation property means a well-trained model can, in principle, extract relative offsets algebraically. In practice, performance degrades beyond the training context length because the model has never seen those absolute position vectors during training; the degradation is softer than with learned embeddings but still substantial.

Learned absolute embeddings are the simplest possible scheme and fast at inference (a single table lookup). The critical failure mode is hard and immediate: any sequence position beyond $L - 1$ maps to an uninitialized or out-of-distribution vector. In GPT-3 with $L = 2048$, a 2049-token sequence causes undefined behavior. There is no graceful degradation.

Relative schemes generalize to unseen lengths because the bias $b ( i - j )$ is indexed by distance, not by absolute index. The trade-off is memory: T5-style relative biases require a tensor of shape $[$batch$,$ heads$, n , n ]$ that is $O ( n^{2} )$ in sequence length. At $n = 4096$, 16 heads, batch size 8, in float32, this tensor is $8 \times 16 \times 4096^{2} \times 4 \approx 8 . 6$ GB - often the single largest activation in the forward pass. Beyond 4k tokens, relative bias memory dominates everything else.

ALiBi (Press et al., 2022) sidesteps the memory issue by making the bias a fixed negative slope times distance: $- m \cdot ( j - i )$, where $m$ is a head-specific constant. No parameters, no $O ( n^{2} )$ tensor, and strong empirical length generalization. The structural assumption, however, is a monotone recency bias: every token always prefers recent context. This hurts tasks that require retrieving information from the beginning of a very long document.

A Worked Example

Take $d = 4$ and compute the sinusoidal vectors for positions $i = 0$ and $i = 1$. There are two dimension pairs, $k = 0$ and $k = 1$:

Pair $( 0 , 1 )$ completes a full oscillation in $2 \pi / 1 . 0 \approx 6 . 3$ tokens; it is highly sensitive to whether two tokens are adjacent or a few positions apart. Pair $( 2 , 3 )$ completes a cycle in $2 \pi / 0 . 01 \approx 628$ tokens; it varies negligibly over short windows but discriminates positions at the scale of hundreds. For a 512-token sequence, the slow pair has advanced only $512 \times 0 . 01 \approx 5$ radians out of $2 \pi$ - a fraction of its dynamic range, preserving sensitivity for sequences well beyond 512.

This exact multi-scale geometry is what motivated RoPE’s paired-dimension rotation design, described in the next subsection.

What You Actually Decide in Practice

• Fixed context, no extrapolation needed. Learned absolute embeddings are the right call. They are trivially fast (one embedding lookup), add a modest $L \times d$ parameter count (for $L = 2048$, $d = 4096$ that is $8$ M parameters - 0.1% of a 7B model), and have no inference overhead.

• Any use case that might require sequences longer than training context. Do not use absolute learned embeddings. The failure is not a smooth performance cliff; it is a hard crash. If you anticipate context growth post-launch, paying the cost of RoPE from day one is always worth it.

• Sinusoidal for new architectures. Generally not worth choosing over RoPE. Sinusoidal requires no learned parameters, which is attractive, but it does not enable any of the context-extension algorithms that RoPE does. There is no upside to picking sinusoidal in 2024 unless you are constrained to not touching the attention layer.

• T5 relative bias at context lengths above 4k. The $O ( n^{2} )$ attention bias tensor becomes the dominant memory cost before the attention matrix itself. Profile memory before committing.

• ALiBi for simple long-context baselines. Zero parameters, strong empirical generalization, and simple to implement. The recency-bias assumption fails for tasks that need long-range retrieval. Fine for summarization or generation; risky for multi-document QA.

How This Shows Up in Interviews

• “Explain the sinusoidal encoding formula. Why use $10000$ as the base?” Strong answer: describe the geometric series of wavelengths from $2 \pi$ to $10000 \cdot 2 \pi$ covering up to $\sim 63000$ tokens, connect high-frequency pairs to short-range discrimination and low-frequency pairs to long-range, and mention the rotation-matrix property that makes relative offsets linearly expressible. Weak answers quote the formula without explaining the frequency intuition.

• “What breaks when you deploy a model trained with absolute learned embeddings on a longer context?” Strong answer: positions beyond $L - 1$ access uninitialized or randomly initialized embedding rows; the model produces incoherent logits, typically NaN or near-uniform distributions. Mention that fine-tuning on longer sequences with the position table extended can fix this, but it requires continued training, not just architecture changes.

• Senior follow-up: “Relative encodings claim to generalize to unseen lengths. Show me where that guarantee breaks down in practice.” The catch: T5 clips relative distances at some maximum offset $k$; beyond $k$ tokens all positions are treated identically, so the model has no signal to distinguish “100 tokens away” from “10000 tokens away.” ALiBi does not clip, but its monotone penalty assumes recency always wins, which is wrong for long-context retrieval. True unlimited-length generalization requires the encoding to remain distinguishable at arbitrary distances - which is the design goal of RoPE’s continuous-frequency rotation.

Key Takeaways

• Transformers have no intrinsic sense of order; position must be injected either into token representations (absolute) or into attention scores (relative).

• Sinusoidal embeddings cover multiple frequency scales and have a linear-combination relative-offset property, but they do not generalize to sequences longer than training context.

• Learned absolute embeddings are fast and simple but fail immediately at any sequence length beyond the declared maximum.

• T5-style relative biases generalize better but impose $O ( n^{2} )$ memory overhead; ALiBi is cheaper but bakes in a recency prior.

• The entire landscape above has been superseded by RoPE in post-2023 production models; understanding these predecessors matters primarily to explain why RoPE won.

3.4.2 Rotary Position Embeddings (RoPE) and Context Extension

Every frontier model released since 2022 - LLaMA 1 through 4, Mistral, Qwen, Gemma, DeepSeek, Falcon - uses Rotary Position Embeddings. If you are asked to design an LLM from scratch, choosing RoPE is almost certainly the right answer. But stating the choice is the easy part; an interviewer will immediately follow up with “explain exactly why the inner product only sees relative position,” and then, if you are interviewing for a senior role, “how would you extend a 4k RoPE model to 128k without retraining from scratch?” Both questions require understanding the mathematics, not just the name.

The design goal is a position encoding $f ( \mathbf{x} , m )$ such that for any two token vectors $\mathbf{q}$ and $\mathbf{k}$ at positions $m$ and $n$:

The inner product must depend on the absolute positions only through their difference $m - n$. This rules out additive embeddings (which introduce cross terms between content and absolute position) and rules out most relative schemes (which modify attention scores but do not satisfy this exact factorization for inner products).

The observation that makes RoPE work is that inner products are invariant under shared rotations: $\langle R \mathbf{u} , R \mathbf{v} \rangle = \langle \mathbf{u} , \mathbf{v} \rangle$ for any rotation matrix $R$. If position $m$ applies rotation $\mathcal{R}_{m}$ and position $n$ applies rotation $\mathcal{R}_{n}$, then:

because $\mathcal{R}_{m}^{T} \mathcal{R}_{n} = \mathcal{R}_{n - m}$ for rotation groups. The right-hand side depends only on $m - n$, satisfying the design goal exactly.

In practice, rotating a $d$-dimensional vector as a single unit requires a $d \times d$ matrix, which is expensive. RoPE reduces this by pairing up dimensions: the vector is sliced into $d / 2$ pairs, and pair $k$ at position $m$ is rotated by angle $m \cdot \theta_{k}$ in the 2D plane of that pair, where:

The rotation for pair $k$ at position $m$ is:

This is applied to queries and keys at every attention layer, immediately after the linear projections for $Q$ and $K$. Values are not rotated; they carry no positional information directly.

RoPE is the only common encoding that provably makes the attention inner product a function of relative distance alone, while requiring zero learnable parameters and no change to token embeddings.

Compared to absolute learned embeddings, RoPE adds zero parameters and produces no out-of-distribution vectors at unseen absolute positions - the rotation angles grow continuously, so position 10001 is just a small step beyond position 10000.

Compared to sinusoidal additive embeddings, RoPE’s rotation acts on the query and key vectors inside the attention layer, which means the dot-product kernel sees only relative angles. Sinusoidal embeddings added to token inputs do not achieve this: the product $( \mathbf{e}_{i} + \mathbf{p}_{i} )^{T} ( \mathbf{e}_{j} + \mathbf{p}_{j} )$ expands into four terms, two of which mix content with absolute position - the encoding is not purely relative.

Compared to T5 relative biases, RoPE requires no $O ( n^{2} )$ memory. The rotation is a per-vector operation (cost $O ( n \cdot d )$), not a per-pair operation. At $n = 32768$ and $d = 4096$, this difference is $32768 \times 4096 \approx 134$ M elements versus $32768^{2} \approx 1$ B elements - roughly an 8$\times$ reduction in the position-related memory footprint.

The geometric spacing of $\theta_{k} = 10000^{- 2 k / d}$ provides the same multi-scale frequency coverage as sinusoidal embeddings: high-$k$ pairs (small $\theta$) oscillate slowly and encode long-range position; low-$k$ pairs oscillate quickly and encode fine-grained adjacency. This design also enables length-extension algorithms discussed below.

A Worked Example

Consider $d = 4$, giving two dimension pairs. The base frequencies are:

Take $\mathbf{q} = [ 1 , 0 , 0 , 1 ]$ at position $m = 3$ and $\mathbf{k} = [ 0 , 1 , 1 , 0 ]$ at position $n = 5$, so the relative gap is $n - m = 2$.

Rotating $\mathbf{q}$ at position 3:

• Pair 0: angle $3 \times 1 . 0 = 3$ rad. $( \cos 3 , \sin 3 ) \approx ( - 0 . 990 , 0 . 141 )$.

• Pair 1: angle $3 \times 0 . 01 = 0 . 03$ rad. Rotating $( 0 , 1 )$: $( - \sin 0 . 03 , \cos 0 . 03 ) \approx ( - 0 . 030 , 1 . 000 )$. So $\mathcal{R}_{3} \mathbf{q} \approx [ - 0 . 990 , 0 . 141 , - 0 . 030 , 1 . 000 ]$.

Rotating $\mathbf{k}$ at position 5:

• Pair 0: angle $5 \times 1 . 0 = 5$ rad. Rotating $( 0 , 1 )$: $( - \sin 5 , \cos 5 ) \approx ( 0 . 959 , 0 . 284 )$.

• Pair 1: angle $5 \times 0 . 01 = 0 . 05$ rad. Rotating $( 1 , 0 )$: $( \cos 0 . 05 , \sin 0 . 05 ) \approx ( 0 . 999 , 0 . 050 )$. So $\mathcal{R}_{5} \mathbf{k} \approx [ 0 . 959 , 0 . 284 , 0 . 999 , 0 . 050 ]$.

Inner product:

Now verify the same $\mathbf{q}$ and $\mathbf{k}$ at absolute positions 0 and 2 (same relative gap of 2): $\mathcal{R}_{0} \mathbf{q} = \mathbf{q} = [ 1 , 0 , 0 , 1 ]$. Rotating $\mathbf{k}$ at position 2: pair 0 gives $( - \sin 2 , \cos 2 ) \approx ( - 0 . 909 , - 0 . 416 )$; pair 1 gives $( \cos 0 . 02 , \sin 0 . 02 ) \approx ( 1 . 000 , 0 . 020 )$.

The inner products are identical. The absolute positions changed from $( 3 , 5 )$ to $( 0 , 2 )$, but the relative gap stayed at 2 and the inner product is unchanged.

What You Actually Decide in Practice

• RoPE base $\theta$. The original paper used $10000$. LLaMA 3 raised it to $500 , 000$. A higher base stretches the period of the slow-frequency pairs, giving the model more resolution at long distances. Increasing the base is the simplest and most reliable way to improve long-context performance when combined with continued pretraining on long documents.

• Context extension without retraining: position interpolation. Chen et al. (2023) showed that dividing all absolute positions by a scale factor $s = L_{\text{new}} / L_{\text{train}}$ before applying RoPE keeps rotation angles within the trained range. The model sees positions that now look like scaled versions of what it trained on. This works but produces a mild perplexity increase; 1000-2000 fine-tuning steps on long-context data almost fully recover quality.

• NTK-aware scaling (Bloc97). Instead of scaling positions, scale the base: replace $10000$ with $10000 \cdot s^{d / ( d - 2 )}$. This redistributes the frequency spectrum so that high-frequency pairs are less affected, preserving short-range sensitivity while stretching the slow pairs. Works better than naive position interpolation out of the box, with no fine-tuning.

• YaRN (Peng et al., 2023). Applies different interpolation factors per frequency group: high-frequency pairs are left unchanged, low-frequency pairs are interpolated, mid-frequency pairs get a temperature adjustment. The most reliable technique for large extension factors ($4 \times$ or more) with minimal fine-tuning. Used in several open-weight long-context derivatives.

• RoPE applies only to $Q$ and $K$, not $V$ . A common implementation bug is accidentally rotating the value projections. Values carry content, not position; rotating them corrupts them.

How This Shows Up in Interviews

• “Why does RoPE make the attention score depend only on relative position?” Strong answer: the rotation $\mathcal{R}_{m}$ applied to $Q$ and $\mathcal{R}_{n}$ applied to $K$ means $\langle \mathcal{R}_{m} \mathbf{q} , \mathcal{R}_{n} \mathbf{k} \rangle = \langle \mathcal{R}_{m - n} \mathbf{q} , \mathbf{k} \rangle$ because the rotations compose as $\mathcal{R}_{m}^{T} \mathcal{R}_{n} = \mathcal{R}_{n - m}$. Work through this in 2D if asked on a whiteboard. Weak answers say “it uses rotation” without the inner product cancellation argument.

• “What breaks when you try to run a RoPE model trained with context 4k on an 8k input?” Strong answer: positions 4001-8000 produce rotation angles that the model has never encountered during training. Unlike absolute embeddings (hard failure), this degrades gradually: perplexity rises as position increases past training context, with the slow-frequency pairs (large $k$) being the first to degrade since they have the least trained coverage. The fix is either NTK-aware base scaling or continued pretraining with longer sequences.

• Senior follow-up: “You have a LLaMA-2 7B checkpoint trained at 4k context and you need 32k context in two weeks. Walk me through what you would actually do.” The strong answer sketches a two-stage recipe: (1) apply NTK-aware base scaling (e.g., raise base to $\approx 128 , 000$ for $8 \times$ extension) with no retraining as an immediate fix; (2) collect a long-document corpus at the target length, run 1000-2000 gradient steps at $32$k sequence length on 8-16 GPUs, using sequence packing and gradient checkpointing to fit memory; (3) evaluate on SCROLLS or LongBench to confirm perplexity recovery. Name the memory cost: at bf16 with a 7B model, a single 32k sequence activates $\approx 32 \times 32^{2} \times 4096 / 10^{9} \approx 4 . 3$ GB of KV activations per layer, so activation checkpointing is not optional.

Key Takeaways

• RoPE rotates $Q$ and $K$ by position-dependent angles so that the attention inner product becomes a function of relative position alone, with zero learnable parameters.

• The dimension pairing with base-10000 frequencies provides multi-scale resolution, just like sinusoidal embeddings, but applied inside the attention operation rather than at the input.

• Virtually all post-2022 production models use RoPE; convergent evidence from 19+ open model releases shows no competing scheme has beaten it at scale.

• Context extension methods - position interpolation, NTK-aware base scaling, YaRN - all work by adjusting how positions map to rotation angles; they differ in how aggressively they disturb the short-range frequency pairs that the model relies on most.

• The simplest production strategy for extending context: increase the RoPE base and fine-tune on long documents; for large extension factors ($> 4 \times$), use YaRN or NTK scaling as the starting point.