PaperlayerView exact revision on arXiv

Source-linked AI summary

Attention Mean Fields Predict Average Representation Dynamics and Reveal Context-Specific Computation

Micah Adler, John W. Byers, Mark Crovella

arXiv:2609.16382v1cs.LGcs.CL

TL;DR

Understanding representation geometry requires tracking how it evolves as a language model processes text, while distinguishing typical behavior from context-specific behavior. The paper uses corpus-averaged attention mean fields to predict average representation evolution and uses deviations to identify context-specific computation, including in-context learning. The authors report accurate predictions of average dynamics, early-training substitutability, and deviations that track reliance on in-context information.

  • Problem

    A useful baseline is needed to distinguish an attention head’s typical behavior from behavior specific to the current context.

    Source: §1 · p. 1

  • Method

    The paper averages trained attention over a corpus to construct mean-field approximations and uses conditional mean fields to characterize average or context-specific behavior.

    Source: §1 · p. 1, §3.1 · p. 5

  • Results

    The authors report accurate predictions of average representation evolution, early-training behavior resembling the mean field, and larger deviations in settings requiring more in-context computation.

    Source: §1 · p. 2, §7 · p. 26

  • Takeaways & Limitations

    Mean-field deviation provides a task-agnostic measure of context-specific computation and yields task residuals that act as transferable function vectors.

    Source: §7 · p. 26

    Mean-field results are conditional on the corpora studied, and different corpora may induce different attention mean fields.

    Source: §7 · p. 26

Abstract

from arXiv · show

A language model's representation geometry is not predetermined; it evolves as the model runs. 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 \emph{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 the model's own MLPs over token representations, never consulting a measured deviation at any layer. The resulting prediction is highly accurate. 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 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.

1 Introduction

The paper constructs a corpus-averaged attention mean field to provide a baseline for typical head behavior and predict average representation dynamics. Its predictions work well for average geometry and early training, while deviations reveal context-specific computation, including in-context learning.

Mean-field approach

  • The corpus mean field averages a trained head’s attention over a corpus, yielding a type-level kernel that summarizes its routing behavior.It is estimated in a single corpus pass without optimization or fitted parameters.Source: §1 · pp. 1–2

Average representation geometry

  • The mean field gives a surprisingly accurate, compact description of how average type representations evolve across model layers.This simplifies the study of average representation geometry.Source: §1 · p. 2

Training dynamics

  • Early in training, mean-field attention can replace attention on frequent query types with almost no loss cost; later, context-specific attention becomes behaviorally important.The transition coincides with a major reorganization in the mean field.Source: §1 · p. 2

Context-specific computation

  • Across controlled induction and few-shot settings, prompts requiring more in-context computation tend to show larger mean-field deviations.Previously documented ICL heads cluster near the top when ranked by average mean-field deviation.Source: §1 · p. 2
  • Mean-field deviation separates context-specific computation from the corpus-average component, framing representations as global structure plus local contextual innovation.The mean field describes global structure; its structured residual measures local variation.Source: §1 · p. 2

2 Related Work

The related work connects representation geometry and model behavior to corpus statistics, while distinguishing static or limiting accounts from this paper’s layerwise analysis of trained models. It also situates mean-field deviation as a diagnostic for context-specific computation rather than a theory of what in-context learning implements.

Corpus statistics and representation geometry

  • Prior co-occurrence accounts derive static representation geometry, whereas this work tracks layerwise transformations in real language models and documents correlation between corpus-level mean fields and co-occurrence.The authors argue this correlation helps explain co-occurrence’s apparent impact on representation geometry.Source: §2 · p. 2

Distributional semantics

  • The paper tests whether a low-order corpus-average operator captures model behavior, finding support for centroid geometry and early training but informative violations for in-context learning.This frames corpus-average statistics as useful but not sufficient for context-specific behavior.Source: §2 · p. 3

Training dynamics

Attention heads, corpus statistics, and contextual algorithms

  • Prior attention analyses connect QK and OV circuits to corpus-average statistics, while this work uses the mean-field residual to identify prompt-specific computation underlying in-context learning.The paper also relates its decomposition to increasing contextuality across upper layers described by Ethayarajh (2019).Source: §2 · p. 3

Input-independent and fixed attention

  • Earlier fixed-attention studies report limited task-performance costs; this work extends the comparison from overall loss to individual layerwise representations.The cited examples include synthetic or random alignment matrices and fixed positional patterns.Source: §2 · p. 3

Mean-field theories of attention

  • Unlike prior attention mean-field work focused on continuum or infinite-token and infinite-depth limits, this work uses finite-sample statistics measured on trained language models.The paper describes its analysis as the first mean-field treatment measured directly on trained language models and validated against their behavior.Source: §2 · p. 3

