Research map / Architectures and efficiency / Attention and transformers
Differentiation and Specialization of Attention Heads via the Refined Local Learning Coefficient
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.
Introduces refined local learning coefficients (weight-, data-, and weight-and-data-refined) to track how attention heads in a two-layer attention-only transformer differentiate and specialize over training, and uses them to find a multigram circuit.
Problem
How attention heads in transformers differentiate into functional roles and specialize to data patterns during training, and how to measure internal structure developmentally using singular learning theory.
Approach
Refine the local learning coefficient to a component (weight subset) and/or an arbitrary data distribution, yielding weight-refined, data-refined, and combined rLLCs that track per-head complexity and specialization over training checkpoints.
Key results
- Weight-refined LLC curves cluster attention heads by functional type, with heads of the same type sharing curve shape and critical points.
- Data-refined LLC for GitHub code separates the two induction heads (1.6 vs 1.7), pulling 1.6 to a significantly higher value, consistent with its specialization to code patterns.
- Identifies a multigram circuit in which layer-0 multigram heads transition to supporting layer-1 multigram heads (e.g., 1.5 for nested Dyck patterns), supported by K-composition scores and path patching.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use weight-refined LLC curves to distinguish attention-head types across training | Two-layer attention-only transformer · Pretraining on a subset of the Pile | Alignment of LLC curves with functional head types | Helped |
| Refine LLC measurements using the GitHub code distribution | Induction heads 1.6 and 1.7 · GitHub code | Induction-pattern specialization reflected by data-refined LLC | Helped |
| Use a one-layer transformer as the generating process for data-refined LLCs | Attention heads in the two-layer transformer · Simple multigrams predicted by a one-layer model | Detection of changes in head information about simple multigrams | Helped |
| Mean-ablate attention head 1.5 (ablation) | Two-layer attention-only transformer · Synthetic sentences with parentheses | Loss | Hurt |
| Mean-ablate layer 0 multigram heads and patch their activations into the input of head 1.5 (ablation) | Two-layer attention-only transformer · Synthetic sentences with nested brackets | Loss | Hurt |
| Mean-ablate layer 0 multigram heads and patch their activations into the input of head 1.5 (ablation) | Two-layer attention-only transformer · Synthetic sentences with unnested parentheses | Performance | Mixed |
Surprising findings
- Layer-0 multigram heads enter the cluster with largest relative decrease in drLLC, despite ablations showing they are primarily involved in multigram prediction throughout training.
Limitations
- LLC estimates presume a local minimum of the population loss, which is unlikely to hold during training checkpoints.
- Absolute accuracy of estimated LLC values is unverified; only ordinality is argued to be reliable.
- Head classification relies on heuristic pattern matching plus manual inspection, and only a two-layer attention-only model is studied.
Experimental setup
Datasets: The Pile GitHub subset of the Pile synthetic parentheses dataset
Metrics: refined local learning coefficient mean-ablation loss increase in-context learning score K-composition score clustering agreement
Baselines: Hessian trace Fisher information matrix trace Hessian rank zero ablation resample ablation
Builds on: local learning coefficient singular learning theory A Mathematical Framework for Transformer Circuits In-context learning and induction heads
Closest papers by meaning
- Causal Head Gating: A Framework for Interpreting Roles of Attention Heads in Transformers (NeurIPS 2025 · similarity 0.88)
- From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics (NeurIPS 2025 · similarity 0.87)
- Circuit Component Reuse Across Tasks in Transformer Language Models (ICLR 2024 · similarity 0.87)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.87)
- Beyond Components: Singular Vector-Based Interpretability of Transformer Circuits (NeurIPS 2025 · similarity 0.87)
- Mechanism and Emergence of Stacked Attention Heads in Multi-Layer Transformers (ICLR 2025 · similarity 0.87)
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.