Attention

The idea#

To understand "it" in a sentence you have to know what it refers to. So you look back and weigh: the noun three words earlier is probably meant, the clause before it probably not.

Attention does exactly that, for every token simultaneously and with numbers instead of intuition.

The three roles#

Every token produces three vectors:

Every query is compared with every key. Where agreement is large, the corresponding value contributes strongly to the result.

By hand#

import numpy as np

def softmax(x):
    x = x - x.max(axis=-1, keepdims=True)     # numerically stable
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

rng = np.random.default_rng(0)
s, d = 6, 16                                   # 6 tokens, 16 dimensions
Q, K, V = (rng.normal(size=(s, d)) for _ in range(3))

scores = Q @ K.T / np.sqrt(d)                  # the scaling is not optional

# Causal mask: tokens must not see the future.
mask = np.triu(np.ones((s, s)), k=1).astype(bool)
scores[mask] = -np.inf

A = softmax(scores)
out = A @ V
print(A.round(2)[3])       # row 3: what token 4 attends to
print(A.sum(axis=-1))      # every row sums to 1

Multiple heads#

Instead of one attention over 4096 dimensions you compute 32 heads of 128 dimensions each. Every head can specialise in a different kind of relationship, syntax, reference, position, topic. The results are concatenated and passed through another matrix.

The formula#

Why the scaling is necessary#

For random vectors with independent unit-variance components, the dot product q·k has variance d_k. At d_k = 128 the standard deviation of the scores is therefore about 11.3. Softmax on values of that magnitude is effectively a maximum: one component takes almost all the weight, all others almost none, and the gradient vanishes. Dividing by √d_k returns the variance to 1.

The cost#

At s = 4096, d = 4096: projections 2.7e11, attention 1.4e11, still comparable. At s = 32768: projections 2.2e12, attention 8.8e12, four times as much. From about s = d the quadratic term dominates completely.

FlashAttention#

The way out is not computing less but never storing the matrix in full. FlashAttention processes blocks, keeps them in the compute unit's fast memory, and updates the softmax result incrementally with a running maximum and a running sum.

The result: memory falls from O(s²) to O(s), and runtime falls despite the same operation count, because slow card memory is no longer the bottleneck. Without it, context lengths above roughly 8,000 tokens would be practically unaffordable. See Context windows in depth.