Self-Attention

Self-attention lets every element of a set exchange information with every other element in one step. Each element, or token, computes a query, a key and a value, and its output is an average of all values weighted by how well its query matches each key. In vision the tokens are usually image patches or feature-map positions, so a patch can draw directly on one at the other side of the image, which a convolutional network reaches only through many layers. It is the core of the transformer and the Vision Transformer (ViT), and appeared in vision earlier in non-local blocks for video.

Definition

Self-attention maps each of $n$ tokens to a query, a key and a value by learned linear projections. The output for token $i$ is a weighted sum of all $n$ values, with weights given by a softmax over the dot products of query $i$ with every key, divided by the square root of the key dimension; Vaswani et al. (2017) call this scaled dot-product attention. It is self-attention because queries, keys and values come from the same tokens. In general attention, queries come from one set and keys and values from another, as in the transformer’s encoder-decoder attention and in the alignment model of Bahdanau, Cho and Bengio (2015), where a translation decoder searches the source sentence for words relevant to its next output.

Intuition

The query describes what a patch is looking for, the key what it contains, and the value what it passes on. A patch on the top edge of a bright square might issue a query that matches the keys of other top edges, and its output blends their values, wherever they are.

Unlike a convolution, the mixing weights are computed from the image content, so they change from image to image, and ignore position: a convolution weights a neighbor by its offset, attention only by what its features look like.

Formal Definition

Stack the tokens as the rows of $X \in \mathbb{R}^{n \times d}$. With weight matrices $W_Q, W_K \in \mathbb{R}^{d \times d_k}$ and $W_V \in \mathbb{R}^{d \times d_v}$,

$$Q = X W_Q, \qquad K = X W_K, \qquad V = X W_V,$$ $$\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V,$$

with the softmax applied to each row of the $n \times n$ score matrix (Vaswani et al., 2017, Eq. 1). Row $i$ of $A = \operatorname{softmax}(QK^\top/\sqrt{d_k})$, nonnegative weights that sum to one, is the attention map of token $i$; the softmax is the one used for class probabilities (see cross-entropy loss).

Why divide by $\sqrt{d_k}$. Vaswani et al. observed, citing earlier experiments, that unscaled dot-product attention did worse than additive attention for larger $d_k$, and suspected that large dot products push the softmax into regions with extremely small gradients. Their footnote gives the reason the products grow: if the components of $\mathbf{q}$ and $\mathbf{k}$ are independent with mean 0 and variance 1, then $\mathbf{q} \cdot \mathbf{k} = \sum_{i=1}^{d_k} q_i k_i$ has mean 0 and variance $d_k$. Dividing by $\sqrt{d_k}$ restores unit variance. The small gradients come from the softmax Jacobian $\operatorname{diag}(\mathbf{p}) - \mathbf{p}\mathbf{p}^\top$, whose entries all vanish as $\mathbf{p}$ approaches a one-hot vector, so a saturated row passes almost no gradient back to the queries and keys in backpropagation.

Multi-head attention. Vaswani et al. run $h$ attention functions in parallel on separate learned projections:

$$\operatorname{MultiHead}(X) = \operatorname{Concat}(\text{head}_1, \dots, \text{head}_h)\, W_O, \qquad \text{head}_i = \operatorname{Attention}(X W_Q^{(i)}, X W_K^{(i)}, X W_V^{(i)}),$$

with $W_O \in \mathbb{R}^{h d_v \times d}$. Each head can attend to different positions in a different representation subspace, which one averaging head cannot, the authors argue. With $d_k = d_v = d/h$ (the base transformer uses $h = 8$, $d_k = 64$, $d = 512$), the cost is close to that of one full-width head, and the four projections hold $4d^2$ weights whatever $h$ is: 1,048,576 for $d = 512$ (1,050,624 with the biases PyTorch’s nn.MultiheadAttention adds by default).

Masks. A mask restricts which keys each query may use by setting the forbidden scores to $-\infty$ before the softmax, so their weights become exactly zero and each row still sums to one.

