Computer Vision MT25, Attention and transformers
Flashcards
Suppose we have:
- $Z \in \mathbb R^{M \times d}$, $M$ “condition” tokens with rows $z _ j$, the encoder output (e.g. tokens in an english sentence we aim to translate to french)
- $X \in \mathbb R^{N \times d}$, $N$ “query” tokens with rows $x _ i$, the decoder output (e.g. a partial translation of a french translation we aim to continue)
- (In self-attention, we take $Z = X, M = N$. Cross-attention has these different).
@State the formula for scaled dot-product attention in this case.
We learn three projection matrices $W^Q, W^K, W^V \in \mathbb R^{d \times d}$ to produce
\[Q = XW^Q \in \mathbb R^{N \times d}, \quad K = ZW^K \in \mathbb R^{M \times d}, \quad V = ZW^V \in \mathbb R^{M \times d},\]In matrix form, we compute
\[\text{Attn}(Q, K, V) = \underbrace{\text{softmax}\left( \frac{QK^\top}{\sqrt d} \right)} _ {=: \, A} V \in \mathbb R^{N \times d},\]with output in $\mathbb R^{N \times d}$.
Per-row breakdown: Let $q _ i, k _ i, v _ i$ be row $i$ of $Q, K, V$ respectively. Then, for each output row $i$, we compute
- Scores: row $i$ of $QK^\top \in \mathbb R^{N \times M}$ has $j$th entry $q _ i k _ j^\top = (x _ i W^Q)(z _ j W^K)^\top$, the similarity between query $i$’s projection from $X$ and condition token $j$’s key projection from $Z$.
- Rescale: divide by $\sqrt d$ to control variance (∆scaled-dot-product-sqrt-d-argument).
- Normalise: row-wise softmax produces attention weights $a _ {i,j}$ with $\sum^M _ {j=1} a _ {i,j} = 1$, i.e. query $i$’s distribution over the $M$ condition tokens
- Mix: Set $y _ i = \sum^M _ {j = 1} a _ {i,j} v _ j$
In scaled dot-product attention, we compute
\[\text{Attn}(Q, K, V) = \text{softmax}\left( \frac{QK^\top}{\sqrt d} \right) V\]
@Justify the choice of dividing by $\sqrt d$ inside the softmax in scaled dot-product attention.
If $q, k \in \mathbb R^d$ have independent entries with mean $0$ and variance $1$, then
\[q^\top k = \sum^d _ {i = 1}q _ i k _ i\]is a sum of $d$ independent zero-mean, unit-variance terms, so $\text{Var}(q^\top k) = d$ and $q^\top k$ has standard deviation $\sqrt d$.
Without rescaling, the entries of $QK^\top$ grow in magnitude as $d$ increases. Dividing by $\sqrt d$ keeps the variance of the scores at $1$ regardless of $d$, which keeps the softmax within its non-saturated regime.
@Define a transformer encoder block, listing its four sub-operations.
A transformer encoder block is a function $\mathbb R^{N \times d} \to \mathbb R^{N \times d}$ takes a sequence of $N$ token representations and produces a refined version of the same shape, via:
- Self-attention: $\text{SelfAttn}(X) = \text{Attn}(X, X)$ (∆scaled-dot-product-attention).
- Add + LayerNorm: a residual connection $X + \text{SelfAttn}(X)$, then a row-wise LayerNorm.
- MLP: a token-wise MLP $\mathbb R^d \to \mathbb R^d$, applied independently to each row.
- Add + LayerNorm again.
@Define a transformer decoder block, listing its six sub-operations.
A transformer decoder block is a function $\mathbb R^{N \times d} \times \mathbb R^{M \times d} \to \mathbb R^{N \times d}$ that takes a decoder-side sequence $X$ of $N$ tokens and an encoder-side condition $Z$ of $M$ tokens, and produces a refined decoder sequence of the same shape as $X$, via:
- Masked self-attention: $\text{SelfAttn}(X)$ with a causal mask applied inside the softmax, so token $i$ can only attend to tokens $1, \ldots, i$ (e.g. as we are typically learning to predict the next token, so this ensures the transformer can’t cheat by looking at the true next encoder token during training).
- Add + LayerNorm: a residual connection $X + \text{SelfAttn}(X)$, then a row-wise LayerNorm
- Cross-attention: $\text{Attn}(X, Z)$ with queries from $X$ (decoder) and keys / values from $Z$ (encoder output from ∆transformer-encoder-block).
- Add + LayerNorm
- MLP: a token-wise MLP $\mathbb R^d \to \mathbb R^d$, applied independently to each row.
- Add + LayerNorm
@State the main variants of the transformer architecture. For each, name a representative model, a task, and what is structurally different about the setup.
Variant Representative model Task Setup … … … Input: …
Architecture: …
Output: …
Architecture: …
Output: …
| Variant | Representative model | Task | Setup |
|---|---|---|---|
| Encoder-only | ViT, BERT, DINO | Classification, embedding | Input: a token sequence (e.g. image patches, class token and positional encoding) Architecture: $L$ blocks of unmasked self-attention Output: Final hidden states for encoding, class token encoding for classification |
| Encoder-decoder | T5, original Transformer, Detr | Translation, summarisation, object detection | Input: source and partial target Architecture: $L$ blocks of self-attention processes source into $Z$, and $L'$ blocks of masked self-attention, cross-attention with $Z$ and FFN process target Output: Decoder hidden states + task head |
| Decoder-only | GPT, Llama, most modern LLMs | Causal language modelling | Input: token sequence Architecture: one stack of masked self-attention + FFN blocks Output: final hidden states + task head (e.g. for next-token prediction) |
| Dual-encoder | CLIP, SigLIP | Image-text retrieval | Input: image $I$ and text $T$ Architecture: two separate encoder-only models, each pooling its final hidden states into a single embedding $e _ I, e _ T$ Output: Similarity $\langle e _ I, e _ T\rangle$, trained via contrastive loss |
| Hybrid (e.g. cross-attention into CNN) | Stable Diffusion UNet | Text-to-image generation | Input: noisy image latent $z _ t$ and a text embedding $Z$ Architecture: a UNet with cross-attention layers inserted between convolutions Output: Predicted noise $\hat \epsilon$ |
| Hybrid (e.g. cross-attention into LM) | Flamingo | Vision-language modelling | Input: interleaved text and image features $z$ Architecture: a decoder-only LLM with cross-attention layers inserted between LM blocks Output: next-token vocab logits per-position |
What does it mean to say that the Transformer architecture is permutation-equivariant?
Permuting the order of the inputs $x _ i$ will permute the order of the outputs $y _ i$ in the same way.
The Transformer architecture is permutation-equivariant, which means that permuting the order of the inputs $x _ i$ will just permutate the order of the outputs $y _ i$ in the same way.
How do transformer-based models where the location of the inputs matters resolve this?
They add a positional encoding
\[x _ i' = x _ i + P(i)\]which is either learned by the network or is hardcoded.
@State the general inputs and outputs of a transformer, and what it means to do “self-attention”, and separately what these would be in an English-French translation task.
- $N$ input tokens $x _ i \in \mathbb R^d$
- English-French: This would be embeddings of English word fragments
- $M$ condition tokens $z _ j \in \mathbb R^d$:
- Self-attention: This would be the input again
- English-French: This would be the predicted translation so far.
- $N$ output tokens $y _ i \in \mathbb R^d$
- English-French: Each vector would be a probability distribution over the words in the vocabulary. During training, each of the output tokens contributes to the loss separately. During inference, you only sample from the next output token.
Single-head self-attention $\text{SAttn}(K, V, Q)$ uses one set of projections to produce one output (∆scaled-dot-product-attention). Multi-head attention runs $H$ such operations in parallel with separate per-head projections, then combines them. This lets the model attend to different patterns simultaneously, analogous to using multiple convolutional filters per conv layer.
@State the multi-head attention formula given an input $x$ and $H$ heads.
Per-head computation: each head $h \in \{1, \ldots, H\}$ has its own learned projection matrices $\mathbf K^{(h)}, \mathbf V^{(h)}, \mathbf Q^{(h)}$, and produces:
\[y^{(h)} = \text{SAttn}(\mathbf K^{(h)} x, \mathbf V^{(h)} x, \mathbf Q^{(h)} x), \quad 1 \le h \le H.\]Combining heads: stack all $H$ outputs and project back via a linear layer $\mathbf W \in \mathbb R^{Hd \times d}$:
\[y = \mathbf W \begin{bmatrix} y^{(1)} \\ \vdots \\ y^{(H)} \end{bmatrix}\]The Vision Transformer (ViT) applies the transformer architecture directly to image classification by treating an image as a sequence of patches. It was the first to show that a pure-transformer model (no convolutions) can match or beat CNNs on ImageNet given enough pre-training data (when there is not enough data, pure CNNs can beat it).
@Describe and @visualise the steps that a vision transformer uses to classify images.
Pipeline:
- Patchification: Split images into non-overlapping fixed-size patches
- Linear patch embedding: A linear layer maps each patch to a vector
- Add positional encoding: Add a 2D positional encoding to each token
- Prepend a CLS token: Add an extra class (“CLS”) token to the sequence. The CLS token’s final representation will be used for classification.
- Transformer encoder: pass the sequence through $L$ layers of standard transformer encoder blocks (∆transformer-encoder-block).
- MLP classification head: Use the encoding of the class token as input to a simple MLP classifier
Bite-sized
Self-attention has complexity $\mathcal{O}(N^2)$ in the number of input tokens, since each of the $N$ queries attends to all $N$ keys to form the $N \times N$ score matrix $QK^\top$ before the softmax.
In cross-attention, where do the queries $Q$, keys $K$, and values $V$ come from, and how does this differ from self-attention?
- Queries $Q$ come from the decoder side (the current decoder hidden states).
- Keys $K$ and values $V$ come from the encoder side (the encoder output $Z$).
So $\text{Attn}(Q, K, V) = \text{softmax}(QK^\top / \sqrt d) V$ with $Q = X W^Q$ for decoder $X$ and $K = Z W^K, V = Z W^V$ for encoder $Z$. Every decoder position thus attends over all encoder positions.
In self-attention by contrast, $Q$, $K$, $V$ are all linear projections of the same sequence ($Z = X$).
The Vision Transformer (ViT) splits an image into non-overlapping patches of size $16 \times 16$ pixels, then treats the resulting sequence of patch embeddings (plus a CLS token and 2D positional encoding) as input tokens to a standard transformer encoder.
@Describe the four normalisation layers (BatchNorm, LayerNorm, InstanceNorm, GroupNorm) by stating which axes they average over for a feature tensor of shape $(N, C, H, W)$.
For batch sample index $n$, channel index $c$, and spatial location $(h, w)$:
- BatchNorm: normalise over $(N, H, W)$ — one $(\mu _ c, \sigma _ c)$ per channel. Depends on batch statistics.
- LayerNorm: normalise over $(C, H, W)$ — one $(\mu _ n, \sigma _ n)$ per sample. Batch-independent.
- InstanceNorm: normalise over $(H, W)$ only — one $(\mu _ {n,c}, \sigma _ {n,c})$ per sample-channel.
- GroupNorm: partition channels into $G$ groups, normalise over $(G _ {\mathrm{group}}, H, W)$ — one $(\mu _ {n,g}, \sigma _ {n,g})$ per sample-group.
Figure: the Other Normalisation Layers slide visualises each by colouring the elements normalised together across the $(N, C, H \cdot W)$ axes (Wu & He, Group Normalization, ECCV 2018).
On the JFT pre-training data-scaling plot from the ViT paper, Vision Transformers only overtake comparable ResNet (BiT) baselines once the pre-training set has roughly 300M images; below that, BiT wins. This is the empirical evidence behind the slogan that ViTs need much more data than CNNs to shine.
BatchNorm normalises a batch $x \in \mathbb R^{B \times d}$ channel-wise. With per-channel mean $\mu _ j = \tfrac{1}{B}\sum _ i x _ {i,j}$ and variance $\sigma _ j^2 = \tfrac{1}{B}\sum _ i (x _ {i,j} - \mu _ j)^2$, the normalised output is
\[x' _ {i,j} = <span class="cloze" tabindex="0">\frac{x _ {i,j} - \mu _ j}{\sqrt{\sigma _ j^2 + \varepsilon</span> }}\]followed by a learnable affine rescaling $y _ {i,j} = \gamma _ j x' _ {i,j} + \beta _ j$ (with $\gamma, \beta \in \mathbb R^d$).
@Describe how BatchNorm behaves differently at training and test time, and why this asymmetry is necessary.
- Training: $\mu, \sigma^2$ are computed from the current mini-batch. Each forward pass through BN therefore depends on the other samples in the batch.
- Test: $\mu, \sigma^2$ are fixed statistics — usually an exponential moving average accumulated over training. The test-time output of any single input is then independent of any other inputs.
Why the asymmetry: at test you may see only one sample, or batches with non-representative statistics, so batch statistics would be noisy or undefined. Using frozen training-set statistics gives deterministic, per-input outputs and lets BN be fused into the previous linear layer for zero test-time overhead.
This different behaviour at train vs. test is flagged in the lecture as “a very common source of bugs”.
@Justify why transformer architectures use LayerNorm (over the channel axis) rather than BatchNorm (over the batch axis).
- Batch-independence: LayerNorm computes $(\mu _ n, \sigma _ n)$ from each sample’s own channels, so the output for one token does not depend on any other token in the batch. This makes inference behaviour identical to training, removes the train/test asymmetry that causes BN bugs, and works fine for batch size 1.
- Variable-length sequences: BN statistics over a batch of variable-length token sequences would mix statistics across very different positions and padding. LayerNorm dodges this entirely by normalising within a single token’s representation.
- Small effective batches: distributed transformer training often uses small per-device batches, where BN’s batch-statistic estimates would be unreliable.
@Justify the use of residual connections $y = f(x) + x$ in transformer (and CNN) blocks.
- Easier optimisation target: writing $y = f(x) + x$ means $f$ only needs to learn the change in the representation rather than reproduce all the information already in $x$. Identity behaviour is achieved at $f \equiv 0$, which is close to the initialisation.
- Gradient flow: $\partial y / \partial x = I + \partial f / \partial x$, so even when $\partial f / \partial x$ is small, gradients still flow back through the identity skip. This is what lets very deep stacks (e.g. ResNet-152, modern transformers with $L \gg 100$) be trained at all.
- Universality: residual blocks now appear in essentially all modern CNN and transformer architectures.
ViT’s positional encoding is learned from scratch (rather than the hand-crafted sin/cos basis used in the original Vaswani transformer). Inspecting the learned embeddings reveals strong local similarity: each patch’s embedding is most similar to its spatial neighbours, which lets the model identify where each token comes from in the 2D image.