Rogue dimensions, attention sinks, and anisotropy

  • The paper distinguishes context-dependent geometry from rogue dimensions and does not correct for those dimensions in body results; appendix analysis finds that omitting them improves predictions.The passage relates this issue to anisotropy, outlier dimensions, and attention sinks.Source: §2 · p. 4

In-context learning

  • The paper does not propose a competing theory of in-context learning; mean-field deviation diagnoses when prediction requires prompt-specific computation and remains compatible with existing accounts.Those accounts include induction heads, implicit gradient descent, and compressed task representations.Source: §2 · p. 4

3 Mean Field Analysis of Attention

The section defines attention contributions and develops mean-field estimates conditioned on query type, query/key types, or an individual context. These estimates replace context-specific routing and values with averaged type-level quantities to approximate representation updates.

3.1 Notation

  • A head updates each query with an attention-weighted sum of value vectors, and mean fields average this contribution over query-type, query/key-type, or single-context sets.Source: §3.1 · pp. 4–5

3.2 The Corpus Mean Field

  • This estimate assumes attention-value independence and negligible curvature; the value-map Jensen gap is small in experiments.Source: §3.2 · p. 6
  • The corpus mean field replaces context-specific attention with a type-level kernel P_st and applies the value map to each type centroid to estimate average head contribution.The kernel is row-stochastic and can be used as an attention pattern in a replacement model.Source: §3.2 · p. 6

3.3 Conditional Mean Fields

  • The (s, t)-conditional field decomposes a contribution using the probability that type t occurs for query s and the average attention and value when it is present.Source: §3.3 · pp. 7–8
  • For one context, the estimate renormalizes corpus attention over present types and uses corpus-average values, assuming value-context independence and proportional attention.Source: §3.3 · pp. 8–9
  • The context attention estimate is not generally unbiased, but observed bias is small and removing it does not change key results.Source: §3.3 · p. 9

4 Mean Field Prediction of Representation Evolution

The paper iterates a corpus-averaged attention kernel and frozen model weights in an open-loop rollout to predict average type-representation evolution. These predictions are accurate across several models, while error analysis identifies the MLP-at-centroid approximation as the leading per-layer source of rollout error.

4.1 Method: Open-Loop Rollout

  • The open-loop rollout starts from input embeddings, applies mean-field attention and each layer’s MLP to type centroids, then feeds those centroids into the next layer without running text.It uses the mean-field kernel and frozen weights alone; applying the MLP once at the average is a simplification of its context-specific executions.Source: §4.1 · pp. 9–10

4.2 Basic Results and Ablations

  • Evaluation covers five models on the 1,000 most frequent types, which account for 64–68% of corpus positions and 67–72% of non-BOS attention.Results also generalize across FineWeb, WikiText, and Pre-1919 books, including when training and testing use different corpora.Source: §4.2 · p. 10
  • Centered cosine similarity exceeds 0.9 through GPT-2’s middle layers and remains above 0.75 at its last; Gemma-2 and Qwen-3 stay well above 0.75 through most layers.Figure 3 compares average per-type centered cosine across the three models.Source: §4.2 · p. 10, Figure 3 · p. 11
  • Query-blind, key-blind, and shuffled kernels fail to match measured attention, while layer-0 and depth-averaged kernels underperform; type-specific distinctions matter increasingly with depth.Figure 4 compares the query-blind and key-blind alternatives.Source: §4.2 · p. 11, Figure 4 · p. 11
  • The measured mean field is most accurate in centered cosine and relative error across all three models; co-occurrence is the closest alternative on GPT-2 and Gemma-2.Qwen-3-14B’s large co-occurrence, uniform, and Pavg relative errors reflect layer-0 rollout instability, not kernel quality.Source: §4.2 · pp. 11–12, Table 1 · p. 12

4.3 Understanding Approximation Errors

  • Teacher-forced head predictions across five models have cross-layer mean centered cosine of 0.78–0.87 and norm ratios of 0.91–1.24.Figure 5 reports centered cosine and centered norm ratio by fractional depth.Source: §4.3 · p. 12, Figure 5 · p. 12
  • The MLP-at-centroid approximation dominates per-layer rollout error, costing about 0.02–0.03 centered-cosine similarity; the value-map and covariance approximations each add at most ~0.005.The covariance term contributes little to centroid prediction but carries context-specific computation.Source: §4.3 · p. 13, Figure 6 · p. 13
  • Llama-3.2-3B’s rollout error comes mostly from the MLP-at-centroid step, especially at layer 1, although its mean-field attention prediction matches the other models.Its centroid rollout therefore requires a Jensen-gap correction to be accurate.Source: §4.3 · p. 13

5 The Mean Field as a Training Waypoint

Replacing attention with its mean field barely affects early Pythia behavior, despite ongoing learning. Around induction, the model departs from that prediction as its attention kernel reorganizes and heads become more selective.