• A causal mask lets position $i$ attend only to positions $j \le i$; the transformer’s decoder uses one to generate a sentence word by word (Vaswani et al., 2017).
• A padding mask hides placeholder tokens added to bring inputs to a common size; Liu et al. (2021) describe padding and masking windows as the simple alternative to their cyclic shift.
• A local or windowed mask limits each query to a neighborhood, as in the restricted self-attention Vaswani et al. suggest for long sequences, or Swin Transformer’s mask that keeps attention inside each window after a cyclic shift.

Properties

Permutation equivariance and position. Every operation above treats the rows of $X$ identically, so $f(PX) = P f(X)$ for any permutation matrix $P$: without a mask, self-attention cannot tell an image from its patches shuffled. Position must be injected separately; the transformer adds a positional encoding to the input tokens (Vaswani et al., 2017). ViT adds learned position embeddings to the input; in its ablation (ViT-B/16, ImageNet 5-shot linear evaluation), removing them lowered accuracy from 64.2% to 61.4%, while 1D, 2D and relative schemes differed little (Dosovitskiy et al., 2021, Table 8). Swin Transformer instead adds a learned relative position bias $B$ to the scores inside each window, $\operatorname{softmax}(QK^\top/\sqrt{d} + B)V$ (Liu et al., 2021).

Cost. Forming $QK^\top$ and multiplying by $V$ each take about $n^2 d$ multiply-adds summed over the heads, and the projections $O(n d^2)$; Vaswani et al. (2017, Table 1) give $O(n^2 \cdot d)$ per layer, with every pair of tokens one step apart. Stored score matrices also take $n^2$ memory per head.

The quadratic term shapes how images are tokenized. A $224 \times 224$ image cut into $16 \times 16$ patches, as in ViT, gives $14 \times 14 = 196$ tokens and 38,416 score entries per head. With one token per pixel it gives 50,176 tokens and about 2.52 billion entries per head and layer, roughly 10 GB in 32-bit floats for one matrix and $256^2 = 65{,}536$ times more. Dosovitskiy et al. (2021) note that pixel-level self-attention does not scale to realistic image sizes, and that sequence length is inversely proportional to the square of the patch size.

Windowed attention. Swin Transformer computes self-attention inside non-overlapping windows of $M \times M$ patches ($M = 7$ by default). For an $h \times w$ patch grid with $C$ channels, the authors give the cost as $4hwC^2 + 2(hw)^2C$ for global and $4hwC^2 + 2M^2hwC$ for window attention, linear in the number of patches for fixed $M$. With Swin’s $4 \times 4$ patches, a $224 \times 224$ input gives $56 \times 56 = 3{,}136$ tokens, and windows of 49 cut the score entries 64-fold. Alternate blocks shift the windows by $\lfloor M/2 \rfloor$ patches so that information crosses window borders.

Examples

Attention over image patches. The code cuts a $32 \times 32$ synthetic image of two bright squares into an $8 \times 8$ grid of $4 \times 4$ patches: 64 tokens of 16 pixels, each with its mean removed so that edges remain. For an interpretable map without training, it ties $W_Q = W_K = 2R$ with $R$ a random orthogonal matrix, so the scaled scores equal the patches’ own dot products (a trained network learns its projections). It prints the attention map of patch (0, 1), on the top edge of the first square.

import numpy as np

rng = np.random.default_rng(0)

# A 32x32 synthetic image: two bright squares on a dark background, light noise.
img = np.zeros((32, 32))
img[2:14, 2:14] = 1.0
img[18:30, 18:30] = 1.0
img += 0.05 * rng.standard_normal(img.shape)

# Cut it into an 8x8 grid of 4x4 patches: 64 tokens of 16 pixels each.
P = 4
X = img.reshape(8, P, 8, P).transpose(0, 2, 1, 3).reshape(64, P * P)
X = X - X.mean(axis=1, keepdims=True)    # remove each patch's mean brightness

def softmax(S):
    E = np.exp(S - S.max(axis=-1, keepdims=True))
    return E / E.sum(axis=-1, keepdims=True)

