Transformer architecture diagram: what every block does

In this article

1. Before you start
2. Why interviewers ask this
3. The transformer architecture from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

Why interviewers ask this

"Draw the transformer" is one of the most common whiteboard requests in machine learning rounds, because a candidate who can draw it and explain each box has usually read the paper rather than a summary. A weak answer names the boxes in order. A strong answer says what each box takes in and gives back, why the residual path and normalisation are there, and which parts of the 2017 design modern models changed. The question sits in the transformers hub learning path and often leads into machine learning system design questions about cost, such as why long inputs are expensive and which part of the model to shrink first.

The transformer architecture from first principles

Transformer layer

A transformer layer is a repeated unit of a transformer that passes every token vector through a self-attention sublayer and then a feed-forward sublayer, each followed by a residual connection and normalisation, and returns vectors of the same size.

Feed-forward sublayer

A feed-forward sublayer is the part of a transformer layer that applies the same small two-layer neural network to each token vector separately, so tokens do not exchange information inside it.

In plain words

Think of a meeting. Every person at the table is a token. In the first round, each person listens to everyone else and writes down what seems relevant to them. That is attention. In the second round, each person goes quietly to a desk and thinks about their own notes alone. That is the feed-forward network. Between rounds, each keeps a copy of what they knew before and adds the new notes to it. That is the residual connection. The analogy breaks here. Real people choose whom to listen to once, but a transformer repeats the two rounds many times, and no one sees the table twice in the same way.

In one precise sentence: a transformer layer lets tokens exchange information, lets each token process what it gathered, and keeps the original signal with an addition.

Step 1: Turn tokens into vectors and add position

The model looks up one vector for each token, then adds information about where the token sits, because attention alone does not know order. In the original paper each vector has 512 numbers, written d_model = 512, and position comes from fixed sine and cosine waves of different wavelengths (Vaswani et al., 2017, sections 3.4 and 3.5). The authors tried learned position vectors instead and report nearly identical results. Many newer models replace both with rotary position embeddings, covered in rotary positional embeddings. Input: token numbers. Output: one 512-number vector per token, with position mixed in. What you now know: everything after this step works on vectors of one fixed size.

Step 2: Let every token look at the others with self-attention

Self-attention gives each token a new vector that is a weighted average of information from all tokens in the sequence, with weights set by how well each pair matches. The original model uses eight attention heads.

Each head projects the full 512-number vector down to 64 numbers with its own learned weights for queries, keys and values, and 8 × 64 = 512, so the cost stays close to one full-size head (Vaswani et al., section 3.2.2). The mechanics of queries, keys and values are in self-attention, and the head mechanism is in multi-head attention. Input: n vectors. Output: n vectors of the same size, each now carrying information from the others. This is the only place tokens exchange information inside a layer.

Step 3: Add the input back and normalise after attention

The attention output is not passed on as it is. The sublayer is wrapped as LayerNorm(x + Sublayer(x)): the input x is added to the sublayer's output, and the sum is normalised (Vaswani et al., 2017, section 3.1). The addition is the residual connection, the dashed path in the diagram. The normalisation keeps the numbers in a stable range; see layer normalisation compared with batch norm for what it computes. The placement matters.

The original "post-norm" placement normalises after the addition, and Xiong et al. (2020) argue it needs a learning-rate warm-up, a short start of training with a very small step size that is then raised, because gradients near the output are large at the start. A gradient is the signal, computed during training, that says how much and in which direction to change each weight. Placing the normalisation inside the residual block, "pre-norm", let them remove warm-up. Input: x and the attention output. Output: one stable vector per token. The same wrapper appears again after the next sublayer.

Step 4: Process each position alone with a feed-forward network

Next, the same small network is applied to every token vector separately. It has two linear layers. A linear layer is a step that multiplies a vector by a learned grid of weights to produce a new vector. They sit around a rectified linear unit (ReLU), a function that sets negative values to zero.

