The Transformer Architecture

Background

Vaswani et al. (2017) "Attention is All You Need" introduced the Transformer architecture, replacing RNNs entirely
Self-attention computes dependencies between all pairs of tokens simultaneously - unlike RNNs which are sequential
BERT, GPT, T5, and all modern LLMs are Transformer-based

The Concept

The Transformer processes sequences by having each token attend to every other token simultaneously via scaled dot-product attention. Multi-head attention runs several attention operations in parallel. Positional encodings add sequence order since there is no recurrence.

How it Links

Attention mechanism is the core operation - see the Attention page for deep dive
BERT uses bidirectional encoder Transformers; GPT uses causal decoder Transformers
Every modern LLM, vision model (ViT), and multimodal model (CLIP) uses Transformer blocks

Concrete Example

Translating "The bank can guarantee deposits" (bank = river bank or financial bank?): the Transformer attends from the word "bank" to "deposits" and "guarantee" simultaneously to resolve the ambiguity. An RNN would need to maintain this context across many sequential steps - Transformers do it in one shot.

Prerequisites

Attention Mechanisms, RNNs, LSTMs, GRUs, Backpropagation & Optimizers, Regularization in Deep Nets

Learning outcomes

Diagram the Vaswani 2017 encoder block, decoder block, and full encoder-decoder, naming every sublayer and residual connection.
Compare sinusoidal, learned, RoPE, and ALiBi positional encodings on length-extrapolation, parameter count, and inference cost.
Explain why pre-norm trains stably at depth and post-norm needs warmup, citing Xiong 2020.
Derive Transformer FLOPs and parameters per block in terms of d_model, d_ff, num_heads, sequence length, and depth.
Compute KV-cache memory and prefill vs decode FLOPs for autoregressive inference, and explain the systems consequences.

TL;DR

A Transformer block is multi-head self-attention plus a token-wise feed-forward network, each wrapped in a residual + layer-norm. The encoder is a stack of these; the decoder adds masked self-attention plus encoder-decoder cross-attention. Position is injected explicitly: Vaswani 2017 used sinusoidal encodings [1], modern LLMs prefer RoPE [2] for clean rotation-based locality and ALiBi [3] for length extrapolation. Pre-norm (LayerNorm inside the residual) trains stably at 100+ layers without learning-rate warmup; post-norm needs warmup but slightly outperforms at small depth [4]. Inference splits into a one-time prefill (full $O(n^2)$ pass over the prompt) and a per-token decode that reads from a KV cache of size $2 L H d n$ floats. The cache is what limits practical context length, not flops. Variants: Transformer-XL [5] for segment recurrence, Sparse Transformer [6] for $O(n \sqrt{n})$ attention, and Mixture-of-Experts [7][8] for sparse parameter scaling.

Intuition

The Transformer is a stack of “attend, then mix” blocks. Attend lets every token gather information from every other token via scaled dot-product self-attention. Mix is a small two-layer MLP applied independently to each token’s vector, which lets the network do non-linear feature combination after the attention has gathered the relevant evidence. The Transformer is the architecture that took deep learning [10] from “wins on vision and speech” to “wins on essentially every sequence-modelling task that has data”.

Two things make the architecture practical at scale. First, every operation is fully parallel across the sequence dimension: there is no recurrence to break, so an entire sequence trains in one matmul-heavy forward and backward pass. Second, residual connections around every sublayer give gradients an unobstructed path to the embedding layer, the same idea ResNet used to make 100+ layer vision networks trainable. Stacking 96 Transformer blocks (GPT-3 scale) is no harder than stacking 6 (the original Vaswani 2017 base model), modulo positional-encoding and norm-placement choices.

Attention gathers, MLP transforms, residual remembers

Every Transformer block does the same three jobs:

Attention is a content-based gather across positions: each token reads a weighted summary of all other tokens.
MLP is a position-wise non-linear transform: each token vector is run through a 2-layer MLP with hidden size $d_{\mathrm{ff}} \approx 4 d$ , in parallel across positions.
Residual + LayerNorm keeps the signal flowing: the block computes a delta added to the residual stream, the layer-norm renormalizes for the next block.

A clean way to read a Transformer is “the residual stream is the working memory; each block reads from it (attention), processes (MLP), and writes back (residual add).” The interpretability literature has formalized this view; treat the residual stream as the bus and each block as a peripheral that reads, transforms, and writes.

The Vaswani 2017 architecture

