Causal Attention and Masking in Transformers

How it works

QKᵀ produces token-pair scores. Scores are divided by √128, future positions are masked, and softmax produces weights to mix Values.

1. Multiply Q × K

   Each cell compares one Query row with one Key row. Click to inspect all multiplications.

2. Sum the dot product

   Four products are summed into one score. Scores are not yet probabilities.

3. Scale

   Dividing by √d_head controls score scale. Here, √4 = 2.

4. Causal mask

   Future positions are unavailable. A score of −∞ becomes exactly zero weight after softmax.

5. Row-wise softmax

   Normalize across Key columns rather than the entire matrix. Every row sums to 1.

6. Mix Values

   Attention weights and Value contributions are different. Click an output to inspect its weighted sum.

7. Concatenate heads

   Two width-4 outputs are concatenated back to width 8.

8. Output projection

   A learned output projection mixes features across heads before residual addition.

What can you observe in the lab?

Lines point only to current and preceding positions. Thickness illustrates weights; other heads may have different patterns.

The Scanner opens notes. The Token tracer selects token positions. The Head lens compares two demo heads. The Time controller pauses signals. The Calibrator changes experiment conditions before you predict and test the outcome.

What to remember

Visuals use illustrative flows and fixtures. GPT-3 175B has 96 blocks, 96 heads per block, width 12,288, and a 2,048-token context. Its paper alternates dense and locally banded sparse attention; this lab examines causal dense attention. The numerical workbench uses a toy model with two blocks, two heads, width 8, and a vocabulary of 16. GPT-3 weights are not loaded, and presets do not generate answers from a trained language model.

Primary sources

• GPT-3 · Language Models are Few-Shot Learners

  175B scale: 96 blocks, width 12,288, 96 heads, context 2,048; alternating dense and locally banded sparse attention.

• GPT-2 · reference implementation

  Pre-norm, learned positions, GELU, residuals, and tied embeddings.

• Attention Is All You Need

  Foundations of scaled dot-product and multi-head attention.

• Word representations · Mikolov et al.

  Background on static word embeddings and geometry.

• KV cache · Hugging Face

  Prefill, decoding, and K/V reuse.