Transformer Model Architecture

Introduction

Transformers were first introduced in the landmark paper Attention Is All You Need by Vaswani et al. in 2017. Unlike previous sequence models such as RNNs and LSTMs, Transformers rely entirely on the attention mechanism to process sequential information. This removes the need for recurrence, enabling highly parallel computation during training and allowing models to capture long-range dependencies much more effectively.

Today, nearly every state-of-the-art language model—including GPT, BERT, Llama, and Gemini—is built upon the Transformer architecture or one of its variants.

Below is the original architecture proposed in the paper.

The full architecture is made of two stacks:

• The Encoder (covered in Part 1) reads the full input sentence and builds a rich contextual representation of it.
• The Decoder (covered in Part 2) generates the output sequence one token at a time, using both what it has generated so far and the encoder's representation of the source sentence.

Part 1 — Encoder

Input and Input Embedding

A Transformer cannot directly process raw text. The input sentence must first be converted into a numerical representation.

Consider the sentence:

"The cat sat on the mat"

The processing pipeline looks like:

Text
 ↓
Tokenization
 ↓
Token IDs
 ↓
Embedding Layer
 ↓
Embedding Vectors

Tokenization

The sentence is split into smaller units called tokens.

["The", "cat", "sat", "on", "the", "mat"]

Each token is mapped to an integer using a vocabulary. Note that the examples given are random numbers.

This produces a vector of token IDs:

For a sequence length \(L = 6\), the token ID tensor has shape:

Or with batching:

Where:

• \(B\) = batch size (number of sequences processed in parallel)
• \(L\) = sequence length (number of tokens in the sequence)

Embedding Layer

Integer IDs carry no semantic information. For example:

cat = 845
dog = 432

This does not imply that "cat" is twice as similar to "dog" as another word.

Instead, each token ID indexes into a learnable embedding matrix \(E\):

Understanding the Matrix Dimensions (\(V \times d_{model}\)):

• Rows (\(V\)): Each row corresponds to a specific vocabulary token ID (from ID 0 to V-1).
• Columns (\(d_{model}\)): Each column represents a learned continuous feature/dimension.

Visual Example of Vocabulary Embedding Matrix \(E\) (\(V \times d_{model}\)):

Given token IDs \(X_{ids} \in \mathbb{R}^{L}\), the lookup extracts the row for each token ID to form the input representation matrix \(X\):

Visual Example of Input Embedding Matrix \(X\) (\(L \times d_{model}\)):

• Rows (\(L = 6\)): Each row represents a word/token in our sentence.
• Columns (\(d_{model} = 512\)): Dense vector features representing word semantics.

Why Embeddings?

Neural networks operate on continuous values. Embeddings allow words with similar meanings to have similar vector representations.

For example, words like cat, dog, lion, and tiger may occupy nearby regions in vector space, whereas cat, airplane, economics, and banana are much farther apart.

Position Encoding (Input)

Unlike RNNs, Transformers process all tokens simultaneously.

Without any notion of position, the Transformer would process "Dog bites man" and "Man bites dog" identically. Therefore, we must inject positional information.

Why Position Encoding?

Attention itself is permutation equivariant (and the overall set-based operation is permutation invariant).

Without position encodings, if you permute (shuffle) the rows of the input matrix \(X\), the self-attention output rows will be permuted in the exact same order, carrying no inherent concept of word order or distance.

Mathematical Proof of Permutation Equivariance

Let \(X \in \mathbb{R}^{L \times d_{model}}\) be the input matrix, and \(W_Q, W_K, W_V\) be the projection weight matrices.

The standard Self-Attention formula is:

Apply a permutation matrix \(P \in \mathbb{R}^{L \times L}\) to shuffle the order of tokens in \(X\). A permutation matrix is an orthogonal matrix where \(P^T P = P P^T = I\).

Let the permuted input be \(X_{\text{perm}} = P X\).

1. Permuted Projections:

1. Permuted Attention Scores:

1. Applying Softmax: Since Softmax operates row-wise, multiplying by \(P\) on the left reorders the rows, and multiplying by \(P^T\) on the right reorders the columns:

1. Multiplying by Permuted Values \(V_{\text{perm}}\):

Since \(P^T P = I\):