The original paper specified three things together: an encoder block, a decoder block, and a sinusoidal positional encoding. We will walk through each.

The encoder block

Inputs: a sequence of token embeddings $X \in \mathbb{R}^{n \times d}$ with positional encoding already added. One encoder block computes

\begin{aligned} Z &= \mathrm{LayerNorm}(X + \mathrm{MHA}(X)), \\ X' &= \mathrm{LayerNorm}(Z + \mathrm{FFN}(Z)), \end{aligned}

where $\mathrm{MHA}$ is multi-head self-attention over $X$ and $\mathrm{FFN}$ is a two-layer MLP applied to each row independently:

\mathrm{FFN}(z) = W_2 \, \mathrm{GeLU}(W_1 z + b_1) + b_2, \qquad W_1 \in \mathbb{R}^{d_{\mathrm{ff}} \times d},\; W_2 \in \mathbb{R}^{d \times d_{\mathrm{ff}}}.

The original used ReLU and $d_{\mathrm{ff}} = 4 d$ ; modern Transformers swap to GeLU or SwiGLU and keep the $4d$ rule of thumb [9]. Note that the layer-norm appears after the residual add: this is post-norm, the original choice. Most modern LLMs flip it (see Pre-norm vs post-norm below).

A full encoder is $L$ identical blocks stacked. There is no causal mask (encoders are bidirectional), so token $i$ can attend to every other token at every layer.

The decoder block

The decoder takes the previously generated tokens (during training: the gold target shifted right; during inference: its own past outputs), embeds them, adds positional encoding, then runs a stack of decoder blocks. One decoder block has three sublayers:

\begin{aligned} Y_1 &= \mathrm{LayerNorm}(Y + \mathrm{MaskedMHA}(Y)), \\ Y_2 &= \mathrm{LayerNorm}(Y_1 + \mathrm{CrossMHA}(Q = Y_1,\, K = X_{\mathrm{enc}},\, V = X_{\mathrm{enc}})), \\ Y' &= \mathrm{LayerNorm}(Y_2 + \mathrm{FFN}(Y_2)). \end{aligned}

The first sublayer is masked self-attention: the upper-triangular mask ensures position $i$ only attends to positions $1, \dots, i$ on the decoder side. The second sublayer is cross-attention: queries come from the decoder, keys and values come from the encoder output. The third is the same FFN as the encoder. Decoder-only models (GPT family) drop the cross-attention and keep just (masked self-attention, FFN); encoder-only models (BERT) drop the decoder entirely. See Encoder, Decoder & Encoder-Decoder Models for the full taxonomy.

Parameter and FLOP accounting

Per encoder block, parameters are dominated by:

$Q, K, V, O$ projections: $4 d^2$ params total.
Two FFN layers: $2 \cdot d \cdot d_{\mathrm{ff}} = 8 d^2$ params if $d_{\mathrm{ff}} = 4 d$ .
LayerNorm scales/biases: $4 d$ (negligible).

Total: $\approx 12 d^2$ params per block. A decoder block adds another set of $4 d^2$ for cross-attention, taking it to $\approx 16 d^2$ . For a base Transformer with $d = 512$ , $L = 6$ , this is about $19$ M params per encoder stack.

FLOPs per token per encoder block (forward only):

Attention: $O(n d)$ for the four projections, $O(n^2 d)$ for the score and value matmul. For the full sequence it is $O(n^2 d)$ .
FFN: $O(n d \cdot d_{\mathrm{ff}}) = O(8 n d^2)$ .

The crossover where attention overtakes the FFN in flops is at $n \approx 8 d$ . For $d = 4096$ (GPT-3-class) that’s $n \approx 32{,}000$ . So at typical training context lengths the FFN still dominates flops; attention starts to dominate only for long-context training.

Positional encodings

Pure self-attention is permutation-equivariant (proof in Attention), so the model has no notion of order until you inject one. The four production choices:

Sinusoidal, learned, RoPE, and ALiBi: the four positional encodings

Sinusoidal (Vaswani 2017)

Add a fixed sinusoidal encoding to the input embeddings before the first block:

\mathrm{PE}_{(\mathrm{pos}, 2i)} = \sin\!\left( \frac{\mathrm{pos}}{10000^{2 i / d}} \right), \qquad \mathrm{PE}_{(\mathrm{pos}, 2 i + 1)} = \cos\!\left( \frac{\mathrm{pos}}{10000^{2 i / d}} \right).

