Transformers and Self-Attention Explained

Before 2017, natural language processing relied on sequential neural networks: Recurrent Neural Networks (RNNs), LSTMs, and GRUs. While groundbreaking for their time, these architectures suffered from a critical bottleneck that made building today's trillion-parameter models impossible: sequential processing.

In December 2017, Vaswani et al. at Google Brain published "Attention Is All You Need", introducing the Transformer architecture. By replacing recurrence entirely with parallelizable Self-Attention, the foundation for modern Generative AI was born.

1. Why RNNs Failed at Scale

In an RNN or LSTM, each hidden state $h_t$ is computed sequentially from the current token $x_t$ and the previous hidden state $h_{t-1}$:

$$h_t = f(h_{t-1}, x_t)$$

flowchart LR
    x1["Token 1<br/>'The'"] --> h1["Hidden 1"]
    h1 --> h2["Hidden 2"]
    x2["Token 2<br/>'data'"] --> h2
    h2 --> h3["Hidden 3"]
    x3["Token 3<br/>'pipeline'"] --> h3
    h3 --> h4["Hidden 4"]
    x4["Token 4<br/>'failed'"] --> h4

    style x1 fill:#64748b,color:#fff
    style x2 fill:#64748b,color:#fff
    style x3 fill:#64748b,color:#fff
    style x4 fill:#64748b,color:#fff
    style h1 fill:#f59e0b,color:#fff
    style h2 fill:#f59e0b,color:#fff
    style h3 fill:#f59e0b,color:#fff
    style h4 fill:#ef4444,color:#fff

The Three Critical Bottlenecks of RNNs:

1. No Parallel Training: Token 100 cannot be processed until Token 99 finishes. This prevents modern GPUs (which rely on massive SIMD parallelism) from saturating their compute clusters.
2. Vanishing / Exploding Gradients: Backpropagating gradients through hundreds of timesteps causes numerical instability.
3. Information Forgetting: By the time an RNN reaches the end of a long paragraph, context from the opening sentence is compressed and degraded.

2. The Core Innovation: Self-Attention

Self-attention eliminates sequential loops. Every token in a sequence looks at every other token simultaneously in $O(1)$ sequential operations (parallelized across matrix tensor cores).

The Library Analogy: Query, Key, and Value

For each token input vector $x_i$, the network projects it into three distinct vector spaces using learned weight matrices:

• Query ($Q$): "What am I looking for?" (like search terms typed into Google).
• Key ($K$): "What information do I possess?" (like the index tags of documents in a database).
• Value ($V$): "What is my actual content?" (the document body itself).

flowchart TD
    subgraph AttentionMechanic["Scaled Dot-Product Attention"]
        X["Input Token Embeddings (X)"]

        X --> WQ["W_Q Matrix"]
        X --> WK["W_K Matrix"]
        X --> WV["W_V Matrix"]

        WQ --> Q["Queries (Q)"]
        WK --> K["Keys (K)"]
        WV --> V["Values (V)"]

        Q --> MM["MatMul (Q × K^T)"]
        K --> MM

        MM --> SCALE["Scale by √d_k"]
        SCALE --> SM["Softmax (Attention Weights)"]

        SM --> OUT["MatMul (Weights × V)"]
        V --> OUT

        OUT --> Z["Context-Aware Representations (Z)"]
    end

    style X fill:#475569,color:#fff
    style Q fill:#2563eb,color:#fff
    style K fill:#059669,color:#fff
    style V fill:#d97706,color:#fff
    style SM fill:#7c3aed,color:#fff
    style Z fill:#0284c7,color:#fff

The Attention Formula:

$$\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^T \in \mathbb{R}^{d_k \times N}$: Transposed Key matrix.
• $QK^T \in \mathbb{R}^{N \times N}$: Raw attention score matrix representing the affinity between every pair of tokens.
• $\sqrt{d_k}$: Scaling factor. As dimensionality $d_k$ grows large, dot products grow exponentially large, pushing softmax into regions with vanishing gradients. Dividing by $\sqrt{d_k}$ stabilizes gradient flow.
• $\text{softmax}(\cdot)$: Normalizes each row into probability weights summing to 1.0.
• Multiply by $V$: Computes a weighted sum of values based on how relevant each token is to the current token.

3. Multi-Head Attention

