Attention Mechanism Guide

TL;DR

• Attention computes weighted relationships between all pairs of tokens in a sequence - determining how much each token should influence every other token.

• Self-attention uses Query, Key, and Value vectors derived from the same input, enabling each position to attend to all other positions in parallel.

• Multi-head attention runs multiple attention operations in parallel, each learning different relationship types (syntax, semantics, coreference, position).

• Attention is O(N²) in sequence length - this quadratic scaling is why context windows have limits and why long prompts are expensive.

• Attention patterns explain LLM behavior - lost-in-the-middle, prompt sensitivity, and in-context learning all stem from how attention weights are computed and applied.

Why This Matters

Attention is not an implementation detail - it is the mechanism that makes LLMs work. Every capability you rely on - understanding context, maintaining coherence across paragraphs, following instructions in system prompts, learning from few-shot examples - is a consequence of attention patterns.

When a model ignores your system prompt, it is because attention weights did not prioritize those tokens. When RAG fails because the model "did not use" retrieved context, attention did not assign sufficient weight to the retrieved passages. When the model exhibits lost-in-the-middle behavior, it is because attention patterns U-shaped across long sequences.

Understanding attention transforms debugging from guesswork into diagnosis. You stop asking "why did the model ignore my instructions?" and start asking "where were my instructions in the context, and what competed for attention?"

The Problem Attention Solves

Sequence models need to relate distant tokens. In the sentence "The cat that chased the mouse across the yard was tired," the word "was" must agree with "cat" (not "mouse" or "yard") despite being separated by six tokens.

RNNs solved this by passing information through sequential hidden states - but distant relationships degraded through many steps (vanishing gradients). Convolutional models used fixed-size windows - but could not reach beyond the window.

Attention solves this directly: any token can access any other token in a single operation, regardless of distance. Token 12 can attend to token 1 with full strength, no degradation.

The insight from Bahdanau et al. (2014) for machine translation: instead of compressing the entire source sentence into one vector, let the decoder attend to different parts of the source at each step. Vaswani et al. (2017) generalized this to self-attention - where a sequence attends to itself - and showed it was sufficient on its own (no recurrence needed).

What Is Attention?

Attention is a mechanism that computes a weighted sum of values, where the weights are determined by the compatibility between a query and a set of keys. In self-attention, queries, keys, and values all come from the same input sequence.

The core operation for a single attention head:

Attention(Q, K, V) = softmax(QK^T / √d_k) × V

Where:

• Q (Query) - "What am I looking for?" - one vector per token

• K (Key) - "What do I contain?" - one vector per token

• V (Value) - "What information do I provide?" - one vector per token

• d_k - dimension of key vectors, used for scaling

Each token produces a query, key, and value by multiplying its embedding by learned weight matrices (W_Q, W_K, W_V). The dot product QK^T measures compatibility between every query-key pair. Softmax converts these to weights that sum to 1. The output is a weighted sum of value vectors.

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

def scaled_dot_product_attention(
    Q: torch.Tensor,  # (seq_len, d_k)
    K: torch.Tensor,  # (seq_len, d_k)
    V: torch.Tensor,  # (seq_len, d_v)
    mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    # scores shape: (seq_len, seq_len) - every query vs every key

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))

    weights = F.softmax(scores, dim=-1)  # Attention weights
    output = torch.matmul(weights, V)     # Weighted sum of values
    return output, weights

How Attention Works

Self-Attention Step by Step

For the input "The cat sat":

Multi-head attention runs several attention functions in parallel, letting the model attend to information from different representation subspaces.

Step 1: Project to Q, K, V

Each token embedding is multiplied by three learned weight matrices:

"The" → Q_the, K_the, V_the
"cat" → Q_cat, K_cat, V_cat
"sat" → Q_sat, K_sat, V_sat

Step 2: Compute attention scores

For the token "sat", compute dot products with all keys:

score(sat, The) = Q_sat · K_the = 0.3
score(sat, cat) = Q_sat · K_cat = 0.8
score(sat, sat) = Q_sat · K_sat = 0.5

Step 3: Softmax to get weights

weight(sat, The) = 0.15
weight(sat, cat) = 0.55  ← "sat" attends most to "cat"
weight(sat, sat) = 0.30

Step 4: Weighted sum of values

output_sat = 0.15 × V_the + 0.55 × V_cat + 0.30 × V_sat

The output for "sat" is now enriched with information from "cat" - the model knows what sat.

Causal (Masked) Attention

Decoder-only LLMs use causal masking - each token can only attend to itself and previous tokens, not future ones:

Attention matrix for "The cat sat" (✓ = allowed, ✗ = masked):

         The  cat  sat
The       ✓   ✗   ✗
cat       ✓   ✓   ✗
sat       ✓   ✓   ✓

This ensures the model cannot "cheat" by looking at future tokens during training or generation.

Multi-Head Attention

Instead of one attention operation, transformers run h parallel heads, each with smaller dimension:

Scaled dot-product attention maps a query and a set of key-value pairs to an output, scaling dot products by 1/√d_k before softmax.