The encoding is deterministic (no parameters) and has an attractive property: the encoding for position $\mathrm{pos} + k$ is a fixed linear function of the encoding at position $\mathrm{pos}$ , so relative offsets are linearly recoverable. Empirically the model can extrapolate moderately past the training length but degrades quickly past 2x.

Learned absolute

Replace the sinusoid with a trainable lookup table $\mathrm{PE} \in \mathbb{R}^{n_{\max} \times d}$ . BERT and GPT-2 both used this. Pros: simple, the model can learn whatever positional structure is useful. Cons: cannot extrapolate past $n_{\max}$ at all; doubling context length requires retraining.

RoPE: rotary position embedding

RoPE [2] applies a 2D rotation to query and key vectors per position, inside the attention computation rather than as an additive embedding. Pair the components of $q$ and $k$ into 2D vectors $(q_{2i}, q_{2i+1})$ . For position $m$ , multiply each pair by the rotation matrix

R_{\theta_i, m} = \begin{pmatrix} \cos(m \theta_i) & -\sin(m \theta_i) \\ \sin(m \theta_i) & \cos(m \theta_i) \end{pmatrix},

with $\theta_i = 10000^{-2 i / d}$ (the same frequencies as the sinusoidal encoding). The crucial identity:

\langle R_{\theta, m}\, q,\; R_{\theta, n}\, k \rangle = \langle q,\; R_{\theta, n - m}\, k \rangle.

The dot product depends only on the relative offset $n - m$ , not the absolute positions. So RoPE-attention is automatically translation-invariant in position, which matches how relative position should affect attention. RoPE is what Llama, Mistral, Qwen, and most open LLMs use.

Cost: zero parameters, but each attention call must apply the rotation to $Q$ and $K$ . A few extra flops, negligible memory.

ALiBi: attention with linear biases

ALiBi [3] goes the opposite direction: leave $Q$ and $K$ alone, and instead add a linear bias to the score matrix:

\mathrm{score}_{ij} = \frac{q_i^\top k_j}{\sqrt{d_k}} - m_h \cdot |i - j|,

where $m_h$ is a per-head fixed slope (the paper gives a geometric schedule $m_h = 2^{-8 h / H}$ for head $h$ of $H$ ). Closer keys get smaller penalty; farther keys get larger penalty. This is a hard-coded “attend to nearby tokens more” prior with zero parameters.

The win that earned ALiBi its keep is length extrapolation: a model trained at $n = 1024$ generates coherent text at $n = 8192$ at inference. RoPE used to extrapolate poorly until the “RoPE scaling” tricks (NTK-aware, YaRN, dynamic) closed most of the gap. As of 2024 RoPE+scaling has won on the open-LLM side, but ALiBi remains competitive in some long-context-from-day-one regimes.

Encoding	Parameters	Where applied	Length extrapolation	Used by
Sinusoidal (Vaswani)	0	Added to input embedding	Modest (degrades past 2x training length)	Original Transformer, T5 (relative variant)
Learned absolute	n_max * d	Added to input embedding	None (hard limit at n_max)	BERT, GPT-2
RoPE	0	Multiplicative rotation on Q, K inside attention	Good with NTK/YaRN scaling tricks	Llama, Mistral, Qwen, GPT-NeoX
ALiBi	0 (per-head fixed slope)	Additive bias to score matrix	Excellent out of the box	BLOOM, MPT
T5 relative	O(num_buckets * h)	Additive bias to score matrix indexed by relative offset bucket	Good	T5, mT5

Pre-norm vs post-norm

LayerNorm placement is the single most important “small” knob in the Transformer.

Why pre-norm trains stably at depth (Xiong 2020)

The original Transformer used post-norm: $\mathrm{LN}(X + \mathrm{Sublayer}(X))$ . The norm sits after the residual add. Xiong et al. [4] showed that for a deep stack, post-norm has gradient magnitude that depends on depth in a way that requires careful learning-rate warmup; without warmup, training diverges past about 12 layers.

Pre-norm moves the norm inside the residual: $X + \mathrm{Sublayer}(\mathrm{LN}(X))$ . The proof in Xiong 2020 shows the gradient at the input is bounded by $1 + O(L^{-1})$ instead of $O(\sqrt{L})$ , so deep stacks have stable gradients without warmup. Empirically pre-norm trains 100+ layer models with no fuss; this is the architecture every modern LLM uses (Llama, GPT, Mistral).

