Self Attention — Module 10.1

How This Lesson Fits the Module & Engineering Practice

Earlier lectures introduced queries, keys, values, and scaled dot-product attention. Self-attention is the special case where those three streams come from the same sequence—every token looks at every other token in that sequence.

This is the computation that lets an encoder or language-model stack build contextual representations without recurrence. Understanding self-attention is the prerequisite for cross-attention, multi-head attention, and the full transformer stack in Module 10.2.

Learning Objectives

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

Define self-attention and state where Q, K, and V come from.
Contrast self-attention with cross-attention using a clear source/target table.
Write the matrix form of self-attention and interpret the attention weight matrix.
Implement a minimal self-attention block in PyTorch.
Explain why self-attention is permutation-equivariant without position information.
Identify common failure modes (full attention cost, missing positions, confusing self vs cross).

Definition

Self-attention is attention where queries, keys, and values are all derived from the same input sequence. Each position produces a weighted combination of value vectors from that same sequence, with weights from similarity between its query and every key.

From Scaled Dot-Product to Self-Attention

Given a sequence matrix \(X \in \mathbb{R}^{n \times d}\) (n tokens, d features), learn three projections:

\[Q = X W_Q,\quad K = X W_K,\quad V = X W_V\]

Then apply scaled dot-product attention:

\[\mathrm{Attention}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V\]

The softmax row for token \(i\) is a distribution over the \(n\) positions in the same sequence. That is the “self” in self-attention: the sequence attends to itself.

Self-Attention vs Cross-Attention

Property	Self-attention	Cross-attention
Query source	Same sequence \(X\)	Target / decoder sequence
Key & value source	Same sequence \(X\)	Source / encoder sequence
Typical role	Contextualize within one stream	Let one stream read another
Where it appears	Encoder layers; decoder self-attn	Decoder–encoder attention

Keep this distinction sharp. Self-attention mixes information inside a sequence; cross-attention lets a sequence look up content from a different sequence (for example, a translation decoder reading encoder states).

Self-attention

Q, K, V from one tensor.
Builds token context for that sequence.
Used in encoders and LM stacks.

Cross-attention

Q from target; K, V from source.
Aligns or retrieves across sequences.
Classic in encoder–decoder models.

What the Attention Matrix Means

The matrix \(A = \mathrm{softmax}(QK^\top / \sqrt{d_k})\) has shape \(n \times n\). Entry \(A_{ij}\) is how much token \(i\) weights token \(j\)’s value. Rows sum to 1. Large \(A_{ij}\) means “when forming the output at \(i\), borrow strongly from \(j\).”

1. Project

Map \(X\) to \(Q\), \(K\), \(V\).

2. Score

Compute scaled similarities \(QK^\top / \sqrt{d_k}\).

3. Softmax

Turn scores into attention weights per query row.

4. Aggregate

Multiply weights by \(V\) to get contextual outputs.

PyTorch: Minimal Self-Attention

import math
import torch
from torch import nn

class SelfAttention(nn.Module):
    def __init__(self, d_model: int, d_k: int | None = None):
        super().__init__()
        d_k = d_k or d_model
        self.d_k = d_k
        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)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        Q, K, V = self.W_q(x), self.W_k(x), self.W_v(x)
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)
        weights = torch.softmax(scores, dim=-1)
        return weights @ V  # (batch, seq_len, d_k)

x = torch.randn(2, 8, 32)
out = SelfAttention(32)(x)
print(out.shape)  # torch.Size([2, 8, 32])

Complexity and Design Notes

Strengths

Direct pairwise interactions in one step.
Highly parallel across positions (unlike RNNs).
Foundation for transformers and modern NLP/CV.

Tradeoffs

Quadratic cost in sequence length \(O(n^2)\).
No built-in order without positional encoding.
Needs residuals and norms for deep stacks.

Common Mistake

Calling every attention layer “self-attention” even when keys and values come from a different tensor. If Q comes from the decoder and K/V from the encoder, that is cross-attention, not self-attention.

Misconception

“Self-attention already knows word order.” Pure self-attention is permutation-equivariant: shuffling tokens (and applying the same shuffle to outputs) yields the shuffled result. Order must be injected via positional encoding (or related position schemes).

Where Self-Attention Appears Next

Production models almost never stop at a single head. Multi-head attention runs several self-attention subspaces in parallel and concatenates them. Stacks also wrap attention with residual connections, layer normalization, and a position-wise FFN.

Sibling topics in this module: Attention, Attention Score, Scaled Dot-Product Attention, Encoder, Decoder.

Knowledge Check

Short Answer: In self-attention, where do Q, K, and V come from? Answer: All three are projected from the same input sequence.
True/False: Self-attention and cross-attention always use the same source for keys and queries. Answer: False. Cross-attention takes keys/values from a different sequence than the queries.
Multiple Choice: The attention weight matrix for a length-n sequence is typically: (a) \(n \times n\), (b) \(d \times d\), (c) \(1 \times n\). Answer: (a).
Short Answer: Why scale by \(\sqrt{d_k}\)? Answer: To keep dot-product magnitudes stable so softmax does not become overly peaked for large \(d_k\).
True/False: Without positional information, self-attention is permutation-equivariant. Answer: True.
Multiple Choice: Which lecture covers Q from one sequence and K/V from another? (a) Self-attention, (b) Cross-attention, (c) LayerNorm. Answer: (b).
Short Answer: Name one computational cost of full self-attention. Answer: Quadratic time/memory in sequence length from the \(n \times n\) score matrix.
True/False: Softmax is applied over the key dimension for each query. Answer: True (each row of scores is normalized).
Multiple Choice: Self-attention is used heavily in: (a) only pooling layers, (b) transformer encoders and many LMs, (c) only BatchNorm. Answer: (b).
Short Answer: What does the next sibling lecture, multi-head attention, add? Answer: Multiple parallel attention heads whose outputs are concatenated and projected.

Key Takeaways

Self-attention derives Q, K, and V from the same sequence and mixes tokens via scaled softmax weights.
It differs from cross-attention, where keys and values come from another sequence.
The \(n \times n\) weight matrix is interpretable but quadratic in length.
Order is not intrinsic; add positional encoding.
Next: Cross Attention, then Multi-Head Attention.

Trainer’s Guide

Hands-on idea: Have learners print the attention weight matrix for a tiny sequence and verify each row sums to ~1.0.

Discussion prompt: When would you prefer sparse or local attention instead of full self-attention?

Recap: Self-attention lets every token in a sequence attend to every other token in that same sequence via shared Q/K/V projections. Continue with Cross Attention.