Each head learns different attention patterns:

• Head 1 might track syntactic relationships (subject-verb agreement)

• Head 2 might track semantic similarity (coreference: "it" → "cat")

• Head 3 might track positional patterns (adjacent tokens)

• Head 4 might track long-range dependencies

Llama 3 8B uses 32 attention heads with head dimension 128. GPT-4 uses ~128 heads. More heads allow finer-grained relationship modeling.

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.W_q = torch.nn.Linear(d_model, d_model)
        self.W_k = torch.nn.Linear(d_model, d_model)
        self.W_v = torch.nn.Linear(d_model, d_model)
        self.W_o = torch.nn.Linear(d_model, d_model)

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
        batch, seq_len, d_model = x.shape

        Q = self.W_q(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2)

        attn_output, weights = scaled_dot_product_attention(Q, K, V, mask)

        attn_output = attn_output.transpose(1, 2).contiguous().view(batch, seq_len, d_model)
        return self.W_o(attn_output), weights

Architecture

Attention appears in three contexts within the transformer ecosystem:

The retriever scores document chunks against the query embedding and returns the top-k passages that ground the generator.

Step-by-Step Flow

Processing one token through the attention mechanism during inference:

1. Input embedding - Token ID mapped to d_model-dimensional vector.

2. Linear projections - Compute Q, K, V by multiplying with W_Q, W_K, W_V.

3. Compute scores - QK^T produces (seq_len × seq_len) compatibility matrix.

4. Scale - Divide by √d_k to prevent softmax saturation with large dimensions.

5. Apply mask - Zero out future positions (causal) or padding positions.

6. Softmax - Convert scores to attention weights (rows sum to 1).

7. Weighted sum - Multiply weights by V to produce attended output.

8. Multi-head concat - Combine outputs from all heads, project with W_O.

9. Residual connection - Add input to attention output: x + attention(x).

10. Feed-forward - Pass through FFN sub-layer (separate from attention).

During generation, steps 3–7 use cached K and V from previous tokens (KV cache) - only the new token's Q is computed against all cached K/V pairs.

Real Production Example

Attention patterns explain why prompt structure matters. Here is a diagnostic tool that analyzes how context placement affects model behavior:

import tiktoken
from openai import OpenAI

client = OpenAI()

def test_attention_sensitivity(
    fact: str,
    question: str,
    distractor_count: int = 10,
    distractor_length: int = 200,
) -> dict:
    """
    Test whether the model retrieves a fact placed at different
    positions in context - revealing attention sensitivity.
    """
    distractor = "This is filler content about unrelated topics. " * (distractor_length // 10)
    results = {}

    positions = {
        "beginning": f"IMPORTANT FACT: {fact}\n\n" + distractor * distractor_count + question,
        "middle": distractor * (distractor_count // 2) + f"\n\nIMPORTANT FACT: {fact}\n\n" + distractor * (distractor_count // 2) + question,
        "end": distractor * distractor_count + f"\n\nIMPORTANT FACT: {fact}\n\n" + question,
    }

for position, prompt in positions.items():
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "Answer based only on the provided context."},
                {"role": "user", "content": prompt},
            ],
            temperature=0,
            max_tokens=100,
        )
        answer = response.choices[0].message.content
        results[position] = {
            "answer": answer,
            "fact_in_answer": fact.lower() in answer.lower(),
            "prompt_tokens": response.usage.prompt_tokens,
        }

    return results

# Usage
results = test_attention_sensitivity(
    fact="The API rate limit is 1000 requests per minute.",
    question="What is the API rate limit?",
)
for pos, data in results.items():
    print(f"{pos}: fact found = {data['fact_in_answer']}, tokens = {data['prompt_tokens']}")

This test typically shows: facts at the beginning and end are retrieved reliably; facts in the middle are missed - direct evidence of U-shaped attention patterns.

Design Decisions

⚠ Common Mistakes

1. Burying instructions in the middle of long context - Attention weights are U-shaped. Instructions at position 50,000 in a 100K context get minimal attention. Place critical instructions at the beginning.

2. Assuming the model "reads" all context equally - Attention is selective, not uniform. The model weights some tokens heavily and ignores others. More context does not mean more attention to your data.

3. Confusing attention with memory - Attention operates within a single forward pass. It is not persistent memory. Each request recomputes attention from scratch. Use external memory systems for cross-session persistence.

4. Ignoring competing attention targets - A system prompt competes with conversation history, retrieved documents, and the user message for attention. Every token added dilutes attention to existing tokens.

5. Expecting precise attention control via prompting - You cannot tell the model "pay attention to paragraph 3." You can only influence attention through position, repetition, and formatting (headers, bold, separation).

Where It Breaks Down

• Quadratic scaling - Attention compute and memory grow as O(N²). A 128K context requires 16 billion attention score computations per head per layer. This is the hard limit on context growth.

• Softmax saturation - When one key matches a query very strongly, softmax assigns nearly all weight to that key, ignoring everything else. This causes the model to fixate on one piece of context and miss contradictory information.

