Add attention-heads playground study note - #45
Merged
Conversation
A TensorFlow-Playground-style interactive note for transformers: train a tiny decoder-only model in the browser and watch nameable attention heads emerge. Complements the existing self-attention note, which covers the mechanism, by showing what the matrices become after training. The model is a real decoder-only transformer (learned token + positional embeddings, pre-LayerNorm blocks, multi-head causal attention, GELU MLP, Adam, ~18.6k params) with a hand-written backward pass. The backward pass was verified against central finite differences in double precision across seven configurations covering every parameter type, 1-2 layers, 1/2/4 heads, and causal/bidirectional masking; worst relative error ~3e-8. Four tasks (induction, reverse, sort, majority), live loss/accuracy chart, per-head attention heatmaps with automatic pattern labelling, per-head ablation, and an editable probe sequence. One training loop, gated by an IntersectionObserver so it pauses off-screen. Every quantitative claim in the prose is measured rather than asserted, and three of them changed after measurement: - The induction task went through three designs. A fixed-period repeat was solvable by a one-layer model via a constant offset; pinning the repeated block near the start let one-layer models score 94% off absolute positions 3-4. A randomly rotated second copy is what finally makes content matching the only strategy that generalises. - The head diagnostics gained a fixed-position detector, which competes with the induction score. Without it, positional heads were being labelled as induction heads. - "Removing positional embeddings breaks reverse" turned out to be false: the causal mask leaks position, and results are bimodal (42-94% over eight seeds). The note reports the spread and explains the leak. Also notes that the clean prev-token to induction circuit appears in only about a third of seeds, while all seeds reach >90% accuracy. Adds a matching index thumbnail sketch for the new slug. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HifqLNmnCY9nin8MDuYtJs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
A TensorFlow-Playground-style interactive note for transformers: train a
tiny decoder-only model in the browser and watch nameable attention heads
emerge. Complements the existing self-attention note, which covers the
mechanism, by showing what the matrices become after training.
The model is a real decoder-only transformer (learned token + positional
embeddings, pre-LayerNorm blocks, multi-head causal attention, GELU MLP,
Adam, ~18.6k params) with a hand-written backward pass. The backward pass
was verified against central finite differences in double precision across
seven configurations covering every parameter type, 1-2 layers, 1/2/4
heads, and causal/bidirectional masking; worst relative error ~3e-8.
Four tasks (induction, reverse, sort, majority), live loss/accuracy chart,
per-head attention heatmaps with automatic pattern labelling, per-head
ablation, and an editable probe sequence. One training loop, gated by an
IntersectionObserver so it pauses off-screen.
Every quantitative claim in the prose is measured rather than asserted,
and three of them changed after measurement:
solvable by a one-layer model via a constant offset; pinning the repeated
block near the start let one-layer models score 94% off absolute
positions 3-4. A randomly rotated second copy is what finally makes
content matching the only strategy that generalises.
with the induction score. Without it, positional heads were being
labelled as induction heads.
the causal mask leaks position, and results are bimodal (42-94% over
eight seeds). The note reports the spread and explains the leak.
Also notes that the clean prev-token to induction circuit appears in only
about a third of seeds, while all seeds reach >90% accuracy.
Adds a matching index thumbnail sketch for the new slug.
Co-Authored-By: Claude Opus 5 noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01HifqLNmnCY9nin8MDuYtJs