Understanding the Transformer Architecture

A Transformer is a network that takes a sequence of tokens and produces a new vector for every token, where each new vector has been enriched with information from the other tokens in the sequence. Stack enough of these enrichment steps, add a final layer that turns vectors into predictions, and you have a model that can classify text, translate it, or generate it one token at a time.

It helps to think of a Transformer as a meeting. Each token arrives with its own notes (its embedding). In each round of the meeting, every token gets to listen to the others, decide who is relevant to it, and update its notes accordingly. That listening is attention. Between rounds, each token also takes a moment to think privately about what it has heard: that is the feed-forward network. After several rounds, each token's notes describe not just itself, but its role in the whole sequence.

Throughout this module, we work with Tiny Shakespeare: about a million characters of Shakespeare's plays in one text file. Our models will read and write it one character at a time, which keeps the vocabulary tiny and lets us train on a CPU.

The file contains 1,115,394 characters, built from just 65 distinct characters: letters, digits, punctuation, the space and the newline. Each character gets a number, its token ID; "First Citizen" becomes [18, 47, 56, 57, 58, 1, 15, 47, 58, 47, 64, 43, 52]. (We look at tokenisation properly in Lesson 6.2.)

Now here is a complete, working Transformer. We build it from PyTorch's ready-made nn.TransformerEncoderLayer, so it fits in a few lines. Don't worry about every argument yet: by the end of the module you will have rebuilt each piece yourself. For now, read it as a list of stages:

1. tok_emb turns each token ID into a vector of 64 numbers.
2. pos_emb adds a vector that encodes each token's position in the sequence.
3. blocks is a stack of 3 Transformer blocks, each containing attention and a feed-forward network.
4. ln_f is a final normalisation step.
5. head turns each token's final vector into a score for every token in the vocabulary: a prediction of what comes next.

Our tiny Transformer has 162,561 parameters. That is minuscule by modern standards, but it has all the essential parts of the models behind today's chat assistants.

Figure 6.1.1. The lesson’s causal decoder-only Transformer adds token and position embeddings, applies three blocks, then normalises and predicts 65 next-character logits per position. The short row illustrates part of a 64-token context.

In Lesson 5.4 we found three problems with recurrent networks:

• A long chain: information, and the learning signal, has to pass through every step between two distant tokens, and fades along the way.
• A fixed-size bottleneck: everything read so far must be squeezed into one hidden state.
• Sequential computation: step 1,000 cannot start until step 999 has finished, so the parallel power of modern hardware is wasted.

The Transformer removes all three. Every token can attend directly to every other token, so the chain has length one. Every token keeps its own vector, so there is no single bottleneck. And because no token has to wait for the previous one, the whole sequence is processed at once. Let's see that last point concretely.

For our 32-character input, a single forward pass of the Transformer produces 32 predictions at once: one at every position, each saying what character should come next at that point. An RNN reading the same input would need 32 steps, one after another.

This matters most during training. With one forward pass, we get a prediction, and therefore a training signal, at every position of every sequence in the batch. That is a key reason Transformers can be trained on enormous amounts of text. We return to this idea, called parallel training, in Lesson 6.6.

The price, as you saw at the end of Lesson 5.4, is that comparing every token with every other token costs work proportional to the square of the sequence length. We measure that cost in Lesson 6.3.

Let's follow one input all the way through our model, and print the shape of the data after each stage. We use the first 32 characters of the play as input.

Read the shapes from top to bottom:

• The input is (1, 32): a batch of 1 sequence, 32 token IDs long.
• After the embeddings it is (1, 32, 64): each of the 32 tokens is now a vector of 64 numbers. This size, called d_model or the model dimension, is the "width" of the Transformer.
• After each of the 3 blocks it is still (1, 32, 64). Blocks never change the shape; they change the content of each token's vector, by mixing in information from other tokens.
• After the output head it is (1, 32, 65): at every one of the 32 positions, one score for each of the 65 characters in the vocabulary. These scores are called logits.

That constant shape inside the stack is one of the Transformer's most useful properties. Because each block takes (batch, tokens, d_model) and returns the same shape, you can stack as many blocks as you like, and the same block design works for any length of input up to the maximum context.

What does the untrained model predict? Let's turn the logits at the last position into probabilities with softmax and look at its best guesses for the character after "proceed".

The untrained model's favourite guess is 'i' with probability 0.065, followed by 'N', 'R', a space and '-'. These are random preferences, not knowledge. The entropy of the distribution, a measure of how spread out it is, is 4.024. The maximum possible, for a perfectly uniform guess among 65 characters, is ln(65) = 4.174. So the untrained model is close to "no idea at all", which is exactly where training will start from in Lesson 6.6.

Key terms