A single attention mechanism might focus only on syntactic relationships (e.g., matching a verb to its subject). Multi-Head Attention projects $Q, K, V$ into $h$ different subspaces (typically 8, 16, or 32 heads), computes self-attention in parallel, and concatenates the results:

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

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

This allows the model to simultaneously attend to:

• Head 1: Syntax & Grammar (subject-verb agreement).
• Head 2: Coreference resolution (linking "it" to "PostgreSQL table").
• Head 3: Long-distance dependencies (connecting function definitions to call sites).

4. The Complete Transformer Block

Each layer of a modern Transformer consists of:

1. Multi-Head Self-Attention
2. Residual Skip Connections: $x + \text{Sublayer}(x)$, preventing vanishing gradients across 80+ layers.
3. Layer Normalization (Pre-LN in modern models like LLaMA): Normalizes activations across features.
4. Feed-Forward Network (FFN): A position-wise two-layer MLP with non-linear activation (SwiGLU or GELU) that stores factual knowledge.

flowchart TD
    In["Token Input x"] --> PreLN1["RMSNorm / LayerNorm"]
    PreLN1 --> MHA["Multi-Head Attention"]
    MHA --> Add1["Add (Residual Connection)"]
    In --> Add1

    Add1 --> PreLN2["RMSNorm / LayerNorm"]
    PreLN2 --> FFN["Feed-Forward Network (SwiGLU)"]
    FFN --> Add2["Add (Residual Connection)"]
    Add1 --> Add2

    Add2 --> Out["Transformer Layer Output"]

    style In fill:#334155,color:#fff
    style MHA fill:#2563eb,color:#fff
    style FFN fill:#7c3aed,color:#fff
    style Out fill:#059669,color:#fff

5. Hands-on: Building Self-Attention in Pure NumPy

Understanding the tensor math demystifies how LLMs operate:

import numpy as np

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Compute Scaled Dot-Product Attention in Pure NumPy.

    Args:
        Q: Query matrix of shape (seq_len, d_k)
        K: Key matrix of shape (seq_len, d_k)
        V: Value matrix of shape (seq_len, d_v)
        mask: Optional causal mask of shape (seq_len, seq_len)
    """
    d_k = Q.shape[-1]

    # Step 1: Raw attention affinity scores (seq_len, seq_len)
    scores = np.matmul(Q, K.T) / np.sqrt(d_k)

    # Step 2: Apply causal mask (used in Decoder LLMs like GPT/LLaMA)
    if mask is not None:
        scores = np.where(mask == 0, -1e9, scores)

    # Step 3: Softmax along the last axis (rows sum to 1.0)
    exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
    attention_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)

    # Step 4: Weighted sum of Value vectors
    output = np.matmul(attention_weights, V)

    return output, attention_weights

# Demonstration: 4 tokens, embedding dimension d k = 8
if __name__ == "__main__":
    np.random.seed(42)
    seq_len = 4
    d_k = 8

    tokens = ["SELECT", "*", "FROM", "customers"]

    # Simulated projections
    X = np.random.randn(seq_len, d_k)
    W_q = np.random.randn(d_k, d_k)
    W_k = np.random.randn(d_k, d_k)
    W_v = np.random.randn(d_k, d_k)

    Q = np.dot(X, W_q)
    K = np.dot(X, W_k)
    V = np.dot(X, W_v)

    # Causal lower-triangular mask (Decoder-only property: cannot peek ahead)
    causal_mask = np.tril(np.ones((seq_len, seq_len)))

    output, weights = scaled_dot_product_attention(Q, K, V, mask=causal_mask)

    print("Attention Weights Matrix (Causal Lower-Triangular):")
    print(np.round(weights, 3))
    print("\nNotice how token 0 ('SELECT') can only look at itself.")
    print("Token 3 ('customers') distributes attention across all prior tokens!")

Key Takeaways

• Transformers replaced sequential RNNs with parallelizable Self-Attention, enabling scaling onto massive GPU clusters.
• Query, Key, Value matrices map relationships between every token in $O(1)$ sequential operations.
• Multi-Head Attention splits queries and keys into multiple subspaces to track grammar, syntax, and long-range semantics independently.
• Decoder models use causal masking (lower-triangular matrices) to prevent tokens from cheating by peeking into the future during next-token prediction.