Lion Cipher Lab

Build scaled dot-product attention in NumPy, and prove the causal mask stops future leakage

• AI Security
• Advanced
• ~90 min
• Free · no account required

Why it matters

Attention is the part of a transformer that decides which earlier tokens each token reads from, and almost every language model a security team touches is built on it. Most people only ever call it through a library, so its two most common bugs stay invisible. A causal mask written slightly wrong lets the model read the future while it trains. Nothing crashes. The training numbers just look better than they should, and the model gets worse when it is used for real. A padding mask written slightly wrong makes a log line's result depend on how long its neighbour in the batch was. In this lab…

What you'll learn

• writing softmax, scaled dot-product attention, multi-head attention, a causal mask and a padding mask in NumPy straight from the equations in Attention Is All You Need, and matching PyTorch's scaled_dot_product_attention to 8.88e-16 in float64
• showing why softmax must subtract the row maximum: the textbook version returned [nan nan nan] for scores of 1000, 1001, 1002, and in float16 it overflows for any score above 11.09
• measuring what the 1/sqrt(d_k) scale is for: at d_k = 1024 the unscaled softmax put 0.952 of its weight on one key, 71.5% of rows were saturated and the median gradient norm fell to 1.77e-04, while the scaled version stayed at 1.79e-01 for every width
• proving a causal mask with a perturbation test: change only the last token and the earlier outputs move by exactly 0.0 with the mask, and by 1.701 without it
• catching four real mask bugs by number: masking after the softmax (leak 0.368), multiplying scores by the mask (0.427), a -1e9 mask cast to float16 (no leak, but 0 of 32 weight rows summed to 1) and left padding with a causal mask (21 of 39 real rows NaN one layer later)
• showing that future leakage makes a next-token model look better than it is: an unmasked layer scored 87.0% train and 81.7% on held-out lines, then 67.8% when it could only see the prefix, below the honest causal layer's 68.4%
• demonstrating with a linear program that attention weights are not explanations: a 40-token query's top token went from weight 0.131 to 0.000 with the output unchanged to 2.22e-15, and saying exactly when that trick is impossible
• measuring attention's O(n^2) cost from n = 128 to 4,096 (peak memory 0.2 MB to 192 MB) and cutting peak memory to 2.1 MB with a tiled online softmax, the idea behind FlashAttention

Steps in this lab

1. The words first: tokens, embeddings, query, key, value, softmax, heads, masks - and what is random here
2. Open a terminal, make a folder, install NumPy and SciPy
3. Write attention.py from the paper's equations, and logdata.py for the synthetic log lines
4. Softmax overflows unless you subtract the maximum - then four sanity checks before trusting the code
5. OPTIONAL: check your attention against PyTorch's scaled_dot_product_attention
6. Why divide by sqrt(d_k): measure saturation and the vanishing gradient with and without it
7. Prove the causal mask: change a future token and measure the past - then catch four mask bugs
8. What future leakage does to a model's numbers: it looks better in training and is worse in use
9. Padding masks on a batch of log lines - and the two places the mask is not enough
10. Honest negative 1: attention weights are not an explanation - two different patterns, one output
11. Honest negative 2: the n x n score matrix - time it from 128 to 4,096 tokens, then tile it
12. Do it better: what a pro checks in any attention code, when to write your own, and clean up

Open the full interactive lab →