5.1 Replacing Attention with its Mean Field

  • For the top 1,000 query types, the intervention replaces each head’s context-specific contribution with its context-conditional mean field on held-out text, leaving other computation unchanged.Remaining query types retain the model’s own computation; the top 1,000 cover 65% of corpus positions.Source: §5.1 · p. 14

5.2 Models Reduce to Their Mean Field, Then Depart from It

  • Through step 102, ΔL is negligible despite >3 nats of loss improvement; from about step 128, ΔL and DMF rise across Pythia sizes, then ΔL keeps rising as DMF plateaus or falls.After induction, larger models have lower true losses while their mean-field losses remain similar.Source: §5.2 · pp. 14–15, Figure 7 · p. 15
  • The same early mean-field effect appears to occur in OLMo-2 at roughly the same training point, although available checkpoints in the key region are limited.Source: §5.2 · p. 15

5.3 Kernel Dynamics Across Training

  • Through step 128, the kernel barely moves even as loss falls; it then reorganizes leading up to and across induction, recovers, and stabilizes, with stronger reorganization in larger models.Early learning is attributed to values, MLPs, and embeddings rather than the attention pattern.Source: §5.3 · p. 16, Figure 8 · p. 16
  • As kernels reorganize, heads become selective, with type-wise attention spread growing from about 1.5× to more than 20×; early attention-value independence then stops holding.Some heads develop large spread while others remain near their initially uniform behavior.Source: §5.3 · p. 16

6 In-Context Learning and Contextual Computation

The section uses context-conditional mean-field deviation to distinguish predictable, uniform attention from context-specific computation and examine its role in in-context learning. It separates mechanisms and head roles in in-context circuits, and explores whether task-specific deviation vectors can transfer task behavior.

6 In-Context Learning and Contextual Computation

  • The section treats deviation from the context-conditional mean field as a principled null for contextual processing, rather than merely as prediction error.The residual captures context-specific computation absent from the mean-field prediction.Source: §6 · pp. 16–17

6.1 Mean-Field Deviation Rises with In-Context Learning

  • Mean-field deviation rises with in-context learning across repeated-copy tokens, prompt classes, and training, emerging in induction heads at their formation transition.Class separation reflects designed differences, not a within-class relationship; cross-head correlation is nonsignificant before transition, rises to 0.37 at step 512 and peaks at 0.71 at the transition, then declines as other contextual heads accumulate.Source: §6.1 · pp. 17, 18, 20, Figure 9 · p. 18, Table 2 · p. 19

6.2 Mean Field Deviation Combines Two Contributors

  • Deviation splits exactly into attention–value covariance and contextualization of attended values, but the dominant mechanism differs by model and covariance alone can misread random controls.Covariance is generally larger, yet random controls show high covariance alongside the lowest contextualization and near-zero permutation loss; GPT-2’s ICL deviation is driven mainly by unusual attention, while Gemma-2’s reflects more contextualized attended tokens.Source: §6.2 · pp. 20–21, Table 3 · p. 21

6.3 Mean Field Deviation Recovers Contextual Heads and Circuitry

  • Ranking GPT-2 heads by corpus deviation surfaces documented ICL heads and other contextual heads, including position movers that transport absolute position and covariance-dominated induction heads.Among the top twenty are three of four identified induction heads and five of eleven name movers; several position movers carry the attended key’s absolute position.Source: §6.3 · p. 22
  • IOI name movers have intermediate induction scores of 0.15–0.45 and are contextualization-dominated, moving contextualized representations of non-repeated names.The passages identify L9H9, L9H6, and L10H0 as name movers, and L10H6 and L10H10 as backup name movers.Source: §6.3 · p. 23
  • For L9H8, all eight top-scoring deviation vectors decode to the document’s topic, suggesting that the head moves topic information into the query.Source: §6.3 · p. 23
  • Other high-deviation content movers track local structure: L5H2 moves abstract completions, L8H3 content in parallel or enumerated frames, and L6H8 content in idiomatic comparative frames.Source: §6.3 · p. 23
  • Duplicate-token and previous-token heads rank near the bottom because their self-attention or one-token-back behavior is type-predictable and already captured by the mean-field kernel.Source: §6.3 · p. 23
  • Induction and IOI circuits straddle the deviation ranking: uniform write-stage heads rank low, while heads that conditionally read and apply in-context bindings rank high.The authors propose that deviation measures where contingency enters a write-then-read computation: corpus statistics capture uniform tagging, but not novel in-context bindings.Source: §6.3 · p. 23

