← Master Index
Vol. 10 Module 10.1 Lecture

Key

Attention Mechanism

How This Lesson Fits the Module & Volume

The Query lecture framed attention as soft retrieval: something asks, something is indexed, something is returned. The key is the index—the vector that says “here is what this position is about, for matching purposes.”

Together with queries, keys produce attention scores. Values (next lecture) carry the payload. Understanding keys clarifies why Q and K share dimension d_k and why encoder memory can be “addressed” by a decoder.

Learning Objectives

By the end of this lesson, students should be able to:

  • Define the key as the matching / addressing vector for each candidate position.
  • Explain the projection k = x WK and why d_k must match the query.
  • Contrast keys in self-attention vs. encoder keys in cross-attention.
  • Describe how padding masks act on key positions (ignore PAD).
  • Implement key projection and score computation shapes in PyTorch.
  • Separate the role of keys from values.
Definition

A key is a vector k = x WK that represents a position in the space where queries look for matches. Compatibility between query q and key k (often a dot product) determines how much that position’s value contributes to the output.

Keys Are Addresses, Not Content

A library card catalog entry tells you where a book is and a short topical label—not the full text. Keys play that catalog role. The full text analogue is the value. Learning separate W_K and W_V lets the model optimize matching independently from what gets copied into the output.

SettingKeys come fromQueries come from
Encoder self-attentionSource positionsSource positions
Decoder self-attentionPast target positionsTarget positions
Cross-attentionEncoder memoryDecoder states

Compatibility with Queries

Same d_k

  • Q and K live in one space.
  • Dot product q · k is well-defined.
  • Per-head dimension is often d_model / h.

Score Matrix

  • scores = Q @ K.transpose(-2, -1)
  • Shape (B, T_q, T_k).
  • High score ⇒ more weight later.

Masking Keys

  • PAD keys get −∞ scores.
  • Causal mask hides future keys.
  • Softmax then ignores them.

Key Projection in PyTorch

import torch from torch import nn B, T_q, T_k, d_model, d_k = 2, 4, 6, 64, 16 decoder_h = torch.randn(B, T_q, d_model) # queries' source encoder_h = torch.randn(B, T_k, d_model) # keys' source (memory) W_Q = nn.Linear(d_model, d_k, bias=False) W_K = nn.Linear(d_model, d_k, bias=False) Q = W_Q(decoder_h) # (B, T_q, d_k) K = W_K(encoder_h) # (B, T_k, d_k) scores = Q @ K.transpose(-2, -1) # (B, T_q, T_k) print(Q.shape, K.shape, scores.shape) # torch.Size([2, 4, 16]) torch.Size([2, 6, 16]) torch.Size([2, 4, 6])

Why Not Reuse Hidden States as Keys?

You could attend with raw hidden states, but learned projections give the model freedom: a position might be easy to match on topic (key) while carrying a different signal to write into the output (value). Multi-head attention further splits keys into specialized subspaces.

Design Upsides

  • Clean interface: match in key space.
  • Works for self- and cross-attention.
  • Masks attach naturally to key indices.

Watch Outs

  • d_k mismatch with Q breaks the matmul.
  • Unmasked PAD keys pollute softmax.
  • Keys ≠ values—do not conflate them.
Common Misconception

“The key is what gets added into the attention output.” No—keys only influence weights. After softmax, the weighted sum is over values. If you mix K into the output by mistake, you are no longer implementing standard attention.

Knowledge Check

  1. Short Answer: What role do keys play in the retrieval metaphor? Answer: Addresses / index entries used for matching.
  2. True/False: Keys and queries must share the same last dimension d_k. Answer: True (for dot-product attention).
  3. Multiple Choice: In cross-attention, keys come from: (a) decoder only, (b) encoder memory, (c) the loss. Answer: (b).
  4. Short Answer: Write the usual score formula involving Q and K. Answer: scores = Q K^T (optionally scaled).
  5. True/False: Softmax is taken over the value dimension, not the key index dimension. Answer: False—softmax is over key positions for each query.
  6. Multiple Choice: PAD positions are handled by: (a) deleting W_K, (b) masking those key scores to −∞, (c) doubling d_k. Answer: (b).
  7. Short Answer: Why learn W_K separately from W_V? Answer: Matching (addressing) can be optimized independently from content (payload).
  8. True/False: In encoder self-attention, keys and queries come from the same sequence. Answer: True.
  9. Multiple Choice: If T_q=4 and T_k=6, scores have shape: (a) (B, 4, 6), (b) (B, 6, 4), (c) (B, 4, 4). Answer: (a).
  10. Short Answer: What lecture covers the payload mixed by attention weights? Answer: Value.

Key Takeaways

  • Keys are matching vectors: k = x WK.
  • Scores compare each query to every key; masks hide invalid keys.
  • Keys address; values (next) provide content for the weighted sum.
  • Continue with Value.
Trainer’s Guide

Hands-on idea: Zero out one key vector and show how that column of scores collapses—linking keys to the score matrix.

Discussion prompt: In cross-attention, should encoder keys be frozen after encoding, or can gradients flow back into the encoder?

Recap: Keys are the addresses queries match against. Continue with Value.