The Transformer Architecture

The previous lesson built the attention mechanism: scaled dot-product attention, multi-head projection, and the positional encoding that restores order to a permutation-equivariant operation. That is the operation. This lesson assembles the architecture: how those operations stack into the encoder–decoder of Attention Is All You Need, how the stack splits into the three families that dominate modern practice, and where its parameters and its quadratic cost actually live.¹

The sublayer wrapper

The atom of the architecture is not a layer but a sublayer wrapper: a residual skip around a function $f$, followed by layer normalization. Every attention block and every feed-forward block in the entire model is wrapped this way, so we fix the convention before stacking anything.

The expansion is not arbitrary. The FFN takes each token's $d$-vector, projects it up into a wider $4d$ space, applies a pointwise nonlinearity there, and projects back down to $d$. The wide hidden layer is where the per-token computation happens: it gives the nonlinearity enough room to compute many independent features of the token before compressing the result. The original paper used $\mathrm{ReLU}$; most modern stacks swap in $\mathrm{GELU}$ (a smooth gate $x\cdot \Phi (x)$, with $\Phi$ the standard-normal CDF) or a gated variant like $\mathrm{SwiGLU}$, which lowers loss at fixed width for reasons that remain partly empirical.

Track the shapes explicitly. A batch of $B$ sequences of length $n$ enters as a tensor of shape $(B,n,d)$. Every operation below preserves the batch and sequence axes and acts only on the last (feature) axis:

The same $W_1,W_2$ apply at every one of the $B\times n$ positions — the FFN carries no positional index and no cross-token interaction, which is what makes it a position-wise map.

Layer normalization is the other half of the wrapper, so its statistics matter for what follows. Given a single token vector $x\in {\mathbb{R}}^d$, LayerNorm computes the mean and variance over the $d$ feature entries of that one token — not over the batch, not over the sequence — and standardizes:

with learned per-feature scale $\gamma \in {\mathbb{R}}^d$ and shift $\beta \in {\mathbb{R}}^d$ (initialized to $1$ and $0$), and $ϵ\approx 10^{-5}$ for numerical safety. Because the reduction is per-token, LayerNorm is independent of batch size and of the other positions — it behaves identically at training and inference and needs no running statistics — the reason it replaced BatchNorm for variable-length sequence models.² The single addition it contributes is $2d$ parameters ($\gamma$ and $\beta$) per LayerNorm.

The two wrapper styles differ only in where the normalization sits relative to the residual sum.

The distinction looks cosmetic but governs whether a deep stack trains at all.²

Post-norm breaks that identity: each residual sum is squeezed back through a LayerNorm, so the gradient path is a product of $N$ normalization Jacobians and can vanish or explode with depth, which is why post-norm Transformers need a learning-rate warmup and pre-norm ones largely do not.

Positional encoding

Attention is permutation-equivariant: shuffle the input tokens and the output set is the same, merely reordered. That is fatal for language, where dog bites man and man bites dog are different sentences built from identical tokens. The attention operator has no notion of position, so order has to be injected into the inputs before the first layer. The original Transformer adds a fixed sinusoidal positional encoding to the token embeddings.

The $10000^{2i/d}$ denominator sets the wavelength of each feature: the lowest feature ($i=0$) oscillates with period $2\pi$ tokens, and the highest ($i=d/2-1$) with period $2\pi \cdot 10000$ tokens. Low-frequency channels change slowly and encode coarse, long-range position; high-frequency channels distinguish adjacent tokens. Together they give each position a unique fingerprint across a geometric band of scales, the same multi-resolution idea a place-value numeral uses to name any integer with a fixed set of digits.

Two properties make the sinusoidal choice more than a lookup table:

• Relative position is linear. For any fixed offset $k$, the encoding at $\text{pos}+k$ is a fixed linear function (a rotation, from the angle-sum identities) of the encoding at $\text{pos}$. A single learned projection inside an attention head can therefore read relative distance, not just absolute index, which is what a language model actually needs.
• Length extrapolation. Because the formula is defined for every real $\text{pos}$, a model can be evaluated on sequences longer than any it saw in training without a new parameter — the encoding simply continues.