6.4 Task-Specific Deviations Are Function Vectors

  • Task-specific deviation vectors average a head’s output over few-shot task prompts at the final answer-cuing token, then subtract its corpus mean-field contribution.The resulting vector is retained for injection rather than reduced to a scalar deviation.Source: §6.4 · p. 24
  • On Gemma-2-2B, neither method consistently performs best, but the best pipeline almost always mixes head selection and injection; on antonym, AIE+MFA reaches 0.82 versus 0.57 for AIE and 0.76 for MFA.The two selection methods have low overlap in Gemma-2-2B; on GPT-J-6B, they select more similar heads and transfer rates tie on each task.Source: Table 5 · p. 26, §6.4 · p. 26
  • On GPT-2’s capital task, deviation-vector injection names a capital city on all 120 held-out queries but almost never the requested one; raw averages produce no task answer.J-lens ranks for seen, held-out, and neither-seen-nor-held-out capitals are 9, 37, and 34, versus about 21,000 for unrelated words, indicating category information with a slight bias toward observed capitals.Source: §6.4 · p. 26

7 Conclusions

The paper presents mean-field attention as a corpus-relative account of average representation dynamics and context-specific computation. Its conclusions are conditional on the corpora studied.

7 Conclusions

  • Mean-field results depend on the corpus used, and selecting an appropriate corpus for a given analysis remains an open direction.The authors report little difference across the corpora they compared.Source: §7 · p. 26
  • Mean-field attention predicts average representation evolution through model layers.The paper describes this as an accurate prediction using the mean-field operator.Source: §7 · p. 26
  • In early training stages, models behave like their mean fields.Source: §7 · p. 26
  • Mean-field deviation measures context-specific computation, recovers known in-context learning machinery, and yields transferable task residuals.Source: §7 · p. 26

Use of AI

The appendix extends conditional mean fields to contexts with repeated token types and examines the bias introduced by context-level renormalization. It finds that the bias is small in the studied data and does not alter the paper’s key results.

Use of AI

  • Claude and ChatGPT assisted with experimental code, analyses, and manuscript editing, while the authors retained responsibility for the research questions, designs, and interpretations.The authors state that they take full responsibility for the paper’s content.Source: Use of AI · p. 27

A Conditional Mean Fields with Repeated Types

  • With repeated token types, the conditional mean field uses expected counts and assumes attention mass grows proportionally with multiplicity.For a fixed context, renormalization weights types by their full multiplicities, not merely by whether each type is present.Source: A Conditional Mean Fields with… · p. 33

B.1 The Kernel ˆA as a Mean Field

  • The context-conditional kernel renormalizes corpus-level attention propensities, a nonlinear step that is not generally unbiased but is approximately unbiased in the studied data.An unbiased replacement can be constructed, and using it leaves key results unchanged.Source: §B.1 · p. 34

B.2 Analysis of Bias in ˆA

  • The bias redistributes attention mass across key types rather than producing a net distortion under the mean-field weighting.Tracked types tend to be under-predicted, while the position-0 bucket is over-predicted.Source: §B.2 · p. 35

B.3 The Measured Bias of ˆA is Small

  • The unweighted cell averages are 0.946 and 0.933, a 5–7% under-prediction for typical pairs, while over 95% of pairs are within 20% of unbiased in both models.The reported estimates use well-sampled head–type pairs.Source: §B.3 · p. 35, Table 6 · p. 35

B.4 Using an Unbiased Estimator of ˆA Does Not Change Key Results

The appendix calibrates the context-conditional kernel to remove its bias and checks the effect on key results. It also explains how mean-field deviation separates into attention-routing and contextualization components.

B.4 Using an Unbiased Estimator of ˆA Does Not Change Key Results

  • The calibration uses iterative proportional fitting to make the kernel’s renormalized attention marginals match the measured marginals.The context-conditional deviation then splits into a covariance contribution from residual routing and a contextualization contribution from in-context values differing from corpus averages.Source: §B.4 · p. 36, §B.5 · pp. 37–38
  • Calibrating the context-conditional kernel leaves the reported training-dynamics, induction, few-shot, and head-ranking findings unchanged.The calibrated head ranking has Spearman correlation 0.9967 with the original, and 19 of the original top 20 heads remain in the top 20.Source: §B.4 · p. 37

C Impact of the Value-Map Mean-Field Approximation

Evaluating the value map at type centroids closely approximates averaging it over occurrences because within-type covariance is small. The appendix also specifies how the mean-field kernel handles untracked types and avoids using measured evaluation centroids during rollout.

C Impact of the Value-Map Mean-Field Approximation

  • The approximation error is the Jensen gap of the input normalizer, governed by within-type covariance; tight clustering makes the centroid value approximate the average value.The value map is linear after its input normalizer, so the normalizer is the source of the discrepancy.Source: C Impact of the Value-Map Mean-Field… · pp. 38–39
  • Type identity explains 0.90 of GPT-2’s z-scored input variance and 0.85 for Gemma-2, declining in deep layers to about 0.25 and 0.15, respectively.Gemma-2’s input residual is exactly degenerate, while GPT-2’s additive positional embeddings already spread occurrences of a type.Source: C Impact of the Value-Map Mean-Field… · p. 39
  • The value-map Jensen gap is negligible at every layer: cosine distance is below 3 × 10−5 for GPT-2 and 3 × 10−4 for Gemma-2, with a 10–28% residual-magnitude gap.The mapped vectors are essentially parallel in both models.Source: C Impact of the Value-Map Mean-Field… · pp. 39–40

