Attention from scratch

Attention from scratch

At 3 a.m., a monitoring model sees the event compressor overheating. The useful clue is not the previous event. It is a voltage spike from 7,000 events ago.

An RNN has been reading those events one at a time. By the time it reaches the alarm, it has compressed the whole history into one hidden vector. That vector has a fixed number of slots. The voltage spike, a maintenance note, and 6,998 less useful events are all competing for those slots.

An LSTM gives the model better gates and a better memory. It does not remove the basic squeeze. To use an old detail, the model must first preserve it in a state that every later time step keeps rewriting.

Attention changes the question. Instead of asking the current step to remember everything, it lets the current step look back through the sequence and retrieve what matches its current need.

That lookup is the important idea. The machinery is ordinary matrix multiplication, dot products, masking, and softmax. Because the lookup is differentiable, backpropagation can learn both what counts as a match and what information to retrieve.

The bottleneck attention was built to remove

A recurrent model updates its state as:

h_t = f(h_{t-1}, x_t)

h_t is a fixed-size summary after seeing the current item. In a sequence of 10,000 items, h_10000 must carry whatever matters from all 10,000.

This causes two problems. Information is compressed: a state with 512 numbers cannot preserve every detail equally well. And the path between distant events is long: information from x_1 must survive 9,999 updates before affecting the output at step 10,000. Gradients travel through the same chain during training. Gating helps, but the path remains long.

Attention creates a direct path. At step t, the model compares the current item with every earlier item and takes a weighted mixture of their information. This replaces a long recurrent path with one attention operation.

It is not free. Attention swaps a compression problem for a comparison problem: compare every item with every other item.

Attention is a learned, differentiable lookup

We will focus on self-attention, where every item can look at other items in the same sequence.

Stack n token or event vectors into X, with shape n × d_model. Each vector might represent a word, image patch, audio frame, or sensor event. Attention makes three learned projections:

• A query is what this position is looking for.
• A key is what this position offers as a possible match.
• A value is the payload returned when this position is selected.

The projections are:

Q = XW_Q

K = XW_K

V = XW_V

Usually, Q and K have width d_k, while V has width d_v. In self-attention, all three start from the same X; their learned projections give them different jobs.

Matching and carrying information are different jobs, so the projections need not be shared. A maintenance event might contain one feature that makes it relevant to a query about overheating and another feature containing the diagnostic information worth returning.

For every query, attention computes a compatibility score with every key:

score(q_i, k_j) = q_i · k_j

The scores become weights with softmax, and the output at position i is the weighted sum of values:

Attention(Q, K, V) = softmax(QKᵀ / sqrt(d_k)) V

QKᵀ contains all query-key dot products and has shape n × n. Softmax runs across each row, so each query’s weights sum to one.

This is a content-addressed lookup: the query supplies the address, the keys decide which addresses match, and the values supply the contents.

Why divide by the square root of d_k?

The scale factor prevents softmax saturation.

Assume query and key components are independent, have mean zero, and each have variance one. Their dot product is a sum of d_k products:

q · k = q_1k_1 + q_2k_2 + ... + q_dk_d

Each product has roughly unit variance, so the sum has variance about d_k and standard deviation about sqrt(d_k).

With d_k = 64, unscaled dot products have a standard deviation near 8. Softmax may then become almost one-hot: one key gets nearly all the weight. Its derivative is tiny near probabilities of zero and one, weakening the learning signal that adjusts matches.

Dividing by sqrt(d_k) brings the score standard deviation back toward one. The variance argument is an initialization approximation—learned weights and correlated features change the exact distribution—but it remains a reliable stabilizer.

A three-token calculation

Take three token positions with d_k = 2 and d_v = 2: voltage, spike, and alarm. Suppose their already-projected queries, keys, and values are:

Q =
[ 1  0 ]
[ 0  1 ]
[ 1  1 ]

K =
[ 1  0 ]
[ 0  1 ]
[ 1  1 ]

V =
[10  0 ]
[ 0 10 ]
[ 5  5 ]

For the first query, the raw dot products with the three keys are [1, 0, 1]. Since sqrt(d_k) = sqrt(2), the scaled scores are approximately [0.707, 0, 0.707].

Softmax gives approximately:

[0.401, 0.198, 0.401]

The output is the weighted sum of the value rows:

0.401[10, 0] + 0.198[0, 10] + 0.401[5, 5]

≈ [6.017, 3.983]

The same calculation happens independently for every query row. The output is a mixture of values, not a copied row. Q and K determine the matching weights; V determines what information those weights retrieve. Since softmax and the matrix operations are differentiable, the loss can increase useful matches, suppress distracting ones, or change the value projection.

Causal masking: do not let the future leak in

For language generation, position 2 must not read position 3. A causal mask blocks future positions. For three positions:

position 1: allowed, blocked, blocked
position 2: allowed, allowed, blocked
position 3: allowed, allowed, allowed

The implementation sets forbidden logits to negative infinity before softmax. Softmax then assigns them zero probability.

For the example above, the first row keeps only its first score, so its weights become [1, 0, 0]. The second row uses [0, 0.707, -inf], giving weights approximately [0.330, 0.670, 0].

Causal masking still permits parallel training: every query is processed at once, with illegal score entries removed. During autoregressive generation, implementations cache previous keys and values instead of recomputing them, but each new query still compares against the growing context.

The cost: every pair is a real pair

For n positions, QKᵀ has n² scores:

• Compute is roughly O(n²d_k).
• Materialized score or weight memory is O(n²).
• Projection work is roughly linear in n for fixed model width.

A single 4096 × 4096 attention matrix contains 16,777,216 numbers. In fp16, that matrix alone occupies 32 MiB. With 32 heads, it would be 1 GiB if all head matrices were materialized together, before gradients and other model state.

FlashAttention computes the same exact result with tiled, memory-aware operations, avoiding materialization of the full matrix in high-bandwidth memory. It can reduce memory dramatically, but it does not make the pairwise work linear.

At inference, a key-value cache grows linearly with the number of past positions. Each new token must still compare its query with all cached keys, so per-token work grows with context length.

A small PyTorch implementation

This single-head implementation exposes the data flow and supports causal masking. Production models commonly add multiple heads, normalization, residual paths, feed-forward layers, and positional information.

import math
import torch


def scaled_dot_product_attention(x, w_q, w_k, w_v, causal=False):
    # x: [batch, tokens, model_width]
    q = x @ w_q  # [batch, tokens, d_k]
    k = x @ w_k  # [batch, tokens, d_k]
    v = x @ w_v  # [batch, tokens, d_v]

    scores = q @ k.transpose(-2, -1)
    scores = scores / math.sqrt(q.size(-1))

    if causal:
        tokens = scores.size(-1)
        future = torch.triu(
            torch.ones(tokens, tokens, dtype=torch.bool, device=x.device),
            diagonal=1,
        )
        scores = scores.masked_fill(future, float("-inf"))

    weights = torch.softmax(scores, dim=-1)
    output = weights @ v
    return output, weights


torch.manual_seed(0)

batch_size, tokens, model_width = 2, 3, 8
d_k, d_v = 4, 4

x = torch.randn(batch_size, tokens, model_width)
w_q = torch.randn(model_width, d_k)
w_k = torch.randn(model_width, d_k)
w_v = torch.randn(model_width, d_v)

output, weights = scaled_dot_product_attention(
    x, w_q, w_k, w_v, causal=True
)

assert output.shape == (batch_size, tokens, d_v)
assert weights.shape == (batch_size, tokens, tokens)
assert torch.allclose(
    weights.sum(dim=-1),
    torch.ones(batch_size, tokens),
)

q @ k.transpose(-2, -1) compares every query token with every key token, producing [batch, query_tokens, key_tokens]. Then weights @ v mixes the value rows. The mask broadcasts across the batch dimension, and the causal diagonal ensures every query has at least one allowed key.

Choosing attention instead of its alternatives

Choose based on dependency structure and affordable sequence length:

Method	Best fit	What it buys you	What you pay
Full self-attention	Short or medium sequences with unpredictable long-range dependencies	Direct retrieval between any positions	Pairwise compute and memory grow as n²
RNN or LSTM	Streaming input or naturally sequential control	Constant-size inference state	Sequential computation and fixed-state compression
1D convolution	Local patterns and high-throughput signals	Parallel computation near O(nk) for kernel width k	Distant interactions require depth, dilation, or pooling
Local or sparse attention	Long sequences with structured dependencies	Content-based lookup with fewer pairs	May miss dependencies outside the chosen pattern

Attention is not automatically right for streaming, enormous contexts, or reliably local signals.

What to remember

• RNNs compress history into a fixed state; attention retrieves directly from the sequence.
• Queries ask, keys match, and values carry the returned information.
• Dividing by sqrt(d_k) keeps softmax in a trainable range.
• Causal masks block future logits before softmax.
• Full attention’s pairwise cost is quadratic; memory-efficient kernels improve implementation, not that fundamental cost.

Quick check

0/3

Q1

Which query vector is producedThe attention weight assigned to the middle valueThe sequence lengthThe key projection width

Check answer

Q2

To make every value vector have unit lengthTo reduce the number of tokens being comparedBecause the dot product variance grows approximately with d_k, which can saturate softmaxBecause causal masks require normalized keys

Check answer

Q3

Full self-attention, because attention is always more expressiveAn RNN or local convolution, because a fixed or local dependency pattern avoids comparing every pairCausal masking alone, because it changes quadratic memory to linear memoryRemove the value projections, because keys contain all needed information

Check answer