The alternative is a learned positional embedding: a trainable table $E_{\text{pos}}\in {\mathbb{R}}^{n_{\max }\times d}$ with one row per position, added just like a word embedding. Learned embeddings match or slightly beat sinusoids on in-distribution data and are the common choice for encoder-only models (BERT) and ViT, but they cap the context at the table's row count $n_{\max }$ and cannot extrapolate past it. Modern decoder-only models mostly use neither and instead rotate the queries and keys inside attention (rotary embeddings, RoPE), which keeps the linear-relative property while removing the additive table entirely.

One encoder layer

An encoder layer is two wrapped sublayers in series: multi-head self-attention, then a position-wise FFN. With pre-norm wrapping and an input $x$,

The self-attention is unrestricted: every position attends to every other, so each output position is a content-weighted summary of the whole input. This is the right behavior for an encoder: the entire sequence is available and there is no causality constraint.

Every tensor in the block has the same shape $(B,n,d)$: the width $d$ is the invariant of the residual stream, so a layer's output plugs into the next layer's input without any reshaping. Trace one block for a concrete configuration $d=512$, $h=8$ heads (so head width $d_k=d/h=64$), $d_{\text{ff}}=2048$, input $(B,n,512)$:

The only place the shape swells is the $(B,8,n,n)$ score tensor in step 3 and the $(B,n,2048)$ FFN hidden in step 6 — the two costs, quadratic-in-$n$ and wide-FFN, that the accounting section makes precise.

One decoder layer

A decoder layer has three wrapped sublayers. It must generate one token at a time without access to tokens it has not yet produced, and it must read the encoder. Two additions handle these requirements.

• Masked self-attention. Before the softmax, the decoder sets the score $S_{ij}=-\infty$ for every $j>i$, forbidding position $i$ from attending to any later position. Generation stays causal: token $t$ depends only on tokens $<t$.
• Cross-attention. A middle sublayer draws its queries from the decoder state but its keys and values from the encoder's output. This is the channel through which a generated token reads the entire source sequence; in translation it is how each output word reads the sentence being translated.³

The full encoder–decoder

The complete Transformer stacks $N$ encoder layers and $N$ decoder layers, feeds both stacks token embeddings plus positional encodings, and routes the encoder's final output into the cross-attention sublayer of every decoder layer. A linear projection and softmax convert the decoder's top state into a distribution over the vocabulary.

The asymmetry is the whole point. The encoder reads bidirectionally because it sees the input all at once; the decoder reads causally because it produces the output left to right. Cross-attention is the only place the two streams meet.

Causal masking in detail

Masking is implemented additively, before the softmax. Define a mask matrix $M\in {\mathbb{R}}^{n\times n}$ with $M_{ij}=0$ for $j\le i$ and $M_{ij}=-\infty$ for $j>i$, and add it to the raw scores:

The $-\infty$ entries become $e^{-\infty }=0$ after exponentiation, so the softmax renormalizes over only the allowed (lower-triangular) positions. The result is row-causal: row $i$ of $A$ has support only on columns $1$ through $i$.

The three families

The encoder and decoder were designed to work as a pair, but each half is a complete model on its own. Modern practice splits the architecture into three families, distinguished by which stack they keep, how they mask attention, and what objective they train against.

The split follows directly from the masking. An encoder-only model sees the whole sequence, so it is built for understanding (classification, retrieval, token tagging) and is trained by corrupting input tokens and reconstructing them. A decoder-only model masks causally, so each position predicts the next token, making it a pure generator; this is the family that scaled into large language models. An encoder–decoder keeps both and is built for transduction: mapping one sequence to another, as in translation or summarization.⁴

The residual stream