D Mean-Field Kernel Details

  • The kernel keeps position 0 and untracked types in separate columns, with untracked types represented by a running mean during rollout.The tracked centroids evolve internally from input embeddings and weights, without consulting measured residuals at each layer.Source: D Mean-Field Kernel Details · pp. 40–41
  • Held-out evaluation estimates kernels and type statistics on one half of OpenWebText and scores predicted centroids against the disjoint other half.Evaluation centroids are not used during estimation.Source: D Mean-Field Kernel Details · p. 41
  • The context-conditional W estimate zeros type pairs observed fewer than 20 times, discarding 53% of pairs but only 3% of Gemma-2’s attention mass.The corresponding GPT-2 figures are 53% of pairs and 3% of attention mass; all context-conditional results use this mask.Source: D Mean-Field Kernel Details · p. 41

E Coverage of the Tracked Type Set

The top 1,000 token types cover a similar share of positions across models, though much less of total attention because position 0 absorbs substantial attention mass. The appendix also reports stable, cross-corpus centroid predictions across the corpora studied.

E Coverage of the Tracked Type Set

  • Top-1,000 types cover 64–68% of corpus positions across models despite a fivefold spread in vocabulary size.The rank-frequency curves are nearly indistinguishable through and beyond the tracked cutoff, so the tracked types select similar text across tokenizers.Source: E Coverage of the Tracked Type Set · pp. 41–42, Figure 14 · p. 42
  • Tracked types receive 18–42% of total attention, but 67–72% of non-sink attention on full-vocabulary models, close to their position share.Position 0 takes 40% of attention in GPT-2 and Gemma-2, 49% in Qwen-3, and 65% in both Llamas.Source: E Coverage of the Tracked Type Set · pp. 42–43, Table 9 · p. 42
  • Untracked types still contribute through bucket values, but their distinct values are compressed into one shared vector.Models with lower tracked attention share have more head behavior compressed into that average.Source: E Coverage of the Tracked Type Set · p. 43

F Mean-Field Predictions Across Corpora

F.1 Corpora Studied

  • The corpora include separately filtered general-web text and more formal, encyclopedic Wikipedia articles.FineWeb is a recent filtered web crawl; Wikipedia’s WikiText-103 consists of featured and good articles.Source: §F.1 · p. 43
  • The four-corpus comparison uses a shared vocabulary of the 1,000 most frequent types in the OpenWebText reference sample.Samples are tokenized with each model’s tokenizer and packed into 512-token windows.Source: §F.1 · p. 43

F.2 Mean Field Predictions Are Stable Across Corpora and Generalize Out of Distribution

Mean-field centroid predictions replicate across corpora and generalize out of distribution, with accuracy generally close to matched-corpus evaluation. The appendix details evaluation metrics and representation anisotropy, then reports how rogue-dimension pinning affects rollout accuracy.

F.2 Mean Field Predictions Are Stable Across Corpora and Generalize Out of Distribution

  • Cross-corpus predictions are generally close to matched-corpus accuracy, and even the worst mismatch beats the context-free baseline in every evaluation column and both models.Generalization degrades with corpus dissimilarity: OWT–FineWeb is strongest, while pre-1919 novels are weakest.Source: §F.2 · p. 44, Figure 15 · p. 44, Figure 16 · p. 45

G Full Mean Field Results

  • Full rollout quality is assessed after centering with directional cosine, cloud-shape RSA, centered relative error, and centered norm ratio.Centered relative error is jointly sensitive to direction and magnitude; the other metrics isolate direction, shape, and magnitude, respectively.Source: G Full Mean Field Results · pp. 45–46

H Anisotropy of Representations

I Impact of Rogue Dimensions on Results

  • A handful of high-magnitude rogue dimensions can perturb every dimension through block normalization, so the appendix pins them to corpus values during rollout.Pinned dimensions are approximately context-invariant and are identified when their per-dimension RMS exceeds 15 times the layer median; they comprise at most 1% of each model’s width.Source: I Impact of Rogue Dimensions on Results · p. 46, §I.1 · p. 49

I.1 Rogue-Dimension Pinning: a Modest Correction

  • Pinning improves final-layer centered cosine for GPT-2, Gemma-2, and Qwen-3-14B, but is neutral to slightly negative for the two Llama models.For example, GPT-2 rises from 0.76 to 0.84 and Gemma-2 from 0.70 to 0.82; the Llama results are 0.72 to 0.71 and 0.60 to 0.58.Source: §I.1 · p. 49

I.2 Rogue Dimensions and the Qwen-3-14B Ablations

