Attention

Attention computes a context-dependent weighted average of value vectors. A query asks what it needs, keys decide which positions match, and values provide the information to mix. This is the central mechanism in transformers, a bridge away from recurrent networks, and a common alignment mechanism in multimodal learning.

How attention works

Every position produces three vectors: a query (what it is looking for), a key (what it offers to be matched against), and a value (the information it will contribute). For each query, attention then

1. scores the query against every key with a dot product,
2. turns those scores into weights with a softmax, so they are non-negative and sum to one, and
3. returns the weighted average of the value vectors.

Scaled dot-product attention writes this as

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

where $Q$, $K$, and $V$ stack the query, key, and value vectors for all positions and $d_k$ is the key dimension. For multi-head attention,

$${\mathrm{head}}_i=\mathrm{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V),$$ $$\mathrm{MHA}(Q,K,V)=\mathrm{Concat}({\mathrm{head}}_1,\dots ,{\mathrm{head}}_h)W^O.$$

The $\sqrt{d_k}$ divisor keeps dot-product logits from growing with key dimension. Masks can forbid future tokens or padded positions, which is essential in language-model architectures.

flowchart TD
  Input[Input sequence] --> Q[Queries Q]
  Input --> K[Keys K]
  Input --> V[Values V]
  Q --> Scores[Score: scaled dot product of Q and K]
  K --> Scores
  Scores --> Weights[Softmax over scores gives attention weights]
  Weights --> Context[Weighted sum of value vectors]
  V --> Context
  Context --> Output[Context vectors]

Worked example

This snippet computes scaled dot-product attention weights from queries and keys, then forms the context vector as the weighted sum of values.

import math, torch
 
Q = torch.tensor([[1., 0.], [0., 1.]])
K = torch.tensor([[1., 0.], [1., 1.], [0., 1.]])
V = torch.tensor([[10., 0.], [0., 5.], [0., 1.]])
scores = Q @ K.T / math.sqrt(2)
weights = scores.softmax(dim=-1)
context = weights @ V
print("weights", torch.round(weights, decimals=3).tolist())
print("context", torch.round(context, decimals=3).tolist())

Observed output:

weights [[0.4009999930858612, 0.4009999930858612, 0.1979999989271164], [0.1979999989271164, 0.4009999930858612, 0.4009999930858612]]
context [[4.011000156402588, 2.203000068664551], [1.9780000448226929, 2.4070000648498535]]

The first query attends most to the first two keys, while the second attends most to the last two. The output vectors are weighted mixtures of the values, not selected tokens.

Caveats

Attention weights are routing weights, not full explanations of a model decision. Full self-attention is $O(n^2)$ in sequence length for its score matrix, so long contexts stress memory and latency. A mask bug changes what information can flow and can silently invalidate evaluation.

References

• Vaswani et al., 2017, Attention Is All You Need
• PyTorch documentation: MultiheadAttention

Nav

Section — Deep Learning

← LSTM and GRU Transformers →