Understanding 'Attention Is All You Need': A Mathematical & Visual Deep Dive into Transformers

Introduction: Why Recurrence Was Replaced

Before 2017, the state of the art in sequence-to-sequence modeling for Natural Language Processing (NLP) relied on Recurrent Neural Networks (RNNs), LSTMs, and GRUs. While groundbreaking for their time, recurrent networks suffered from two fundamental architectural limits:

flowchart LR
    subgraph RNN["Traditional RNN: Strictly Sequential O(N) Processing"]
        x1["x1: The"] --> h1["h1"]
        x2["x2: animal"] --> h2["h2"]
        x3["x3: didn't"] --> h3["h3"]
        x4["x4: cross"] --> h4["h4"]
        h1 -->|Sequential bottleneck| h2
        h2 -->|Information decay| h3
        h3 -->|Cannot parallelize| h4
    end

1. Sequential Computation Bottleneck ($O(N)$ Path): The hidden state at time-step $t$ depends strictly on $t-1$: $h_t = \sigma(W_{hh} h_{t-1} + W_{xh} x_t + b)$ Because tokens must enter the GPU sequentially, we cannot compute token representations across long paragraphs simultaneously.
2. Long-Range Context Decay: When connecting words separated by 50 or 100 tokens, gradients and semantic signals decay over the chain of recurrent matrix multiplications.

In June 2017, Vaswani et al. published “Attention Is All You Need”, proposing the Transformer—an architecture that dispenses entirely with recurrence and convolutions, computing direct all-to-all relationships across all tokens in parallel in a single matrix multiplication step ($O(1)$ sequential operations).

flowchart TD
    subgraph Transformer["Transformer: Parallel All-to-All Self-Attention O(1)"]
        Tokens["Tokens: The, animal, didn't, cross, the, street, because, it, was, tired"]
        Matrix["All-to-All Self-Attention Matrix (Q · Kᵀ)"]
        Output["Context-Enriched Dynamic Embeddings"]
        Tokens ==> Matrix
        Matrix ==> Output
    end

---

1. The Core Intuition: Queries, Keys, and Values

To understand self-attention, consider a search engine database analogy:

Component	Database Analogy	Transformer Role
Query ($Q$)	What you type into the search bar	What the current token is searching for
Key ($K$)	Video / webpage titles and tags	What other tokens advertise about themselves
Value ($V$)	The actual video content you watch	The semantic information contributed to the context

flowchart LR
    Q["Query Q: 'it' (Searching for antecedent)"] 
    K1["Key K1: 'street' (Physical location)"]
    K2["Key K2: 'animal' (Animate entity)"]
    
    Q -->|"Score: Low (0.05)"| K1
    Q -->|"Score: High (0.92)"| K2
    
    K2 ==>|"Extract Value Vector"| V2["Value V2: Semantics of 'animal'"]
    V2 ==> Output["Enriched embedding of 'it' with animal semantics"]

---

2. Scaled Dot-Product Attention

The core attention formula is defined as:

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

Where:

• $Q \in \mathbb{R}^{N \times d_k}$: Query matrix for $N$ tokens.
• $K \in \mathbb{R}^{M \times d_k}$: Key matrix for $M$ tokens.
• $V \in \mathbb{R}^{M \times d_v}$: Value matrix containing token semantics.
• $d_k$: Dimension of keys and queries (e.g., $d_k = 64$).
• $\sqrt{d_k}$: Scaling factor.

flowchart TD
    Q[Query Matrix Q] --> MatMul1["1. Matrix Multiply: S = Q · Kᵀ"]
    K[Key Matrix K] --> MatMul1
    MatMul1 --> Scale["2. Scale: S ÷ √dₖ"]
    Scale --> Mask["3. Optional Causal / Padding Mask"]
    Mask --> Softmax["4. Softmax across rows: W = softmax(S)"]
    Softmax --> MatMul2["5. Weighted Sum: Output = W · V"]
    V[Value Matrix V] --> MatMul2
    MatMul2 --> Result[Final Contextualized Output]

Why Scale by the Square Root of Key Dimension?

Assume the query and key components $q_i$ and $k_i$ are independent random variables with mean $0$ and variance $1$. The dot product is:

$q \cdot k = \sum_{i=1}^{d_k} q_i k_i$

• Mean: $\mathbb{E}[q \cdot k] = \sum \mathbb{E}[q_i]\mathbb{E}[k_i] = 0$
• Variance: $\text{Var}(q \cdot k) = \sum_{i=1}^{d_k} \text{Var}(q_i k_i) = d_k$

For $d_k = 64$, the variance is $64$ and standard deviation is $\sigma = 8$. Large values in magnitude push the softmax function into exponential saturation where gradients approach zero:

$\frac{\partial \text{softmax}(z)_i}{\partial z_j} \approx 0 \quad \text{for large } |z|$

Dividing by $\sqrt{d_k}$ normalizes the variance back to $1$, ensuring healthy gradient flow during backpropagation.

---

3. Concrete Numerical Walkthrough with Visual Heatmap

Let us trace attention for a sequence of 3 tokens: "AI", "builds", "future".

Suppose $d_k = 2$ and we project our tokens into $Q, K, V$:

flowchart LR
    Token1["Token 1: AI"] --> Q1["Q₁ = [1, 2]"] & K1["K₁ = [1, 2]"] & V1["V₁ = [10, 0]"]
    Token2["Token 2: builds"] --> Q2["Q₂ = [0, 1]"] & K2["K₂ = [0, 1]"] & V2["V₂ = [0, 20]"]
    Token3["Token 3: future"] --> Q3["Q₃ = [2, 1]"] & K3["K₃ = [2, 1]"] & V3["V₃ = [5, 15]"]

Step 1: Compute Similarity Scores ($S = QK^T$)

$S = \begin{bmatrix} 1 & 2 \\ 0 & 1 \\ 2 & 1 \end{bmatrix} \begin{bmatrix} 1 & 0 & 2 \\ 2 & 1 & 1 \end{bmatrix} = \begin{bmatrix} 5 & 2 & 4 \\ 2 & 1 & 1 \\ 4 & 1 & 5 \end{bmatrix}$

Step 2: Scale ($S / \sqrt{2} \approx S / 1.414$)

$\text{Scaled } S = \begin{bmatrix} 3.53 & 1.41 & 2.82 \\ 1.41 & 0.70 & 0.70 \\ 2.82 & 0.70 & 3.53 \end{bmatrix}$

Step 3: Softmax Probabilities (Row-wise $\sum = 1.0$)

$\text{Attention Weights } W = \begin{bmatrix} 0.61 & 0.07 & 0.32 \\ 0.50 & 0.25 & 0.25 \\ 0.32 & 0.07 & 0.61 \end{bmatrix}$

Visualizing Attention Weight Heatmap:

Query Token	Attends to “AI”	Attends to “builds”	Attends to “future"
"AI”	61% (self)	7%	32% (contextual connection)
“builds”	50% (subject)	25%	25% (object)
“future”	32% (connected)	7%	61% (self)

Step 4: Multiply by Value Matrix ($Output = W \times V$)

$\text{Output} = \begin{bmatrix} 0.61 & 0.07 & 0.32 \\ 0.50 & 0.25 & 0.25 \\ 0.32 & 0.07 & 0.61 \end{bmatrix} \begin{bmatrix} 10 & 0 \\ 0 & 20 \\ 5 & 15 \end{bmatrix} = \begin{bmatrix} 7.70 & 6.20 \\ 6.25 & 8.75 \\ 6.25 & 10.55 \end{bmatrix}$

Each token vector now contains a contextual blend of its own meaning plus relevant related tokens!

---

4. Multi-Head Attention: Parallel Semantic Views

A single attention calculation averages all token relationships into one distribution. To capture different linguistic patterns simultaneously, Multi-Head Attention projects queries, keys, and values $h$ times into distinct lower-dimensional subspaces:

flowchart TD
    Input[Input Embeddings] --> SplitQ[Project Q into 8 Heads] & SplitK[Project K into 8 Heads] & SplitV[Project V into 8 Heads]
    
    SplitQ & SplitK & SplitV --> H1["Head 1: Syntactic Dependencies (Subject - Verb)"]
    SplitQ & SplitK & SplitV --> H2["Head 2: Coreference Resolution (it to animal)"]
    SplitQ & SplitK & SplitV --> H3["Head 3: Positional Proximity (Adjacent tokens)"]
    SplitQ & SplitK & SplitV --> H8["Head 8: Semantic Tone & Mood"]
    
    H1 & H2 & H3 & H8 --> Concat["Concatenate all 8 Heads (8 × 64 = 512)"]
    Concat --> OutputLinear["Linear Projection Wᵒ (512 -> 512)"]
    OutputLinear --> FinalOut[Multi-Head Output]

Mathematical Formulation

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O$

$\text{where } \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$

With standard parameters:

• Model Dimension $d_{\text{model}} = 512$
• Number of Heads $h = 8$
• Head Dimension $d_k = d_v = d_{\text{model}} / h = 64$
• Total parameters for projections equal single-head attention with full dimensionality.

---

5. Positional Encoding: Encoding Sequence Order

Because self-attention is a set operation (permutation invariant), the model has zero inherent awareness of token position. Vaswani et al. added fixed Sinusoidal Positional Encodings directly to input token embeddings:

$\text{PE}_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$

$\text{PE}_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)$

flowchart LR
    Word["Word Embedding: 'Deep' (d_model = 512)"] 
    Pos["Positional Vector: pos = 0 (d_model = 512)"]
    Word --> Add["+ Element-wise Addition"]
    Pos --> Add
    Add --> Out["Position-Aware Input Vector to Encoder"]

Mathematical Property: Linear Relative Offsets

For any fixed offset $k$, there exists a linear transformation matrix $M_k$ such that:

$\text{PE}_{pos + k} = M_k \cdot \text{PE}_{pos}$

Using the trigonometric identity: $\sin(\alpha + \beta) = \sin(\alpha)\cos(\beta) + \cos(\alpha)\sin(\beta)$ $\cos(\alpha + \beta) = \cos(\alpha)\cos(\beta) - \sin(\alpha)\sin(\beta)$

This allows the model to learn relative position distances seamlessly.

---

6. The Complete Encoder-Decoder Architecture

flowchart TB
    subgraph EncoderStack["Encoder (N = 6 Stacked Identical Layers)"]
        Inp[Inputs] --> Emb1[Input Embedding + Positional Encoding]
        Emb1 --> MHA1[Multi-Head Self-Attention]
        MHA1 --> AddNorm1[Add & LayerNorm]
        AddNorm1 --> FFN1[Position-Wise Feed Forward Network]
        FFN1 --> AddNorm2[Add & LayerNorm]
    end

    subgraph DecoderStack["Decoder (N = 6 Stacked Identical Layers)"]
        Outp[Outputs Shifted Right] --> Emb2[Output Embedding + Positional Encoding]
        Emb2 --> MaskMHA[Masked Multi-Head Self-Attention]
        MaskMHA --> DecNorm1[Add & LayerNorm]
        DecNorm1 --> CrossMHA[Cross-Attention: Queries from Decoder, Keys/Values from Encoder]
        AddNorm2 -.->|K, V Memory Keys| CrossMHA
        CrossMHA --> DecNorm2[Add & LayerNorm]
        DecNorm2 --> DecFFN[Position-Wise Feed Forward]
        DecFFN --> DecNorm3[Add & LayerNorm]
        DecNorm3 --> Lin[Linear Classifier]
        Lin --> Smax[Softmax Vocabulary Probabilities]
    end

Key Architectural Mechanisms:

1. Masked Decoder Self-Attention (Lookahead / Causal Mask): When predicting token $t$, the decoder must not peek at tokens $t+1, t+2, \dots$. A lower-triangular mask sets upper values to $-\infty$: $\text{Mask} = \begin{bmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{bmatrix} \xrightarrow{\text{softmax}} \begin{bmatrix} 1.0 & 0 & 0 \\ 0.5 & 0.5 & 0 \\ 0.33 & 0.33 & 0.33 \end{bmatrix}$
2. Cross-Attention: Queries ($Q$) originate from the previous decoder layer, while Keys ($K$) and Values ($V$) originate from the final Encoder output, enabling English-to-French translation alignment.
3. Position-Wise Feed-Forward Network (FFN): Applied to each position separately and identically with an expansion factor of 4 ($d_{ff} = 2048$): $\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2$

---

7. Clean PyTorch Implementation

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int = 512, num_heads: int = 8):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads # 64
        
        # Linear projections for Queries, Keys, and Values
        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)
        
        # 1. Project & Reshape: [Batch, SeqLen, Heads, d_k] -> Transpose to [Batch, Heads, SeqLen, d_k]
        Q = self.w_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.w_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.w_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 2. Scaled Dot-Product: (Q @ K.T) / sqrt(d_k)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        # 3. Apply Causal or Padding Mask
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
            
        attn_weights = F.softmax(scores, dim=-1)
        
        # 4. Context Matrix Multiply: Weights @ V
        context = torch.matmul(attn_weights, V)
        
        # 5. Concatenate heads and project output
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        return self.w_o(context), attn_weights

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int = 512, max_len: int = 5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

---

8. Summary Comparison

Metric	Recurrent Networks (LSTM)	Transformer (Self-Attention)
Sequential Operations	$O(N)$ (slow serial processing)	$O(1)$ (fully parallel GPU tensor ops)
Max Information Path	$O(N)$ (vanishing context)	$O(1)$ (direct token-to-token link)
Complexity per Layer	$O(N \cdot d^2)$	$O(N^2 \cdot d)$
Hardware Scalability	Poor across multi-GPU nodes	Optimal (Powers modern LLMs & GPT-4)

Vaswani et al.’s breakthrough established that attention is indeed all you need—replacing recurrent bottlenecks with parallel linear algebra and unlocking the modern era of Generative AI.