Attention Mechanism

An attention mechanism is a neural network operation that builds each output as a weighted average of input features, with the weights computed from how relevant each input is to the element being processed. Relevance is scored by comparing a query vector with key vectors, and the resulting scores weight the matching value vectors. Attention was introduced in 2014 for machine translation, where it let a recurrent neural network (RNN) decoder look back at every source word instead of squeezing the whole sentence into one fixed-length vector. The 2017 transformer built an entire architecture from attention, which is now the core operation of large language models, vision transformers, and many object detection models.

How Attention Works#

Attention uses three learned linear projections of its input. The query (Q) describes what the current token is looking for, each key (K) describes what a token offers, and each value (V) carries the content that gets passed on. The model compares the query with every key, turns the comparison scores into weights, and returns the weighted sum of the values. When a translation model produces "noir" for "The black cat", the query for that step scores highest against the key for "black", even though French places the adjective after the noun.

Multi-head attention runs several of these computations in parallel, each on its own lower-dimensional projections, and concatenates the results. Different heads can then specialize in different relationships within the same layer.

Scaled Dot-Product Attention#

The transformer's attention is written as Attention(Q, K, V) = softmax(QKᵀ / √d_k) V. Each query is compared with every key by a dot product, and the scores are divided by the square root of the key dimension d_k so that they do not grow with vector size and saturate the softmax. Softmax turns each query's scores into weights that sum to 1, and the output is the weighted sum of the value vectors.

import torch
import torch.nn.functional as F

torch.manual_seed(0)

# Queries, keys, and values for a sequence of 4 tokens, each 8-dimensional
q, k, v = torch.randn(3, 4, 8).unbind(0)

# Score every query against every key, scale by sqrt(d_k), and normalize each row with softmax
weights = (q @ k.T / k.shape[-1] ** 0.5).softmax(dim=-1)

# Each output token is a weighted average of the value vectors
output = weights @ v

print("Attention weights of token 0:", [round(w, 3) for w in weights[0].tolist()])
print("Matches PyTorch SDPA:", torch.allclose(output, F.scaled_dot_product_attention(q, k, v), atol=1e-6))

It prints [0.121, 0.741, 0.106, 0.032] and True: token 0 takes 74% of its output from token 1's value vector, its four weights sum to 1, and the manual result matches PyTorch's fused scaled_dot_product_attention, which dispatches to memory-efficient kernels such as FlashAttention on supported GPUs.

Types of Attention#

• Self-attention: Queries, keys, and values all come from the same sequence, so every token gathers context from every other token. It is the attention inside transformer encoders and ViTs.
• Cross-attention: Queries come from one sequence and keys and values from another. A translation decoder attends to the source sentence, a text-to-image diffusion model attends to the prompt, and the object queries of a DETR decoder attend to image features.
• Causal attention: A mask blocks each token from attending to later tokens, so a language model generates text left to right, and the keys and values of earlier tokens can be stored in a KV cache.
• Efficient attention: Sparse, sliding-window, and linear attention reduce the quadratic cost of comparing every token with every other, and deformable attention samples only a few learned points around each reference location.
• Channel attention in CNNs: Squeeze-and-Excitation networks reweight the channels of a feature map using its global average, a cheap form of attention that formed the basis of the winning ILSVRC 2017 classification entry.

Attention vs. Convolution#

A convolution applies the same small filter at every position, so each output sees a fixed local window and its weights do not depend on the image. Attention computes its weights from the content and can link any two positions within one layer. The price is a cost that grows with the square of the token count. A 640×640 image at stride 32 gives a 20×20 grid of 400 tokens and 160,000 query-key pairs, and the stride-8 grid of the same image has 6,400 tokens and nearly 41 million pairs. YOLO26 places its attention blocks accordingly: the C2PSA block that ends the backbone and one attention-enabled C3k2 block in the neck both run multi-head self-attention on the 20×20 stride-32 feature map, and every higher-resolution layer stays convolutional.