It widens each vector from 512 numbers to 2,048 and back (Vaswani et al., section 3.3). Its output is wrapped in add and norm exactly as in Step 3. Counting weights with the paper's dimensions, attention in an encoder layer has 1,048,576 weights and the feed-forward network has 2,097,152, so the feed-forward network holds two thirds of the layer.¹ That answers the hook. Attention decides where to look; the feed-forward network holds most of what the layer can compute with what it found. Input: n vectors. Output: n vectors, each changed independently of the others, then added to their input and normalised.

Step 5: Repeat the layer

The stack is not a separate stage after the layer; it is the layer repeated, and the original uses six encoder layers and six decoder layers (Vaswani et al., section 3.1). An encoder layer has two sublayers, attention and feed-forward. A decoder layer has three: it adds a second attention block that reads the encoder's output, and its first attention block is masked so a position cannot see later positions, which keeps generation left to right. Both details are explained in causal masks and cross-attention.

Each stacked layer takes vectors in and gives vectors out, so layers chain without any change of shape. Counting every sublayer, the model has 30 residual connections: 12 in the encoder and 18 in the decoder.¹

Text description of the diagram

• The original model has six encoder layers and six decoder layers.
• A decoder layer adds masked self-attention and cross-attention to the encoder's output.

Step 6: Turn the last vectors into next-token probabilities

A final linear layer projects each output vector onto the vocabulary, one score per token, and softmax turns them into probabilities. Softmax is a function that turns a list of scores into positive numbers that sum to one.

The original model also shares the same weight matrix between the embedding step and this projection (Vaswani et al., section 3.4). Training compares the probabilities with the true next token; generation picks or samples one and feeds it back. A decoder-only model is a transformer that keeps only the decoder stack and removes the encoder and the cross-attention block, the design behind most chat models (the transformers hub explains how it is used). The reuse of earlier work during generation is the topic of the key-value (KV) cache.

1 · Embed2 · Attend3 · Add and norm4 · Feed forward5 · Repeat x66 · OutputReplay

Step-by-step description

1. Tokens become vectors, with position added.
2. Tokens gather information from each other.
3. Attention output plus its input x, then normalised.
4. Each token alone, then wrapped in add and norm again.
5. The whole layer repeats; the original uses six.
6. Linear layer and softmax give probabilities.

• The input step turns tokens into 512-number vectors and adds position.
• Self-attention is the only place tokens exchange information inside a layer.
• Each sublayer is wrapped as LayerNorm(x + Sublayer(x)).
• The feed-forward network holds two thirds of an encoder layer's weights.

Worked example

We counted the weights in one layer of the base model from the dimensions printed in the paper: d_model = 512, d_ff = 2,048, eight heads of 64 numbers, six encoder and six decoder layers.¹ An attention block has four weight matrices (queries, keys, values and the output projection), each 512 by 512, so 4 × 512 × 512 = 1,048,576 weights. The feed-forward network has two matrices, 512 by 2,048 and back, so 2 × 512 × 2,048 = 2,097,152. Biases and normalisation weights add 6,656 to an encoder layer, 0.2% of its weights, so they are ignored in the split. A decoder layer has two attention blocks and one feed-forward network, so its feed-forward share is one half.

Text description of the diagram

• An encoder layer has 1.05 million attention weights and 2.10 million feed-forward weights.
• A decoder layer has two attention blocks, 2.10 million weights, and 2.10 million feed-forward weights.

The paper's own parameter count for this model is 65 million in Table 3, which also includes the embedding table that this layer-by-layer count leaves out. The two stacks together hold 44,040,192 weights by our count.¹

Trade-offs and edge cases

The design of 2017 has been changed in four places by the models that followed. The Llama 2 paper describes its architecture as the standard transformer with pre-normalisation using root mean square normalisation (RMSNorm), the SwiGLU activation function and rotary position embeddings, and adds grouped-query attention (GQA) in its 34-billion and 70-billion parameter models, with the context window doubled from 2,048 to 4,096 tokens (Touvron et al., 2023).

SwiGLU is a gated activation: one branch of the feed-forward network multiplies another, replacing the plain ReLU. GQA lets several query heads share one key and value head, which shrinks the key-value cache that generation must keep in memory (see the KV cache). RMSNorm drops the re-centring step of layer normalisation and, its authors report, cuts running time by about 7% to 64% depending on the model at comparable quality (Zhang and Sennrich, 2019).