The trade-off: pre-norm sometimes gives slightly worse final loss at small depth (12 layers and below), because the un-normalized residual stream can grow large and the model has less direct control over the sublayer input scale. For LLM scale (depth $\ge 24$ ) pre-norm is strictly better.

Property	Post-norm (original)	Pre-norm (modern)
Formula	X_{l+1} = LN(X_l + Sublayer(X_l))	X_{l+1} = X_l + Sublayer(LN(X_l))
Gradient at input as a function of depth L	O(sqrt(L))	1 + O(1/L)
Warmup required at depth?	Yes (Vaswani used 4000 steps)	No
Trainable past 24 layers without tricks?	No	Yes
Final loss at small depth (6-12 layers)	Slightly better	Slightly worse
Final loss at LLM depth (24+ layers)	Diverges or needs careful tricks	Strictly better
Used by	Original Transformer, BERT, T5	GPT-2, GPT-3, Llama, Mistral, most modern LLMs

KV cache and inference math

Training and inference look almost identical for a Transformer, but inference has one structural difference: autoregressive decoding generates one token at a time, and naively re-running the full forward pass for each token would cost $O(L n^2 d)$ per token. The KV cache fixes this.

Prefill vs decode and KV-cache memory

Two phases of inference

Generating a response from a prompt of length $P$ followed by $T$ generated tokens splits into:

Prefill. Run the full Transformer forward on the $P$ -token prompt in one shot. Compute the keys and values for every position at every layer and every head, and cache them in the KV cache. Cost: $O(L \cdot P^2 \cdot d)$ FLOPs. Done once.
Decode. For each generated token $t = 1, \dots, T$ : project the new token to $q_t, k_t, v_t$ , append $k_t, v_t$ to the cache, compute attention by reading the entire cache, and produce the next token. Cost per generated token: $O(L \cdot (P + t) \cdot d)$ FLOPs (linear in current sequence length). Done $T$ times.

So total FLOPs are roughly $L (P^2 + (P + T)^2 - P^2) d = L (P + T)^2 d$ , but with a very different memory access pattern. Prefill is matmul-heavy (parallel across $P$ ). Decode is memory-bandwidth-heavy: each generated token reads the entire cache, does a few small matmuls, writes one row back.

KV cache size

Per layer, per head, each token cached holds $d_k$ floats for the key and $d_v$ floats for the value. With $L$ layers, $H$ heads, $d_k = d_v = d / H$ , and sequence length $n$ , the cache size is

\mathrm{KV\,cache\,bytes} = 2 \cdot L \cdot H \cdot \frac{d}{H} \cdot n \cdot \mathrm{bytes\,per\,float} = 2 L d n \cdot \mathrm{bytes\,per\,float}.

In bf16 (2 bytes), a 70B Llama-3 with $L = 80$ , $d = 8192$ at $n = 8192$ has

2 \cdot 80 \cdot 8192 \cdot 8192 \cdot 2 \approx 21 \text{ GB}

per request, just for the KV cache. That is why batched inference is hard: each user gets their own cache, and aggregating ten users instantly fills an 80 GB H100. Grouped-query attention (GQA) and multi-query attention (MQA) cut this by sharing $K, V$ across groups of heads; Llama-3-70B uses GQA with 8 KV-heads instead of 64, cutting cache to about 2.6 GB.

Decode latency is bandwidth-bound

Generating one token reads the entire KV cache from HBM. At 21 GB cache and 1 TB/s HBM, that is 21 ms minimum per token, before any compute, just from memory transfer. This is why long-context decoding is slow per token even though compute is low: you cannot escape the bandwidth read. FlashDecoding, paged-attention (vLLM), and speculative decoding are all systems-level fixes for this fundamental constraint.

Modern variants

A small zoo of Transformer variants has appeared since 2017. The four worth knowing for interview depth:

2019Transformer-XLpaper
Segment-level recurrence: cache hidden states from the previous segment as additional context, with relative position encoding. Lets the model carry context across segments without retraining on longer sequences.
2019Sparse Transformer (Child et al.)efficient
Replace full attention with strided + local sparse patterns. Cuts attention cost from O(n^2) to O(n sqrt(n)) while retaining most quality on long-context tasks.
0.9980207817911925Mixture-of-Experts (Shazeer 2017, Switch Transformer 2021)architecture
Replace the FFN with a router over E expert FFNs; each token activates the top-1 or top-2 experts. Total parameters scale with E, FLOPs per token stay constant.
2022FlashAttentionsystem
Exact attention with O(n) memory via tiled softmax in SRAM. Drop-in replacement, ubiquitous in modern LLM training and inference. Detail: see Attention page.