• Transformer: a network built from stacked blocks of attention and feed-forward layers that processes all tokens of a sequence in parallel.
• Model dimension (d_model): the length of each token's vector inside the Transformer.
• Transformer block: one round of attention plus feed-forward processing; input and output have the same shape.
• Logits: the raw scores the output head produces, one per vocabulary entry; softmax turns them into probabilities.
• Entropy: a measure of how uncertain a probability distribution is; highest when every outcome is equally likely.

Each stage of the pipeline has one job. Here they are in the order the data meets them:

• Tokens. The text is split into units and each unit gets an ID. The model only ever sees IDs. (Lesson 6.2.)
• Token embeddings. A lookup table turns each ID into a vector, exactly as in Lesson 5.3. This vector describes the token on its own, with no context.
• Positional information. Attention, as you will see in Lesson 6.3, has no built-in sense of order. So we add a second vector that encodes where each token sits. (Lessons 6.2 and 6.3.)
• Attention. Each token gathers information from other tokens, weighted by relevance. This is the only place where tokens exchange information. (Lesson 6.3.)
• Feed-forward network (FFN). A small two-layer MLP applied to each token's vector separately. It processes what attention gathered. (Lesson 6.4.)
• Output head. A linear layer that turns each final vector into one score per vocabulary entry. (Lesson 6.5.)

A useful slogan: attention moves information between tokens; the feed-forward network processes information within each token.

Where do the parameters live? Let's count them by component.

The feed-forward networks hold the most parameters: 99,264, or 61.1% of the model. Attention holds 49,920 (30.7%). The embeddings, output head and layer norms together hold well under a tenth. The roughly two-to-one ratio between feed-forward and attention follows from a standard design choice: each feed-forward network expands every vector to 4 × d_model numbers and back, which needs about twice as many weights as attention's four d_model × d_model projections. Large language models keep a similar balance, which is one reason researchers think of the feed-forward layers as the place where much of a model's "knowledge" is stored.

All Transformers share the blocks we have just seen. What distinguishes the three main families is which tokens each token is allowed to look at, and that is controlled by a mask: a grid that says, for every pair of positions, whether attention between them is allowed.

The three grids show the three patterns. Each row is a position doing the looking; each column is a position being looked at; blue means "allowed".

Encoder-only (bidirectional). Every token may look at every other token, before and after it (left grid). This gives the richest possible understanding of a complete input, so encoders are used to understand text: classify it, extract information from it, or produce embeddings for search. The best-known example is BERT. An ordinary bidirectional encoder is not trained for the left-to-right autoregressive generation loop we build here. It can predict masked tokens, and encoder architectures can be adapted for other generation procedures; "encoder" does not mean that text generation is fundamentally impossible.

Decoder-only (causal). Each token may look only at itself and the tokens before it (middle grid). This is called a causal mask, because the output at each position can only depend on what came earlier, just as a writer can only build on what has already been written. This is exactly what you need to generate text one token at a time. GPT-style language models, including the chat assistants you may have used, are decoder-only. So is the model we build in this module.

Encoder–decoder. An encoder reads the whole input bidirectionally; a decoder generates the output causally, and at every layer it also looks at the encoder's output through cross-attention (right grid: every output position may look at every input position). This suits tasks that turn one sequence into another, such as translation or summarisation. The original Transformer was of this kind; T5 is a well-known modern example.

We can verify the difference between bidirectional and causal attention directly. Below, two inputs differ only in their last character: "the king is dead" and "the king is deaf". We run both through our model with and without the causal mask, and check whether the output at the first position changes.

With bidirectional attention, changing the last character changes the output at position 0 (by up to 0.095371): the first token can "see" the end of the sentence. With the causal mask, the change is exactly 0.000000: the first token cannot see anything after it. The mask does exactly what it promises.

For completeness, here is PyTorch's full encoder–decoder nn.Transformer, with a source sequence of 10 tokens (say, an English sentence) and the first 7 tokens of a target sequence (its translation, generated so far):

The decoder produces one vector per target position, (1, 7, 64), each informed by all 10 source tokens through cross-attention. To generate the 8th target token, a head would turn the last of these vectors into scores over the target vocabulary.

Key terms

• Mask: a grid that allows or blocks attention between each pair of positions.
• Bidirectional attention: every token can attend to every other token (encoders).
• Causal attention: each token can attend only to itself and earlier tokens (decoders).
• Cross-attention: attention from the decoder's tokens to the encoder's output.
• Encoder-only / decoder-only / encoder–decoder: the three Transformer families, used respectively for understanding, generating, and transforming one sequence into another.

There is a powerful way to picture what happens inside the stack of blocks. Think of each token's vector as a stream that flows upward through the model. This is often called the residual stream. Each block does not replace the stream; it adds something to it: first whatever attention gathered from other tokens, then whatever the feed-forward network computed. (This "adding" is the residual connection, which you will meet properly in Lesson 6.4.)