The pattern has also held into 2024-25. As of October 2026, the DeepSeek-V3 report (December 2024) describes a mixture-of-experts (MoE) language model with 671 billion total parameters, 37 billion of them activated for each token, that adopts Multi-head Latent Attention and DeepSeekMoE (DeepSeek-AI, 2024). A mixture of experts replaces the one feed-forward network with many smaller ones and sends each token to only a few. So the layer still alternates attention and feed-forward blocks with residual paths, and the feed-forward part became the experts; this also means the two-thirds split above does not carry over to such models.

Two edge cases matter in an interview. Attention compares every token with every other token, so its cost grows with the square of the sequence length, which is why long contexts are expensive. And the weight split above is for the 2017 base model; models with gated activations, grouped-query attention or experts have different splits, and even in the 2017 design it depends on the width: with d_ff = 2 × d_model = 1,024 the feed-forward share of an encoder layer would be 50%.¹ Check a model's configuration before quoting a fraction.

How to answer it in an interview

Say it in this order: input, the repeated layer, why the wrapping exists, the output, then what changed since 2017.

Common mistakes

• "Attention is where most of the parameters are" (wrong). In the original base model the feed-forward network has twice the weights of an attention block (it depends on d_ff and, in newer models, on experts). (Right: "attention routes information; the feed-forward network holds most of the weights.")
• "The residual connection is an extra layer" (wrong). It adds no weights; it adds the input to the output. (Right: "it is an addition that lets a sublayer learn a change.")
• "Transformers know word order from attention" (wrong). Attention is blind to order without position information. (Right: "position is added to the input, or rotated into the queries and keys.")
• "Normalisation always comes after the addition" (wrong). The 2017 paper does that, but Llama 2 normalises before each sublayer. (Right: "post-norm in the paper, pre-norm in Llama 2 and many models since.")

Follow-up questions

What is the difference between an encoder and a decoder in a transformer?

An encoder layer has self-attention and a feed-forward network and reads the whole input at once. A decoder layer adds masked self-attention, so a position cannot see later positions, and a cross-attention block that reads the encoder's output, giving it three sublayers. Causal masks and cross-attention covers both.

What is a decoder-only transformer architecture?

It keeps only the decoder stack: no encoder and no cross-attention block, just masked self-attention and a feed-forward network in each layer, predicting the next token from the tokens before it. Most chat models use this design.

Who invented the transformer architecture, and when?

Vaswani and colleagues at Google introduced it in the paper "Attention Is All You Need", posted to arXiv on 12 June 2017 (arXiv 1706.03762).

Why does a transformer need positional information?

Self-attention treats its input as an unordered set, so swapping two tokens would not change what attention computes without a position signal. The paper adds sine and cosine waves to the embeddings.

Where does the layer normalisation go?

In the original paper after each residual addition. Xiong et al. (2020) show that placing it inside the residual block removes the need for learning-rate warm-up, and Llama 2 normalises before each sublayer. What the normalisation computes is in batch norm versus layer norm.

What if the feed-forward width were only twice the model width?

With d_ff = 2 × d_model = 1,024, the feed-forward network has 1,048,576 weights, equal to attention's 1,048,576, so its share of an encoder layer would be 50% instead of two thirds.¹

Sources

1. Attention Is All You Need — arXiv (Vaswani et al.), v7 revised 2023-08-02, 12 Jun 2017
2. On Layer Normalization in the Transformer Architecture — arXiv (Xiong et al.), ICML 2020, v2 revised 2020-06-29, 12 Feb 2020
3. Root Mean Square Layer Normalization — arXiv (Zhang and Sennrich), NeurIPS 2019, 16 Oct 2019
4. DeepSeek-V3 Technical Report — arXiv (DeepSeek-AI), v2 revised 2025-02-18, 27 Dec 2024
5. Llama 2 technical paper: open foundation and chat models — arXiv (Touvron et al.), 18 Jul 2023