How Transformers Work: The Architecture Powering Modern AI
Introduced in the foundational 2017 paper "Attention Is All You Need" by Ashish Vaswani and colleagues at Google Brain and Google Research, the Transformer fundamentally transformed natural language processing, computer vision, audio synthesis, and reinforcement learning. Every dominant foundation model todayâincluding GPT-4, Llama 3, Claude 3.5, Gemini, and Mistralârelies on variants of this self-attention neural architecture.
To understand how transformers work, one must first recognize the fundamental computational limitation they overcame: the sequential bottleneck of recurrent neural networks.
1. The Problem With Sequential Architectures (RNNs and LSTMs)
Before transformers, state-of-the-art sequence transduction relied on Recurrent Neural Networks (RNNs), Gated Recurrent Units (GRUs), and Long Short-Term Memory networks (LSTMs). These models processed sequences step-by-step: to compute the hidden state $h_t$ at time step $t$, the network strictly required the preceding hidden state $h_{t-1}$:
This recurrent dependency imposed two fatal bottlenecks:
- No Parallelization Across Time: Training cannot compute word 50 until words 1 through 49 are sequentially computed. This prevented modern graphics processing units (GPUs) and tensor processing units (TPUs) from utilizing their massive parallel SIMD (Single Instruction, Multiple Data) computing cores.
- Vanishing Gradients and Catastrophic Forgetting: Information from the beginning of a long sequence had to survive dozens of non-linear matrix multiplications to reach the end. Despite the gating mechanisms of LSTMs, signals degraded over long token contexts, causing models to lose track of early dependencies.
The Transformer dissolved recurrence completely. Instead of feeding information sequentially through an internal memory tape, it feeds all tokens into the GPU simultaneously and calculates direct pairwise relationships across the entire sequence using Self-Attention.
2. Scaled Dot-Product Attention: The Core Engine
At the heart of the transformer is scaled dot-product attention. To compare tokens dynamically, the model projects each input vector into three distinct learned functional representations:
- Query ($Q$): What a given token is looking for in other tokens (e.g., the pronoun "it" searching for the antecedent noun it replaces).
- Key ($K$): What a given token offers or identifies as (e.g., the noun "animal" identifying its semantic properties).
- Value ($V$): The actual substantive semantic payload or content passed forward if a Query and Key match.
Given an input sequence matrix $X \in \mathbb{R}^{N \times d_{\text{model}}}$, where $N$ is sequence length and $d_{\text{model}}$ is hidden dimension, projections are computed via learned weight matrices $W_Q, W_K, W_V \in \mathbb{R}^{d_{\text{model}} \times d_k}$:
The Four Calculation Steps of Attention
- Dot-Product Similarity Matrix ($Q K^T$): The model multiplies the Query matrix by the transpose of the Key matrix. For $N$ tokens, this generates an $N \times N$ matrix where entry $(i, j)$ represents the raw unnormalized compatibility score between token $i$ and token $j$.
- Scaling by $\sqrt{d_k}$ (Square Root of Key Dimension): For large projection dimensions (such as $d_k = 64$ or $128$), the dot products grow large in magnitude: $$\operatorname{Var}\left(\sum_{k=1}^{d_k} q_k k_k\right) = d_k$$ Dividing by $\sqrt{d_k}$ normalizes the variance back to 1. Without this scaling, huge values would push the subsequent softmax function into regions with near-zero gradients, stalling backpropagation.
- Softmax Normalization: Softmax is applied row-wise along the scaled scores: $$A_{i, j} = \frac{\exp(Q_i K_j^T / \sqrt{d_k})}{\sum_{m=1}^N \exp(Q_i K_m^T / \sqrt{d_k})}$$ Every row now sums to exactly 1.0, turning arbitrary dot products into a legitimate probability distribution of attention weights.
- Weighted Aggregation of Values ($A \times V$): The attention matrix $A$ multiplies the Value matrix $V$. Each output token vector is a linear combination of all Value vectors, weighted by how relevant each token was to that position.
3. Why Multi-Head Attention (MHA)?
A single attention mechanism tends to average together conflicting relational objectives. A word often has multiple simultaneous relationships:
- Grammatical / Syntactic Role: Linking a verb to its direct object or auxiliary verb.
- Coreference / Semantic Binding: Connecting pronouns, demonstratives, and named entities.
- Positional Proximity: Attending to immediate predecessor and successor words.
Rather than performing a single attention function with dimension $d_{\text{model}}$, Multi-Head Attention projects Queries, Keys, and Values $h$ separate times with different learned parameter matrices:
In a model with $d_{\text{model}} = 512$ and $h = 8$ heads, each head operates on vectors of dimension $d_k = 512 / 8 = 64$. Because the dimension of each head is reduced, the total computational cost is identical to single-head attention with full dimensionality, yet the network gains the ability to jointly attend to information from disparate representation subspaces at disparate positions.
4. Positional Encoding: Giving Order to Unordered Sets
Because self-attention computes dot products pairwise between all tokens simultaneously, it is inherently permutation-invariant. To a pure self-attention layer, the phrase "dog bites man" produces the exact same set of attention scores as "man bites dog".
To inject sequence order, transformers add a Positional Encoding vector to each token's input embedding before the first attention block:
Each frequency corresponds to a sinusoid. For any fixed offset $k$, $PE_{pos+k}$ can be represented as a linear function of $PE_{pos}$, allowing the model to easily learn relative position. Modern LLMs frequently adopt advanced alternatives such as RoPE (Rotary Position Embeddings), which rotates query and key vectors in the complex plane to encode relative distance directly into the dot product.
5. Deep Transformer Block Anatomy: Residuals, Norms, and MLPs
A transformer model does not consist of attention layers alone. Stacking attention layers without non-linear transformation and stabilization would collapse gradient flow. Each transformer layer block contains:
| Sub-Component | Mathematical Operation | Primary Purpose |
|---|---|---|
| Multi-Head Attention | $y = \operatorname{MHA}(x)$ | Enables tokens to exchange context and information dynamically across positions. |
| Residual Connection | $x + \operatorname{Sublayer}(x)$ | Prevents vanishing gradients; creates an unimpeded gradient highway across hundreds of layers. |
| Layer Normalization | $\operatorname{LayerNorm}(z) = \frac{z - \mu}{\sigma} \odot \gamma + \beta$ | Stabilizes activation distributions across features within each token vector. |
| Feed-Forward Network (FFN) | $\operatorname{FFN}(z) = \operatorname{GELU}(z W_1 + b_1) W_2 + b_2$ | Position-wise MLP with dimension expansion (typically $4 \times d_{\text{model}}$) storing factual knowledge. |
Pre-LN vs. Post-LN Architecture
The original 2017 transformer used Post-LN ($\text{LayerNorm}(x + \text{Sublayer}(x))$). Modern production models (GPT-3/4, Llama, PaLM) almost universally use Pre-LN:
Pre-LN allows direct unnormalized residual skip connections from the initial layer to the final layer, eliminating the need for delicate warm-up learning rate schedules and preventing training instabilities in deep 100+ layer architectures.
6. Computational Complexity: The $O(N^2)$ Scaling Challenge
The dominant computational characteristic of standard self-attention is its quadratic scaling with sequence length:
Because the $Q K^T$ matrix has dimensions $N \times N$, doubling the context window from 4,000 tokens to 8,000 tokens quadruples the memory required for attention weights. At 128,000 or 1,000,000 tokens, standard attention would exhaust GPU High-Bandwidth Memory (HBM). Modern systems deploy three key architectural innovations to address this:
- FlashAttention (Dao et al.): Fuses attention matrix computation by computing softmax tiling in fast on-chip SRAM without ever writing the massive $N \times N$ matrix to GPU HBM, achieving 2â4Ã speedups.
- Multi-Query & Grouped-Query Attention (MQA / GQA): Shares Key and Value projection heads across multiple Query heads (e.g., Llama 3 uses 8 KV heads for 32 Query heads), slashing KV-cache memory footprints during autoregressive generation by up to 75%.
- KV Caching: In autoregressive generation (inference), previously computed Key and Value vectors are cached in GPU memory so each new token generated only computes a single $1 \times N$ query vector instead of recalculating the entire history.