In Qwen-3-14B, pinning changes how kernel-substitution results are interpreted: rogue dimensions amplify some already-drifted substitute rollouts, while the mean-field rollout remains stable. The damage is concentrated in layer 0, and pinning is not uniformly beneficial.

I.2 Rogue Dimensions and the Qwen-3-14B Ablations

  • Pinning rescues several substitute-kernel rollouts in Qwen-3-14B, but barely affects the mean-field rollout; the rogue dimensions act as a conditional amplifier.Some substitute kernels drift far enough to trigger a blow-up through the shared RMSNorm scale; pinning removes that blow-up.Source: §I.2 · p. 49, Table 13 · p. 50
  • Layer 0 accounts for 0.279 of the 0.327 total damage from substituting co-occurrence kernels across all layers, an 85% share.Substituting across the other 39 layers incurs only 0.065 damage; the two shares sum to 105%.Source: §I.2 · p. 49, Table 14 · p. 50
  • Pinning can improve centered cosine while worsening the mean field’s relative error from 0.640 to 0.721, so it is not a uniform accuracy gain.The uniform kernel remains broken when pinned.Source: §I.2 · p. 50

J The Llama Models: Accurate Per-Layer Prediction, Failed Rollout

Llama mean-field predictions remain accurate layer by layer, but open-loop rollout fails sharply for Llama-3.2-3B after a single MLP amplification step. The Llama-3.2-3B failure is traced to a layer-1 MLP curvature effect, while the MLP-at-centroid approximation is a leading per-step error source across models.

J The Llama Models: Accurate Per-Layer Prediction, Failed Rollout

  • Teacher-forced per-layer predictions remain accurate, with cross-layer mean centered cosine of 0.858 for Llama-3.2-3B and 0.837 for Llama-3.1-8B.Both models show a mid-stack dip followed by recovery; the respective minima are 0.69 and 0.62.Source: J The Llama Models: Accurate Per-Layer… · pp. 50–51, Figure 21 · p. 52
  • Llama-3.2-3B rollout cosine collapses to 0.30 by layer 2 and averages 0.499, while predicted centered magnitude reaches 4.9× the true value.Llama-3.1-8B instead has a mean rollout cosine of 0.719 and final-layer cosine of 0.72.Source: J The Llama Models: Accurate Per-Layer… · p. 50, Figure 19 · p. 51
  • A single layer-1 MLP step amplifies Llama-3.2-3B rollout error 17×, while attention expands it only 1.06× and later steps mostly preserve it.The incoming error is 36% of the centroid norm; the corresponding amplification is about 1× at every layer in the other tested models.Source: J The Llama Models: Accurate Per-Layer… · p. 51
  • Approximately 25 first-MLP gates account for the amplification, with the top 25 recovering 94.7% of nonlinear excess versus 0.1% for random gates.The error is heavy-tailed over tokens: the ten most-amplified carry 53% of output-error energy, and 43 of the top 50 are closed-class tokens.Source: J The Llama Models: Accurate Per-Layer… · pp. 51–52
  • Excluding structural tokens leaves divergence essentially unchanged, while tested sink-magnet, massive-activation, and positional explanations do not account for the amplification.Boundary tokens receive 1.4–1.7× more attention per occurrence, but their magnet strength does not predict error amplification; their massive-activation medians are 0.71 versus 0.73 for content tokens, and capitalization pairs show no systematic positional effect.Source: J The Llama Models: Accurate Per-Layer… · pp. 52–53
  • The remaining explanation is a model-specific curvature effect in Llama-3.2-3B’s layer-1 MLP: concentrated centroid error pushes a few SwiGLU gates across their nonlinearity.Why this model’s gates lie near the nonlinear knee while others do not remains open.Source: J The Llama Models: Accurate Per-Layer… · p. 53

K Impact of the MLP on Mean-Field Approximation

  • The MLP-at-centroid approximation is near-exact in early layers but weaker mid-network, with layer-averaged centered cosine of 0.77 for GPT-2 and 0.86 for Gemma-2.The MLP term dominates the per-step error decomposition, but its disagreement costs less in the full centroid prediction because attention output is also included and much of the difference is a shared magnitude offset.Source: K Impact of the MLP on Mean-Field… · p. 53

L Replication on OLMo-2, and Evidence of a Shared Mean-Field Ceiling

The early mean-field substitution pattern also appears in OLMo-2-7B, while later model losses separate from a narrow mean-field loss band. Increasing top-N coverage has little effect on trajectory shape at matched estimation quality, and kernel similarity falls during training reorganization before recovering.