Conclusion: Permuting the input sequence by \(P\) simply permutes the output sequence by \(P\). The operation has no intrinsic preference for sequence order, which is why positional encodings must be explicitly added.

Sinusoidal Position Encoding

The original paper introduced fixed sinusoidal encodings. For position \(pos\) and embedding dimension \(i\):

Visual Example of Position Encoding Matrix \(PE\) (\(L \times d_{model}\)):

Adding Position Information

We perform element-wise addition of the input embedding \(X\) and position encoding \(PE\):

Visual Example of Final Encoder Input Matrix \(X'\) (\(L \times d_{model}\)):

Encoder Block

Each encoder block consists of:

1. Multi-Head Self-Attention
2. Add & Layer Normalization
3. Position-wise Feed-Forward Network
4. Add & Layer Normalization

Query, Key, and Value

The positional input \(X'\) is projected into Query (\(Q\)), Key (\(K\)), and Value (\(V\)) matrices using learnable weight parameters \(W_Q, W_K, W_V\).

Let \(d_k = d_v = 64\) for a single attention head.

Weight Matrix \(W_Q \in \mathbb{R}^{d_{model} \times d_k}\) (\(512 \times 64\)):

• Rows (\(d_{model} = 512\)): Maps each feature of the input embedding.
• Columns (\(d_k = 64\)): Projects into the smaller query subspace.

Computation of \(Q\), \(K\), and \(V\):

Visual Example of Query Matrix \(Q\) (\(L \times d_k\)):

Visual Example of Key Matrix \(K\) (\(L \times d_k\)):

Visual Example of Value Matrix \(V\) (\(L \times d_v\)):

Each row represents the value vector (content payload) of that word to be extracted during attention weighting.

Scaled Dot-Product Attention

1. Raw Attention Score Matrix \(QK^T\) (\(L \times L\)):

Multiplying Query (\(6 \times 64\)) by Transposed Key (\(64 \times 6\)) yields dot-product similarities between all pairs of tokens.

• Row \(i\): Token paying attention.
• Column \(j\): Token being attended to.

2. Scaling and Softmax Matrix \(A = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)\) (\(L \times L\)):

Divide by \(\sqrt{64} = 8\) and apply Softmax row-wise so each row sums to \(1.0\):

Interpretation: In Row 1 ("cat"), the model places 68% attention on itself ("cat") and 25% attention on "mat".

3. Attention Output Matrix \(\text{Attention}(Q,K,V) = A V\) (\(L \times d_v\)):

Multiply Softmax probabilities (\(6 \times 6\)) by Values (\(6 \times 64\)):

Each row is a weighted sum of all value vectors in the sentence, encoding context.

Multiple Heads

With \(h = 8\) heads, each head generates an output matrix of size \((L \times d_v) = (6 \times 64)\).

Concatenation

Concatenating 8 heads side-by-side along columns produces a single wide matrix (\(L \times h\cdot d_v\)):

Visual Example of Concatenated Heads Matrix (\(L \times 512\)):

Final Linear Projection

The output projection weight matrix \(W_O \in \mathbb{R}^{512 \times 512}\) mixes features across all heads.

Visual Example of Output Weight Matrix \(W_O\) (\(512 \times 512\)):

Final Output Computation:

Visual Example of Final Multi-Head Attention Matrix Output (\(L \times d_{model}\)):

• Rows (\(L = 6\)): Context-aware representations for each token in the sentence.
• Columns (\(d_{model} = 512\)): Unified embedding space ready for the Feed-Forward Network and residual addition.

Why Multi-Head Attention Works Better

Using multiple heads allows the model to attend to different aspects of the sequence simultaneously. One head might focus on syntactic structure, another on semantic similarity, another on resolving pronouns, and yet another on long-range dependencies.

Because each head has its own learned projection matrices (\(W_Q^i, W_K^i, W_V^i\)), they learn complementary views of the same input. The final projection matrix \(W_O\) fuses these diverse representations into a single contextual embedding that is richer than what any individual head could produce.

The output of the final encoder block, \(\text{EncOut} \in \mathbb{R}^{6 \times 512}\), is held fixed and fed into every decoder block's cross-attention layer, described next.

Part 2 — Decoder

Introduction

While the Encoder reads the full input sentence and builds a rich contextual representation of it, the Decoder's job is to generate an output sequence one token at a time, using both what it has generated so far and the encoder's representation of the source sentence.

This is the architecture used for tasks like machine translation (source language → target language) and, in a simplified single-stack form, for autoregressive language models like GPT.

Continuing our running example, suppose we are translating:

Source (English): "The cat sat on the mat"
Target (French):  "Le chat s'est assis sur le tapis"

The Decoder consists of N stacked decoder blocks (6 in the original paper), each containing:

1. Masked Multi-Head Self-Attention (looks only at previous target tokens)
2. Add & Layer Normalization
3. Multi-Head Cross-Attention (queries the encoder's output)
4. Add & Layer Normalization
5. Position-wise Feed-Forward Network
6. Add & Layer Normalization
7. Final Linear + Softmax layer (only after the last decoder block) to produce output probabilities

Output Embedding and Position Encoding

Just like the encoder, the decoder cannot process raw text directly. The target sequence generated so far is embedded and given positional information exactly the same way as the encoder input.

At training time (via teacher forcing), the whole target sentence is fed in at once, shifted right by one position and prefixed with a special start token:

Target tokens (shifted right):
["<start>", "Le", "chat", "s'est", "assis", "sur", "le", "tapis"]

Why "shifted right"?

The decoder predicts token \(t\) using only tokens \(0 \dots t-1\). Shifting right and prepending <start> means: at position 0 the decoder sees <start> and must predict "Le"; at position 1 it sees <start>, "Le" and must predict "chat"; and so on.

Tokenization and Embedding

Exactly as in the encoder:

where \(E_{tgt}\) is the (target-language) embedding matrix — in many implementations this is tied (shared weights) with the output projection layer and sometimes with the encoder's embedding matrix too.

Adding Position Encoding

using the same sinusoidal formula:

For our example, \(L_{tgt} = 8\) (including <start>), so \(Y' \in \mathbb{R}^{8 \times 512}\).

Masked Multi-Head Self-Attention

Intuition

This is nearly identical to the encoder's self-attention — Q, K, V are all projected from the same matrix \(Y'\) — with one crucial difference: a token must never be allowed to attend to future tokens.

If "Le" could peek at "chat" while predicting "Le" itself, the model would be cheating during training (it would just learn to copy the answer) and would break completely at inference time, when future tokens don't exist yet.

To prevent this, we apply a look-ahead mask (also called a causal mask) before the softmax step.

Computing Q, K, V

using a separate, independently learned set of weight matrices from the encoder's self-attention (and from the cross-attention weights described later).

Raw Attention Scores \(QK^T\) (\(L_{tgt} \times L_{tgt}\))

For our 8-token shifted target sequence:

Applying the Look-Ahead Mask

We construct a mask matrix \(M \in \mathbb{R}^{L_{tgt} \times L_{tgt}}\) that is upper-triangular with \(-\infty\) above the diagonal, and \(0\) elsewhere:

The masked attention scores become:

Since \(e^{-\infty} = 0\), softmax assigns exactly zero probability to any future position:

Visual Example of Masked Softmax Matrix \(A_{masked}\):

Notice the matrix is lower-triangular: each row only distributes probability mass across itself and earlier tokens, and every row still sums to \(1.0\).

Interpretation: When generating "assis" (row 5), the model can only distribute attention over <start>, Le, chat, s'est, assis — it has literally zero access to sur, le, tapis, which haven't been generated yet.

Masked Attention Output

This is then split across \(h=8\) heads, concatenated, and projected with \(W_O\), exactly as in the encoder, producing \(\text{MaskedMultiHead}(Y') \in \mathbb{R}^{8 \times 512}\).

Add & Norm

The residual connection (\(Y' + \cdot\)) helps gradients flow through the deep stack, and Layer Normalization stabilizes the scale of activations across the feature dimension.

Multi-Head Cross-Attention (Encoder–Decoder Attention)

Intuition

This is where the decoder actually looks at the source sentence. It is structurally identical to self-attention, but with one key change:

• Queries (\(Q\)) come from the decoder's own representation \(Z_1\) (i.e., "what am I currently trying to generate?")
• Keys (\(K\)) and Values (\(V\)) come from the encoder's final output \(\text{EncOut} \in \mathbb{R}^{L_{src} \times d_{model}}\) (i.e., "what information exists in the source sentence?")

There is no masking here — the decoder is always allowed to look at the entire source sentence, since the whole input is available up front. Only the target side needs to be masked.

Computing Q, K, V for Cross-Attention

Note the shapes: \(Q_{cross}\) has \(L_{tgt}=8\) rows (one per target token) but \(K_{cross}, V_{cross}\) have \(L_{src}=6\) rows (one per source token). This is why the attention matrix below is not square.

Cross-Attention Score Matrix (\(L_{tgt} \times L_{src}\))

Interpretation: Row 2 ("chat" is being generated) places 72% attention on the source word "cat" — this is the model learning word-level alignment between languages, entirely analogous to the alignment weights in older attention-based seq2seq models (Bahdanau et al., 2014), except now computed via scaled dot-product attention inside a fully parallel architecture.

Cross-Attention Output

Concatenating across \(h=8\) heads and projecting with \(W_O^{cross}\):

Add & Norm

Position-wise Feed-Forward Network

Identical in structure to the encoder's FFN — applied independently and identically to every row (token position) of \(Z_2\):

where:

• \(W_1 \in \mathbb{R}^{d_{model} \times d_{ff}}\), typically \(d_{ff} = 2048\)
• \(W_2 \in \mathbb{R}^{d_{ff} \times d_{model}}\), projecting back down to \(512\)

This gives each token's representation a chance to be transformed through a nonlinearity, independently of the other tokens — all the mixing between tokens has already happened via the two attention sub-layers.

Add & Norm

\(Z_3 \in \mathbb{R}^{8 \times 512}\) is the output of one full decoder block. This becomes the input to the next decoder block (there are \(N=6\) stacked in the original paper), with each block repeating: masked self-attention → cross-attention → FFN, each block's cross-attention re-reading the same fixed \(\text{EncOut}\) from the encoder.

Final Linear Layer and Softmax

After the last (\(N^{th}\)) decoder block, we have \(Z_3^{(N)} \in \mathbb{R}^{L_{tgt} \times d_{model}}\). To turn this back into a probability distribution over the vocabulary:

Linear Projection

where \(W_{vocab} \in \mathbb{R}^{d_{model} \times V}\) projects each token's \(512\)-dim representation up to the size of the vocabulary \(V\) (e.g. \(30{,}000\) subword tokens). This matrix is frequently weight-tied with the target embedding matrix \(E_{tgt}\) to reduce parameter count and improve generalization.

Softmax

Visual Example: Predicting the token after "Le" (row 2 of Logits → Softmax)

The model assigns the highest probability (\(0.81\)) to "chat" — correctly predicting the next French word — since the cross-attention layer had already routed information from the source word "cat" into this position.

Training vs. Inference

Why the Decoder Needs Two Attention Sub-Layers

Without masked self-attention, the decoder couldn't build coherent internal structure in the target language (e.g., subject-verb agreement across several generated words). Without cross-attention, the decoder would be generating text with no connection to the source sentence at all — effectively an unconditional language model.

Full Decoder Data Flow Summary

Target tokens (shifted right)
 ↓
Output Embedding + Position Encoding        → Y'
 ↓
[ Masked Multi-Head Self-Attention ]        → uses Y' for Q, K, V (causal mask)
 ↓
Add & Norm                                   → Z1
 ↓
[ Multi-Head Cross-Attention ]              → Q from Z1, K & V from Encoder Output
 ↓
Add & Norm                                   → Z2
 ↓
[ Position-wise Feed-Forward Network ]
 ↓
Add & Norm                                   → Z3   (repeat block × N)
 ↓
Linear Projection → Vocabulary logits
 ↓
Softmax → Next-token probability distribution

References

1. Vaswani, A., et al. (2017). Attention Is All You Need. arXiv:1706.03762. https://arxiv.org/abs/1706.03762
2. Alammar, J (2018). The Illustrated Transformer [Blog post]. Retrieved from https://jalammar.github.io/illustrated-transformer/
3. https://nlp.seas.harvard.edu/annotated-transformer/
4. https://blog.timodenk.com/linear-relationships-in-the-transformers-positional-encoding/
5. https://kazemnejad.com/blog/transformer_architecture_positional_encoding/
6. https://krypticmouse.hashnode.dev/attention-is-all-you-need
7. https://www.youtube.com/watch?v=bCz4OMemCcA