Attention Is All You Need

Before this paper, the strongest sequence models read text one token at a time with recurrence. This walkthrough rebuilds the paper's central mechanism, scaled dot-product attention, from scratch, so the equation stops being abstract and becomes something you can run.

The problem

A recurrent network processes a sentence left to right, carrying a hidden state forward. Two things hurt:

• It is sequential. Token t cannot be computed until token t − 1 is done, so training cannot fully use parallel hardware.
• Long-range links fade. Information from the first word has to survive many update steps to influence the last word, and often it does not.

We want a layer where any position can draw directly on any other position, in one parallel step.

The idea

Represent every token as a vector, then let each token ask a question about all the others and collect a weighted blend of their content. Three learned projections give each token a query, a key, and a value. A token's query is compared against every key to produce weights; those weights mix the values. No recurrence is involved, the whole sentence is processed at once.

The architecture around it

Attention is the core, but a Transformer block wraps it in a few standard parts:

• Multi-head attention, run several attention operations in parallel, each with its own projections, then concatenate. Different heads can specialise (one tracks syntax, another tracks a referent).
• Positional encoding, attention alone is order-blind, so a position signal is added to the token vectors.
• Feed-forward network, a small per-token MLP after attention.
• Residual connections and layer normalization, around each sub-layer, so gradients flow and training stays stable.

The mathematics

For queries $Q$, keys $K$, and values $V$ (each a matrix whose rows are tokens), scaled dot-product attention is:

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

$QK^\top$ scores every query against every key. Dividing by $\sqrt{d_k}$ (the key dimension) keeps the scores from growing large enough to push softmax into a near–one-hot regime where gradients vanish. The softmax turns each row of scores into weights that sum to 1, and multiplying by $V$ blends the values.

A minimal implementation

Just NumPy, one function, plus a row-wise softmax:

import numpy as np

def softmax(x):
    # subtract row max for numerical stability
    x = x - x.max(axis=-1, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V):
    d_k = K.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)   # (n_queries, n_keys)
    weights = softmax(scores)         # each row sums to 1
    return weights @ V, weights

# Three tokens, 4-dimensional. Here Q = K = V (self-attention).
x = np.array([
    [1.0, 0.0, 1.0, 0.0],   # token A
    [0.0, 2.0, 0.0, 2.0],   # token B
    [1.0, 1.0, 1.0, 1.0],   # token C
])

out, weights = attention(x, x, x)
print(np.round(weights, 3))
print(np.round(out, 3))

What to observe

• Every row of weights sums to 1. Each token spends a fixed budget of attention across the others.
• Similar tokens attend to each other. Token C, whose vector overlaps both A and B, spreads its weight; A and B, being near-orthogonal, attend mostly to themselves.
• Remove the / np.sqrt(d_k) scaling and scale x up (say x * 10): the softmax collapses toward one-hot and the blend becomes a hard pick. That saturation is exactly what the scaling factor exists to prevent.

Limitations

• The score matrix is $n \times n$ in the sequence length, so cost grows quadratically, the main reason long-context models need extra tricks.
• This is one attention operation, not a trained Transformer: no multi-head projections, positional encoding, feed-forward layer, or optimisation loop.
• Attention weights are often shown as an "explanation," but high weight does not reliably mean causal importance, read such visualisations with care.

Further reading

• Attention Is All You Need (arXiv:1706.03762)
• The Illustrated Transformer