Transformer Architecture: The Engine Behind Modern AI

Introduction

In June 2017, eight researchers at Google published a paper titled “Attention Is All You Need.” It was 11 pages long. It proposed a neural network architecture that dispensed with recurrence and convolution entirely, replacing them with a mechanism called attention. The paper was not immediately recognised as a turning point, but within five years the Transformer had become the single most consequential architecture in the history of machine learning.

Every major language model you have heard of — GPT-4, Claude, Llama 3, Gemini, Mistral, DeepSeek — is a Transformer. So are the vision models (ViT, DALL-E 3), the audio models (Whisper, AudioLM), and the multimodal models that combine all of the above. The Transformer is not a language model architecture; it is a general-purpose computational primitive that happens to work remarkably well for sequence data.

This post explains the Transformer from first principles. We walk through every component — self-attention, multi-head attention, positional encoding, feed-forward networks, normalisation, and residual connections — then examine how the architecture evolved from the original encoder-decoder design to the decoder-only giants that power today's AI, including modern innovations like Flash Attention, Grouped Query Attention, Rotary Position Embeddings, and Mixture-of-Experts.

The “Attention Is All You Need” Breakthrough

Before 2017, the dominant approach for sequence modelling was the recurrent neural network (RNN) — specifically LSTMs and GRUs with attention mechanisms bolted on. These models processed tokens one at a time, maintaining a hidden state that accumulated information over time. The sequential nature of recurrence made training slow (no parallelisation across the sequence) and made it difficult to capture long-range dependencies (the vanishing gradient problem).

The Transformer solved both problems in one stroke. By removing recurrence entirely and replacing it with a pure attention mechanism, the architecture could process all tokens in a sequence simultaneously. Training time dropped from weeks to hours. Long-range dependencies — relationships between tokens hundreds or thousands of positions apart — could be captured in a single forward pass.

The original Transformer used an encoder-decoder structure inspired by sequence-to-sequence models in machine translation. The encoder reads the input sequence and produces a representation. The decoder generates the output sequence, attending to both the encoder's representation and its own previously generated tokens. This design remains foundational, but the industry has since converged on simplified variants.

High-Level Architecture

The Transformer is a stack of identical layers. Each layer has two sub-layers: a multi-head self-attention mechanism and a position-wise feed-forward network. Around each sub-layer there is a residual connection followed by layer normalisation.

In equation form, the output of each sub-layer is:

output = LayerNorm(x + Sublayer(x))

This “pre-norm” formulation — where normalisation is applied before the sub-layer — is the convention in modern Transformers. The original paper applied “post-norm” (normalise after the residual addition), but pre-norm has been shown to produce more stable training dynamics at scale.

Figure: The original Transformer architecture from Vaswani et al. 2017. The encoder (left) processes the input sequence. The decoder (right) generates the output sequence. (Wikimedia Commons, CC BY-SA 4.0)

Self-Attention: The Core Innovation

Self-attention is the mechanism that lets each token in a sequence “look at” every other token and decide how much to incorporate information from each one. It is the reason Transformers can capture long-range dependencies without the sequential bottleneck of recurrence.

Q, K, V: Query, Key, Value

The intuition comes from information retrieval. Imagine you are searching a database. You have a query (what you are looking for), keys (labels on each database entry), and values (the actual content of each entry). You compare the query against every key to find the best matches, then retrieve the corresponding values.

Self-attention does exactly this, but every token plays all three roles simultaneously. For each token in the sequence, the model computes three vectors by multiplying the token embedding against three learned weight matrices:

• Query (Q): What information is this token looking for?
• Key (K): What information does this token have to offer?
• Value (V): The actual content this token can contribute if matched.

The attention mechanism computes a compatibility score between every Q and every K, converts those scores into a probability distribution via softmax, and then computes a weighted sum of the V vectors. The result is a new representation for each token that incorporates context from the entire sequence.

Scaled Dot-Product Attention

The specific operation is called scaled dot-product attention. Given matrices Q, K, and V, the attention output is:

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

