Transformer Architecture Overview

A common interview question is deceptively small: What happens to one token inside one transformer layer? Candidates who answer with attention happens stop at the buzzword. The answer that lands walks through the tensor: shape in, projections, scaling, heads, residual, layer norm, feed-forward, shape out.

Start with the input. A sequence of 16 tokens, each embedded to d_model = 512, gives a matrix X of shape (16, 512). Every encoder block then does exactly two kinds of work in a fixed order: mix information across positions with multi-head self-attention, and transform each position independently with a feed-forward network. Nothing else in the block moves information between positions.

Multi-head attention is the mechanism from Attention Mechanisms with three learned projections turning X into queries, keys, and values. With 8 heads, each head projects down to d_k = 512 / 8 = 64 dimensions, and one head computes:

import math
import torch.nn as nn

class ScaledDotProductAttention(nn.Module):
    def forward(self, q, k, v, mask=None):
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.size(-1))
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        return torch.softmax(scores, dim=-1) @ v

The division by sqrt(d_k) is not decoration. If q and k have roughly unit-variance components, the variance of their dot product grows with d_k, and by d_k = 1000 the logits spread wide enough that softmax collapses to a near-one-hot vector whose gradients vanish. Scaling keeps the logits in the range where softmax still differentiates.

The heads each emit (16, 64); the eight are concatenated back to (16, 512) and one final projection W_O mixes them. Each head is a separate weighted average over its own subspace, which is why one layer can carry several relations at once: one head tracking the previous token, another a subject-verb link, another punctuation. A single softmax could only produce one convex blend per position.

Why 8 heads instead of 1 or 128? The product h x d_k must equal d_model, so head count trades width for variety: 64 dimensions per head is wide enough to learn a meaningful compatibility function, and 8 heads is enough to cover several relation types simultaneously. At 128 heads of 4 dimensions each every score becomes a near-random similarity; at 1 head of 512 you are back to a single averaged blend. The default lands at 8 or 16 heads of 64, and the feed-forward inner width is conventionally 4 x d_model because that expansion is where each position's transformation is stored - halving it saves parameters but measurably costs quality on language benchmarks.

After attention the block adds the residual x = x + Head(LayerNorm(x)), then repeats the pattern with a position-wise feed-forward: two linear layers with GELU between them, widening 512 to 2048 and back. The feed-forward never sees other positions; every cross-position mixing already happened in attention, and probing studies keep locating much of a model's factual knowledge in those feed-forward weights, which is why they dominate parameter counts.

Residuals and layer normalization are what make depth trainable. The identity path lets gradients reach early layers without passing through every matrix, and layer norm keeps activations in a stable range as the sums accumulate. Modern stacks (GPT-2 onward, Llama) place layer norm before each sublayer - pre-LN - with one extra norm at the end; the original 2017 design placed it after, and deep post-LN models needed warm-up tricks to avoid divergence. When an interviewer asks how anyone stacks 96 layers, the honest answer is residual connections plus normalization, not attention itself.

A useful follow-up they may push: what breaks if you delete the residual add? Gradients reaching the early layers then have to multiply through every intervening matrix, so the signal decays roughly exponentially with depth and deep stacks stop training at all. The residual turns that chain of products into a chain of sums, which is why adding layers to a modern transformer degrades gracefully instead of catastrophically.

Feed the same eight tokens in a different order and self-attention produces the same outputs, just permuted. That should bother you. Attention scores positions purely by content: each score is a dot product between a query from one token and a key from another, and a softmax mixes the values. Shuffle the input rows and the output rows shuffle identically - self-attention is permutation-equivariant, so dog bites man and man bites dog yield the same multiset of representations. Recurrence hid this by consuming tokens one at a time in a fixed order; the step index was the position signal.

You can prove it in five lines:

import torch, torch.nn as nn
attn = nn.MultiheadAttention(512, 8, batch_first=True)
x = torch.randn(2, 8, 512)               # batch 2, 8 tokens, d_model 512
perm = torch.tensor([3, 0, 7, 1, 6, 2, 5, 4])
y, _ = attn(x, x, x)
yp, _ = attn(x[:, perm], x[:, perm], x[:, perm])
print(torch.allclose(yp, y[:, perm], atol=1e-5))   # True - order is invisible

The fix is to add a position signal to the embeddings before the first layer. The original transformer used sinusoids:

PE(pos, 2i)     = sin(pos / 10000^(2i / d_model))
PE(pos, 2i + 1) = cos(pos / 10000^(2i / d_model))

Each frequency pair gives one dimension a wavelength on a geometric ladder from 2*pi up to roughly 62,832, so every position gets a unique smooth pattern and any fixed offset becomes a fixed rotation of the vector. The model can therefore learn attend three positions back as an addition in angle space, and sinusoids extrapolate a little past training length, which a learned table cannot. Building the table takes one vectorized pass:

