Super logo Super

How Transformers Work: Interactive Attention Lab

Preset: Head:
Stage:
INPUT SEQUENCE TOKENS (Click to focus query) 10 tokens
Attention Weight Matrix $A = \text{softmax}(QK^T / \sqrt{d_k})$
Head 1/4 (d_k = 64)
Weight: 0.0
1.0
Hover on any cell to inspect mathematical dot product
Stage 3: Scaled Dot-Product Attention QK^T Calculation
$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$$
Dot-product scores determine how strongly query token "it" attends to key token "animal" versus "street".
Query [q]:
Key [k]:
Raw q·k:
+14.28
÷ √d_k (8):
+1.785
Softmax wt:
0.842 (84.2%)
TRANSFORMER BLOCK ARCHITECTURE
1. Token + Positional Encoding E + PE ∈ ℝ^(N×d)
2. Multi-Head Self-Attention + ResAdd LayerNorm(x + MHA(x))
3. Position-Wise Feed-Forward (MLP) ReLU / GeLU(xW₁ + b₁)W₂
4. Final Norm & Linear Logits Output Softmax(Linear(x))
Query token selected: "it" (Index 6) O(N²) Parallel Matrix

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}$:

h_t = \tanh(W_{hh} h_{t-1} + W_{xh} x_t + b_h)

This recurrent dependency imposed two fatal bottlenecks:

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:

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}$:

Q = X W_Q, \quad K = X W_K, \quad V = X W_V

The Four Calculation Steps of Attention

  1. 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$.
  2. 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.
  3. 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.
  4. 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.
The Winograd Schema Example in Our Simulator: In the sentence "The animal didn't cross the street because it was too tired", humans effortlessly understand that "it" refers to the animal. In the interactive visualizer above, select Head 1 and hover over the row for "it". You will see that the Query vector for "it" yields an attention score above 80% with the Key for "animal", resolving the pronoun antecedent in a single feed-forward pass.

3. Why Multi-Head Attention (MHA)?

A single attention mechanism tends to average together conflicting relational objectives. A word often has multiple simultaneous relationships:

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:

\operatorname{MultiHead}(Q, K, V) = \operatorname{Concat}(\text{head}_1, \dots, \text{head}_h) W^O \\ \text{where } \text{head}_i = \operatorname{Attention}(Q W_i^Q, K W_i^K, V W_i^V)

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:

PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i / d_{\text{model}}}}\right), \quad PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i / d_{\text{model}}}}\right)

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:

x_{t+1} = x_t + \operatorname{MHA}(\operatorname{LayerNorm}(x_t))

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:

\text{Time Complexity} = O(N^2 \cdot d), \quad \text{Memory Complexity} = O(N^2 \cdot h + N \cdot d)

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:

Frequently Asked Questions

What is self-attention in a transformer model? +
Self-attention is a mathematical mechanism that allows each word or token in a sequence to score and incorporate information from every other token in the sequence simultaneously, dynamically weighting relationships regardless of word distance.
Why do transformers divide Query-Key dot products by the square root of d_k? +
Dividing by the square root of the key dimension (sqrt(d_k)) prevents the dot products from growing excessively large for high dimensions. Without scaling, large dot products push the softmax activation function into regions with extremely small gradients, causing vanishing gradients during training.
What is the purpose of having multiple attention heads? +
Multi-head attention projects Queries, Keys, and Values into multiple lower-dimensional representation subspaces. This enables the model to simultaneously attend to information from different representation aspects, such as syntactic dependencies in one head and semantic references or pronoun binding in another.
Why do transformers need positional encodings? +
Unlike Recurrent Neural Networks (RNNs) that process text sequentially step-by-step, self-attention operations are permutation-invariant. Positional encodings add deterministic or learned spatial order signals to token embeddings so the model knows word order.