Mixture-of-Experts in one paragraph

Replace the FFN sublayer with $E$ “expert” FFNs and a learned router: a small linear that produces routing logits $g(x) \in \mathbb{R}^E$ per token. Top- $k$ routing picks the $k$ experts with the largest logits and runs only those. The output is the (renormalized) weighted sum of expert outputs. Switch Transformer [8] uses $k = 1$ (one expert per token, simplest possible routing). The result is a model with $E \cdot$ FFN parameters but the same FLOPs per token as a dense model, since each token activates only one expert. Mixtral-8x7B is 8 experts of 7B-FFN each; total params 47B, active params 13B. The catch is load balancing: without an auxiliary loss, the router collapses to a few favorite experts and the rest go untrained.

A from-scratch encoder block

A complete encoder block in about 80 lines of PyTorch:

python

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

class MultiHeadSelfAttention(nn.Module):
  def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0):
      super().__init__()
      assert d_model % n_heads == 0
      self.n_heads = n_heads
      self.d_k = d_model // n_heads
      # Fused QKV projection: one big linear instead of three.
      self.qkv = nn.Linear(d_model, 3 * d_model)
      self.proj = nn.Linear(d_model, d_model)
      self.dropout = nn.Dropout(dropout)

  def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
      B, N, D = x.shape
      qkv = self.qkv(x).reshape(B, N, 3, self.n_heads, self.d_k)
      q, k, v = qkv.unbind(dim=2)              # each (B, N, H, d_k)
      q, k, v = (t.transpose(1, 2) for t in (q, k, v))  # (B, H, N, d_k)
      # PyTorch dispatches to FlashAttention on supported GPUs.
      out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
      out = out.transpose(1, 2).reshape(B, N, D)
      return self.dropout(self.proj(out))


class TransformerEncoderBlock(nn.Module):
  """Pre-norm encoder block: x + Sublayer(LN(x))."""

  def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.0):
      super().__init__()
      self.ln1 = nn.LayerNorm(d_model)
      self.attn = MultiHeadSelfAttention(d_model, n_heads, dropout)
      self.ln2 = nn.LayerNorm(d_model)
      self.ffn = nn.Sequential(
          nn.Linear(d_model, d_ff),
          nn.GELU(),
          nn.Linear(d_ff, d_model),
      )
      self.dropout = nn.Dropout(dropout)

  def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
      x = x + self.attn(self.ln1(x), mask)
      x = x + self.dropout(self.ffn(self.ln2(x)))
      return x


# Sanity: forward a dummy batch through a 6-layer encoder.
torch.manual_seed(0)
x = torch.randn(4, 16, 512)
blocks = nn.Sequential(*[TransformerEncoderBlock(512, 8, 2048) for _ in range(6)])
y = blocks(x)
print(y.shape)   # (4, 16, 512)

Pre-norm pattern. The fused QKV linear shaves one kernel launch off the hot path. F.scaled_dot_product_attention is the production exact-attention kernel and dispatches to FlashAttention on supported GPUs.

Transformer vs RNN vs CNN for sequence modelling

Property	RNN / LSTM	CNN (1D)	Transformer
Parallelism across time	No (sequential h_t depends on h_{t-1})	Yes (each conv is parallel)	Yes (full sequence in one matmul)
Long-range dependency	Limited (vanishing gradient even with LSTM gating)	Receptive field grows linearly with depth	Direct any-to-any in one layer
Per-layer compute	O(n d^2)	O(n k d^2)	O(n^2 d) attn + O(n d^2) FFN
Memory at training	O(n d) activations	O(n d) activations	O(n^2) attention activations (or O(n) with FlashAttention)
Inductive bias	Recency, sequential dependence	Locality, translation equivariance	None (data-driven; needs positional encoding)
Inference autoregressive cost	O(d^2) per token (RNN state update)	O(k d^2) per token	O(n d) per token (KV cache read)
Parameter count for fixed budget	Smallest	Mid	Largest (FFN dominates)
Why it eventually won	Dominated NLP through 2016	Dominated NLP-via-CNN through 2017 (ConvS2S, ByteNet)	Parallelism + scale wins (Vaswani 2017 onward)

