Attention: mixing information from context

Attention lets a token’s representation draw information from other positions. The network learns which information is useful through projections called queries, keys and values: Q, K and V.

Score, normalize, combine

A query is compared with the available keys. In scaled dot-product attention, these comparisons are dot products divided by the square root of the query dimension. Softmax turns the scores into non-negative weights that sum to one. The output is the weighted combination of the value vectors.

For a token referring to something earlier in a sentence, this provides a route for information from that earlier position. Different attention heads can construct different mixtures. Later operations combine and transform their results.

Inside one causal head

Here x contains the token vectors and the three weight matrices are learned parameters. NumPy makes the calculation visible:

import numpy as np


def causal_attention(x, wq, wk, wv):
    q, k, v = x @ wq, x @ wk, x @ wv
    scores = q @ k.T / np.sqrt(q.shape[-1])
    scores[np.triu_indices(len(x), k=1)] = -np.inf
    scores -= scores.max(axis=-1, keepdims=True)
    weights = np.exp(scores)
    weights /= weights.sum(axis=-1, keepdims=True)
    return weights @ v

The triangular mask prevents a position from using future tokens. This allows a whole training sequence to be processed in parallel without revealing the answers to earlier positions. Subtracting each row’s maximum keeps the softmax computation numerically stable.

The returned vectors can feed the next operations in a Transformer block. Full models also have residual connections, normalization and feed-forward networks around their attention operations.

For n positions, full attention considers roughly n² query–key pairs. Optimized implementations can avoid storing that entire score matrix. A KV cache saves repeated projections during generation, while each new query still reads the retained context.

Related: embeddings, Transformer, our model implementation.