• Attention dilution - With many tokens in context, attention weights spread thin. Each individual token gets less weight, making it harder for the model to extract specific facts from large contexts.

• No selective attention control - Unlike human attention, you cannot programmatically tell the model which tokens to attend to (without fine-tuning or architectural modifications like sparse attention).

• Inductive bias limitations - Attention learns positional relationships from data, not from built-in structure. It must learn that adjacent tokens are related, that subjects agree with verbs, etc. - all from training examples.

Running in Production

Best Practice

✅ Best Practices - Instrument every stage, version embedding models, enforce access control at retrieval time, and evaluate on a fixed golden set before shipping changes.

Ecosystem

• Flash Attention (Dao et al.) - IO-aware exact attention algorithm. Standard in all production inference engines.

• Multi-Query Attention (MQA) - Share K/V heads across Q heads. Used in PaLM, Falcon. Reduces KV cache memory.

• Grouped-Query Attention (GQA) - Intermediate between MHA and MQA. Used in Llama 2/3. Best balance of quality and efficiency.

• Sliding Window Attention - Mistral uses attention limited to a local window + global tokens. Reduces O(N²) to O(N × W).

• Attention visualization tools - BertViz, exBert - visualize attention patterns in open models (not available for API models).

Related Technologies

• Transformers - Attention is the core operation inside every transformer block.

• Large Language Models - LLM capabilities - context understanding, few-shot learning, instruction following - are attention patterns at scale.

• Context Windows - Context limits exist because attention scales quadratically with sequence length.

• Tokens - Attention operates on token sequences. Token boundaries affect what the model can attend to.

• Embeddings - Token embeddings are projected to Q, K, V - the input to the attention mechanism.

• Prompt Engineering - Prompt structure directly influences attention patterns and model behavior.

Learning Path

Prerequisites: Transformers · Embeddings · Tokens

Next topics: Large Language Models · Context Windows

Estimated time: 55 min · Difficulty: Advanced

FAQs

What is the difference between attention and self-attention?

Attention generally refers to one sequence attending to another (cross-attention). Self-attention means a sequence attends to itself - each token attends to all other tokens in the same sequence. LLMs use self-attention exclusively.

Why divide by √d_k in scaled dot-product attention?

Without scaling, dot products grow large with dimension, pushing softmax into regions with tiny gradients. Dividing by √d_k keeps variance stable, ensuring useful gradient flow during training.

How many attention heads should a model have?

More heads allow finer-grained relationship modeling but increase compute. Typical ratio: d_model / 64 or d_model / 128 = num_heads. Llama 3 8B uses 32 heads with d_model=4096 (head_dim=128).

Can I visualize attention for GPT-4?

Not directly - OpenAI does not expose attention weights. You can visualize attention for open models (Llama, Mistral) using tools like BertViz. For API models, use behavioral tests (fact retrieval by position) as a proxy.

What is cross-attention?

In encoder-decoder models, cross-attention lets the decoder attend to encoder outputs. The decoder's queries attend to the encoder's keys and values. This is how T5 connects input understanding to output generation.

Does attention explain in-context learning?

Partially. Research suggests attention heads implement algorithms (induction heads) that copy patterns from context - enabling few-shot learning without weight updates. This is an active area of research.

What is Flash Attention?

An exact attention algorithm that computes in tiles, avoiding materialization of the full N×N attention matrix in GPU HBM. Provides 2–4x speedup and enables longer contexts with the same memory. Used by default in vLLM, TGI, and all modern inference engines.

Why do models ignore my system prompt?

System prompts compete with all other tokens for attention. In long contexts, system prompt tokens at the beginning may receive lower attention weights as the model focuses on recent tokens. Keep system prompts concise and repeat critical instructions near the user message.

What is KV cache and how does it relate to attention?

During generation, K and V vectors from previous tokens are cached. Each new token only computes its Q and attends to all cached K/V pairs. This avoids recomputing attention for previous tokens - critical for generation speed.

How does grouped-query attention (GQA) work?

GQA uses fewer K/V head groups than Q heads. Multiple Q heads share the same K/V head, reducing KV cache memory by 4–8x with minimal quality loss. Llama 3 uses GQA with 8 KV groups and 32 Q heads.

References

• Attention Is All You Need (Vaswani et al., 2017)
• Hugging Face Transformer Architecture
• PyTorch Transformer Tutorial

Further Reading

• Google Research Blog: Transformer
• TensorFlow Transformer Guide
• The Illustrated Transformer (Jay Alammar)

Summary

• Attention computes weighted relationships between tokens - the core mechanism enabling language understanding and generation. - Self-attention with Q, K, V projections lets every token directly access every other token in parallel. - Multi-head attention learns diverse relationship patterns across parallel heads. - O(N²) scaling is the fundamental constraint on context window size and inference cost.

• Attention patterns explain LLM behavior: lost-in-the-middle, prompt sensitivity, and context competition. - Prompt engineering is attention engineering - placement, structure, and length directly affect what the model focuses on.