Attention and Transformers: The Mathematics of Modern AI

The chatbots, translators and code assistants that dominate modern AI are almost all built on one architecture: the transformer. At its heart is a surprisingly compact formula called scaled dot-product attention, which lets every word in a sentence look at every other word and decide what matters. In this article we unpack that formula with matrices small enough to compute by hand.

Why Attention Was Needed

Earlier language models, such as recurrent neural networks, read text one word at a time and squeezed everything they had seen into a single hidden vector. In a long sentence like "The book that my teacher in Lahore recommended last year was excellent", the model must remember "book" across many words to know what "was excellent" refers to. Information faded over long distances, and the sequential processing could not easily be parallelised.

Attention solves this by letting each position directly consult every other position, with learned weights saying how much to borrow from each. The 2017 paper "Attention Is All You Need" by Vaswani and colleagues at Google showed that a model built from attention alone, with no recurrence, could outperform previous translation systems while training much faster.

Tokens as Vectors

Text is first split into tokens (words or word pieces), and each token is mapped to an embedding vector of dimension \( d_{\text{model}} \), often 512 or more. A sentence of \( n \) tokens becomes a matrix \( X \in \mathbb{R}^{n \times d_{\text{model}}} \), one row per token.

Queries, Keys and Values

Each token produces three vectors by multiplying its embedding by learned weight matrices:

$$ Q = X W_Q, \qquad K = X W_K, \qquad V = X W_V . $$

A helpful analogy is a library search:

• The query is what this token is looking for.
• The key is what each token advertises about itself.
• The value is the information each token actually hands over.

A token compares its query with every key. High similarity means "pay attention to that token", and the output is a weighted average of the values.

Scaled Dot-Product Attention

The full operation is

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

where \( d_k \) is the length of each key vector and softmax is applied to each row:

$$ \operatorname{softmax}(z)_j = \frac{e^{z_j}}{\sum_{l} e^{z_l}} . $$

The matrix \( QK^\top \) has shape \( n \times n \); entry \( (i, j) \) is the dot product of query \( i \) with key \( j \). Softmax turns each row into non-negative weights that add up to 1.

Why Divide by the Square Root of d_k?

If the components of \( \mathbf{q} \) and \( \mathbf{k} \) are independent with mean 0 and variance 1, then \( \mathbf{q} \cdot \mathbf{k} = \sum_{i=1}^{d_k} q_i k_i \) is a sum of \( d_k \) terms each with variance 1, so its variance is \( d_k \). For \( d_k = 64 \) the scores would have a standard deviation of 8, and softmax of such large numbers is almost one-hot, giving tiny gradients. Dividing by \( \sqrt{d_k} \) brings the standard deviation back to 1.

Worked Example 1: Attention on Three Tokens

Take \( n = 3 \) tokens with \( d_k = 2 \) and

