TinyStories Suppresses Induction Heads in Language Models
On the TinyStories dataset with 473 million tokens, Previous Token Heads form reliably in the early layers of GPT-2 Small (124 million parameters), reaching a score of 0.20 for L0H3. Induction Heads emerge weakly—maxing out at 0.05 for L6H4 with a peak at 16,000 steps followed by a decline. A Sparse Autoencoder (SAE) on layer 6 confirms no pure Induction features.
The methodology isolates the impact of data complexity: the model was trained from scratch in PyTorch, with the SAE expanding 768 dimensions to 8192 using an L1 coefficient of 3.9 over 25 epochs. Metrics compute_induction_score and compute_previous_token_score were implemented from scratch.
Formation of Previous Token Heads
Previous Token Heads concentrate in layers 0–4. Head L0H3 shows a steady plateau at 0.20 from 14,000 steps onward. Dynamics stabilize by the end of training (20,000 steps, ~5.5 epochs).
Distribution across layers and heads:
- L0H3: 0.20 (maximum)
- Layers 0–4: dominant values
- Later layers: minimal contribution
This points to a fundamental mechanism essential for next-token prediction, regardless of data complexity.
Analysis of Induction Heads
Induction Heads don't fully develop: peak of 0.0537 at L6H4 around 15–16,000 steps, then a drop-off. No head exceeds 0.05, aligning with the low repetition statistics in TinyStories.
Key observations:
- Peak in a mid-layer (L6).
- No stable plateau—indicating instability.
- Comparison to Previous Token: 4x lower score.
Hypothesis: The dataset's simplicity makes Induction mechanisms redundant for minimizing loss.
Sparse Autoencoder Results
SAE trained on activations from layer 6 (L6H4). Active features per token: 20–35. Dead features: ~3% (228 out of 8192) thanks to L1 regularization.
Two dominant features:
- F#8107: syntactically predictable context, activating on first tokens (difference from repeats: 0.04–0.12).
- F#635: general contextual pattern at predictable positions.
No pure Induction features detected—confirming the weak score.
Dataset Impact on Mechanisms
TinyStories doesn't encourage In-Context Learning: the model relies on Previous Token Heads for low loss. On complex datasets (e.g., OpenWebText), Induction mechanisms should grow due to frequent repeats.
Next test plans:
- Models of varying sizes on TPU.
- TinyStories vs. OpenWebText comparison.
- Measuring dependence on data statistics.
Key Takeaways
- Previous Token Heads form universally (score ~0.20 in early layers).
- Induction Heads are weak on simple data (max 0.05, unstable).
- SAE is effective: 20–35 active features, 3% dead features.
- No pure Induction features—F#8107/635 tied to syntax/context.
- Hypothesis: Dataset complexity drives ICL development.
— Editorial Team
No comments yet.