Attention Mechanism

The Self-Attention Mechanism is the mathematical engine at the core of the Transformer architecture.

It gives language models the ability to dynamically route information across a sequence—allowing every token to weigh, reference, and extract context from every other relevant token in the prompt.

#Why does this matter?

Human language is deeply interconnected and context-dependent. Pronouns, adjectives, and verbs depend on references scattered across paragraphs:

Consider these two sentences:

Sentence A: "The animal didn't cross the street because it was too tired."

Sentence B: "The animal didn't cross the street because it was too wide."

What does the word "it" refer to in each sentence?

• In Sentence A, because of the word "tired", "it" refers to the animal.
• In Sentence B, because of the word "wide", "it" refers to the street.

A human resolves this ambiguity effortlessly. The self-attention mechanism allows a neural network to calculate this exact relationship: when computing the vector for "it", the model places heavy mathematical "attention" on "animal" in Sentence A, and on "street" in Sentence B.

#The Query, Key, and Value Intuition

To compute attention, the Transformer borrows a concept from database systems and search retrieval: Queries, Keys, and Values ($Q, K, V$).

For every token in the sequence, the model creates three distinct vectors by multiplying the token's embedding by three learned weight matrices ($W_Q, W_K, W_V$):

┌─────────────────────────────────────────────────────────────┐
│ 1. Query (Q): What is this token searching for?             │
│    (e.g., "I am a pronoun looking for my subject")          │
├─────────────────────────────────────────────────────────────┤
│ 2. Key (K): What information does this token offer?         │
│    (e.g., "I am a noun describing a physical animal")       │
├─────────────────────────────────────────────────────────────┤
│ 3. Value (V): What is the actual semantic content payload?  │
│    (e.g., The rich descriptive features of the animal)      │
└─────────────────────────────────────────────────────────────┘

#The Self-Attention Formula Step-by-Step

The complete scaled dot-product attention formula is:

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

Let's break down this calculation intuitively into four straightforward operations:

Step 1: Calculate Compatibility (Dot Product)
        Token Query (Q)  •  All Token Keys (Kᵀ) ──► Raw Similarity Scores

Step 2: Scale Factor (Stability)
        Divide by √d_k to prevent numerical gradients from exploding

Step 3: Normalize to Probabilities (Softmax)
        Convert raw scores into percentages (weights) that sum to 100% (1.0)

Step 4: Weighted Sum of Values (V)
        Multiply each token's Value vector by its attention score and sum them up

The output is a new, enriched contextual embedding for the token that incorporates relevant details from every other word in the prompt.

#Multi-Head Attention: Multiple Perspectives

A single attention calculation can only focus on one type of relationship at a time (e.g., pronoun resolution).

In practice, human language contains dozens of simultaneous structural relationships:

• Subject-verb agreements
• Direct object relationships
• Temporal relationships (verb tenses)
• Factual and entity associations
• Punctuation and clause boundaries

Multi-Head Attention solves this by splitting the embedding into multiple parallel "heads" (often 32 to 128 heads in production models):

Input Vector ──┬──► Head 1: Focuses on grammatical syntax
               ├──► Head 2: Focuses on coreference (pronoun tracking)
               ├──► Head 3: Focuses on factual entity connections
               └──► Head N: Focuses on long-range dependencies
                      │
                      ▼
               Concatenate all Heads ──► Linear Projection ──► Output

Each head learns to specialize in a different linguistic or logical pattern independently.

#Causal Attention (Masked Attention)

In generative decoder-only models (like GPT-4 or Llama), models must generate text autoregressively—predicting one token into the future based only on past tokens.

If a token at position 3 could look ahead at position 5 during training, the model would simply "cheat" by reading the answer rather than learning to predict it.

To prevent looking into the future, decoder models apply a Causal Mask (an upper-triangular matrix of $-\infty$ values):

                  Position Attended To:
                  Token 1   Token 2   Token 3   Token 4
Token 1 ("The")   [  ✓         ✗         ✗         ✗    ]
Token 2 ("AI")    [  ✓         ✓         ✗         ✗    ]
Token 3 ("Rishi") [  ✓         ✓         ✓         ✗    ]
Token 4 ("is")    [  ✓         ✓         ✓         ✓    ]

Every token can attend to itself and all preceding tokens, but future tokens are masked out completely.

#The Computational Cost: Quadratic Scaling $O(N^2)$

Because every token calculates a dot product with every other token in the prompt, the computational cost of standard self-attention grows quadratically with sequence length:

$$\text{Attention Operations} \propto N^2$$

• A prompt of $1,000\text{ tokens} \implies 1,000,000\text{ attention calculations}$.
• A prompt of $100,000\text{ tokens} \implies 10,000,000,000\text{ attention calculations}$.

This quadratic scaling is why ultra-long context windows require immense GPU memory (VRAM) and specialized hardware kernels like FlashAttention to execute efficiently in production.

#Common Misconceptions

• Misconception: Attention means the model is conscious and paying attention like a person. Reality: Attention is a mathematical matrix dot-product that calculates statistical correlation weights between vector representations.
• Misconception: Higher attention scores always mean factual importance. Reality: Attention heads often place high attention weights on syntax delimiters (like commas, periods, or the first token) simply as numerical anchors for the network.

#Key Takeaways

[!IMPORTANT]

• Self-attention allows tokens to dynamically share information across the entire sequence.
• The Query ($Q$), Key ($K$), Value ($V$) system functions like a soft database lookup: $Q$ searches, $K$ indexes, and $V$ provides the content payload.
• Multi-Head Attention allows the model to simultaneously track grammar, facts, coreferences, and syntax.
• Causal masking prevents generative models from looking ahead at future tokens.
• Standard attention scales quadratically ($O(N^2)$) with context length.

#What comes next?

We now understand the architecture of a Transformer and how self-attention operates. But how do the billions of weights in these matrices actually get their values?

In Lesson 06, we explore Training an LLM—from self-supervised pre-training to instruction tuning and RLHF alignment.