Anatomy of a transformer: the brain behind generative AI

So far we have walked the path from text to numbers: tokenization → IDs → embeddings (vectors with meaning). But those vectors, on their own, do not generate text. They need a processor, a machine that reads them, interprets them in context, and decides which token comes next. That machine is the transformer.

The transformer has a reputation for being intimidating — and it is true that production models have hundreds of billions of parameters. But the core design, the architecture, fits on a napkin. There are few pieces, and the most important one (the attention block) repeats over and over. Understanding these pieces is understanding 90% of what makes generative AI work.

The overall structure: an end-to-end pipeline

A transformer is, at its core, a processing pipeline: vectors go in (the token embeddings) and a prediction comes out (a probability distribution over the vocabulary). The full path:

1. Tokenize and embed — turn text into IDs and then into vectors (position comes later, inside attention).
2. Pass through a stack of blocks — where the real “brain” lives: each block refines the representations.
3. Normalize and project — a final normalization and an output head that assigns scores to every token in the vocabulary.
4. Turn them into probabilities — apply softmax so the scores sum to 1 and can be sampled.

The most important thing: from the embeddings to the output, the tensor shape never changes. Each token travels through the model as a fixed-size vector (the model width). The transformer is a transformer of representations: it takes a sequence of vectors and produces another sequence of vectors of the same length and dimensionality, but richer in context each time.

The full journey, with an example — take the sentence “the cat sat on the mat”. After tokenization, we have 6 tokens:

Notice something crucial: the shape does not change across the blocks. The transformer does not compress or expand the sequence; it transforms it while keeping the structure. Each block is a “refinement layer” that improves the representations.

There is no domain-specialized machinery — the same design works for text, code, images (with adapters), and more. Leading labs add improvements (more efficient attention, mixture of experts), but the central architecture stays the same. What separates the leaders from the rest is not so much the architectural blueprint as the data, the compute, and post-training.

Inside a block: the heart of the transformer

Each block has two main sublayers:

• Causal multi-head attention — where tokens “talk” to each other, exchanging information.
• MLP (Multi-Layer Perceptron) — where each token “thinks” individually about what it has heard.

And two structural elements that wrap these sublayers: residual connections (they let information flow straight through the block) and pre-norm normalization (it stabilizes training).

The block runs in this order (modern pre-norm variant):

1. x = x + Attention(Norm(x))     # tokens talk to each other
2. x = x + MLP(Norm(x))           # each token thinks on its own

The residual connections (the additions) are highways that let information from earlier layers reach later layers intact. Without them, training deep networks would be nearly impossible.

Causal multi-head attention: the mechanism that changed everything

Attention is the star ingredient — the mechanism that lets each token look at other tokens and decide how much weight to give them. The intuition: not every word in a sentence is equally important for understanding a given word.

Analogy: imagine you are reading a sentence and you hit the word “she”. To understand who it refers to, you need to look backward and find the most recent feminine noun. Attention does exactly that: each token can “ask” the others for relevant information, and receive weighted answers.

How does it work technically? Each token generates three vectors from its embedding:

• Query (Q): “What am I looking for?” — the question the token asks.
• Key (K): “What do I offer?” — the label the token offers to the others.
• Value (V): “What information do I contribute?” — the content it shares if selected.

The process: each token computes its attention score toward every other token (including itself) via the dot product of its query with each one’s key — producing a score matrix of size (T × T), where T is the sequence length. Softmax is applied to those scores (row by row) to turn them into probabilities that sum to 1. That distribution weights the other tokens’ values — the result is a weighted sum: each token receives a mix of information from the others, weighted by relevance.

Causal attention: in generative models (GPT, Llama, and so on), tokens can only look backward — at earlier positions or at themselves, never at the future. It is implemented with a mask: in the score matrix, future-position cells are set to -infinity, so softmax gives them probability zero. This is what lets you train the model by predicting the next token in parallel: in a single pass, each position predicts based only on the past.