Pitfalls

Forgetting positional encoding

A vanilla Transformer block has no notion of order. If you swap the input from sinusoidal-PE BERT to RoPE Llama by changing only the embeddings, the rest of the architecture has nothing to anchor positions. Always check that every attention call has a positional mechanism applied: either to the input embedding (sinusoidal, learned) or inside attention (RoPE, ALiBi, T5 relative). The bug is silent: training looks normal at first, then plateaus.

Weight tying: same matrix, different roles

The original Transformer ties the input embedding matrix to the output projection (the “softmax weights”): both are $W \in \mathbb{R}^{V \times d}$ . This halves the parameter count and ties the geometry of the input and output token spaces. Most modern LLMs untie them because at scale the savings are marginal but the constraint hurts (Llama 3 is untied; PaLM is untied; Gemma is tied for the 2B and untied for the 7B+). When loading a pretrained model, check the tied/untied state in the config; passing tied weights to an untied model is a silent shape mismatch.

KV cache OOM is the real long-context limit

You will hit OOM on the KV cache long before you hit OOM on weights. A 13B model in bf16 takes 26 GB for weights, and at $n = 8192$ takes another $\sim 4$ GB for the cache; at $n = 32768$ it is $\sim 16$ GB. Multiply by batch size for serving. Standard fixes: grouped-query attention (Llama-3), paged-attention (vLLM, SGLang) which fragments the cache into pages so different requests can share blocks, and quantizing the cache itself to int8 or int4. Bandwidth bound at decode time means even a “free” cache hit costs latency.

Pre-norm if you can, post-norm only with warmup

If you are training a Transformer from scratch at depth $\ge 12$ , use pre-norm and skip the warmup schedule. Pre-norm + a constant or cosine learning rate is the modern default and trains fine to hundreds of layers. Post-norm + 4000-step warmup matches the original recipe and is fine at small depth, but it is a fragile combo in production: a warmup-step misconfigure on a checkpoint resume will diverge silently [4].

Pop quiz

Check yourself

Score: 0 / 5(0 answered)

1. A Transformer encoder block has approximately how many parameters in terms of d_model, with d_ff = 4 * d_model?
2. Why does pre-norm train more stably than post-norm at depth above 24 layers?
3. What is the KV cache memory for a Transformer with L layers, d model dimension, sequence length n, in bf16?
4. RoPE applies position information by:
5. Mixture-of-Experts (Switch Transformer) replaces the FFN with E experts and top-1 routing. The effect on parameters and FLOPs per token is:

Interview Q&A

Derive Transformer FLOPs per training step in terms of d_model, d_ff, sequence length n, depth L, and batch size B.onsite

Per layer, per token, forward. Two contributions:

Attention. Q, K, V projections: $3 \cdot d^2$ FLOPs per token (each is a $d \times d$ matmul). Output projection: $d^2$ . Score matmul $QK^\top$ : $n \cdot d$ FLOPs per token. Softmax: $n$ FLOPs per token (cheap). Multiply by V: $n \cdot d$ per token. Total attention: $4 d^2 + 2 n d$ FLOPs per token.
FFN. Two matmuls: $d \cdot d_{\mathrm{ff}}$ each, and a typical $d_{\mathrm{ff}} = 4 d$ , so $8 d^2$ FLOPs per token.

Per token, per layer, forward: $\approx 12 d^2 + 2 n d$ FLOPs. Backward is $\approx 2\times$ forward (one re-evaluate plus the gradient matmul), so per token, per layer, total: $\approx 36 d^2 + 6 n d$ FLOPs.

Ignoring the $n d$ attention-score term (which dominates only for $n \gg 8 d$ ):

\mathrm{FLOPs\,per\,training\,step} \approx 36 \cdot B \cdot n \cdot L \cdot d^2.

A useful approximation when scoring estimates: a rule of thumb is $\approx 6 N$ FLOPs per token where $N$ is total parameters (factor 6 = 2 forward + 4 backward in fused form). For a 175B model on 300B tokens, $6 \cdot 175 \mathrm{e}9 \cdot 3 \mathrm{e}{11} \approx 3 \mathrm{e}{23}$ FLOPs total, matching the GPT-3 published number.

The key takeaways are: (i) FFN dominates compute at typical $n$ , (ii) attention overtakes FFN above $n \approx 8 d$ , (iii) doubling depth doubles FLOPs at fixed $d$ , and (iv) doubling $d$ quadruples FLOPs, which is why scaling-laws fits care so much about the relative growth of $d$ and $n$ .