def attention(Q, K, V, mask=None):
    """Scaled dot-product attention; mask[i, j] = False forbids query i -> key j."""
    S = Q @ K.T / np.sqrt(K.shape[1])
    if mask is not None:
        S = np.where(mask, S, -np.inf)
    A = softmax(S)
    return A @ V, A

# Tie W_Q = W_K = s R with R orthogonal, so the logits are s^2 times patch dot products.
d = X.shape[1]
R, _ = np.linalg.qr(rng.standard_normal((d, d)))
s = 2.0
W_q = W_k = s * R
W_v = rng.standard_normal((d, d)) / np.sqrt(d)
Y, A = attention(X @ W_q, X @ W_k, X @ W_v)

q = 1                                    # patch (row 0, col 1): top edge of the first square
print("attention of patch (0,1) over the 8x8 patch grid:")
print(np.array2string(A[q].reshape(8, 8), precision=2, suppress_small=True))
top = np.argsort(A[q])[::-1][:6]
print("top patches:", [divmod(int(j), 8) for j in top])
print("weight on the second square's top edge (row 4, cols 5-6):",
      round(A[q, 4 * 8 + 5] + A[q, 4 * 8 + 6], 3))

attention of patch (0,1) over the 8x8 patch grid:
[[0.02 0.17 0.19 0.03 0.   0.   0.   0.  ]
 [0.01 0.   0.   0.   0.   0.   0.   0.  ]
 [0.   0.   0.   0.   0.   0.   0.   0.  ]
 [0.   0.   0.   0.   0.   0.   0.   0.  ]
 [0.   0.   0.   0.   0.02 0.15 0.21 0.03]
 [0.   0.   0.   0.   0.   0.   0.   0.  ]
 [0.   0.   0.   0.   0.   0.   0.   0.  ]
 [0.   0.   0.   0.   0.   0.   0.   0.  ]]
top patches: [(4, 6), (0, 2), (0, 1), (4, 5), (4, 7), (0, 3)]
weight on the second square's top edge (row 4, cols 5-6): 0.357

The query attends to the top edges of both squares, the far one as much as its neighbors; bottom edges (negative dot product) and flat patches get almost none.

Multi-head attention, permutation, and a window mask. Continuing with random projections: four heads of width 4, a shuffle of the patches, and a Swin-style mask with four windows of $4 \times 4$ patches.

def multi_head(X, W_q, W_k, W_v, W_o, h, mask=None):
    """h heads of width d/h; head outputs are concatenated and projected by W_o."""
    n, d = X.shape
    dk = d // h
    heads, maps = [], []
    for i in range(h):
        cols = slice(i * dk, (i + 1) * dk)
        Yi, Ai = attention(X @ W_q[:, cols], X @ W_k[:, cols], X @ W_v[:, cols], mask)
        heads.append(Yi)
        maps.append(Ai)
    return np.concatenate(heads, axis=1) @ W_o, np.stack(maps)

h = 4
Wq, Wk, Wv, Wo = (rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(4))
Y, A = multi_head(X, Wq, Wk, Wv, Wo, h)
print("output", Y.shape, "attention maps", A.shape, "parameters", 4 * d * d)

# Permutation equivariance: shuffling the patches shuffles the outputs the same way.
perm = rng.permutation(64)
Y_perm, _ = multi_head(X[perm], Wq, Wk, Wv, Wo, h)
print("f(X[perm]) == f(X)[perm]:", np.allclose(Y_perm, Y[perm], rtol=0, atol=1e-12))