$$ Q = K = \begin{pmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{pmatrix}, \qquad V = \begin{pmatrix} 1 & 0 \\ 0 & 2 \\ 1 & 1 \end{pmatrix}. $$

We compute the output for token 1, whose query is \( \mathbf{q}_1 = (1, 0) \).

Step 1. Dot products with each key. \( \mathbf{q}_1 \cdot \mathbf{k}_1 = 1 \), \( \mathbf{q}_1 \cdot \mathbf{k}_2 = 0 \), \( \mathbf{q}_1 \cdot \mathbf{k}_3 = 1 \).

Step 2. Scale by \( \sqrt{2} \approx 1.4142 \). Scores become \( 0.7071, 0, 0.7071 \).

Step 3. Exponentiate. \( e^{0.7071} \approx 2.0281 \) and \( e^{0} = 1 \), giving \( 2.0281, 1, 2.0281 \).

Step 4. Normalise. The sum is \( 5.0562 \), so the weights are

$$ \alpha_1 = (0.4011,\ 0.1978,\ 0.4011). $$

Step 5. Weighted sum of values.

$$ \mathbf{z}_1 = 0.4011(1, 0) + 0.1978(0, 2) + 0.4011(1, 1) = (0.8022,\ 0.7967). $$

Token 1 pays most attention to itself and to token 3, whose keys align with its query, and less to token 2, whose key is perpendicular.

Repeating for token 3 with \( \mathbf{q}_3 = (1, 1) \): dot products \( 1, 1, 2 \), scaled to \( 0.7071, 0.7071, 1.4142 \), exponentials \( 2.0281, 2.0281, 4.1133 \), sum \( 8.1695 \), weights \( (0.2483, 0.2483, 0.5035) \). The output is

$$ \mathbf{z}_3 = 0.2483(1,0) + 0.2483(0,2) + 0.5035(1,1) = (0.7517,\ 1.0000). $$

(With the rounded weights the components come out as \( 0.7518 \) and \( 1.0001 \); unrounded arithmetic gives \( 0.7517 \) and exactly 1, because the three weights must sum to 1.)

Multi-Head Attention

One attention pattern cannot capture everything: one relationship might be grammatical (subject and verb), another about meaning (pronoun and noun). Transformers therefore run \( h \) attention "heads" in parallel, each with its own \( W_Q, W_K, W_V \) of width \( d_k = d_{\text{model}}/h \). The head outputs are concatenated and multiplied by an output matrix \( W_O \):

$$ \text{MultiHead}(X) = \operatorname{Concat}(\text{head}_1, \dots, \text{head}_h)\, W_O . $$

Worked Example 2: Counting Parameters and Cost

Use the original base transformer settings: \( d_{\text{model}} = 512 \) and \( h = 8 \) heads.

Step 1. Head width. \( d_k = 512/8 = 64 \).

Step 2. Projection weights. Across all heads, \( W_Q \), \( W_K \) and \( W_V \) are each effectively \( 512 \times 512 \), and \( W_O \) is also \( 512 \times 512 \). Each has \( 512^2 = 262{,}144 \) weights.

Step 3. Total weights. \( 4 \times 262{,}144 = 1{,}048{,}576 \), plus \( 4 \times 512 = 2048 \) biases, giving \( 1{,}050{,}624 \) parameters for one attention block.

Step 4. Attention matrix size. For a sequence of \( n = 1000 \) tokens, each head computes an \( n \times n \) score matrix with \( 1000^2 = 1{,}000{,}000 \) entries. With 8 heads that is 8 million scores per layer.

Notice that the parameter count does not depend on \( n \), but the computation grows as \( n^2 \). Doubling the context length quadruples the attention cost, which is why long-context models are an active research area.

Positional Encoding

Attention treats its input as a set: shuffle the tokens and each output is shuffled the same way but otherwise unchanged. Word order must therefore be injected explicitly. The original transformer added sinusoidal vectors to the embeddings:

$$ PE(\text{pos}, 2i) = \sin\!\left(\frac{\text{pos}}{10000^{2i/d_{\text{model}}}}\right), \qquad PE(\text{pos}, 2i+1) = \cos\!\left(\frac{\text{pos}}{10000^{2i/d_{\text{model}}}}\right). $$

For position 0, every sine component is 0 and every cosine component is 1. Different frequencies let the model distinguish both nearby and distant positions. Many newer models use learned or rotary position encodings instead.

The Complete Transformer Block

A single block applies these steps in order:

1. Multi-head self-attention, with a residual connection (input added back) and layer normalisation.
2. A position-wise feed-forward network, typically expanding to about four times \( d_{\text{model}} \) and back, again with a residual connection and normalisation.

Large language models stack dozens of such blocks. In text generation, a causal mask sets scores for future positions to \( -\infty \) before softmax, so each token can only attend to earlier ones.

Did You Know?

• Every row of an attention matrix is a probability distribution, so researchers often visualise these weights to see what a model "looks at".
• Softmax is unchanged if you subtract the same constant from all scores, a trick programmers use to avoid overflow when exponentiating.
• The same transformer architecture now powers models for images, protein structures and audio, not just text.

Key Takeaways

• Attention computes \( \operatorname{softmax}(QK^\top/\sqrt{d_k})V \): similarity scores, normalised weights, then a weighted average of values.
• Scaling by \( \sqrt{d_k} \) keeps scores at unit variance so softmax does not saturate.
• Multi-head attention runs several attention patterns in parallel and combines them.
• Parameters are independent of sequence length, but compute grows with \( n^2 \).
• Positional encodings supply word order, which attention alone cannot see.

Need help with a transformer assignment, a linear algebra derivation or an AI project report? Amathify's experts and PhD scholars offer step-by-step solutions via WhatsApp or the Ask a Question form.