Self-Attention and QKV Computation

🎯 Learning Objectives

• Explain why Q, K, V matrices are necessary and how they differ
• Derive Q, K, V from input embeddings through linear projections
• Interpret attention weight matrices and visualize attention patterns
• Distinguish between self-attention, cross-attention, and bidirectional attention
• Compute attention outputs manually for small examples

📋 Prerequisites

• Transformer Architecture (previous topic) — high-level understanding of the architecture
• Matrix multiplication — required for QKV projections
• Softmax function — normalizing attention weights

1. 📖 Core Content

1.1 Intuition: Why Q, K, V?

• Query (Q): Your search query — "What I'm looking for"
• Key (K): The titles/index of each document — "What each document is about"
• Value (V): The actual content of each document — "The information itself"
• Attention: The process of matching your query against all keys to retrieve relevant values In self-attention, every token simultaneously acts as a query, a key, and a value. Each token asks: "What should I pay attention to?" (Query), each token announces: "Here's what I contain" (Key), and each token offers: "Here's my contribution" (Value). Why can't we use the embeddings directly as Q, K, V? Because the model needs to learn different representations for these three roles. The same word used as a query might need different features than when used as a key.

1.2 Linear Projections: Computing Q, K, V

• $W^Q$, $W^K$, $W^V$ are learned weight matrices (different for each head)
• $d_k$ is typically $d_{model} / h$ (e.g., 512/8 = 64)
• In the original Transformer, $d_k = d_v$ (both = 64 per head)

1.2.1 Worked Example: QKV Projection

1.3 Self-Attention vs Cross-Attention

1.4 Attention Pattern Analysis

1. Diagonal attention: Tokens attend mostly to themselves and immediate neighbors
2. Vertical attention: Certain tokens (like [CLS] or [SEP] in BERT) attend to many tokens
3. Block diagonal: Within-phrase attention (tokens attend within their own clause)
4. Long-range: Attending to distant but semantically related tokens

1.4.1 Visualizing Attention Patterns

python
Run Inline
# runnable
import numpy as np
import matplotlib.pyplot as plt
def visualize_attention(attention_weights, tokens):
    """
    Visualize attention weights as a heatmap
    Args:
        attention_weights: (n_tokens, n_tokens) matrix
        tokens: List of token strings
    """
    fig, ax = plt.subplots(figsize=(8, 6))
    im = ax.imshow(attention_weights, cmap='Blues')
    # Labels
    ax.set_xticks(range(len(tokens)))
    ax.set_yticks(range(len(tokens)))
    ax.set_xticklabels(tokens, rotation=45)
    ax.set_yticklabels(tokens)
    ax.set_xlabel('Keys')
    ax.set_ylabel('Queries')
    # Colorbar
    plt.colorbar(im)
    plt.title('Attention Weights')
    plt.tight_layout()
    return fig
# Example: Sentence with self-attention pattern
tokens = ['The', 'cat', 'sat', 'on', 'the', 'mat']
# Simulated attention: diagonal + long-range (cat→mat)
np.random.seed(42)
attn = np.random.rand(6, 6)
# Make it more interpretable
attn = attn / attn.sum(axis=-1, keepdims=True)  # Normalize rows
print("Simulated attention weights (6 tokens):")
print(np.round(attn, 2))
# Each row sums to 1.0
print("\nRow sums (should be all 1.0):", attn.sum(axis=-1))

1.5 Attention in Encoder vs Decoder

Encoder Self-Attention

• Bidirectional: Each token can attend to ALL tokens (before and after)
• No masking needed
• Produces contextualized representations of the input

Decoder Masked Self-Attention

• Unidirectional: Each token can only attend to itself and previous tokens
• Uses a causal mask (upper triangular matrix of $-\infty$)
• Prevents "cheating" during generation

Decoder Cross-Attention

• Q comes from decoder, K and V come from encoder
• Each decoder token can attend to all encoder tokens
• Allows the decoder to "look at" the input

1.5.1 Causal Masking

python
Run Inline
# runnable
import numpy as np
def create_causal_mask(seq_len):
    """
    Create a causal attention mask (lower triangular)
    """
    mask = np.triu(np.ones((seq_len, seq_len)) * -1e9, k=1)
    return mask
seq_len = 5
mask = create_causal_mask(seq_len)
print("Causal Mask (showing -1e9 as -inf):")
print(np.where(mask < -1e8, float('-inf'), 0))

1.5.2 Worked Example: Masked Attention

1.6 Attention Output Interpretation

• "I went to the bank to deposit money" → "bank" attends to "deposit", "money" → financial institution
• "I sat on the river bank" → "bank" attends to "river" → river bank The same word gets different representations based on what it attends to.

1.7 Why This Matters