# Windowed attention: each patch attends only within its 4x4-patch window.
r, c = np.divmod(np.arange(64), 8)
win = (r // 4) * 2 + c // 4                       # window index 0..3
mask = win[:, None] == win[None, :]
_, A_win = multi_head(X, Wq, Wk, Wv, Wo, h, mask)
print("patch (0,1): attention mass outside its window:", A_win[:, q, ~mask[q]].sum())
print("score entries per head: global", mask.size, "windowed", mask.sum())

output (64, 16) attention maps (4, 64, 64) parameters 1024
f(X[perm]) == f(X)[perm]: True
patch (0,1): attention mass outside its window: 0.0
score entries per head: global 4096 windowed 1024

Shuffling the patches only shuffles the outputs; the window mask removes all weight outside the query’s window.

The effect of $\sqrt{d_k}$. Continuing, the tokens are now the sixteen $2 \times 2$ patches of each of 300 scikit-learn $8 \times 8$ digit images, randomly projected to queries and keys with roughly unit-variance entries, as the footnote assumes. For each $d_k$ the code reports the spread of raw scores, the mean largest weight, and the mean softmax Jacobian size, unscaled and scaled:

from sklearn.datasets import load_digits

# Tokens: the 16 non-overlapping 2x2 patches of 300 digit images, standardized per pixel.
imgs = load_digits().images[:300]
T = imgs.reshape(300, 4, 2, 4, 2).transpose(0, 1, 3, 2, 4).reshape(300, 16, 4)
T = (T - T.mean(axis=(0, 1))) / T.std(axis=(0, 1))

def softmax_jacobian_norm(p):                     # || diag(p) - p p^T ||_F, averaged over rows
    J = p[..., :, None] * (np.eye(p.shape[-1]) - p[..., None, :])
    return np.sqrt((J ** 2).sum(axis=(-2, -1))).mean()

print(" d_k  logit std  max weight (raw / scaled)  Jacobian norm (raw / scaled)")
for dk in [4, 16, 64, 256, 1024]:
    Gq, Gk = (rng.standard_normal((300, 4, dk)) / 2.0 for _ in range(2))  # fresh weights per image
    Q, K = T @ Gq, T @ Gk                         # entries have roughly unit variance
    S = Q @ K.transpose(0, 2, 1)                  # 300 x 16 x 16 raw logits
    p_raw, p_scl = softmax(S), softmax(S / np.sqrt(dk))
    print(f"{dk:5d} {S.std():9.1f}      {p_raw.max(-1).mean():.2f} / {p_scl.max(-1).mean():.2f}"
          f"                {softmax_jacobian_norm(p_raw):.3f} / {softmax_jacobian_norm(p_scl):.3f}")

 d_k  logit std  max weight (raw / scaled)  Jacobian norm (raw / scaled)
    4       2.0      0.28 / 0.17                0.313 / 0.285
   16       4.1      0.45 / 0.18                0.318 / 0.289
   64       8.2      0.60 / 0.18                0.269 / 0.288
  256      16.3      0.74 / 0.18                0.200 / 0.290
 1024      32.1      0.81 / 0.17                0.142 / 0.289

The raw scores’ standard deviation doubles each time $d_k$ quadruples, as predicted. Unscaled, attention concentrates on one patch and the Jacobian falls to less than half its size by $d_k = 1024$; scaled, both stay constant. Other seeds give the same pattern.

Checking against PyTorch. Continuing (run with PyTorch 2.14.1), the same weights go into torch.nn.MultiheadAttention, and head 0 with the window mask is recomputed by torch.nn.functional.scaled_dot_product_attention. Their documented boolean masks are opposite: True forbids a position in MultiheadAttention and allows it in scaled_dot_product_attention.

import torch
import torch.nn.functional as F

mha = torch.nn.MultiheadAttention(embed_dim=d, num_heads=h, bias=False, batch_first=True).double()
with torch.no_grad():                    # PyTorch computes x @ W.T, so load the transposes
    mha.in_proj_weight.copy_(torch.from_numpy(np.concatenate([Wq.T, Wk.T, Wv.T])))
    mha.out_proj.weight.copy_(torch.from_numpy(Wo.T))
x = torch.from_numpy(X)[None]            # batch of one image: (1, 64, 16)

with torch.no_grad():
    Y_t, A_t = mha(x, x, x, average_attn_weights=False)
    # Opposite boolean conventions: MultiheadAttention's True blocks, SDPA's True allows.
    Y_tw, _ = mha(x, x, x, attn_mask=torch.from_numpy(~mask))
    Q, K, V = (x @ torch.from_numpy(W[:, :4]) for W in (Wq, Wk, Wv))   # head 0
    Y_sdpa = F.scaled_dot_product_attention(Q, K, V, attn_mask=torch.from_numpy(mask))

Y_w, _ = multi_head(X, Wq, Wk, Wv, Wo, h, mask)
Y_h0, _ = attention(X @ Wq[:, :4], X @ Wk[:, :4], X @ Wv[:, :4], mask)
print(f"torch {torch.__version__}")
print(f"MultiheadAttention vs NumPy, output:   {np.abs(Y_t[0].numpy() - Y).max():.1e}")
print(f"MultiheadAttention vs NumPy, maps:     {np.abs(A_t[0].numpy() - A).max():.1e}")
print(f"windowed MultiheadAttention vs NumPy:  {np.abs(Y_tw[0].numpy() - Y_w).max():.1e}")
print(f"windowed SDPA (head 0) vs NumPy:       {np.abs(Y_sdpa[0].numpy() - Y_h0).max():.1e}")
print("parameters:", sum(p.numel() for p in mha.parameters()))

torch 2.14.1
MultiheadAttention vs NumPy, output:   3.1e-17
MultiheadAttention vs NumPy, maps:     1.0e-17
windowed MultiheadAttention vs NumPy:  6.9e-17
windowed SDPA (head 0) vs NumPy:       6.9e-17
parameters: 1024

Outputs and per-head maps agree to rounding error.

Common Misconceptions

• “Self-attention knows where each patch is.” It is permutation equivariant; layout enters only through positional encodings, position biases, or masks.
• “Masking means zeroing weights after the softmax.” That leaves rows that no longer sum to one. Masks set scores to $-\infty$ before the softmax, so the remaining weights renormalize.
• “The $\sqrt{d_k}$ is a tuned temperature.” It is a fixed constant that cancels the growth of dot products with the key dimension.
• “An attention map explains the prediction.” It shows how one head of one layer mixed values that are themselves projections, before further heads, layers and residual paths. On NLP tasks, Jain and Wallace (2019) found attention weights often uncorrelated with gradient-based importance, and very different attention distributions giving equivalent predictions.

Where It Is Used

Non-local blocks for video. Wang et al. (2018) defined a non-local operation, inspired by classical non-local means, that computes the response at a position as a normalized weighted sum of the features at all positions. Its “embedded Gaussian” version, a softmax over dot products of embedded features, has the self-attention form; the authors describe the translation model’s self-attention as a special case. They applied it across all $T \times H \times W$ positions of a video feature map. With five blocks in a C2D ResNet-50 baseline on Kinetics, they reported top-1 accuracy of 71.8% without them, 72.9% with attention over space only, 73.1% over time only, and 73.8% over space-time. Versions without the softmax worked about as well, so they judged the attentional behavior not essential in their applications.

Image generation. SAGAN (Zhang et al., 2019) added self-attention over feature-map locations to the generator and discriminator of a convolutional GAN, scaled by a learnable factor initialized at zero so that the network first relies on local cues.

Transformers for vision. The transformer stacks multi-head self-attention with position-wise fully connected layers (Vaswani et al., 2017). ViT applies it to image patches for classification; Swin Transformer builds a hierarchical backbone for detection and segmentation as well.

Matching and motion. In feature matching, SuperGlue and LoFTR alternate self-attention within each image with cross-attention between the two. In optical flow, GMA adds a transformer-based aggregation module to RAFT. RAFT’s own all-pairs correlation volume is related: it holds the dot products between every feature vector of one frame and every one of the next, like a score matrix $QK^\top$ between two images (see cost volume), and grows quadratically with the number of pixels for the same reason. The comparison is ours and has limits: RAFT compares outputs of one shared encoder rather than separate query and key projections, applies no softmax and takes no weighted sum of values; its correlations are looked up as matching evidence.