Multiple heads: instead of a single attention, the model has several heads in parallel (8, 12, 96…). Each head learns its own Q, K, V projections, and can focus on different kinds of relationships — one head might learn syntactic dependencies, another semantic ones, another long-range ones. The outputs of all heads are concatenated and projected back to the original width.

The attention formula, with mask:

Where QKᵀ is the product of queries and keys (the score matrix), √d_k scales the scores to avoid vanishingly small gradients in the softmax, M is the causal mask (0 where looking is allowed, -∞ where it is not), and V are the values being weighted.

In practice, this operation runs as a fused kernel (scaled dot-product attention), which computes everything without materializing the huge QKᵀ matrix in memory — dramatically faster.

# Project X into Q, K, V (a single projection produces all three)
q, k, v = mx.split(self.c_attn(x), 3, axis=-1)   # (B, T, 3*C)

# Reshape for multiple heads: (B, n_head, T, head_dim)
B, T, C = x.shape
hd = C // self.n_head
q = q.reshape(B, T, self.n_head, hd).transpose(0, 2, 1, 3)
k = k.reshape(B, T, self.n_head, hd).transpose(0, 2, 1, 3)
v = v.reshape(B, T, self.n_head, hd).transpose(0, 2, 1, 3)

# Apply RoPE (rotary positions)
q = self.rope(q)
k = self.rope(k)

# Attention with a causal mask
y = mx.fast.scaled_dot_product_attention(
    q, k, v, scale=1.0 / math.sqrt(hd), mask="causal"
)

# Reshape back
y = y.transpose(0, 2, 1, 3).reshape(B, T, C)

The MLP: when each token thinks for itself

If attention is the conversation (where tokens exchange information), the MLP is the moment of individual reflection — it processes that information locally, with no communication with other tokens.

Analogy: imagine a team meeting. First, everyone talks and shares information (attention). Then each person goes back to their desk, processes what they heard, and takes notes (MLP). Finally, they return to the meeting with new ideas (the next block).

The MLP is a two-layer network with a nonlinearity in the middle. Without the nonlinearity, stacking linear layers would be equivalent to a single layer — you would gain no capacity. The nonlinearity is what lets the MLP learn complex functions.

Modern structure (SwiGLU): instead of the classic “expand to 4x, apply GELU, compress”, modern models use SwiGLU. The recipe: expand (project the vector into a larger space, ~2.67× the original width), and compress (project back to the original width, with a gate that decides what information gets through — that gate is the input passed through SiLU, multiplied by the main branch). This gives three linear matrices (up, gate, down) instead of two, but it improves performance.

def __call__(self, x):
    # First, attention with normalization and residual
    x = x + self.attn(self.ln_1(x))

    # Then, MLP with normalization and residual
    h = self.ln_2(x)
    # SwiGLU: gate (with SiLU) * up, then down
    x = x + self.down(nn.silu(self.gate(h)) * self.up(h))
    return x

The MLP takes roughly two thirds of a block’s parameters (attention takes the remaining third) — that is why, when models scale, they grow mainly by adding more MLP. Mnemonic: attention moves information between positions; the MLP transforms information inside each position.

The residual highway and normalization: the glue that holds it together

Residual connections: each sublayer (attention and MLP) does not replace the input vector — it adds its output to the original vector: output = input + sublayer(input). This creates an information highway that runs straight from the block’s input to its output, without going through the sublayer. In deep networks, this is essential because it lets the gradient flow without fading during training. Before residual connections became popular (2015-2016), training networks deeper than a couple dozen layers was nearly impossible.

Normalization (pre-norm): re-scales the vectors so they have a consistent magnitude, keeping the numbers from “drifting” in deep layers. Modern models use RMSNorm: it normalizes by the square root of the mean of the squares (without subtracting the mean) — cheaper computationally. The pre-norm variant (normalize before the sublayer, not after) became popular with GPT-2 and is what every modern model uses — it makes optimization much easier because the residual connections stay “clean”.

