Self-Attention in Vision Transformers

A Vision Transformer does not run attention directly on pixels. It first turns image patches into a sequence of vectors, adds position information, and then lets each sequence row construct a new representation by mixing information from every other row. Self-attention is the learned, content-dependent mixing rule inside that update.

The word “global” describes the available connectivity: in a standard ViT encoder block, any query row can score any key row in one layer. It does not mean every learned attention row is uniform, semantic, or globally spread.

Follow one patch row through attention

Play the full mechanism or select any stage. The image field shows one query row over its original spatial positions rather than replacing the image with an abstract matrix.

What reaches the attention layer

For an image of height H, width W, and square patch size P, non-overlapping patchification produces (H/P) × (W/P) patch rows. A learned patch projection maps each flattened patch to model width D. Position information is then added or otherwise supplied so that identical-looking patches at different locations need not remain indistinguishable.

The original ViT classification architecture prepends a learned CLS row. For a 224 × 224 image with 16 × 16 patches, that gives 14² + 1 = 197 rows. Other vision transformers may omit CLS, pool patch rows, use relative position mechanisms, or change the spatial hierarchy.

This page starts after patch embedding and position injection. See positional embeddings in ViT and the CLS token for those separate mechanisms.

Q, K, and V are learned projections

In a pre-normalized ViT block, let U = LN(X) be the normalized input sequence. For one attention head of width dₕ:

The three roles are operational rather than semantic labels:

• Query: the vector used by one row to score all keys,
• Key: the vector each source row exposes for matching,
• Value: the vector each source row contributes if it receives weight.

A patch does not have one permanent “importance.” Its contribution depends on the querying row, head, layer, image, learned weights, and any mask.

Scores, scaling, and row-wise softmax

One head writes a score matrix

with shape N × N. Entry Sᵢⱼ is the scaled dot product between query row i and key row j. Division by √dₕ controls logit scale as head width grows; without it, dot products tend to grow in magnitude and can push softmax toward saturated distributions.

Softmax runs independently across each query row:

Therefore Aᵢⱼ > 0 and Σⱼ Aᵢⱼ = 1 for each unmasked row. Scaling changes the distribution’s sharpness but does not change score ordering because √dₕ is positive.

Values carry the mixed content

The selected row’s output is a weighted sum of value vectors:

Keys participate in deciding the weights; values carry the content that is mixed. Replacing V with K is not the standard operation.

A ViT uses several heads. Each head has its own Q/K/V projections and probability matrix. Their outputs are concatenated and mapped back to model width:

The attention sublayer then writes a residual update:

A separate normalized MLP sublayer follows: Z = Y + MLP(LN(Y)). The figure models one head with compatible toy widths so the residual arithmetic remains visible; it does not claim to be a complete production block.

What an attention map does—and does not—show

Selecting one patch and displaying Aᵢ,: shows one query row from one head at one layer. Bright cells received more weight in that particular value mixture. This is useful for inspecting the mechanism, but several boundaries matter:

• an attention row is not the same as feature magnitude or class evidence,
• high weight does not prove a source patch caused the final prediction,
• residual paths, value vectors, other heads, later layers, and the MLP also affect the result,
• averaging heads can hide distinct behavior,
• deterministic teaching scores are not learned checkpoint activations.

Causal claims need interventions or attribution methods appropriate to the question, not attention weights alone.

Global attention versus local convolution

The clean comparison is between operators, not universal claims about model families.

This does not establish that every ViT is more data-hungry, slower, or less parameter-efficient than every CNN. Pretraining, augmentations, architecture, resolution, kernels, and hardware all matter. Swin restricts attention to local windows; ConvNeXt remains convolutional; architectures such as CoAtNet intentionally combine convolution and attention.

The quadratic pair term

For N sequence rows, model width D, H heads, and dₕ = D/H, the all-head score product and value mixture each contribute roughly N²D multiply-accumulates. Q/K/V projection contributes 3ND², and the output projection contributes ND², excluding biases and elementwise work.

The logical attention matrices contain H N² score or probability slots. ViT-B/16 at 224² has N = 197 including CLS and H = 12, so one layer has 12 × 197² = 465,708 logical score slots. Fused attention kernels may avoid materializing the complete matrix in high-bandwidth memory, but they do not change the dense all-pairs dependency or arithmetic result.

Shape-faithful implementation

import math
import torch


def attention_head(x, wq, wk, wv, mask=None):
    # x: [batch, tokens, model_width]
    # w*: [model_width, head_width]
    q = x @ wq
    k = x @ wk
    v = x @ wv

    scores = q @ k.transpose(-2, -1)
    scores = scores / math.sqrt(q.shape[-1])
    if mask is not None:  # boolean mask, broadcastable to score shape
        scores = scores.masked_fill(~mask, float('-inf'))

    weights = torch.softmax(scores, dim=-1)
    output = weights @ v
    return output, weights

For a classification ViT encoder, the attention mask is usually non-causal: patch and CLS rows can interact in both directions. Padding or architecture-specific masks may still apply. Decoder-style causal attention is a different connectivity rule.

Common implementation errors

• Scaling by model width instead of head width: use √dₕ for each head.
• Normalizing over the wrong axis: softmax must normalize the key axis separately for every query row and head.
• Mixing keys instead of values: the normalized matrix multiplies V.
• Dropping position information: bare self-attention is permutation-equivariant; ViT must supply spatial position through an explicit mechanism.
• Showing only the selected row but charging only N scores: the full dense operation writes N rows, so its pair term is N² per head.
• Reading attention as causal attribution: weights describe a value-mixing route, not the entire prediction computation.

Primary sources

• An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale
• Attention Is All You Need

Related concepts