Transformer attention from first principles

Attention is easier to understand when you start with its purpose. A token needs information from other tokens to represent its meaning in context.

A language model starts with a vector for each token. It retrieves that vector from an embedding table. The initial vector is the same each time the token appears. The token "bank" therefore starts with the same vector in a sentence about a river or a loan.

Recurrent models read tokens in sequence and maintain a running summary. That summary may retain "river" when the model reaches "bank." A fixed-size summary must represent the whole sequence. Information from earlier tokens can become harder to retain.

Attention lets each token compare itself with other tokens directly. It assigns weights to those comparisons and uses them to build a new vector. The token "bank" can receive information from "river." Each permitted pair has a direct connection, rather than a chain of recurrent steps.

This post explains the calculation in stages:

1. Calculate dot products between token vectors.
2. Convert scores into weights with softmax.
3. Introduce trainable query, key and value projections.
4. Apply a causal mask for next-token prediction.
5. Combine several attention heads.

The examples use four tokens, so you can calculate the results by hand and compare them with code.

Where the vectors come from

Attention receives vectors and returns vectors. Token IDs, embeddings and position information determine its input. The context window limits the sequence length.

Tokens. A tokeniser splits text into vocabulary entries and assigns an integer ID to each entry. GPT-2's vocabulary contains 50,257 entries. Common words can correspond to single tokens. Less common words can require several tokens. Sequence length counts tokens, rather than words.

Embeddings. Each token ID selects a row from an embedding table with shape $\text{vocab} \times d_{\text{model}}$. GPT-2 small uses a $50{,}257 \times 768$ table. Row 3 supplies the vector for token ID 3.

Training adjusts the table through backpropagation. Tokens with similar uses can develop related vectors. However, the initial table supplies one vector per token, regardless of the current sentence. Attention adds information from that sentence.

Position. The basic attention formula does not identify a token's position. If you reorder the input rows, the output rows follow that order. Language needs position information because word order can change meaning.

The original Transformer adds sinusoidal position encodings to token embeddings. Other approaches include learned position embeddings and rotary encodings. This example omits position information to keep the arithmetic small. The matrix $X$ represents the vectors that enter attention.

Context window. The context window limits the number of tokens a model can process in one sequence. GPT-2 uses a window of 1,024 tokens. Its learned position table cannot supply a position embedding beyond that range. Other models use different limits and position methods.

The symbol $n$ denotes the sequence length in this post. The attention score matrix has shape $n \times n$. The final section explains the cost of increasing $n$.

The input matrix $X$ contains $n$ token vectors. Each vector has $d_{\text{model}}$ dimensions, so $X$ has shape $n \times d_{\text{model}}$. This example uses $n = 4$ and $d_{\text{model}} = 4$.

The one idea

Attention is a weighted average. Each token's new vector is

For each token, the weights $w_{ij}$ sum to 1. The remaining steps explain how to calculate those weights.

They come from dot products. For two vectors, the dot product is

Multiply the matching coordinates. Sum the products. For $a = [0.9, 0.1]$ and $b = [0.8, 0.2]$, the result is $0.72 + 0.02 = 0.74$. For $a$ and $c = [-0.3, 0.9]$, the result is $-0.27 + 0.09 = -0.18$.

A dot product measures alignment and also depends on vector magnitude. Aligned vectors give positive products. Orthogonal vectors give zero. Opposing directions give negative products.

The calculation needs one multiply-add per dimension. Matrix multiplication, or matmul, calculates many dot products together. This makes dot products useful as attention scores.

The initial example has no trainable projections. It follows three steps:

1. Calculate each token's dot product with every token.
2. Convert the scores into weights that sum to 1.
3. Use those weights to average the token vectors.

A sentence small enough to finish on paper

The example uses four tokens with four dimensions each:

This example uses invented vectors. The vocabulary contains four token IDs, from 0 to 3. Stacking the vectors produces $X$ with shape $4 \times 4$. Each row represents a token. Each column represents a dimension.

Simple self-attention

Calculate the dot product of "saw" with each token, including itself:

These numbers are scores. In this example, "saw" has its highest score with itself and a negative score with "cat". Calculating every row produces $X X^{\top}$ with shape $4 \times 4$. Entry $(i, j)$ contains the dot product between tokens $i$ and $j$.

The scores do not yet form a weighted average. They include a negative number and do not sum to 1. Softmax converts them into positive weights that sum to 1:

Exponentiate each score. Divide each result by the sum of all results in its row. For "saw", the exponentials are $e^{0.04} = 1.04$, $e^{1.09} = 2.97$, $e^{0.21} = 1.23$ and $e^{-0.74} = 0.48$.

The sum is approximately 5.73. Dividing each value by the sum gives weights of approximately [0.18, 0.52, 0.22, 0.08].

The weighted average uses 18% of "she", 52% of "saw", 22% of "the" and 8% of "cat":

The result is a context-aware vector for "saw." Its original vector was [0.20, 0.80, 0.40, -0.50]. The weighted average combines information from its neighbours. For all four rows, the calculation is:

What is missing

In this example, each row's highest weight is on its own token. Simple attention therefore retains much of each original vector. It also mixes information from nearby vectors in embedding space. This pattern follows from these chosen vectors, rather than a rule that every self-dot-product must be largest.

This version has no trainable projection weights. The input embeddings determine its attention scores. It cannot adjust a separate comparison to represent a subject-verb relationship.

Trainable projections change the comparison. Project each token into query and key vectors. Compare queries with keys instead of comparing the original token vectors directly.

Trainable attention

This example uses three matrices with shape $d_{\text{model}} \times d_k$. The dimensions are $d_{\text{model}} = 4$ and $d_k = 2$:

Each projection is a linear map without a bias term, as in $y = mx$ from the backpropagation article.

Multiply $X$ by $W_Q$ to calculate queries. Multiply $X$ by $W_K$ to calculate keys. Multiply $X$ by $W_V$ to calculate values:

A query describes what a token seeks. A key supplies information against which queries compare. A value supplies the information that the weighted average combines.

Scores compare queries with keys. The output averages values. Separate projection weights let training adjust each role independently.

The query for "saw" is [1.15, -1.29]. The key for "she" is [1.18, -1.13]. These vectors align because of the chosen projection weights, rather than the original embeddings.

The key for "saw" is [0.20, 0.06]. It receives a much smaller score from the "saw" query. A token's own position no longer needs to receive the highest score.

Scores, now with a divide

The score for "saw" against "she" is $1.15 \times 1.18 + (-1.29)(-1.13) = 2.81$. The score against "the" is $-2.29$. Divide each score by $\sqrt{2} = 1.414$. The scaled "saw" row is [1.99, 0.11, -1.62, 0.16].

The scale factor controls the size of dot products. With independent, zero-mean entries of variance 1, the dot product has variance $d_k$. The following example measures variance over twenty thousand random pairs:

At $d_k = 512$, raw scores vary by tens. Softmax of [2, 1, 0, -1] gives [0.64, 0.24, 0.09, 0.03]. Multiplying those scores by 8 gives weights that round to [1.00, 0.00, 0.00, 0.00].

Almost all weight goes to one token. Gradients through softmax become small in this region. Dividing by $\sqrt{d_k}$ keeps score variance near 1 under the stated assumptions. The original attention paper explains this scaling.

Weights and output

Apply softmax to each row of $S$. Use the resulting weights to average the values:

Without the causal mask, "saw" assigns about 75% of its weight to "she" and 2% to "the". The token "cat" assigns about 67% to "the". The chosen projections produce relationships that differ from the simple embedding comparison.

The output for "saw" is:

$0.75 \, v_{\text{she}} + 0.11 \, v_{\text{saw}} + 0.02 \, v_{\text{the}} + 0.12 \, v_{\text{cat}} = [0.55, -0.12]$.

The output has two dimensions because $d_k = 2$. The multi-head section restores $d_{\text{model}}$ dimensions with an output projection.

The complete attention calculation is:

Calculate the scores, scale them, apply softmax and average the values. The projection matrices $W_Q$, $W_K$ and $W_V$ contain the trainable parameters in this example. Backpropagation supplies their gradients. The calculation has no recurrent dependency between positions.

Causal attention

Unmasked attention lets "she" use information from "cat." This can suit a model that receives the complete sentence. A next-token model must not use the token it must predict.

During training, a next-token model predicts the token after each input position. For example, "she saw the" supplies the input for predicting "cat." Training can process these positions together.

If "the" can attend to "cat", the input reveals the expected answer. The model can exploit that leak instead of learning the prediction task. At inference, the future token is unavailable.

Token $i$ may therefore attend only to positions $j \le i$. A lower-triangular mask enforces this restriction. Apply the mask before softmax.

Removing weights after softmax would make their sum less than 1. You would need another normalization step. Instead, set future-position scores to $-\infty$ before softmax.

Because $e^{-\infty} = 0$, those positions receive zero weight. Softmax normalizes the remaining positions automatically.

The unmasked scaled scores for "the" are [0.53, -0.02, -0.40, 0.12]. Mask the final position. The remaining exponentials are approximately [1.70, 0.98, 0.67], with a sum of 3.35. The resulting weights are [0.51, 0.29, 0.20, 0]. The final zero prevents access to "cat".

The last row does not change when the causal mask applies. The final input token, "cat", can already attend to every input position. Its output supplies the context for predicting the next token.

During training, attention dropout removes a random fraction of weights and rescales the remaining weights. It reduces reliance on particular attention connections. Inference disables dropout.

Raschka's chapter 3 includes notebook examples for masking and dropout.

Multi-head attention

A single head produces one weight distribution and one output vector for each token. The token "cat" may need information about both its determiner and its verb. One weighted average must combine those relationships.

Multi-head attention uses separate query, key and value projections for each head. Different heads can learn different attention patterns. The second set of example weights produces another pattern:

Each head returns a $4 \times 2$ matrix. Concatenating two heads along the feature axis gives a $4 \times 4$ matrix. An output projection, $W_O$, then mixes information from the heads:

For "cat", head 1 returns [-0.10, 0.39] and head 2 returns [-0.02, 0.33]. Concatenation gives [-0.10, 0.39, -0.02, 0.33].

The output projection is $W_O$ = [[0.5, -0.1, 0.3, 0.2], [0.1, 0.6, -0.2, 0.4], [-0.3, 0.2, 0.5, 0.1], [0.2, 0.3, 0.1, -0.4]].

Applying it gives [0.06, 0.34, -0.09, 0.00]. The resulting four-dimensional vector combines information from both heads.

Implementations can combine the projections for all heads into larger matrices. One matrix multiplication produces the projected vectors. Reshaping separates the final dimension into $(h, d_{\text{head}})$. The following code uses this approach.

Code that matches the hand numbers

import numpy as np
 
x = np.array([
    [0.9, 0.1, -0.3, 0.2],
    [0.2, 0.8, 0.4, -0.5],
    [-0.4, 0.1, 0.9, 0.3],
    [0.3, -0.6, 0.2, 0.8],
])
 
WQ = np.array([[0.1, 0.3], [0.7, -0.9], [0.3, -0.7], [-0.9, 0.7]])
WK = np.array([[0.9, -0.9], [0.6, -0.3], [-0.9, 0.7], [0.2, -0.4]])
WV = np.array([[0.5, -0.2], [0.1, 0.6], [-0.3, 0.4], [0.7, 0.1]])
 
 
def softmax(s):
    s = s - s.max(axis=-1, keepdims=True)
    e = np.exp(s)
    return e / e.sum(axis=-1, keepdims=True)
 
 
Q, K, V = x @ WQ, x @ WK, x @ WV
scores = Q @ K.T / np.sqrt(K.shape[-1])
 
mask = np.tril(np.ones_like(scores, dtype=bool))
scores = np.where(mask, scores, -np.inf)
 
weights = softmax(scores)
out = weights @ V
 
print(weights[3])   # [0.0298 0.1493 0.6738 0.1471]
print(out[1])       # [0.5605 -0.1183]

weights[3] contains the "cat" row from the causal example. out[1] contains the "saw" output. Subtracting the row maximum before exponentiation leaves softmax unchanged because the shared factor cancels. It also prevents exponential overflow for large scores.

The following PyTorch class combines projections and separates heads through reshaping:

import torch
import torch.nn as nn
import torch.nn.functional as F
 
 
class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)
 
    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=-1)
 
        # (B, T, C) -> (B, heads, T, d_head)
        q = q.view(B, T, self.n_heads, -1).transpose(1, 2)
        k = k.view(B, T, self.n_heads, -1).transpose(1, 2)
        v = v.view(B, T, self.n_heads, -1).transpose(1, 2)
 
        scores = q @ k.transpose(-2, -1) / (k.size(-1) ** 0.5)
        mask = torch.tril(torch.ones(T, T, dtype=torch.bool))
        scores = scores.masked_fill(~mask, float("-inf"))
        weights = F.softmax(scores, dim=-1)
 
        y = weights @ v                                 # (B, heads, T, d_head)
        y = y.transpose(1, 2).reshape(B, T, C)          # concat heads
        return self.out(y)

view separates the head dimension, and transpose reorders the axes. These operations can share the original tensor storage. After attention, transpose and reshape concatenate the heads. That final reshape can require a copy, depending on the tensor's strides.

Check that the output shape is (B, T, C). An incorrect head axis can produce the wrong layout.

The original example compares the hand-picked weights with NumPy to 1e-16. It also compares the result with F.scaled_dot_product_attention(q, k, v, is_causal=True).

PyTorch can select FlashAttention when the device and inputs support that backend. The fused calculation avoids storing the complete score matrix. Karpathy's nanoGPT includes explicit and fused attention paths.

What it costs as the sequence grows

Four tokens produce a $4 \times 4$ score matrix. A sequence of $n$ tokens produces an $n \times n$ matrix. At 128k tokens, this means about 16 billion scores per head.

Standard dense attention has quadratic cost in sequence length. Projections and feed-forward layers have linear sequence-length cost for fixed model dimensions. The score matrix therefore creates a major cost for long sequences.

Two implementation techniques address different parts of that cost.

FlashAttention calculates attention in tiles. It avoids writing the complete $n \times n$ score matrix to memory. It improves memory use and data movement without removing dense attention's quadratic arithmetic cost.

A key-value cache, or KV cache, stores keys and values during token-by-token generation. The new query compares with earlier keys and combines earlier values. A causal model can reuse those keys and values instead of recalculating them.

Each step adds a new query, key and value. It calculates one new attention row rather than the complete grid. For the standard multi-head layout here, cache storage is $2 \times n \times d_{\text{model}}$ floats per layer. Longer sequences therefore need more cache memory.

In practice, use F.scaled_dot_product_attention when it suits the model. The hand calculations help you inspect incorrect attention weights, mask errors and memory costs. Queries, keys, values, softmax and the output projection each have a distinct role.