L Replication on OLMo-2, and Evidence of a Shared Mean-Field Ceiling

  • Early in training, all three models fail to improve over their mean fields: ∆L is −0.001 and −0.006 for Pythia at 0.13B tokens and +0.003 for OLMo at 0.63B.OLMo’s early deviation is 0.0713, close to Pythia-160m’s 0.0720 at initialization; OLMo has no published checkpoints before 0.63B tokens.Source: L Replication on OLMo-2, and Evidence… · p. 53
  • Across Pythia-1.4B, Pythia-12B, and OLMo-2-7B, model losses fall while mean-field losses flatten at 4.93–5.45 nats, leaving late gaps of 2.49, 2.81, and 2.66 nats.The models differ eightfold in parameter count, and OLMo also uses different training data; the authors suggest the mean-field loss ceiling may be set by corpus statistics.Source: L Replication on OLMo-2, and Evidence… · pp. 53–54, Figure 22 · p. 54

M Sensitivity to the Top-N Truncation

  • At matched estimation quality, top-N truncation has little effect: ∆L stays within 0.001 nats across N values, and normalized curve shapes agree within 0.01 across a fourfold change in N.The 40M-token extraction controls estimation quality across arms; position coverage rises from 65% to 72% and 80% at N = 2000 and 4000.Source: M Sensitivity to the Top-N Truncation · p. 55, Figure 23 · p. 55
  • The three truncation arms share the model-loss curve, and their mean-field loss stays within 0.01 nats through step 64 before detaching between steps 128 and 256.The mean-field loss reaches a minimum at step 2000, a late maximum at step 32000, and then falls back by step 143000.Source: M Sensitivity to the Top-N Truncation · p. 55, Figure 24 · p. 56

N Kernel Comparison: Estimation Details

  • Adjacent-checkpoint kernel similarity stays at 0.99–1.00 during the frozen stretch, then drops through reorganization to a minimum of 0.40 in the 12B model.The two smaller models recover to about 0.9 by step 16000, while 12B recovers more slowly to 0.75.Source: N Kernel Comparison: Estimation Details · pp. 55–56
  • Row spread starts near 0.40 for 160m and 0.38 for 1.4B, dips around the deviation minimum, then increases as heads differentiate, with some remaining near-uniform.Figure 25 plots the mean and interquartile band as well as per-head spread grouped by layer.Source: N Kernel Comparison: Estimation Details · p. 56

O How Reliably Each Deviation Metric Separates the Prompt Classes

The appendix tests how mean-field deviation components separate designed prompt classes from a random control and qualifies what that separation can establish. It also documents head-ranking lenses and measurement procedures used to interpret deviations.

O How Reliably Each Deviation Metric Separates the Prompt Classes

  • For induction and few-shot prompts, DMF separates classes from the random control in both models, while covariance and contextualization components divide the separation differently by model.Induction and few-shot Hedges’ g values are +3.60 and +2.06 for GPT-2 and +2.93 and +2.21 for Gemma-2; GPT-2 covariance separates induction and contextualization separates few-shot, whereas Gemma-2 contextualization separates every class.Source: O How Reliably Each Deviation Metric… · p. 57, Table 16 · p. 57
  • The separation is between prompt classes, not individual prompts: within-class metrics add at most 0.004 R2 in predicting ∆Lperm beyond class membership.GPT-2 DMF does not distinguish natural text from control (g = +0.04, p = 0.84), while Gemma-2 does (g = +1.06).Source: O How Reliably Each Deviation Metric… · p. 57

P Provenance of the Heads’ Roles and Ranking

P.1 Lenses for Mean Field Analysis

  • The covariance/contextualization plane splits each head’s deviation by attention routing versus the values carried by attended tokens.Its signed shares sum to per-occurrence deviation, and both axes use residual-stream units.Source: §P.1 · p. 58
  • Top-deviation context mining keeps each head’s largest per-occurrence deviations with surrounding text, attended keys, and the next token, filtering truncation artifacts.Vocabulary-space decoding projects each deviation vector through a logit lens with a J-lens refinement for middle-layer mediation.Source: §P.1 · p. 58

P.2 Measurement Protocol

  • Repeated-text readings present a real text chunk twice and sample queries only from the second copy, where an exact earlier context copy is available.The measurement uses 300,000 tokens at stride 8 and ranks all 144 heads.Source: §P.2 · p. 58
  • The vocabulary-space and occurrence-mining procedures filter occurrences with substantial out-of-vocabulary attention and decode each deviation separately to avoid mixing entities or topics.The J-lens uses the model’s average linear account of how layer content reaches the output; head rankings are stable across document-interleaved corpus halves (ρ = 0.999, sharing 18 of the top 20).Source: §P.2 · pp. 58–59

P.3 What the Lenses Show

The section uses mean-field deviation and related validations to examine context-specific head behavior, including repetition, position re-encoding, induction-related structure, and selection of task-relevant heads. It also tests whether head-selection verdicts remain consistent across two injection protocols.