The dot product QK^T produces a matrix of raw attention scores. Each entry (i, j) represents how much token i should attend to token j. The scaling factor 1/sqrt(d_k) prevents the dot products from growing too large in magnitude as the key dimension d_k increases — without scaling, the softmax would saturate to near-one-hot distributions, producing rigid attention patterns.

The softmax ensures that each row of the attention matrix sums to 1, producing a valid probability distribution over the sequence. The final multiplication with V produces the weighted context vector for each token.

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    attention_weights = F.softmax(scores, dim=-1)
    return torch.matmul(attention_weights, V)

The computational cost of this operation is O(n^2 * d_k) in both time and memory, where n is the sequence length. This quadratic scaling is the fundamental constraint on Transformer context windows and the primary target of every efficiency improvement discussed later.

Multi-Head Attention

Single-head attention computes one weighted combination per token. That is powerful, but it captures only one relationship pattern at a time. In practice, tokens need to attend to other tokens for multiple reasons simultaneously — syntactic structure, semantic similarity, positional proximity, and so on.

Multi-head attention solves this by running multiple attention operations in parallel, each with different learned projections. The Q, K, V matrices are each projected h times into lower-dimensional subspaces (typically d_k = d_model / h). Each projection produces a different attention head, and the heads are concatenated and projected back to the model dimension.

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W_o

The original paper used h = 8 heads with d_model = 512, giving each head d_k = 64. Each head learns to focus on different types of relationships. In practice, some heads learn syntactic relationships (subject-verb agreement), some learn positional relationships (adjacent tokens), and some learn semantic relationships (coreference resolution).

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch_size = x.size(0)
        Q = self.W_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        attn = scaled_dot_product_attention(Q, K, V)
        attn = attn.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
        return self.W_o(attn)

Modern implementations rarely use h = 8. GPT-3 uses 96 heads. Llama 3 70B uses 64 heads. The trend is towards more heads with smaller per-head dimensions, but this interacts with other design choices like Grouped Query Attention, which we cover later.

Positional Encoding

Self-attention is permutation-invariant: it processes the tokens as a set, not a sequence. If you shuffle the tokens, the attention output is identical. This is desirable for many properties but catastrophic for language, where word order determines meaning.

Positional encoding injects information about token position into the model. The original paper used sinusoidal encodings — fixed sine and cosine functions of different frequencies that produce a unique vector for each position:

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

Sinusoidal encodings have the advantage that the model can extrapolate to sequence lengths longer than those seen during training, since the functions are defined for all positions. However, they are fixed and do not adapt to the data. Learned positional embeddings — where each position is assigned a learnable vector — became the standard in early BERT and GPT models.

Rotary Position Embeddings (RoPE)

The most widely adopted positional encoding in modern Transformers is Rotary Position Embeddings (RoPE), introduced by Su et al. in 2021. RoPE encodes position by rotating the query and key vectors by an angle proportional to the position. Concretely, it applies a rotation matrix to each pair of dimensions:

R_theta, pos * X = X_rotated by pos * theta_i

RoPE has two crucial properties. First, the dot product between a query at position m and a key at position n depends only on their relative position (m - n), not their absolute positions. Second, as the distance between tokens grows, the attention score decays naturally — a useful inductive bias for language. These properties make RoPE the default choice for virtually all recent open-source models, including Llama, Mistral, Qwen, and DeepSeek.

Feed-Forward Networks and Activation Functions

Between attention sub-layers, each token passes through an identical feed-forward network (FFN) — also called a multilayer perceptron (MLP). The FFN consists of two linear transformations with a non-linear activation in between:

FFN(x) = W_2 * activation(W_1 * x + b_1) + b_2

The inner dimension is typically 2x to 4x larger than d_model. In the original Transformer, d_ff = 2048 with d_model = 512 (a 4x expansion). In GPT-3, the expansion ratio is about 2.7x. In Llama 3 8B, the intermediate dimension is 14,336 with d_model = 4096 (a 3.5x expansion).

