Attention Without a Library

Goal

You implement attention without a library. You use neither numpy nor torch — only the standard library (math).

It is slow. That is why you use small dimensions (d=8–64, n=4–8). The goal is not speed but being able to see inside.

What to build

Define the following in /root/work/tf/model.py.

Function	Contract
softmax(xs)	Sums to 1. Does not blow up on large inputs
attention(Q, K, V, mask=None)	Returns (out, weights). Divides by √d_k
causal_mask(n)	If mask[i][j] is true, i can see j
multi_head(Q, K, V, h, mask=None)	Splits the last dimension into h equal parts, runs attention on each, and concatenates
pos_encoding(n, d)	A vector in the [-1,1] range that differs at every position
layer_norm(row, eps=1e-5)	Mean 0, variance 1
block(X, h, mask=None)	X + multi_head(LN(X), ...)

All matrices are Python lists of lists ([[float]]). Positions × dimensions.

Check

cd /root/work/tf
python3 -c "import model; print(model.softmax([1,2,3]))"

Steps

1. softmax — no overflow
2. attention + measure √d → 02-scale.txt
3. causal_mask
4. multi_head (with h=1 it is the same as attention)
5. pos_encoding + permutation experiment → 05-perm.txt
6. layer_norm
7. block (pre-LN + residual)
8. Summary → 08-notes.md

Notes

The grader compares against a reference implementation within 1e-6. Attention is an operation that gives the same numbers however it is implemented, so if your values are off, it is usually because the scaling or the order of the mask is wrong.

A softmax that does not blow up

Create softmax(xs) in /root/work/tf/model.py. The result must sum to 1, and it must work without inf or nan even on large inputs such as [1000, 1001, 1002].

Run mkdir -p /root/work/tf. Use only the standard library (import math). Subtract the maximum before exp — exp(x - max) is mathematically the same value, but it does not overflow. Without this one line, it blows up the moment a large value comes in.

The three lines of attention, and √d

Create attention(Q, K, V, mask=None) and have it return (출력, 가중치) (the placeholders are the output and the weights). Always divide the scores by √d_k. Then measure how much the maximum weight differs with and without the division, and record it in 02-scale.txt.

Q, K and V are [[float]] (positions × dimensions). score[i][j] = dot(Q[i], K[j]) / sqrt(len(K[j])), w[i] = softmax(score[i]), out[i] = Σ w[i][j] * V[j]. Do the measurement with random vectors of d=64 — without the division, the maximum weight sticks to 1.0. That is the state where it has become "pick one" instead of a weighted average.

Hide the future

Create causal_mask(n). If mask[i][j] is true, it means the i-th query can see j. When attention receives this mask, the weight at masked positions must be exactly 0.

If i >= j, it can see. When you implement it, set the score to -inf (or a very small value) before the softmax. If you multiply by 0 after the softmax, the remaining weights no longer sum to 1 — that is the common bug.

Split and join again

Create multi_head(Q, K, V, h, mask=None). Split the last dimension into h equal parts, run attention on each, and concatenate them again. The output shape must be the same as the input.

The dimension of each head is d // h, and the scaling must also be based on that smaller dimension. With h=1 you must get exactly the same result as attention — that is the best self-check.

Prove that attention does not know order

Create pos_encoding(n, d). Then confirm that when you shuffle the input order, without positional encoding the output is shuffled right along with it, and with it added, that is not the case, and record it in 05-perm.txt.

Sine/cosine or any other method is fine. The values must be in the range [-1,1] and differ at every position. How to prove it: compare whether the result of running attention on X' (a shuffled X) equals the original result shuffled in the same order. This is permutation equivariance.

LayerNorm along the feature axis

Create layer_norm(row, eps=1e-5). It sets one token vector to mean 0 and variance 1.

Compute the mean and variance within that vector, not across the batch. That is why it is not affected by batch size or sequence length — the decisive difference from BatchNorm. Use eps so that you do not divide by zero even if a constant vector comes in.

Assemble the block

Create block(X, h, mask=None). The structure is pre-LN + residual: X + multi_head(LN(X), ...). The output shape must be the same as the input.

Do rows = [layer_norm(r) for r in X], run multi-head on them, and then add the original X. Without the residual, the gradient cannot reach the input when you stack many layers. Models today are pre-LN rather than post-LN because they train without warmup.

Sum up the three things in numbers

Write at least three lines in 08-notes.md: the two maximum weights you measured in step 2, the property step 5 showed, and how many times larger the score matrix becomes when you double the sequence length.

The text must contain 스케일링, 순서 and 제곱 (the Korean words for "scaling", "order" and "square"). The third is what you run into most often in practice — raising the context from 4k to 8k makes it 4 times as large.