So a token's vector starts as "this is the character d at position 8", and after each block it carries a little more about its surroundings. Let's watch this happen. We take two inputs with the same ending, "good lord" and "evil lord". Their final character, d, has the same ID and sits at the same position in both, so before any block their vectors must be identical. Do they stay identical?

Before any block, the cosine similarity of the two d vectors is 1.000: identical, as expected. After block 1 it is 0.997, after block 2 0.989, and after block 3 0.978. Because each block mixes in information from earlier tokens, the d in "good lord" and the d in "evil lord" gradually become different vectors. This is a contextual representation, the thing that static embeddings could not provide in Lesson 5.3.

The effect is small here because the model is untrained: its attention is essentially random, so it mixes in context weakly and without purpose. After training, the context will matter far more.

Figure 6.1.2. A schematic residual stream: the same character can acquire different context-dependent representations. The coloured updates are an analogy, not measured semantic coordinates or guaranteed attention behaviour.

This section is included in class. It connects our teaching model to refinements used in research papers and model documentation.

The Transformer in "Attention Is All You Need" has been refined in several ways since it was published, mostly to make large models train more reliably:

• Where normalisation goes. The original placed layer normalisation after each attention and feed-forward step ("post-norm"). Most modern models place it before each step ("pre-norm"), which trains more stably in deep stacks. You will measure this difference yourself in Lesson 6.4.
• Activation function. The original feed-forward network used ReLU; modern models usually use smoother functions such as GELU or gated variants.
• Positional information. The original added fixed sine-and-cosine patterns. Many models instead learn position vectors, as ours does, and many large models use methods such as rotary position embeddings (RoPE) that encode relative positions inside attention itself.
• Families. The original was an encoder–decoder for translation. The most prominent models today are decoder-only.

These are refinements, not a different architecture. PyTorch builds both styles from the same layer class, with exactly the same number of parameters:

Both layers have 49,984 parameters. The only differences are the order of operations (norm_first) and the activation function: the same components, arranged slightly differently.

Nothing in a Transformer block is specific to text. A block only needs a sequence of vectors. So if you can turn something into a sequence of vectors, you can process it with a Transformer.

The Vision Transformer (ViT) does exactly this for images. It cuts an image into a grid of small square patches, flattens each patch into a list of pixel values, and passes each through a linear layer to make a vector. Those patch vectors are then treated just like word embeddings. Here is the idea applied to a 28×28 Fashion-MNIST image, cut into 7×7 patches:

The 28×28 image becomes a 4×4 grid of patches, (4, 4, 7, 7), which we flatten into a sequence of 16 "tokens" of 49 numbers each, (16, 49). A linear layer turns each into a 64-number vector, giving (1, 16, 64): exactly the shape of a 16-token text sequence. We can then pass it through the same kind of Transformer block we used for Shakespeare, and get back (1, 16, 64). (In a real ViT, a position embedding is added too, so the model knows where each patch came from.)

Notice how this connects back to Module 5. A CNN builds in the assumption that nearby pixels belong together. A ViT does not: from its first layer, any patch can attend to any other patch. With enough training data, this flexibility pays off, which is why many modern vision models are Transformers.

The same trick powers multimodal models. Turn image patches, audio frames or video frames into vectors, put them in the same sequence as text tokens, and one Transformer can attend across all of them. That is how assistants that "see" an uploaded image work: the image becomes a set of vectors placed alongside the words of your question.

A Transformer is a pipeline. Text is split into tokens and converted to IDs. Token embeddings and position embeddings turn each ID into a vector of size d_model. A stack of Transformer blocks then refines every token's vector: in each block, attention moves information between tokens and a feed-forward network processes each token on its own, and both add their results to the token's residual stream. Finally, an output head turns each vector into one score per vocabulary entry.

In our small model, a 32-character input flowed through as (1, 32) → (1, 32, 64) → three blocks that kept that shape → (1, 32, 65) logits: 32 next-character predictions from a single forward pass. Of its 162,561 parameters, 61.1% sit in the feed-forward networks and 30.7% in attention.

What separates the Transformer families is the mask. Encoders let every token see every token, which is ideal for understanding. Decoders use a causal mask, so each token sees only the past, which is needed for generation; we confirmed that a change at the end of the input left the first position's output exactly unchanged. Encoder–decoders combine the two with cross-attention. And because blocks only need a sequence of vectors, the same design works for image patches and mixed modalities.

Figure 6.1.3. Follow the same system from inputs (6.2), through attention (6.3), the block (6.4), the full architecture (6.5), and training and generation (6.6). All Module 6 material is covered in class.