Tiny Attention Lab

Question and hypothesis

Can a minimal example make causal attention’s shape and masking rules observable? Hypothesis: after masking logits before softmax, each row assigns zero probability to future positions, remaining probabilities sum to one, and changing a future value cannot alter an earlier output.

Procedure

1. Create a fixed tensor with batch 1, sequence 4, one head, and head dimension 2.
2. Use explicit small projection matrices, or set (Q=K=V=X) for the first pass.
3. Compute (QK^T / \sqrt{2}); write the expected [1, 1, 4, 4] score shape before running.
4. Fill positions above the causal diagonal with negative infinity, then apply softmax.
5. Multiply probabilities by (V) and confirm output shape [1, 1, 4, 2].
6. Change only the final value vector. Assert that outputs for positions 0–2 are unchanged.
7. Repeat without scaling and compare probability concentration.

Use deterministic values and assertions rather than relying on visual inspection.

Expected observations

The first token can attend only to itself. Later rows can distribute probability over progressively longer prefixes. Removing the scale can make logits and probabilities more extreme, although the tiny example is not evidence about training stability.

Connections

Read Attention and the KV Cache before extending the lab. The lab is a checkpoint in Transformers.

Limitations

This exercise excludes learned projections, multiple heads, dropout, numerical edge cases, optimized kernels, gradients, and KV-cache layout. It validates semantics for chosen inputs, not model quality or production performance.

Self-test

1. Which axis does softmax normalize, and why?
2. Why must masking occur before softmax?
3. What assertion detects information leakage from a future value?
4. How would you extend the test to a batch with padding?