GELU and SwiGLU

The original Transformer used ReLU activation. In 2016, Hendrycks and Gimpel proposed the Gaussian Error Linear Unit (GELU), which smoothly approximates ReLU with a small but beneficial gradient for negative values. GELU became the default for BERT and GPT-2.

Today, the most common activation in top-performing models is SwiGLU — a gated variant introduced by Noam Shazeer in 2020. SwiGLU combines the outputs of two linear projections through a element-wise multiplication with a sigmoid-weighted gate:

SwiGLU(x) = (W_1 * x) * sigmoid(W_2 * x) * (W_3 * x)

SwiGLU uses three weight matrices instead of two, which increases parameter count. To compensate, models using SwiGLU typically reduce the hidden dimension by 2/3 so that the total parameters match. Despite the added complexity, SwiGLU consistently outperforms ReLU and GELU on language modelling benchmarks, and is used in Llama 2, Llama 3, Mistral, Gemma, and DeepSeek.

Layer Normalisation and Residual Connections

Deep networks suffer from training instability as gradients either vanish or explode through the layers. Transformers address this with two mechanisms: residual connections (skip connections) and layer normalisation.

Residual connections, introduced by He et al. in 2016 for ResNet, allow gradients to flow directly through the network by adding the input of a sub-layer to its output. This creates a “gradient highway” that prevents signal decay in deep stacks.

Layer normalisation (LayerNorm), introduced by Ba et al. in 2016, normalises the activations across the feature dimension for each token independently:

LayerNorm(x) = gamma * (x - mu) / sigma + beta

The original Transformer applied LayerNorm after the residual addition (post-norm). Modern practice applies it before the sub-layer (pre-norm), which provides more stable gradients during the initial stages of training. Llama 3 uses RMSNorm — a simplified variant that removes the mean-centering step and only normalises by the root mean square — as a computational optimisation.

Encoder-Decoder vs Encoder-Only vs Decoder-Only

The original Transformer was an encoder-decoder architecture. The encoder processes the input sequence bidirectionally (each token attends to all tokens in both directions), producing a rich contextual representation. The decoder generates the output sequence auto-regressively, attending to both the encoder output and previously generated tokens via causal masking.

Encoder-Only: BERT and Friends

In 2018, Devlin et al. introduced BERT, an encoder-only Transformer trained with masked language modelling (predicting randomly masked tokens). BERT's bidirectional nature makes it ideal for understanding tasks: classification, named entity recognition, question answering, and semantic similarity. The encoder-only design is still the dominant choice for embeddings and retrieval.

Decoder-Only: The GPT Family

In 2018, Radford et al. introduced GPT, a decoder-only Transformer trained with auto-regressive language modelling (predicting the next token). The key difference from encoder-only is causal masking: each token can only attend to tokens before it, not after. This makes the model inherently generative.

The decoder-only architecture has become the dominant paradigm for large language models. GPT-3, GPT-4, Claude, Llama, Mistral, Qwen, DeepSeek, and Gemini are all decoder-only. The reasons are practical: decoder-only models are simpler to train (single objective), simpler to deploy (single model), and scale more predictably. The emergence of few-shot and zero-shot capabilities showed that generative pre-training alone is sufficient for most understanding tasks — the need for a separate encoder has largely disappeared.

Scaling Laws: Why Transformers Scale So Well

In 2020, Kaplan et al. (and later Hoffmann et al. in 2022 as part of the Chinchilla study) published scaling laws for Transformer language models. These papers demonstrated that model performance follows a predictable power-law relationship with three resources: model parameters, dataset size, and compute budget.

The key insight is that increasing model size, training data, and compute all improve performance with no diminishing returns visible at the scales tested — provided they are increased together at the correct ratio. The Chinchilla study showed that most large models were undertrained: given a fixed compute budget, the optimal model is smaller and trained on more data than was common practice.

