Research map / Architectures and efficiency / Attention and transformers
Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Finds sparse continuous-sparsification subnetworks (~10M active non-embedding params, <0.2% in Pythia 1B) that reproduce bigram next-token predictions in GPT-2 and Pythia models, concentrated in the first MLP layer and overlapping with optimal pruning subnetworks.
Problem
Identifying a minimal, interpretable circuit that is sufficient and necessary for naive bigram next-token prediction in Transformer language models, as a foundation for building up more complex circuits.
Approach
Uses continuous sparsification to learn binary masks over frozen pretrained weights so that the masked model matches the bigram distribution P(w_i|w_{i-1}), then analyzes where these subnetworks sit, how they transform the residual stream, and whether ablating them harms language modeling.
Key results
- Bigram subnetworks reach bigram surprisal correlation r>0.95 with roughly 10M active non-embedding parameters across model sizes; Pythia 1B subnetwork with 0.17% of non-embedding parameters reaches r=0.959.
- Ablating the Pythia 1B bigram subnetwork degrades language modeling similarly to ablating an optimal subnetwork of comparable size, while random subnetworks of the same size have almost no effect.
- The Pythia 1B bigram subnetwork and optimal subnetwork overlap 15.3x more than chance, with 38.0% of the bigram subnetwork contained in the optimal subnetwork.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use the identified bigram subnetwork instead of the full model | Pythia and GPT-2 · OSCAR bigram surprisal | Pearson correlation: 0.690 -> 0.964 (Pythia 160M); 0.632 -> 0.987 (Pythia 1B); 0.680 -> 0.987 (GPT-2 small*) | Helped |
| Ablate the bigram subnetwork (ablation) | Fully trained language models · OSCAR held-out language modeling | evaluation loss | Hurt |
| Ablate a same-size random subnetwork (ablation) | Fully trained language models · OSCAR held-out language modeling | evaluation loss | No clear effect |
| Patch trained token embeddings and unembeddings into a randomly initialized model (ablation) | Pythia 1B subnetwork · OSCAR bigram surprisal | Pearson correlation: 0.896 -> 0.986 | Helped |
| Keep only embeddings and unembeddings with an empty subnetwork (ablation) | Pythia 1B · OSCAR bigram surprisal | Pearson correlation: 0.987 -> 0.628 | Hurt |
| Add a learned linear transform between embeddings and unembeddings (ablation) | Pythia 1B · OSCAR bigram surprisal | Pearson correlation: 0.628 -> 0.936 | Helped |
| Use 500K active parameters at checkpoint 143K rather than checkpoint 4K | Pythia 1B bigram subnetwork · OSCAR bigram surprisal | Pearson correlation: 0.949 -> 0.919 | Hurt |
Surprising findings
- Bigram subnetworks can be found in randomly initialized models if trained token embeddings and unembeddings are patched in (Pythia 1B r=0.986).
- Embeddings with an empty subnetwork do not reproduce bigram predictions, but a learned linear transform from embeddings can reach r>0.90 for Pythia 1B and GPT-2 small, questioning tuned lens interpretations.
Limitations
- Bigram subnetworks are not shown to be complete or unique; other redundant parameters could also recreate bigram predictions.
- Analyses focus on GPT-2 and Pythia models up to 1B parameters only.
- Unclear whether the bigram subnetwork corresponds to genuine intermediate bigram predictions inside full models.
Experimental setup
Datasets: OSCAR
Benchmarks: OSCAR held-out subset (1.2M tokens) language modeling loss
Metrics: Pearson correlation of bigram surprisal language modeling loss median rotation of activations covariance cross-layer similarity parameter overlap ratio
Baselines: random subnetwork optimal pruning subnetwork empty subnetwork embeddings with linear transform
Builds on: continuous sparsification Pythia GPT-2 tuned lens
Closest papers by meaning
- Transformer Block Coupling and its Correlation with Generalization in LLMs (ICLR 2025 · similarity 0.91)
- Distilling Sequential Computation in Transformer Language Models (NeurIPS 2026 · similarity 0.90)
- Transformers need glasses! Information over-squashing in language tasks (NeurIPS 2024 · similarity 0.90)
- Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models (NeurIPS 2024 · similarity 0.89)
- Understanding Transformers via N-Gram Statistics (NeurIPS 2024 · similarity 0.89)
- Separations in the Representational Capabilities of Transformers and Recurrent Architectures (NeurIPS 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.