Trophic cascade Part B: TinyStories models trained with and without dropout
Eighteen small GPT-2-architecture language models trained from scratch on TinyStories to test whether
dropout causes self-repair, the restoration by later heads of an effect removed by ablating a head.
There are three training regimes and six seeds of each. Seed k of every regime starts from the same
initialisation and sees the same training sequences in the same order, so models are compared in pairs
by seed.
| Regime | Embedding, attention, residual dropout | Head dropout |
|---|---|---|
none |
0 | 0 |
standard |
0.1 | 0 |
drophead |
0 | 0.1 |
Head dropout zeroes each head's output independently, per sequence, with probability 0.1 during
training, and scales the remaining heads by 1/0.9. It is not part of GPT2LMHeadModel, so the
drophead configs load as plain models with all dropout at 0. The standard configs keep their
dropout of 0.1, which applies only in training mode.
Architecture: GPT2LMHeadModel, 6 layers, 8 heads of dimension 48 (d_model 384), MLP width 1,536,
context 256, tied embeddings, about 30 million parameters, GPT-2 tokenizer. Training: 200 million
tokens, AdamW, peak learning rate 1e-3 with 2% linear warm-up and cosine decay to 10%, weight decay 0.1,
batch 128 sequences, on a Colab A100.
Loading
Each model is in its own folder:
from transformers import AutoTokenizer, GPT2LMHeadModel
tok = AutoTokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("kvenanzi/trophic-cascade-partb", subfolder="drophead-s0").eval()
Models
Validation loss is on the TinyStories validation split. The pooled self-repair fraction is the share
of a head's removed direct effect on the correct-token logit that the rest of the model restores,
summed over all 48 heads (resample ablation, 256 held-out sequences, the 2% of positions where each
head's direct effect is largest), measured in evaluation mode. The run IDs refer to the original
Weights & Biases records, which are no longer public; the same records are in runs.json.
| Folder | Regime | Seed | Run | Validation loss | Pooled self-repair |
|---|---|---|---|---|---|
none-s0 |
none | 0 | ryqgyaml |
1.614 | 0.360 |
none-s1 |
none | 1 | b08p7lt5 |
1.610 | 0.353 |
none-s2 |
none | 2 | r9l9uae6 |
1.612 | 0.382 |
none-s3 |
none | 3 | 3tfis5fz |
1.610 | 0.372 |
none-s4 |
none | 4 | evackiyh |
1.612 | 0.298 |
none-s5 |
none | 5 | 7j3yqh5d |
1.613 | 0.388 |
standard-s0 |
standard | 0 | zsyblp32 |
1.656 | 0.623 |
standard-s1 |
standard | 1 | z3jlfpko |
1.659 | 0.609 |
standard-s2 |
standard | 2 | z7oqo7od |
1.659 | 0.604 |
standard-s3 |
standard | 3 | 7b5nmdi5 |
1.661 | 0.630 |
standard-s4 |
standard | 4 | aoja09w6 |
1.660 | 0.592 |
standard-s5 |
standard | 5 | wgxi6q2o |
1.661 | 0.632 |
drophead-s0 |
drophead | 0 | jlnsa7xj |
1.626 | 0.683 |
drophead-s1 |
drophead | 1 | 5o6hsruz |
1.626 | 0.653 |
drophead-s2 |
drophead | 2 | 6j67iu7u |
1.625 | 0.671 |
drophead-s3 |
drophead | 3 | k6zpon6c |
1.626 | 0.691 |
drophead-s4 |
drophead | 4 | dpxouw3m |
1.623 | 0.642 |
drophead-s5 |
drophead | 5 | d59he38q |
1.627 | 0.679 |
The means over seeds are 0.359 (none), 0.615 (standard), and 0.670 (drophead) for self-repair,
and 1.612, 1.659, and 1.626 for validation loss. These are small models trained on one corpus at one
dropout rate; the write-up gives the intervals and the limitations.
Write-up: https://withinnoise.dev/blog/posts/transformer-trophic-cascades/ (section 5.8) Code: https://github.com/kvenanzi/trophic-cascade