Attention — Module 10.1

How This Lesson Fits the Module & Volume

Module 10.1 began with the encoder / decoder frame, then named the three actors: Query, Key, and Value. This lecture is the synthesis: attention is the mechanism that turns Q, K, V into contextual outputs.

Historically, attention rescued RNN seq2seq from the single-vector bottleneck (Volume 08). Transformers then made attention the main layer. Everything that follows—scores, scaled dot-product, self-/cross-/multi-head variants—is a refinement of this core idea.

Learning Objectives

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

Define attention as soft, content-based selection over a set of values.
Write the generic pipeline: project Q/K/V → score → softmax → weighted sum.
Contrast attention with fixed pooling and with hard argmax retrieval.
Locate attention inside encoder and decoder stacks.
Implement a complete single-head attention step in PyTorch.
Preview why raw dot products need scaling (next lectures).

Definition

Attention computes, for each query, a distribution over key positions and returns the corresponding weighted average of values. Formally: Attention(Q, K, V) = softmax(score(Q, K)) V. It is differentiable soft lookup.

The Attention Pipeline

1. Project

Build Q, K, V from hidden states.

2. Score

Compare every query to every key.

3. Softmax

Turn scores into weights α.

4. Mix

Output = α V (contextual vector).

Why Attention Beats Fixed Summaries

Method	How it summarizes	Limitation
Mean / max pool	Same recipe for every query	Not content-adaptive
Final RNN state	Compress all history into one vector	Early tokens fade
Hard attention	Sample one position	High variance / non-diff (often)
Soft attention	Weighted blend of all values	Cost O(T_q · T_k)

Where It Appears

Encoder

Self-attention over source.
Builds contextual memory.
Bidirectional within the source.

Decoder (self)

Causal self-attention on targets.
Models language prefix.
No future leakage.

Decoder (cross)

Cross-attention to memory.
Aligns generation to source.
Q from decoder; K, V from encoder.

End-to-End Single-Head Attention

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

class SingleHeadAttention(nn.Module):
    def __init__(self, d_model, d_k):
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
        self.scale = d_k ** -0.5

    def forward(self, x_q, x_kv, mask=None):
        Q = self.W_Q(x_q)
        K = self.W_K(x_kv)
        V = self.W_V(x_kv)
        scores = (Q @ K.transpose(-2, -1)) * self.scale
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        return weights @ V, weights

attn = SingleHeadAttention(d_model=64, d_k=16)
x = torch.randn(2, 7, 64)
out, w = attn(x, x)          # self-attention
print(out.shape, w.shape)    # (2, 7, 16), (2, 7, 7)

Interpretability Bonus

The weight matrix w is often visualized as an alignment heatmap: which source tokens a decoder step looked at. Treat heatmaps as hypotheses, not proofs—models can use attention in non-obvious ways—but they remain a valuable teaching and debugging tool.

Strengths

Adaptive, content-based focus.
Differentiable end-to-end.
Short paths for long-range deps.

Tradeoffs

Quadratic memory/time in length.
Needs masks for PAD and causality.
Deep stacks need residuals/norm.

Common Misconception

“Attention is only for translation alignment.” Alignment was the first famous use, but attention is now a general layer: vision (ViT), speech, multimodal models, and any setting where soft selection over a set beats fixed pooling.

Knowledge Check

Short Answer: Write the schematic formula for Attention(Q, K, V). Answer: softmax(score(Q, K)) V.
True/False: Attention is a hard argmax over positions by default. Answer: False—standard neural attention is soft (softmax).
Multiple Choice: Soft attention returns: (a) one discrete index, (b) a weighted sum of values, (c) only the keys. Answer: (b).
Short Answer: Name the four pipeline stages of attention. Answer: Project Q/K/V, score, softmax, weighted sum of V.
True/False: Encoder self-attention is typically bidirectional. Answer: True.
Multiple Choice: Cross-attention uses queries from: (a) the encoder, (b) the decoder, (c) random noise. Answer: (b).
Short Answer: Why did attention help classic RNN seq2seq? Answer: The decoder can look at all encoder states instead of one final vector.
True/False: Attention weight matrices are always faithful explanations of model reasoning. Answer: False—useful but not guaranteed explanations.
Multiple Choice: Dense attention cost in sequence length n scales roughly: (a) O(n), (b) O(n log n) always, (c) O(n²). Answer: (c).
Short Answer: What does the next lecture focus on specifically? Answer: Attention scores (how Q and K are compared).

Key Takeaways

Attention = soft lookup: score keys with queries, softmax, mix values.
It powers encoder self-attention and decoder self-/cross-attention.
It replaces brittle fixed summaries with content-adaptive focus.
Next: Attention Score details the comparison step.

Trainer’s Guide

Hands-on idea: Plot w[0] as a heatmap for a short self-attention run on random inputs, then on repeated tokens—discuss structure.

Discussion prompt: Is mean-pooling a special case of attention? (Yes—uniform weights.)

Recap: Attention turns Q, K, V into contextual outputs via soft weighting. Continue with Attention Score.