• Debugging: When attention patterns go wrong, model behavior becomes unpredictable
• Interpretability: Attention weights can be extracted to understand what the model focuses on
• Architecture design: Many improvements (sparse attention, linear attention, Flash Attention) are modifications to the QKV computation
• Fine-tuning: Methods like PEFT modify how attention is computed

4. 📐 Key Formulas / Concepts

5. ⚠️ Common Pitfalls

Pitfall 1: Thinking Q, K, V are arbitrary names

Pitfall 2: Confusing attention dimension with model dimension

Pitfall 3: Assuming attention weights sum to 1.0 for each column

• Row i (query i): $\sum_j A_{ij} = 1.0$
• Column j (key j): $\sum_i A_{ij}$ can be anything
• Each query distributes 1 unit of attention across all keys

6. 📝 Practice Questions

Q1: What is the shape of Q, K, V for a batch of 8 sequences, each of length 20, d_model=768, and 12 heads?

• d_k = d_model / heads = 768 / 12 = 64
• Q shape: (8, 20, 768) × (768, 64) = (8, 20, 64) per head → after splitting: (8, 12, 20, 64)
• K shape: same as Q: (8, 12, 20, 64)
• V shape: same: (8, 12, 20, 64)
• Attention weights: (8, 12, 20, 20)
• Output: (8, 12, 20, 64) → combined: (8, 20, 768) Q2: If d_model=256, h=8, what is d_k? What if d_model=1024, h=16?

Case 1: d_k = 256/8 = 32 Case 2: d_k = 1024/16 = 64

The per-head dimension typically stays between 32-128 in most architectures. Q3: Why can't the decoder use bidirectional attention during training?

If the decoder could attend to future tokens during training, it would "cheat" by looking at the correct next token instead of learning to predict it. For example, training to predict the third word "sat" in "The cat sat":

• Without mask: The model sees "sat" as a key and uses its value directly → trivial prediction
• With causal mask: The model must predict "sat" using only "The" and "cat" → actual learning

During inference, future tokens don't exist yet, so the mask naturally prevents looking at non-existent tokens. Q4: In cross-attention, why does Q come from the decoder while K and V come from the encoder?

Q represents "what the decoder needs to know next." The decoder asks: "Given what I've generated so far, what part of the input should I look at to generate the next token?"

K and V from the encoder represent "what the input contains." The encoder has already processed the full input and created rich representations.

If the roles were reversed (Q from encoder, K, V from decoder), the encoder would be trying to "query" the decoder, which doesn't make sense since the decoder is being built incrementally. Q5: Compute Q, K, V for X = [[2, 1], [0, 3]] with W^Q = [[1, 0], [0, 1]], W^K = [[0, 1], [1, 0]], W^V = [[0.5, 0.5], [0.5, 0.5]].

Q = X·W^Q = [[2·1+1·0, 2·0+1·1], [0·1+3·0, 0·0+3·1]] = [[2, 1], [0, 3]] K = X·W^K = [[2·0+1·1, 2·1+1·0], [0·0+3·1, 0·1+3·0]] = [[1, 2], [3, 0]] V = X·W^V = [[2·0.5+1·0.5, 2·0.5+1·0.5], [0·0.5+3·0.5, 0·0.5+3·0.5]] = [[1.5, 1.5], [1.5, 1.5]] Q6: What happens to attention weights if we add a constant C to all scores before softmax?

Nothing. Adding a constant to all elements of a row before softmax doesn't change the result because:

$\text{softmax}(x_i + C) = \frac{e^{x_i + C}}{\sum_j e^{x_j + C}} = \frac{e^{x_i} \cdot e^C}{e^C \sum_j e^{x_j}} = \frac{e^{x_i}}{\sum_j e^{x_j}} = \text{softmax}(x_i)$

But adding a constant to a single element (while keeping others unchanged) does change the distribution. Q7: For a sequence of length n, how many pairwise attention comparisons are made in self-attention?

n² comparisons. Every token (as query) compares against every token (as key). For n=100, that's 10,000 comparisons. This quadratic growth is the reason Transformers are computationally expensive for long sequences. Q8: What is the difference between "content-based" attention and "position-based" attention?

Content-based attention: The attention score depends on what the tokens represent (the QK dot product captures semantic similarity). "bank" attends to "money" because their representations are similar.

Position-based attention: The attention score depends on where tokens are in the sequence. The model might learn to always attend to the token at position i-1 (previous token) regardless of content.

Multi-head attention can capture both: some heads learn content-based patterns while others learn position-based patterns. The learned projections determine which type each head specializes in.

7. 🔗 Cross-References

• Next topic: Multi-Head Attention Deep Dive (Week 1)
• Previous topic: Transformer Architecture Introduction
• Video: Week 1 Lectures 2-4 in BSDA5004 transcripts Join Discord PreviousTransformer ArchitectureNextMulti-Head Attention