Transformer Architecture

AI & Machine Learning Foundations beginner How AI Works 12 min read

Overview

The Transformer is the neural network architecture that powers virtually every modern generative AI system — from OpenAI’s GPT series to Google’s Gemini to Anthropic’s Claude. Introduced in the landmark 2017 paper “Attention Is All You Need” by Vaswani et al. at Google Brain, the Transformer replaced recurrent neural networks (RNNs) and convolutional neural networks (CNNs) as the dominant architecture for sequence processing tasks.

What made the Transformer revolutionary was a single insight: you don’t need recurrence or convolution to model sequences. Instead, the architecture relies entirely on a mechanism called self-attention, which allows every token in a sequence to directly attend to every other token in a single computation step. This eliminates the sequential bottleneck of RNNs — where each word must be processed one after another — and enables massive parallelization during training.

For AI engineers, understanding the Transformer is foundational. Every concept in this training program — from embeddings and tokenization to RAG pipelines and agentic systems — ultimately runs on top of this architecture. Whether you’re tuning inference parameters, debugging retrieval quality, or estimating API costs, your decisions trace back to how Transformers process information.

How It Works

The Encoder-Decoder Structure

The original Transformer uses an encoder-decoder design. The encoder reads the entire input sequence and produces a rich contextual representation. The decoder then generates the output sequence one token at a time, attending to both its own previously generated tokens and the encoder’s output.

The original paper stacked 6 encoder layers and 6 decoder layers. Each encoder layer contains two sub-layers: a multi-head self-attention mechanism and a position-wise feed-forward network. Each decoder layer adds a third sub-layer: cross-attention over the encoder output.

Key Concept

Modern LLMs like GPT and Claude use a decoder-only architecture — they drop the encoder entirely and use only the decoder stack with causal (masked) self-attention. This is because language generation is fundamentally an autoregressive task: predict the next token given all previous tokens.

Input Embeddings & Positional Encoding

Before any attention computation, raw text is converted into numerical representations through two steps:

  1. Token Embeddings: Each token (subword unit) is mapped to a dense vector, typically of dimension d_model = 512 or 768 or larger. These vectors are learned during training.
  2. Positional Encoding: Since the Transformer processes all tokens simultaneously (no inherent notion of order), positional information must be explicitly injected. The original paper uses sinusoidal functions at different frequencies:
PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

The token embedding and positional encoding are added element-wise, giving each token both semantic meaning and position awareness. Modern models often use learned positional embeddings or relative position encodings like RoPE (Rotary Position Embedding) instead.

Multi-Head Self-Attention

Self-attention is the core mechanism. For each token, the model computes three vectors — Query (Q), Key (K), and Value (V) — by multiplying the input embedding by learned weight matrices. The attention output is computed as:

Attention(Q, K, V) = softmax(Q · K^T / √d_k) · V

The dot product Q · K^T produces a score matrix indicating how much each token should attend to every other token. Dividing by √d_k prevents the dot products from growing too large, which would push the softmax into regions with extremely small gradients. The softmax normalizes scores into attention weights (summing to 1), and these weights are used to create a weighted sum of the Value vectors.

Multi-head attention runs this process h times in parallel (typically 8 or 16 heads), each with different learned projections. This allows the model to simultaneously capture different types of relationships — one head might learn syntactic dependencies, another semantic similarity, another positional patterns. The outputs of all heads are concatenated and projected through a final linear layer.

Feed-Forward Network

After the attention sub-layer, each token’s representation passes through a position-wise feed-forward network (FFN) — the same two-layer MLP applied independently to each position:

FFN(x) = max(0, x · W1 + b1) · W2 + b2

The inner dimension is typically 4× the model dimension (e.g., d_model = 512 → inner dim = 2048). This is where much of the model’s “knowledge” is stored — the FFN layers act as key-value memory banks encoding learned facts and patterns.

Residual Connections & Layer Normalization

Each sub-layer (attention and FFN) is wrapped with a residual connection and layer normalization:

output = LayerNorm(x + Sublayer(x))

Residual connections allow gradients to flow directly through the network during backpropagation, preventing the vanishing gradient problem in deep stacks. Layer normalization stabilizes training by normalizing activations across the feature dimension.

Output: Linear Layer & Softmax

In the final layer, the decoder’s output is projected through a linear layer to a vector of size equal to the vocabulary (often 32,000–100,000+ tokens), followed by a softmax to produce a probability distribution over the next token. The token with the highest probability (or one sampled from the distribution) becomes the output.

Practical Applications

Why Transformers Won

The Transformer’s dominance comes down to one word: parallelism. RNNs process sequences step-by-step — token 5 cannot be computed until tokens 1-4 are done. Transformers compute all token relationships simultaneously, making them dramatically faster to train on modern GPU/TPU hardware. This parallelism enabled scaling to billions and then trillions of parameters.

Architectural Variants in Production

Scale in Practice

Modern Transformer-based models are massive. GPT-4 is estimated at ~1.8 trillion parameters across a mixture-of-experts architecture. Claude and Gemini operate at similar scales. Training these models requires thousands of GPUs running for weeks to months, costing tens of millions of dollars. Understanding the architecture helps you reason about why inference costs what it does and why context windows have the limits they do.

Key Takeaways

Further Reading