Attention, computed step by step

Attention lets every position build its next representation from a weighted average of all previous positions, with the weights computed from the content itself. It is a learned, content-addressed lookup.

The equation

Attention(Q, K, V) = softmax( (Q Kᵀ) / √d_k + M ) V

Reading it term by term:

• Q, K, V are three linear projections of the same input: Q = XW_Q, K = XW_K, V = XW_V. Interpret a query as "what am I looking for", a key as "what do I offer", a value as "what I will contribute if selected".
• Q Kᵀ gives an n × n matrix of raw compatibility scores between every pair of positions.
• ÷ √d_k is not cosmetic. With unit-variance components, a dot product over d_k dimensions has variance d_k; without rescaling the softmax saturates and gradients vanish. This one division is what makes the mechanism trainable at scale.
• + M is the causal mask: −∞ above the diagonal, so position i cannot attend to anything after it. This is what makes the model autoregressive, and it is why one training pass can supply a learning signal at every position simultaneously.
• softmax turns scores into a distribution over positions; multiplying by V takes the weighted average.

import numpy as np

def attention(X, Wq, Wk, Wv):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)
    n = scores.shape[0]
    scores += np.triu(np.full((n, n), -np.inf), k=1)      # causal mask
    weights = np.exp(scores - scores.max(-1, keepdims=True))
    weights /= weights.sum(-1, keepdims=True)
    return weights @ V

Multi-head, and why

One attention operation produces one weighted average — a single relationship per position. Splitting d_model into h heads of width d_model/h and running attention independently in each lets different heads specialise: some track syntactic dependency, some carry positional patterns, some perform the induction behaviour ("this pattern appeared earlier, copy what followed it") that underlies in-context learning. Outputs are concatenated and projected back.

The quadratic problem

Time is O(n² · d). Naively, memory is also O(n²) because the score matrix is materialised — which is what actually blocked long contexts for years, before compute did.

The KV cache

During generation, token t+1 needs the keys and values of every earlier token — which are unchanged. Recomputing them each step would make generation O(n²) per token. Caching them makes each step O(n).

The cost is memory, and it is substantial:

KV cache bytes ≈ 2 × layers × heads_kv × d_head × seq_len × batch × bytes_per_element

For a large model at long context and reasonable batch size this reaches tens of gigabytes — often exceeding the weights themselves. Serving throughput is usually limited by KV cache memory, not by parameters, which is the fact that drives the whole next generation of attention variants.

MQA and GQA

Since the cache scales with the number of key/value heads, reduce them. Multi-query attention keeps many query heads but a single K/V head — a large cache reduction with some quality loss. Grouped-query attention is the compromise now standard in production models: groups of query heads share one K/V head, capturing most of the memory saving with almost no quality cost.

Try it yourself

Run the code above on a short sequence and print the weight matrix. Confirm it is lower-triangular and each row sums to 1. Then remove the √d_k division and inspect the weights with d_k = 512: you will see them collapse to nearly one-hot, which is exactly the vanishing-gradient failure the scaling prevents.