Why do Transformers scale so well? Three reasons. First, the architecture has no inductive bottlenecks — attention can theoretically capture any relationship between any tokens, so adding more parameters directly increases representational capacity. Second, training at scale on diverse data produces emergent abilities that do not exist in smaller models, creating a virtuous cycle where larger models become more capable in unexpected ways. Third, the hardware utilisation of Transformers — dominated by matrix multiplications on GPUs — also scales well, with throughput remaining high even at enormous parameter counts.

For an operational perspective on the infrastructure that makes this scaling possible, see our post on AI infrastructure in 2026 and parallel processing with GPUs.

Efficient Attention: Breaking the O(n^2) Barrier

The quadratic memory and compute cost of standard attention — O(n^2) for a sequence of length n — is the fundamental constraint on context window size. A 128K-token sequence requires over 16 billion attention scores, consuming hundreds of gigabytes of GPU memory with standard attention. This has driven intense research into efficient attention mechanisms.

Flash Attention

Flash Attention, introduced by Dao et al. in 2022, is the most impactful efficiency innovation in the Transformer ecosystem. Rather than reducing the computational complexity (it still computes O(n^2) dot products), Flash Attention eliminates the memory bottleneck by never materialising the full NxN attention matrix in HBM (high-bandwidth memory).

The key insight is that the attention computation can be tiled: split the Q, K, V matrices into blocks, compute partial attention scores on-chip in SRAM (which is orders of magnitude faster than HBM), and accumulate results incrementally. This reduces HBM reads/writes from O(n^2 + n*d) to O(n^2 / M) where M is the SRAM size — typically a 10-50x speedup in practice.

from flash_attn import flash_attn_func

# Flash Attention 3 — fused kernel
# No need to materialize the full NxN attention matrix
attn_output = flash_attn_func(
    q, k, v,
    dropout_p=0.0,
    softmax_scale=None,
    causal=True
)

Flash Attention v3, released in 2025, extends the approach to FP8 precision and Hopper architecture GPUs, achieving over 2 petaFLOPS on a single H100 GPU for attention operations. It is now the default in virtually all major training and inference frameworks.

Sparse Attention and Linear Attention

Sparse attention mechanisms reduce the O(n^2) cost by only computing a subset of the attention pairs. Approaches include sliding window attention (each token attends only to nearby tokens), dilated attention (skip patterns), and global+local patterns (a small set of tokens attend globally while most attend locally). Mistral uses sliding window attention with a window of 4,096 tokens. Longformer and BigBird use combination patterns that scale to hundreds of thousands of tokens.

Linear attention (Katharopoulos et al., 2020) reformulates the softmax attention as a kernel feature map, reducing the complexity from O(n^2) to O(n). However, linear attention has not matched the quality of standard attention at scale, and most production models still prefer Flash Attention with sliding window approximations for long contexts.

Modern Innovations: GQA, RMSNorm, and RoPE

The Transformer architecture has evolved significantly since 2017. Three innovations in particular have become standard in nearly every modern model.

Grouped Query Attention (GQA)

In standard multi-head attention, each head has its own Q, K, and V projections. With 96 heads in GPT-3, this means 96 separate K and V projections per layer — a significant memory and compute cost during inference, since every head's K and V must be cached in the KV cache.

Grouped Query Attention, introduced by Shazeer in 2019, reduces the number of key-value heads while keeping the full number of query heads. For example, Llama 3 70B uses 64 query heads but only 8 key-value heads — the query heads are grouped, with each group sharing a single key-value head. This reduces the KV cache size by 8x while preserving most of the representational capacity of full MHA. GQA has become the standard in all large models because KV cache size is the dominant memory bottleneck at inference time for long sequences.

RMSNorm

RMSNorm (Root Mean Square Layer Normalisation), introduced by Zhang and Sennrich in 2019, is a simplified variant of LayerNorm that removes the mean-centering step:

RMSNorm(x) = x / sqrt(mean(x^2) + epsilon) * gamma

This reduces the normalisation computation by roughly 10-15% while producing nearly identical training dynamics. RMSNorm is used in Llama 2, Llama 3, Mistral, Gemma, and most other open-source models.