What a block does, in four steps: normalize, attend, add to the highway; normalize, run the MLP, add again. There is no special final block — the model gets its power from repetition.

Back to words: the output head

After passing through all the blocks, the vector at the last position is the final representation the model has of the sequence up to that point. The output head is a linear projection: it multiplies that vector by a matrix of size (vocabulary × dim), producing a score for each token in the vocabulary. Softmax turns those scores into probabilities that sum to 1.

Weight tying (sharing weights): in many models, the output-head matrix is the same as the input embedding table (transposed) — a token’s embedding and its output score use the same vectors. Three advantages: it saves parameters (the embedding table is not duplicated), it improves generalization (shared geometry between input and output), and stability (you do not have to learn two separate representations). It is a common design decision in the GPT line and many other modern models.

Sizing the model: the key hyperparameters

A transformer is defined by five main hyperparameters:

Concrete example: a model with 16 blocks, width 768, 32 heads, context 2.1M, and vocabulary 10,000 (the same vocabulary size we used in this chapter’s manual counting example) has approximately 120.9 million parameters. Typical split: Embeddings (~6%), Attention (~31%), MLP (~62%).

hidden = 64 * ((int(8 * n_embd / 3) + 63) // 64)   # MLP width (SwiGLU), rounded to 64
per_block = 4 * n_embd * n_embd + 3 * n_embd * hidden
emb = vocab_size * n_embd
total = emb + n_layer * per_block

In small models, embeddings can be a significant fraction (up to a third, as in our 120.9M example). In large models, the block term (12d² × layers) completely dominates, because d² grows fast. That is why giant models are, overwhelmingly, made of repeated transformer blocks.

Where this design comes from: a brief history

The transformer was introduced in the paper “Attention Is All You Need” (Vaswani et al., 2017), originally designed for machine translation (hence the original encoder-decoder). But it was the decoder-only (causal) variant that changed the game:

• 2018 — GPT-1 (Radford et al.): first large-scale pretrained generative transformer.
• 2019 — GPT-2: 12 layers, width 768, context 1,024 → 124 million parameters. Introduced the pre-norm variant and popularized weight tying.
• 2020 — GPT-3: 96 layers, width 12,288, 96 heads → 175 billion parameters. The architecture barely changed; they just scaled.

Since then, the central architecture has held, with incremental improvements: RoPE replaces absolute positions, SwiGLU replaces GELU in the MLP, and mixture of experts (MoE) replaces some dense MLPs with multiple experts and a router — giving two size metrics: total parameters vs. active parameters (a recent model has 975B total parameters but activates only 41B per token).

One fact worth underlining: in the jump from GPT-2 to GPT-3, the architecture barely changed — the biggest difference was the volume of data and the number of parameters. That underscores an important lesson: architecture matters, but data and compute are just as crucial.

Chapter summary

The transformer is a sequence processor that:

• Tokenizes and embeds, turning text into vectors.
• Stacks blocks, each with causal multi-head attention (tokens exchange information looking only backward) and an MLP with SwiGLU (each token individually processes what it received), wrapped in residual connections and pre-norm normalization that make it possible to train deep networks.
• Projects and applies softmax to get a probability distribution over the vocabulary.

The model has no external memory and no persistent state: every prediction is based solely on the current token sequence. All of the model’s “knowledge” lives in the parameters — attention and MLP weights, and the embedding table — which are learned during training.

In the Training data chapter we will see how a model is actually trained: the training data, the loss function, and gradient descent at massive scale. Understanding the architecture is only half; the other half is how to turn mountains of text into useful parameters.

If at any point you feel overwhelmed by the complexity, come back to this central idea: a transformer block is just two sublayers (attention and MLP) joined by residual additions and normalization. The power does not come from a magic piece, but from repetition and scale. It is like a craft workshop: each artisan (block) makes a small refinement, and after many artisans and many iterations, the piece is perfect. The transformer is not magic, it is engineering.