Attention from the Math Up

After this lesson, you will be able to: Derive and implement scaled dot-product attention: queries, keys, values, the scaling factor, the softmax, and causal masking.

Attention is the heart of the transformer and the one equation every NLP researcher must know cold. This lesson builds it from the intuition to the formula to a few lines of code.

Prerequisites:Embeddings and Positional Encoding

The intuition: query, key, value

Each token produces three vectors by multiplying its representation by learned matrices: a query (what am I looking for), a key (what do I offer), and a value (what I will contribute if attended to). A token's new representation is a weighted average of all tokens' values, where the weight from token i to token j is how well i's query matches j's key. 'It' attends strongly to the noun it refers to because their query/key dot product is high. The model learns the matrices that make these matches meaningful.

Scaled dot-product attention, the formula and the code

Attention(Q,K,V) = softmax(Q Kᵀ / sqrt(d_k)) V. A few lines of NumPy make it concrete.

python

import numpy as np

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

def attention(Q, K, V, mask=None):
    # Q,K,V: (seq_len, d_k). One head.
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)      # (seq, seq): every query dot every key
    if mask is not None:
        scores = np.where(mask, scores, -1e9)   # block disallowed positions
    weights = softmax(scores, axis=-1)   # each row sums to 1
    return weights @ V                   # weighted average of values

# Causal mask for a decoder LM: token i may only attend to j <= i
seq = 4
causal = np.tril(np.ones((seq, seq))).astype(bool)

Why divide by sqrt(d_k)

The dot product of two d_k-dimensional vectors grows with d_k, so for large d_k the raw scores get large, pushing softmax into regions where one weight is near 1 and the rest near 0 and the gradients vanish. Dividing by sqrt(d_k) keeps the variance of the scores roughly constant regardless of dimension, so softmax stays in a useful range and training is stable. This is the kind of detail reviewers expect you to know if your method touches attention.

Common mistakes only experienced researchers catch

Forgetting the scaling factor and wondering why training is unstable. Applying the causal mask after softmax instead of before (set disallowed scores to a large negative number, then softmax). Confusing the attention weights (which sum to 1 per query) with the output. Computing attention in float16 without care, where the unmasked positions can overflow. Assuming attention is O(n) when it is O(n squared) in sequence length, which is exactly why efficient-attention is a whole research subfield (FlashAttention, sparse/linear attention).

Quick Check

Why is the dot product scaled by 1/sqrt(d_k) before the softmax?

Pick the correct explanation.

←Embeddings and Positional Encoding

Back to Transformers and NLP Foundations

Multi-Head Attention and the Transformer Block→