P.3 What the Lenses Show

  • Under repeated text, four induction heads show the network’s largest deviation increases (+0.48 to +0.64), while positional heads become more predictable, including L2H7’s drop from 0.395 to 0.166.The induction heads rank low on natural text because their trigger is rare there; other head groups respond inconsistently to repetition.Source: §P.3 · p. 59
  • The lowest-deviation heads include both important mean-field carriers and duplicate-token machinery; four of Gemma’s function-vector core heads also rank among its bottom fifteen.In GPT-2, four of the six most predictable heads are among the eight highest-impact heads; duplicate-token detection is expressible in the mean field, making it invisible to deviation.Source: §P.3 · p. 60, Table 21 · p. 62

P.4 Layer-1 Positional Transport: Re-encoding the Query’s Position

  • GPT-2’s positional cluster writes a deterministic linear image of the query’s absolute-position embedding, with BOS-attending heads gating a fixed BOS value through position-dependent attention.The cluster’s writes occupy largely new directions rather than amplifying the embedding code; removing their position subspace raises loss by +0.0117 nats, 38× the controls, especially for numeric predictions (+0.028).Source: §P.4 · pp. 62–64
  • The five GPT-2 positional heads encode position in a smooth, low-dimensional code spread across partially overlapping channels, and downstream attention changes concentrate in layer 2 when that code is ablated.L1H2 and L1H10 overlap at 0.61, while L1H8 and L2H7 overlap at 0.02; the combined code subspace captures 0.88 of the output variance.Source: §P.4 · pp. 63–64

P.5 Gemma-2-2B Head Survey

  • In Gemma-2-2B, induction heads rank highly by deviation: L15H0 ranks third, and deviation correlates with induction score across 208 heads (ρ = 0.45).L15H0 and L4H4 lean toward covariance; the top 20 heads lie in layers 7–17, and their mean split favors covariance (0.43) over contextualization (0.32).Source: §P.5 · pp. 64–65, Table 22 · p. 65

Q Function-Vector Extraction: Baselines, Controls, and Results

  • The function-vector comparison selects top-10 heads from Gemma-2-2B, with MFA selection sharing 7–10 heads across its three splits while AIE selections are fixed per task.The appendix compares MFA task-residual selection with Todd et al.’s AIE selection and reports task-by-task results and injected-vector analyses.Source: §Q.1 · p. 65, Q Function-Vector Extraction:… · p. 65

Q.1 Validations

  • AIE reproduces Todd et al.’s GPT-J selection, while removing the null baseline makes task-head selection essential: raw-magnitude transfer on capital is 0.00 at every scale.AIE recovers 7, 8, and 7 of the ten published heads on the antonym, capital, and past tasks; raw-magnitude and MFA rankings agree on only 5–6 of ten heads.Source: §Q.1 · pp. 65–66
  • Per-head injection scored 0.570 against 0.442 for summed injection across 72 paired cells (Wilcoxon p < 10−4), but did not reorder the selections.The penalty was similar for AIE (−0.125) and MFA (−0.130; Mann–Whitney p = 0.83), despite their different layer ranges.Source: §Q.1 · p. 66

Q.2 Detailed Results

The detailed results show that head selection and vector construction trade off across models and tasks, while their combined pipelines transfer task information. Decoding further distinguishes category content from answer binding and reveals different function-vector behavior in GPT-2 and Gemma-2.

Q.2 Detailed Results

  • AIE selection wins capital and currency on Gemma-2, MFA wins past, language, and translation, and the methods tie on antonym.The two selections overlap by only 2–3 of ten heads and favor different depth bands.Source: §Q.2 · p. 66
  • 10 of 12 per-head selection–task cells match or favor the MFA deviation vector over the raw average, with the two exceptions within the 0.02 tie margin.Under summed injection, the deviation advantage grows for AIE-selected heads but reverses for MFA-selected heads.Source: §Q.2 · pp. 66–67
  • 7/10, 5/10, and 4/10 of Todd et al.’s heads are recovered by MFA selection on GPT-J antonym, capital, and past tasks, respectively, with transfer tied on all three.The apparent antonym gap in smaller pools does not persist at n = 120.Source: §Q.2 · p. 67, Table 24 · p. 67
  • 0.60–0.97 transfer covers all six tasks through at least one pipeline combination, while random ten-head controls transfer at 0.00–0.01.On GPT-J, the two selections are statistically indistinguishable on every task despite often choosing different heads.Source: §Q.2 · p. 67, Table 24 · p. 67

Q.3 Decoding Category Vectors

  • Per-prompt J-Lens decodes distinguish query-category writers from answer-family writers in GPT-2 capital prompts, while no antonym residual decodes the answer well.For past tense, decoded residuals can express the category without binding it to the correct verb.Source: §Q.3 · p. 67
  • GPT-2 function vectors favor seen answers, whereas Gemma-2 shows no seen-answer preference, supporting a readout-binding rather than category-content deficit for GPT-2.Gemma-2’s vectors carry the task without retaining which exemplars taught it.Source: §Q.3 · p. 68

Related papers

Loading 2609.16382v1…