Rotary Position Embeddings (RoPE)

As discussed in the positional encoding section, RoPE has become the de facto standard for position encoding. Its relative-position bias, natural decay over distance, and support for context extension through frequency scaling (e.g., YaRN, NTK-aware scaling) make it the most practical choice for modern models. Llama 3 uses RoPE with a base frequency of 500,000 and supports context extension from 8K to 128K tokens through YaRN interpolation.

Mixture-of-Experts and Sparse Transformers

Standard Transformers use all their parameters for every token — each FFN applies the full set of weights regardless of whether the token needs that capacity. Mixture-of-Experts (MoE) replaces each FFN with multiple “expert” FFNs and a learned routing mechanism that activates only a subset of experts per token.

The Mixtral 8x7B model (Mistral, 2024) demonstrated the effectiveness of MoE at scale: it uses 8 experts per layer, routes each token to 2 experts, and achieves the performance of a 60B-parameter dense model while using only 13B parameters per forward pass. DeepSeek-V2 and V3 pushed MoE further with 256 experts and a shared expert, routing to only 6 experts per token.

MoE introduces engineering challenges: load balancing across experts (ensuring all experts receive a similar number of tokens), memory management (all expert weights must fit in GPU memory even when only a subset is used), and communication overhead in distributed training. But the compute-quality trade-off is so favourable that MoE is now the default architecture for frontier models.

From BERT to GPT to Llama: A Short History

Tracing the architectural lineage helps contextualise where we are today.

2017 — Transformer (Vaswani et al.): The original encoder-decoder architecture with sinusoidal positional encodings, ReLU activation, and post-norm LayerNorm.

2018 — BERT (Devlin et al.): Encoder-only, trained with masked language modelling and next-sentence prediction. Bidirectional context. The foundational model for NLP understanding tasks.

2018 — GPT (Radford et al.): Decoder-only, trained with auto-regressive language modelling. Lacked the encoder cross-attention of the original Transformer. Introduced the paradigm of generative pre-training followed by fine-tuning.

2020 — GPT-3 (Brown et al.): Scaled the decoder-only architecture to 175B parameters. Demonstrated few-shot learning without fine-tuning. Established scaling laws and the in-context learning paradigm.

2022 — Chinchilla (Hoffmann et al.): Showed that most models were undertrained. Introduced the compute-optimal training paradigm: smaller models trained on more data outperform larger models trained on less data for the same compute budget.

2023 — Llama (Touvron et al.): Decoder-only with a few critical changes: pre-norm with RMSNorm, SwiGLU activation, RoPE positional encodings, and no bias terms in linear layers. This recipe became the blueprint for open-source LLMs.

2024-2025 — Llama 3, DeepSeek-V2/V3, Qwen 2.5: Expanded vocabularies, grouped query attention, MoE variants, longer context windows (128K+), and FP8 training. The architecture stabilised around the Llama recipe with expert-specific optimisations.

Understanding this lineage helps when choosing a base model for fine-tuning. For a practical guide, see our LLM fine-tuning guide.

The Mathematical Intuition Behind Transformers

For readers who want to understand the maths without the notation overload: a Transformer is a differentiable program that repeatedly applies two operations — association and transformation.

Association (attention):Each token broadcasts “I contain information about X” (key) and “I am looking for information about Y” (query). Tokens that match exchange information (value). This operation is linear: it is a weighted sum of inputs. The weighting (attention scores) is itself learned through dot products that measure compatibility.

Transformation (FFN):After exchanging information, each token independently applies a non-linear function to its updated representation. This is where the model performs “thinking” — pattern matching, rule application, and feature interaction that cannot be expressed as simple weighted sums.

Stack 96 layers of this alternating pattern, and you have a system capable of representing complex reasoning, world knowledge, and linguistic structure — all learned from predicting the next word.

For more on the underlying linear algebra, see our post on vectors, tensors, and scalars in AI.