BibTeX Entry |
@techreport{AdlerEtAl:arxiv26,
author = {Adler, Micah and Byers, John and Crovella, Mark},
title = {Attention Mean Fields Predict Average Representation Dynamics and Reveal Context-Specific Computation},
year = {2026},
URL = {https://arxiv.org/abs/2609.16382},
doi = {10.48550/arXiv.2609.16382},
abstract = {A language model's representation geometry is not predetermined; it evolves as the model runs, with each layer reshaping it through attention that depends upon the surrounding context. A faithful account of that geometry must capture that dynamic process, and so cannot be based solely on model-independent statistics such as co-occurrence. Here we introduce a mean-field analysis of attention. The average attention from one token to another defines a kernel that carries representations layer to layer and can be iterated through the network to model how the geometry is transformed. We condition this average two ways. Conditioned on a whole corpus, the kernel predicts the average-case evolution of representation geometry. Conditioned instead on a single context, it predicts the expected geometry for that context. A head's departure from that prediction, its mean-field deviation, isolates the context-specific computation that the mean field misses. Under the corpus-conditional reading, the kernel yields an open-loop model: from the input embeddings and the frozen weights alone, we can iterate the kernel and consulting a measured deviation at any layer. With no free parameters, the rollout reproduces the trajectory of each token's mean representation across models from GPT-2 to Qwen-3-14B. We show that several alternatives, including one built from co-occurrence, are all poorer predictors. In early training the model and its corpus mean field are indistinguishable. Replace every attention head with its mean field, and across three scales of Pythia - 160m, 1.4B, and 12B - the substitution leaves the loss on real text unchanged. Around the onset of induction, the two diverge, and the gap widens as representations become contextualized and attention becomes correlated with the values it transports. Recent theory characterizes early-training weights as closed-form compositions of corpus statistics; we give a mean-field account of the computation, under which the young model's behavior reduces to its corpus mean field until in-context computation emerges. Under the context-conditional reading, deviation from the mean field is a task-agnostic measure of context-specific computation. The residual decomposes additively into unusual attention routing and contextualization of the transported values. Across controlled induction and few-shot settings, greater deviation tracks greater reliance on in-context information. Ranking heads by corpus-level deviation recovers documented in-context machinery while also revealing positional and content-tracking heads. On ICL tasks, task-conditioned residuals act as function vectors; we show that selecting function-vector heads by mean-field deviation achieves transfer accuracy comparable to causal-mediation selection, and only requires a single forward pass.},
note = {Also appeared in 2026 New England Mechanistic Interpretability Workshop (NEMI-2026)}
}