import torch

def sinusoidal_position(n: int, d_model: int) -> torch.Tensor:
    pos = torch.arange(n).unsqueeze(1)                 # (n, 1)
    i = torch.arange(0, d_model, 2).unsqueeze(0)       # (1, d_model / 2)
    angle = pos / (10000 ** (i / d_model))
    pe = torch.zeros(n, d_model)
    pe[:, 0::2] = torch.sin(angle)
    pe[:, 1::2] = torch.cos(angle)
    return pe                                           # (n, d_model), added to embeddings

Note that the original paper also scales token embeddings by sqrt(d_model) before the addition, so position and content arrive at comparable magnitudes; if the embedding norm dwarfs the position vector, the first layers effectively ignore order.

Variants you should be able to name on sight:

RoPE is the one that draws follow-ups: positions m and k are rotated by angles mtheta and ktheta, and the dot product between them depends only on m minus k, so relative distance falls out of the scores themselves instead of being bolted on afterward.

Two details interviews probe. First, position is added, not concatenated, so the tensor keeps d_model = 512 through every layer and content cannot be cleanly separated from position - that entanglement is intentional, not sloppiness. Second, batching interacts with position: padding tokens are given no meaningful position and are excluded by the attention mask, otherwise their rows leak probability mass into real positions.

So the sharp version of transformers learn syntax is a two-part claim: attention supplies content-based matching, positional encodings supply order, and the stack learns the interaction. Drop either and you have a bag of words, or a model that gives the same answer for every permutation.

You are choosing an architecture and the interviewer asks which transformer to use. The reliable answer comes from masks and wiring, not from brand names.

The causal mask is one boolean matrix. Position i may attend only to positions at or below i, so everything above the diagonal becomes -inf before the softmax and each prediction is built strictly from left context. That is what makes teacher-forced training valid: without the mask, position 5 could read the answer at position 6 and the loss would be meaningless.

import torch

def causal_mask(n: int) -> torch.Tensor:
    return torch.tril(torch.ones(n, n, dtype=torch.bool))   # True = attend

mask = causal_mask(4)
scores = scores.masked_fill(~mask, float("-inf"))           # keep only the past

Padding masks run alongside: positions holding PAD tokens are marked False so the softmax does not spread probability over meaningless rows.

Encoder-decoder models add one more mechanism, cross-attention, where queries come from the target side and keys and values come from encoder outputs. The decoder therefore reads the whole source with no causality required there while generating the target left to right. That split is why seq2seq work was the original transformer's job, and why Attention Mechanisms flows straight into this architecture: same scoring function, three different wiring patterns.

Why did the industry converge on decoder-only? Simplification rather than superiority - one stack, one objective, one inference path, and a KV cache that makes autoregressive decoding cheap, with classification folded in through prompting. Encoder-only still wins when you need bidirectional understanding cheaply: reranking a candidate list or embedding documents for search, one forward pass, no generation.

How a stack is pre-trained predicts what it can do. Masked language modeling lets every position see both sides, so the representations are symmetric in context and well suited to understanding tasks - but nothing in the objective teaches the model to emit tokens one after another, so you cannot simply generate from BERT. Next-token prediction bakes left-to-right emission into the weights from day one, and everything generative follows from that. Encoder-decoder spends capacity on both sides: full source comprehension plus autoregressive decoding, which is why translation and summarization benchmarks favored it for years before decoder-only models caught up by scale and prompting.

Then the cost. Every position scores every position, so n positions produce n squared score pairs per head per layer: attention is O(n^2 * d) in time and memory. At n = 8192 that is about 67 million score pairs per head, per layer, before gradients are stored. Production long-context systems answer with FlashAttention, which tiles the computation so the n-by-n matrix is never materialized; sliding-window or sparse attention, which bounds each position's neighborhood; and at inference a KV cache so each new token attends to stored keys instead of recomputing the prefix, which is why decoding becomes memory-bandwidth-bound rather than compute-bound.

That inference path splits cleanly in two. Prefill reads the whole prompt in one parallel pass and fills the cache - compute-bound and fast. Decode then emits one token per step, each new token reading the cached keys and values - bandwidth-bound, because every step streams the whole cache to produce a single token. It is why a chat model's responsiveness is quoted in tokens per second, and why a 500-token answer costs far more than the prefill of the question that asked for it.

The follow-up cluster is predictable: why quadratic?, what breaks when I double the context?, why is generation slow when the prefill was fast? - the mask explains the first two and the cache explains the third.

Evaluation of everything you build from here starts at Cross-Validation Strategies, where the question changes from how does it work? to how do I know it works?