Derive RoPE: show that rotating Q and K by R(theta, pos) makes the attention score depend only on the relative position.onsite

Setup. Pair the components of $q, k \in \mathbb{R}^d$ into $d/2$ 2D vectors. For each pair index $i$ and each token position $m$ , define the rotation

R_{\theta_i, m} = \begin{pmatrix} \cos(m \theta_i) & -\sin(m \theta_i) \\ \sin(m \theta_i) & \cos(m \theta_i) \end{pmatrix},

with frequency $\theta_i = 10000^{-2 i / d}$ . RoPE rotates each pair of $q, k$ by this matrix at its own position before computing the dot product.

Key identity. A 2D rotation is multiplicative in angle: $R_\alpha R_\beta = R_{\alpha + \beta}$ and $R_\alpha^\top = R_{-\alpha}$ . So for query at position $m$ and key at position $n$ :

\langle R_{\theta, m}\, q,\; R_{\theta, n}\, k \rangle = q^\top R_{\theta, m}^\top R_{\theta, n}\, k = q^\top R_{\theta, n - m}\, k.

The right-hand side depends only on the relative offset $n - m$ , not on the absolute positions $m$ and $n$ separately. So translating both query and key by the same offset $\delta$ leaves the dot product unchanged: RoPE-attention is translation-invariant in position.

Why it generalizes. Independence of absolute position is exactly the property a model should have for natural language: “the dog ran” should look the same to attention whether it appears at position 5 or position 5000. With RoPE, any pair of tokens at offset $k$ produces the same dot product regardless of where they sit in the sequence. With absolute (sinusoidal or learned) position, translating the sentence shifts every dot product by some position-dependent amount.

Length extrapolation. RoPE’s dependence on relative offset, plus the smooth decay of high-frequency components past the training length, means a RoPE-trained model has a smooth degradation past $n_{\max}$ rather than a hard cliff. Modern “RoPE scaling” tricks (NTK-aware, YaRN, dynamic) reshape the frequency spectrum so the model extrapolates 4-32x past training length with fine-tuning of just a few hundred steps [2].

Why pre-norm wins for deep nets, mathematically.research

Setup. A pre-norm Transformer applies LayerNorm inside the residual: $x_{l+1} = x_l + f_l(\mathrm{LN}(x_l))$ . A post-norm Transformer applies it outside: $x_{l+1} = \mathrm{LN}(x_l + f_l(x_l))$ .

Forward variance. In pre-norm, the residual stream $x_l$ is never re-normalized between blocks, so its variance can grow: at layer $l$ , $\mathrm{Var}(x_l)$ is approximately $l \cdot \sigma_f^2$ where $\sigma_f^2$ is the per-block sublayer-output variance (independent residual increments). The LN inside the sublayer rescales the input to $f_l$ to unit variance. So sublayer inputs stay normalized while the residual stream grows.

In post-norm, the residual is normalized at every layer, so $\mathrm{Var}(x_l) = O(1)$ . But the sublayer input is the previous post-norm output, which is unit-variance, and the residual + sublayer output then gets re-normalized.

Backward gradient. Differentiate the recurrence. In pre-norm, $\partial x_{l+1} / \partial x_l = I + \mathrm{Jac}_f \cdot \mathrm{Jac}_{\mathrm{LN}}(x_l)$ . The LN Jacobian shrinks the second term as $1 / \lVert x_l \rVert$ , which grows like $\sqrt{l}$ . So the Jacobian product over the full stack is $\prod_l (I + O(1/\sqrt{l}))$ , which Xiong shows is $1 + O(\log L / \sqrt{L})$ at depth $L$ . Bounded.

In post-norm, the LN is on the residual + sublayer output, and its Jacobian factor multiplies the whole residual, not just the sublayer increment. The chain ends up multiplying $L$ independent LN Jacobians of magnitude $1/\sqrt{d}$ , giving a total gradient norm of order $L / \sqrt{d}$ at the input layer. This is why training without warmup diverges past $L \approx 12$ : the gradient at the input is too large at small step counts [4].

Practical consequence. Pre-norm trains stably out of the box at any depth; post-norm needs a warmup of $\sqrt{L}$ steps to compensate. For depth 100+ models (modern LLMs), pre-norm is the only reasonable choice. Empirically pre-norm slightly underfits at small depth (6-12 layers), which is why some legacy NMT and BERT-style models still use post-norm [1].

