Attention Mean Fields Predict Average Representation Dynamics

Paper: 2609.16382 Authors: Micah Adler, John W. Byers, Mark Crovella Categories: cs.LG, cs.CL

The Gap

People who study transformer representations often reach for model-independent statistics such as co-occurrence. The paper’s opening claim is that this cannot be enough: a language model’s representation geometry is not predetermined; it evolves as the model runs. A snapshot statistic misses the process.

What is missing is a faithful account of that geometry that stays cheap enough to iterate. Probing every layer on every context is expensive and hard to separate into “what the weights already do on average” versus “what this particular context is doing.”

   SNAPSHOT STATS VS A DYNAMIC ACCOUNT

   co-occurrence / static geometry stats
     fixed before the forward pass
     cannot see how geometry is TRANSFORMED layer to layer
         |
         v
   representation geometry ACTUALLY EVOLVES
     not predetermined
     changes as the model runs
         |
         v
   need: iterate an average transformation
         condition on corpus OR on one context
         residual = context-specific computation

The Increment

One sentence: Before this paper, average attention had not been turned into an iterated kernel that predicts representation dynamics; after it, a mean-field model predicts average-case geometry accurately, and deviation from that mean field is a task-agnostic measure of context-specific computation.

Core Mechanism

The construction is a mean-field analysis of attention. The average attention from one token to another defines a kernel that carries representations from layer to layer. That kernel can be iterated through the network to model how geometry is transformed.

Two conditionings give two readings. 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 is named mean-field deviation, and it isolates the context-specific computation that the mean field misses.

Under the corpus-conditional reading, the model is open-loop: from input embeddings and frozen weights alone, iterate the kernel and the model’s own MLPs, never consulting a measured deviation at any layer. The resulting prediction is highly accurate. Two training-phase facts make the story sharp. In early training the model and its corpus mean field are indistinguishable. Replace every attention head with its mean field, and 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.

Under the context-conditional reading, deviation is a measure of context-specific computation, and the residual decomposes additively into unusual attention routing and contextualization of the transported values. Greater deviation tracks controlled induction and few-shot settings.

   MEAN-FIELD KERNEL AS A CARRIER

   input embeddings + frozen weights
     |
     v
   average attention kernel  -->  carry reps layer to layer
     |                            (iterate with MLPs)
     +-- corpus-conditional  --> average-case geometry
     |
     +-- context-conditional --> expected geometry for THIS context
     |
     v
   measured head attention
     |
     v
   mean-field deviation = residual
     -> unusual routing
     -> contextualized values
     -> context-specific computation

Think of it as crowd flow through a station versus one passenger’s path. The mean field is the crowd: given the gates and the timetable, you can predict how dense each corridor will be on average, and most early training looks exactly like that crowd. Individual passengers who cut against the flow — induction heads, few-shot copying — are the deviation. The paper’s engineering gift is that you can forecast the crowd without watching each passenger, then spend analysis budget only on those who leave it.

Key Concepts

• Attention as an iterated kernel: average attention is not just a diagnostic; it is a linear carrier you can compose through depth.
• Mean-field deviation: measured attention minus the predicted average — a residual that names context-specific computation.
• Induction as the divergence point: early training is mean-field-like; the interesting specialization appears when model and mean field separate.

Framework Shift

Before (static interpretability):     After (mean-field dynamics):
  co-occurrence / probing snapshots   iterated average-attention kernel
  geometry treated as given           geometry treated as evolving
  every head analyzed separately      average vs deviation decomposition
  cost scales with contexts           corpus kernel is open-loop cheap

From reading representations as fixed objects, to predicting how they move on average and measuring only what departs, the core shift is dynamics over snapshots.

Expert Assessment

Problem choice: Excellent. Interpreting attention head-by-head does not scale; a theory of the average transformation that can be iterated is the right missing tool.

Method maturity: The open-loop prediction without measured deviations is the strong result — it is a statement about weights, not about a fit to a probe. An unexamined assumption to name: that “average attention” is the right coarse variable. If important structure lives in the variance of routing at fixed average, the mean field will look predictive while missing the computation that matters.

Experimental integrity: The strongest controls are the substitution test (mean-field heads leave early loss unchanged) and the timing of divergence around induction. Those are falsifiable and reported as negative-turned-positive facts. The limitation is scope: controlled induction and few-shot are clean instruments; natural long-context editing may need a richer residual.

Writing quality: The two conditionings (corpus vs context) are the paper’s spine and are stated clearly. A denser worked example of one head’s deviation in a real prompt would help practitioners.

Verdict: weak accept — a clean, useful lens with a real open-loop predictive result; not yet a complete account of context-specific computation.

Takeaways

• If you need a cheap “what should this layer do on average” baseline, iterate an average-attention kernel instead of only probing.
• Treat deviation from the average as the analysis object when hunting specialized heads — not raw attention patterns.
• Watch training time: mean-field collapse early and mean-field divergence around induction is a diagnostic you can reuse.