Research map / Architectures and efficiency / Attention and transformers
Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers
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.
Develops a finite-width statistical-mechanics (BPKR) theory of Bayesian learning in a linear-value multi-head multi-layer self-attention network, showing predictor statistics are a U-weighted sum of path-path kernels and validating on synthetic HMM and one-shot Omniglot/FashionMNIST tasks with head pruning.
Problem
Transformers lack theoretical characterization of how their multiple attention paths (sequences of heads across layers) interact and affect generalization; infinite-width GP theory loses this interplay.
Approach
Exact finite-width thermodynamic-limit equations (N,P→∞, P/N fixed) for a deep linear-in-values multi-head attention network show the mean predictor is a kernel regression with a kernel summing path-path kernels weighted by an order parameter U that aligns the kernel with task labels; this task-relevant kernel combination is absent in the GP limit.
Key results
- In the HMC task with N=10 the finite-width network clearly outperforms the GP limit, with the order parameter suppressing adversarial paths and anti-correlating good and denoising paths.
- On one-shot Omniglot classification the finite-width regime beats the GP limit, and the gain is preserved OOD on FashionMNIST.
- Pruning the two lowest-scoring heads (25% of parameters) from a gradient-descent-trained network causes only marginal in-distribution accuracy loss, and OOD accuracy even improves.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use finite-width learning with task-relevant attention-path kernel combination instead of the GP limit | Transformer-like attention network · HMC classification | classification accuracy | Helped |
| combine the good and denoising attention paths using cross-path kernels | Transformer-like attention network · HMC classification | classification accuracy | Helped |
| increase perpendicular noise to σ⊥=5 while using the good and denoising paths | Transformer-like attention network · HMC classification with stronger perpendicular noise | classification accuracy | Helped |
| use finite-width learning instead of the GP limit | Transformer-like attention network · One-shot Omniglot classification | classification accuracy | Helped |
| use finite-width learning instead of the GP limit | Transformer-like attention network · One-shot FashionMNIST classification (OOD) | classification accuracy | Helped |
| prune the two lowest-scoring heads (ablation) | Gradient-descent-trained one-shot image classifier · One-shot Omniglot classification (ID) | classification accuracy: 25% | Hurt |
| prune the two lowest-scoring heads (ablation) | Gradient-descent-trained one-shot image classifier · One-shot FashionMNIST classification (OOD) | classification accuracy: 25% | Helped |
| use only the designed good attention path (ablation) | Transformer-like attention network · HMC classification | classification accuracy: A∼94% | Helped |
| use each of the three other paths in isolation rather than the good path (ablation) | Transformer-like attention network · HMC classification | classification accuracy: A∼94% to A∼50% | Hurt |
| include adversarial paths rather than suppressing their associated kernels (ablation) | Transformer-like attention network in the GP limit · HMC classification | classification accuracy | Hurt |
| combine the good and denoising paths instead of using the good path alone (ablation) | Transformer-like attention network in the GP limit · HMC classification | classification accuracy | Helped |
Surprising findings
- Pruning low-scoring heads from the gradient-descent-trained network slightly improves OOD accuracy on FashionMNIST.
- Theory agrees well with sampled statistics even at N=10, below the number of independent order-parameter entries (136).
Limitations
- Attention depends only on bare input and the model is linear in value weights (no MLP blocks or skip connections).
- Query and key weights are fixed and not learned by the theory; learning them is an open challenge.
- Theory characterizes the inductive bias of the architecture under Bayesian posterior, not the bias of gradient descent training.
Experimental setup
Datasets: Omniglot MNIST FashionMNIST synthetic hidden Markov chain data
Benchmarks: hidden Markov chain classification Omniglot one-shot classification FashionMNIST OOD one-shot classification MNIST one-shot classification
Metrics: classification accuracy order parameter structure kernel-label PC overlap head pruning accuracy
Baselines: GP limit (infinite-width) predictor gradient descent-trained network
Builds on: back-propagating kernel renormalization Gaussian process limit of attention networks globally gated deep linear networks
Closest papers by meaning
- Hierarchical Self-Attention: Generalizing Neural Attention Mechanics to Multi-Scale Problems (NeurIPS 2025 · similarity 0.89)
- Bayes optimal learning of attention-indexed models (NeurIPS 2025 · similarity 0.88)
- Phase Transitions in Attention: A Bayesian Theory of Copy Head Emergence (NeurIPS 2026 · similarity 0.88)
- Mind the Gap: a Spectral Analysis of Rank Collapse and Signal Propagation in Attention Layers (ICML 2025 · similarity 0.87)
- High-Dimensional Learning Dynamics of Attention-Indexed Models (NeurIPS 2026 · similarity 0.87)
- On the Role of Attention Masks and LayerNorm in Transformers (NeurIPS 2024 · 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.