A productive way to read the stack is as a residual stream: a fixed-width $d$-dimensional bus running from the embeddings to the output head. Each sublayer reads a normalized copy of the stream, computes an update, and writes that update back by addition. No sublayer overwrites the stream; they only accumulate into it.

This linear-accumulation view explains why Transformers are so editable: because contributions add, an attention head and an FFN can read each other's writes, and the final representation is literally the embedding plus every increment the sublayers added.

Where the parameters and the compute live

Two accounting questions decide how a Transformer scales: where do its weights sit, and what dominates its runtime. Both have clean closed forms.⁵

Parameters. Per layer, attention contributes the four projection matrices $W_Q,W_K,W_V,W_O$, each $d\times d$, for $4d^2$. The FFN contributes $W_1\in {\mathbb{R}}^{d\times d_{\text{ff}}}$ and $W_2\in {\mathbb{R}}^{d_{\text{ff}}\times d}$; with $d_{\text{ff}}=4d$ that is $2\cdot 4d^2=8d^2$. So per layer,

Two-thirds of a Transformer's weights live in the feed-forward networks, not the attention — a fact that motivates much of the work on sparsifying or scaling the FFN (mixture-of-experts) rather than attention.

The $12d^2$ count ignores the small additive terms — the biases ($\sim 5d$ per layer) and the two LayerNorms' $\gamma ,\beta$ ($4d$ per layer) — which are linear in $d$ and vanish against the quadratic weights for any real width. Put $d=512$, $d_{\text{ff}}=2048$ (the base configuration) into the table to see the actual per-layer budget:

A $6$-layer encoder of this width holds $\approx 19\text{M}$ weights in its blocks; the token embedding and output projection add $2|V|d$ on top, which for a $30\text{k}$ vocabulary is another $\approx 31\text{M}$ — a reminder that at small depth the vocabulary table, not the blocks, can dominate the count.

Compute. The cost split is reversed. The FFN is linear in sequence length — $O(n\cdot d^2)$ per layer — but the attention matrix is quadratic.

The crossover is at $n\approx d$: for $n<d$ the per-layer cost is dominated by the $O(n{d}^2)$ projection and FFN work, but once $n\gtrsim d$ the $O(n^{2}d)$ attention term takes over and grows quadratically. This single $n^2$ is the reason context windows are expensive and the target of nearly every efficient-attention variant.

That is the architecture and its cost model. The same design turns out to be indifferent to what a token is — image patches, DNA bases, or program tokens all work — and the decoder-only half of it scaled into the large language models that dominate current practice. Those applications, and the empirical scaling laws that make scale pay off so reliably, continue in Transformers in Practice.

Takeaways

• The Transformer is attention plus position-wise feed-forward layers, stacked with residual + LayerNorm wrappers, with no recurrence or convolution anywhere.
• Pre-norm ($x+f(\mathrm{LN}(x))$) leaves a clean identity path up the stack, so deep Transformers train stably; post-norm ($\mathrm{LN}(x+f(x))$) squeezes every residual through a norm and needs warmup.
• Sinusoidal positional encoding gives each position a unique fingerprint across a geometric band of frequencies, makes relative offsets a linear map, and extrapolates past the training length; learned tables and rotary (RoPE) encodings are the modern alternatives.
• The encoder attends bidirectionally; the decoder adds masked self-attention (causal, lower-triangular) and cross-attention to the encoder output, and reads out through linear + softmax.
• The three families specialize the design: encoder-only (BERT, masked LM, understanding), decoder-only (GPT, causal LM, generation), encoder–decoder (T5, seq2seq, transduction).
• Causal masking realizes the autoregressive factorization $p(x)={\prod }_{t}p(x_t\mid x_{<t})$ in a single parallel forward pass.
• The residual stream view reads the stack as a fixed-width bus each sublayer reads from and adds back into, never overwriting — which is why contributions compose so cleanly.
• Two-thirds of the parameters live in the FFN ($8d^2$ vs $4d^2$ per layer), while compute is dominated by the $O(n^{2}d)$ attention term — the quadratic cost that bounds context length.