Compute the KV-cache memory and decode latency for a 70B Llama-3-class model at n = 8192 in bf16. How does grouped-query attention change it?staff

Llama-3-70B params. $L = 80$ , $d_{\mathrm{model}} = 8192$ , $H = 64$ , $d_k = 128$ . With multi-head attention, each layer caches $K$ and $V$ for every head, so per token, per layer the cache is $2 \cdot H \cdot d_k = 2 \cdot 64 \cdot 128 = 16{,}384$ floats. Times 80 layers, 8192 tokens, 2 bytes (bf16):

80 \cdot 8192 \cdot 16{,}384 \cdot 2 \approx 21.5 \text{ GB}.

That is bigger than half the model weights (140 GB in bf16) per request. At batch size 4, $\approx 86$ GB, which already exceeds an 80 GB H100.

With grouped-query attention. Llama-3-70B uses GQA with $H_{\mathrm{kv}} = 8$ KV heads (the 64 query heads are grouped into 8 groups of 8 that share K, V). Cache per layer per token: $2 \cdot 8 \cdot 128 = 2048$ floats. Total at $n = 8192$ :

80 \cdot 8192 \cdot 2048 \cdot 2 \approx 2.7 \text{ GB}.

A 8x reduction, with negligible quality loss. This is what makes serving 70B-class models with multi-thousand-token context realistic on a single 80 GB GPU.

Decode latency. Generating one token reads the entire cache from HBM. With $\sim 1$ TB/s effective bandwidth, the floor is

\frac{2.7 \text{ GB}}{1 \text{ TB/s}} \approx 2.7 \text{ ms / token}

just from KV-cache memory transfer, before any compute. At 21.5 GB without GQA, the floor is 21 ms / token. This is why bandwidth, not flops, dominates LLM decode latency, and why systems work (vLLM paged-attention, FlashDecoding, speculative decoding) targets the read pattern rather than the compute.

A useful cross-check: dense compute per generated token at $n = 8192$ is $\approx 2 \cdot 70 \mathrm{e}9 = 1.4 \mathrm{e}{11}$ FLOPs, which an H100 (about 800 TFLOPs effective at bf16) does in $0.18$ ms. So compute is 10x faster than memory bandwidth at the relevant scale, confirming the bound is bandwidth.

Build it yourself

A weekend-scale project that takes the Transformer from formulas to a working tiny LM:

Implement the from-scratch encoder block above. Verify forward shapes match nn.TransformerEncoderLayer on a fixed seed.
Add a learned absolute positional embedding and train a tiny BERT (4 layers, $d = 128$ , 2 heads) on $10^4$ Wikipedia sentences with masked-language modelling. Get the training loss curve and confirm it goes below $\log V$ .
Swap to RoPE, retrain, and read the loss differential. RoPE typically wins by a small but consistent margin.
Build a 12-layer decoder-only model with KV cache and confirm that adding model.eval() plus the cache cuts per-token decode latency by roughly $L \cdot n / 2$ vs naive recompute-everything decoding.
Stretch: replace the FFN with a 4-expert MoE, train, and observe the parameter / flop trade-off in practice.

References

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., and Liu, Y. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv preprint. arXiv:2104.09864
Press, O., Smith, N. A., and Lewis, M. (2021). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR 2022 (arXiv preprint 2021). arXiv:2108.12409
Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., and Liu, T.-Y. (2020). On Layer Normalization in the Transformer Architecture. ICML 2020. arXiv:2002.04745
Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q. V., and Salakhutdinov, R. (2019). Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context. ACL 2019. arXiv:1901.02860
Child, R., Gray, S., Radford, A., and Sutskever, I. (2019). Generating Long Sequences with Sparse Transformers. arXiv preprint. arXiv:1904.10509
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR 2017. arXiv:1701.06538
Fedus, W., Zoph, B., and Shazeer, N. (2021). Switch Transformer: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. JMLR 2022 (arXiv preprint 2021). arXiv:2101.03961
Goodfellow, I., Bengio, Y., and Courville, A. (2016). Deep Learning. MIT Press. https://www.deeplearningbook.org
LeCun, Y., Bengio, Y., and Hinton, G. (2015). Deep learning. Nature 521, 436-444. doi:10.1038/nature14539

Attention Mechanisms Tokenization & Embeddings