Conceptio › Archive › arXiv CS

arXiv CSopen access

Layerwise Dynamics for In-Context Classification in Transformers

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Patrick Lutz 1 Themistoklis Haris 1 Arjun Chandra 1 Aditya Gangrade 1 Venkatesh Saligrama 1

arXiv:2604.11613v1 [cs.LG] 13 Apr 2026

Abstract

The obstacle: symmetry breaking hides the computation. Classification problems possess inherent symmetries: permuting examples, permuting feature coordinates, or relabeling classes should not change the decision rule. While self-attention is invariant to token order, standard training does not enforce feature- or label-permutation symmetry. As a consequence, models achieve high accuracy, but encode spurious asymmetries in their weights. These asymmetries hinder interpretability by obscuring the underlying algorithmic structure. To move from models that predict well to discovering the underlying algorithm, we must break through this barrier. We do so by forcing the internal computations of the model to respect the task’s symmetries.

Transformers can perform in-context classification from a few labeled examples, yet the inference-time algorithm remains opaque. We study multi-class linear classification in the hard no-margin regime and make the computation identifiable by enforcing feature- and labelpermutation equivariance at every layer. This enables interpretability while maintaining functional equivalence and yields highly structured weights. From these models we extract an explicit depth-indexed recursion—an end-to-end identified, emergent update rule inside a softmax transformer, to our knowledge the first of its kind. Attention matrices formed from mixed feature–label Gram structure drive coupled updates of training points, labels, and the test probe. The resulting dynamics implement a geometry-driven algorithmic motif, which can provably amplify class separation and yields robust expected class alignment.

The method: enforcing symmetry layer by layer. We enforce symmetry by construction. At each layer, we conjugate the attention block with a random block permutation over feature and label coordinates (apply P , run attention, then apply P ⊤ ). This forces each layer to implement the same computation in every symmetric coordinate system. We verify that the learned inference rule is preserved by checking output agreement and query/context sensitivity fingerprints. Empirically, behavior is unchanged; instead, the constraint acts as a structural denoiser, collapsing the learned weights onto a low-dimensional, interpretable manifold.

Introduction

Transformers trained for in-context learning (ICL) routinely solve learning problems using only a forward pass. Yet, the core mechanistic question remains open:

Result 1: Coupled mean-shift dynamics. With symmetry enforced, the forward pass becomes algebraically readable: for linear multiclass classification we recover an explicit layerwise recursion (layer ℓ = iteration ℓ). Let Xℓ ∈ Rn×d be the layer-ℓ feature vectors of the n training tokens, and Yℓ ∈ Rn×K their label vectors. Each iteration first builds an affinity between training tokens by combining feature similarity and label agreement, Aℓ = softmax αXℓ Xℓ⊤ + γYℓ Yℓ⊤ ,

What algorithm is the model running at inference time? A common hypothesis is that transformers implicitly simulate a generic optimizer—most often gradient descent— within the residual stream. In this paper, we show that for in-context classification, this is a misleading abstraction. Instead, we show that once natural task symmetries are enforced layer by layer, trained transformers instantiate a shared dynamical template: a label-aware, geometrydriven algorithmic motif. This motif is not a loose analogy: it yields a closed-form layerwise recursion (depth ↔ iterations), supports theory with testable predictions, and reappears when we retrain transformers on classification tasks with different distributional and geometric structures.

and then applies this same weighting to propagate both training features and labels, Xℓ+1 = Xℓ + α′ Aℓ Xℓ ,

Yℓ+1 = Yℓ + γ ′ Aℓ Yℓ .

The query token is updated analogously by attending to the training set, accumulating class evidence across iterations. Unlike GD-style ICL abstractions that update a classifier on fixed features, this is a representation-shaping procedure in which feature and label geometry co-evolve. We refer to

Boston University, Departments of Computer Science and Electrical & Computer Engineering. Correspondence to: Patrick Lutz <[email protected]>. Preprint. April 14, 2026.

Layerwise Dynamics for In-Context Classification in Transformers Symmetry constraint

Interpretable weights

Emergent algorithm

Permutation sandwich

Structural denoising

Coupled feature-label dynamics

Attention

without permutations

Same inference rule query prediction

Same predictions

Same response to query perturbations

sample query influence

with permutations

depth

iterations

Feature geometry evolves

Canonical weight structure

prediction

Same influential examples

Figure 1. Symmetry reveals coupled mean-shift inference in transformers. Left: Enforcing feature/label permutation symmetry layer by layer preserves the inference rule: we match predictions, response to query perturbations (local decision rule), and influential examples (context sensitivity). Middle: The symmetry constraint denoises weights into a canonical low-dimensional structure. Right: The resulting structure makes the forward pass readable as a closed-form iterative classifier (depth ↔ iterations) in which feature and label geometry co-evolve, yielding a label-aware mean-shift motif.

2. Closed-form inference dynamics. We extract a concise layerwise recursion for in-context classification (depth ↔ iterations): a coupled update dynamic in which feature and label geometry co-evolve. 3. A Shared Algorithmic Motif. The extracted dynamics allows for theory and testable predictions that hold beyond the base setting; retraining transformers on semi-supervised, label-noise, and prototype tasks recovers the same weight structure and the same coupled algorithmic motif.

this template as a coupled feature-label dynamic. Why this matters: Predictions, Theory, and Robustness. With a closed-form layerwise recursion in hand, we can analyze the transformer’s inference-time computation and derive testable predictions about behavior and performance. In a label-driven regime, the dynamics preserve class structure while monotonically amplifying the margin between class clusters. This yields a simple mechanism: a small number of labeled points anchors the classes, after which performance is driven primarily by geometric separation. A direct corollary is semi-supervised ICL: adding unlabeled points sharpens the geometry and improves accuracy even at a fixed label budget. More broadly, because the mechanism is geometry-driven, it should extend beyond linear separators when clusters are well separated, which we test in prototype (Voronoi cell) classification settings.

1.1

Related Work

ICL as algorithm learning and implicit optimization. Work in stylized in-context learning (ICL) settings trains transformers on sequences of examples sampled from a task family and analyzes the resulting inference-time procedure on held-out tasks (Garg et al., 2022; Zhang et al., 2024; Ahuja & Lopez-Paz, 2023). Several papers argue that, in linear regression and related regimes, attention and residual connections can realize closed-form estimators like gradient-descent updates (von Oswald et al., 2023; Akyürek et al., 2023; Dai et al., 2023; Lutz et al., 2025). Complementary theory formalizes this viewpoint as algorithm learning and studies learnability, generalization, and stability of the learned solver (Li et al., 2023; Wies et al., 2023). At the same time, empirical correspondences between ICL and explicit gradient descent can be evaluation-dependent (Deutch et al., 2024). Our work adopts the algorithmic lens but focuses on multiclass classification in a hard regime; after enforcing the task symmetries layer-by-layer, the extracted computation is best described as a representation-space dynamical system rather than parameter-space empirical risk minimization.

Result 2: Dynamics predict behavior across tasks. The analysis above predicts that the same coupled feature-label dynamic should succeed in semi-supervised, label-noise, and prototype (Voronoi) settings. Strikingly, when we retrain transformers on these tasks, this prediction carries over to the learned models: the trained transformers repeatedly recover the same symmetry-aligned weight structure and implement the same underlying coupled mean-shift recursion (with task-specific parameter schedules). In other words, the closed-form iteration is not just an explanatory fit after training—it predicts the inference-time mechanism that re-emerges across tasks. Contributions. In summary, we make three contributions: 1. Symmetry as an identifiability protocol. We enforce feature/label symmetries layer by layer and verify the inference rule is preserved (output alignment and query/context sensitivity), removing an important interpretability barrier and yielding structured, interpretable weights.

Symmetry and equivariance. Symmetry-respecting architectures are a long-standing theme, including permutationinvariant set models (Zaheer et al., 2017; Lee et al., 2019) 2

Layerwise Dynamics for In-Context Classification in Transformers

Semi-Supervised Prediction and ICL. Later in the paper, we also investigate semi-supervised prediction (SSP). This augments nlab labeled examples of the prediction task with u , and demands the same sonunlab unlabeled points (xj )nj=1 lution. In semi-supervised ICL, the prompts are expanded to add these unlabeled examples coupled with a ‘null’ label.

and more general group-equivariant networks that impose structured parameter sharing (Cohen & Welling, 2016; Maron et al., 2019; Satorras et al., 2021). Recent work studies permutation equivariance properties within transformer architectures and their consequences (Xu et al., 2024). We use equivariance differently: we enforce feature- and labelpermutation symmetry inside each layer (a permutationconjugation “sandwich”), not primarily to improve accuracy but to remove symmetry-breaking degrees of freedom that obscure identifiability, enabling closed-form extraction and cross-task motif comparisons.

Transformers. We follow common conventions in the ICL literature and encode the input data (i.e., the prompt) ⊤ X̃0 = x1 . . . xn xtest ∈ R(n+1)×d ⊤ Ỹ0 = y1 . . . yn 0K ∈ R(n+1)×K . Z0 = X̃0 Ỹ0 ∈ R(n+1)×(d+K) .

Mean-shift, diffusion, and semi-supervised learning. Mean-shift is a classical mode-seeking and clustering method based on iterative kernel-weighted averaging (Comaniciu & Meer, 2002). Graph diffusion and label propagation methods likewise exploit unlabeled geometry by iterating a stochastic similarity operator (Zhu & Ghahramani, 2002). Our extracted recursion is naturally situated in this family: attention produces a learned stochastic similarity matrix and depth implements repeated diffusion steps, with label-dependent terms coupling representation updates to class information. This connection clarifies why behavior can be primarily margin-driven (geometry-first) and why similar dynamical templates can remain effective across perturbations that preserve the underlying symmetry structure. Related theoretical work studies self-attention as an interacting particle system and proves clustering phenomena consistent with mean-shift-like dynamics in deep attention stacks (Geshkovski et al., 2023; Rigollet, 2025).

We consider single-head, attention-only transformers (Vaswani et al., 2017), which transform an input Z0 ∈ R(n+1)×(d+K) over L layers via Qℓ = Zℓ WQ,ℓ , Kℓ = Zℓ WK,ℓ , Vℓ = Zℓ WV,ℓ , Qℓ K⊤ ℓ Attn(Zℓ ) = σ √ + M Vℓ WP,ℓ , (2) d+K Zℓ+1 = Zℓ + Attn(Zℓ ), −∞1n+1 is an additive masking where M = 0n+1 0⊤ n bias that reproduces the causal constraint that earlier tokens cannot use query information and σ is the row-wise softmax operator whereby for a matrix A = (Aij )i,j∈[n+1] we have (σ(A))ij = exp(Aij )/

Further discussion. Appx. A provides additional context on (i) how ICL emerges from pretraining distributions; (ii) Bayesian/kernel interpretations; and (iii) mechanistic circuit analyses.

(1)

exp(Aim ).

(3)

Architectural scope. End-to-end algorithms have so far mainly been identified for linear transformer models (von Oswald et al., 2023; Ahn et al., 2023). We take the next step by studying attention-only transformers with softmax, which capture data-dependent weighting while remaining tractable. To keep the mechanism identifiable, we omit MLPs, layer norm, and multi-head attention, and use this controlled setting to isolate how attention couples feature geometry and labels across layers.

Setup

Supervised Prediction Task and In-Context Prediction. Let F be a class of functions f : X → Y from features to labels. We study prediction problems where, given examples, the value of a map in F at a test point must be predicted. Specifically, given n pairs (xi , yi )ni=1 with yi = f (xi ) for some unknown f ∈ F, and a test point xtest , we would like to approximate f (xtest ). We call this the prediction task induced by F. We set X = Rd and Y = RK throughout.

In-Context Prediction. In in-context learning, the data ((xi , yi )ni=1 , xtest ) is formatted to a prompt p that is fed a transformer Tθ . We use P to denote the set of prompts induced by F, and we say that a prompt is induced by f ∈ F if for all i, yi = f (xi ). A prediction function is a map P → Y, and the weights of the transformer encode a specific implementation of one such function. We will use xtest (p) as notation for the test point in p.

Algorithm Extraction in In-Context Learning via Symmetries

In this section, we (i) formalize the natural symmetries that prediction tasks satify, and (ii) describe how we use these symmetries as a diagnostic tool to uncover the core inference procedure a trained transformer relies on. Concretely, we do this by enforcing the symmetries inside the network at every layer, thus improving interpretability. Symmetries of a Prediction Task. We begin by formally setting up a notion of symmetries of a prediction task. 3

Layerwise Dynamics for In-Context Classification in Transformers

Figure 2. Learned weight matrices for the unconstrained transformer (left) and the symmetry-preserving transformer (right) trained on in-context linear classification. The unconstrained model exhibits little visible structure, whereas enforcing the task symmetries produces a more regular pattern that is easier to interpret. The top and bottom rows show WQK,ℓ and WV P,ℓ , respectively.

Definition 1. Let G be a set of transformations g : P → P where if p is induced by some f ∈ F and p′ = g(p), then p′ is induced by some f ′ ∈ F. We say G captures prediction invariance if for all such (p, f, p′ , f ′ ),

internal computations. Even if the target algorithm is equivariant, the network may implement it using intermediate representations defined in arbitrary, rotated coordinate systems. Consequently, the same input-output behavior can arise from obscure hidden-state dynamics whose representations do not reflect the problem’s natural symmetries.

f ′ (xtest (p′ )) = f (xtest (p)). If moreover the action g extends to Y, we say G captures prediction equivariance if for all such (p, f, p′ , f ′ ), f ′ (xtest (p′ )) = g f (xtest (p)) .

Example. Gradient descent (GD) for standard linear models is permutation equivariant: permuting the input feature coordinates simply results in a corresponding permutation of the weight updates. However, a transformer could simulate this algorithm while encoding the iterates in a basis that rotates layer-by-layer. While the final predictions remain correct (and equivariant), the internal activations would appear disordered, effectively masking the simple permutation structure of the simulated algorithm.

In words, a transformation g modifies the prompt in a structured way. If we apply g, the correct answer for the transformed prompt is either unchanged (invariance) or changes in the corresponding structured way (equivariance). Symmetries commonly arising in common ICL task families (e.g., regression/classification) include:

Enforcing symmetry via randomization. To resolve the ambiguity of arbitrary internal coordinates, we explicitly enforce symmetries layer-by-layer. We introduce a stochastic permutation operation inside each attention block. Crucially, if we applied a fixed permutation P , the transformation would be vacuous: it would act as a mere coordinate reparameterization (gauge choice) that the layer’s linear weights could algebraically absorb. However, by randomizing Pℓ at every forward pass, we convert this gauge freedom into an active constraint. This prevents the model from relying on fixed coordinate choices and restricts the hypothesis class to functions that are permutation-equivariant:

• Feature symmetries. The transformations g apply an identical permutation to the feature coordinates of every xi in the prompt and to the test point. This only relabels the coordinates, so the correct test prediction is unchanged. • Label symmetries. The transformations g relabel target classes by a fixed permutation (e.g., 0 7→ 1 and 1 7→ 0) applied to every yi in the prompt. This only renames the labels, so the correct test label is permuted in the same way. • Permuting examples. The transformations g reorder the context pairs (xi , yi ). This does not change the underlying task, so the correct test prediction is unchanged.

Zℓ+1 = Zℓ +Attn(Zℓ Pℓ ) P⊤ ℓ ; Pℓ ∼ Unif(Sd ×Sk ). (4)

Role of Symmetries in Prediction. Symmetries transform the prompt without changing the underlying prediction task. It is therefore natural to ask whether a predictor respects these symmetries, i.e., whether applying g to the prompt leaves its prediction unchanged (invariance) or transforms it by g (equivariance). For example, empirical risk minimization respects them provided g preserves the model class and the loss: applying g maps any loss-minimizer for p to a corresponding loss-minimizer for g(p).

Here Zℓ is the representation matrix. Pℓ independently permutes the subspaces corresponding to the d input features and k output dimensions. Because the permutation changes constantly, the attention block is forced to process information based on relationships rather than absolute positions. Interpretation via behavioral alignment. We validate this intervention by comparing the original transformer to its symmetrized counterpart across a wide range of problem instances (varying dimension, sample size, and margin). Crucially, we observe a tight match between the two mod-

Why symmetries help interpretability. A trained transformer can implement a specific function via many distinct 4

Layerwise Dynamics for In-Context Classification in Transformers

Figure 3. Transformer symmetrization preserves the learned algorithm. Left: The symmetrized transformer (S) matches the unconstrained (U): on the same prompt, U–S fingerprint matches the U–U baseline across query (local decision rule at the test point), context (which demonstrations drive the logits), and predictions, while the different-prompt control is near zero (mean over 2,048 tasks; averaged over 10 training runs; ±1 s.d. across runs). Right: After averaging predictions across 5 training runs, U and S are functionally identical: ground-truth-class probabilities agree with R2 > 0.96 over 2,048 tasks, yielding the same decision boundaries.

ture (2) on i.i.d. instances with xi ∼ N (0, Id ) and each wk ∼ Unif(Sd−1 ), where Sd−1 is the unit sphere in Rd . Because {wk } vary across instances and are never revealed, achieving low loss requires the model to infer the induced linear decision rules from the in-context examples. See Appx. B.1 for additional training detail.

els not only in final predictive error but also in their input sensitivities. This strong behavioral alignment confirms that enforcing symmetry has not forced the model to switch strategies; rather, it has simply stripped away the arbitrary, symmetry-breaking coordinate choices found in the unconstrained model. Consequently, the symmetrized network serves as a faithful proxy for the original: it implements the same underlying algorithm, but in a canonical representation that is accessible to reverse engineering.

4.1

Step 1: Train a transformer. We first verify that the unconstrained architecture in (2) solves the task. We fix K = 3, d = 7 and n = 64, and train transformers with depths L ∈ {2, 4, 8}. Figure 4 compares their performance to logistic regression and SVM across context lengths. The transformer is consistently competitive and improves with depth, indicating that solving this task requires L > 1 layers.

Case Study: Linear Classification

In this section, we study transformers trained for in-context multi-class linear classification through the lens of symmetries. We apply the symmetry-based interpretability probe in Section 4.1 to extract the transformer’s underlying update rule. We then study the emerging class of classifiers in Section 4.2: we prove convergence in the label-dominated regime, validate the mechanism in simulation, and, guided by the theory, discuss how the same classifier can extend to label noise, semi-supervised prompts, and selected nonlinear tasks.

Step 2: Make the weights easier to interpret. Figure 2 compares a standard transformer trained without additional constraints (2) to a symmetry-enforced transformer trained under our approach (4). The symmetry-enforced model exhibits a markedly more regular weight structure, which makes the learned computation easier to interpret.

Problem setup. Each ICL instance is defined by K latent unit-norm vectors W ≡ (w1 , . . . , wK ) ∈ Rd . Given n + 1 feature vectors x1 , . . . , xn , xtest ∈ Rd , we assign labels by

Behavioral alignment. Crucially, symmetrization changes the parameterization, not the learned algorithm: it removes redundant degrees of freedom while preserving the model’s input–output map and local decision rule. We test this by comparing unconstrained (U) and symmetrized (S) transformers with three “fingerprints,” in the spirit of outputand gradient-level matching (e.g., Akyürek et al., 2023; von Oswald et al., 2023). We compare (i) the predicted groundtruth probability softmax(s(p))ctest , (ii) the query Jacobian ∇xtest s(p) ∈ RK×d capturing the local decision rule around the test input, and (iii) per-example context influence scores S ∈ RK×n where Sc,i = ∥∇xi sc (p)∥2 indicating which demonstrations drive the logits. See Appx. B.2 for details.

ci = arg max ⟨wk , xi ⟩, and k∈[K]

Extracting the Emergent Algorithm

(5)

yi = eci ∈ {0, 1}K . where e1 , . . . , ek are the standard basis vectors in RK , so yi is the one-hot embedding of class ci . In words, each wk is a class-specific direction, and xi is assigned to the class whose direction is most aligned with it. Given the prompt built from examples (xi , yi )i∈[n] and a test feature xtest , the goal is to predict ytest , i.e., determine which class direction is most aligned with xtest , without directly observing the {wk }. We denote s(p) ∈ Rk the transformer’s predicted logit vector for the query point.

Our experiments show that symmetrization preserves the learned algorithm: it leaves not only the predictions, but also the local decision rule and context usage essentially

Transformer training. We train the transformer architec5

Layerwise Dynamics for In-Context Classification in Transformers

Figure 4. Transformers match baselines. Test performance vs. context length (mean over 10,000 tasks; averaged over 5 training runs; ±1 s.d. across runs).

Figure 5. Label-driven mean-shift increases class separation. Left: Simulated dynamics with fixed parameters (α, γ, α′ , γ ′ ) = (1, 5, 0.08, 0.1), showing class centroids drifting apart.

unchanged. In Fig. 3, on the same prompt, U–S is as similar as U–U: the query Jacobian and context-influence fingerprints align strongly (Spearman’s correlation > 0.95 and ≈ 0.8), whereas the different-prompt control has near-zero alignment. Predicted probabilities also match instance-byinstance (R2 > 0.96 after averaging over training seeds; Fig. 3, right). Further, additional experiments show that, on the aggregate level, overall test accuracy agrees across a broad sweep of task parameters (Appx. B.2).

4.2

In the previous section, we identified a family of attentiondriven update rules (7) that the transformer instantiates. We now show that this family contains parameterizations whose induced dynamics implement a classification rule. Label-Driven Regime. To make this concrete, we consider a fixed schedule α, γ, α′ , γ ′ with positive parameters and analyze the corresponding iterations in the regime γ/α ≫ 1. In this regime, the attention matrices are dominated by the term γ Ỹ Y ⊤ , and the resulting dynamics are effectively ‘label-driven.’

Step 3: Extract the algorithm. Using the structured weights from Step 2, we project each learned matrix onto a two-parameter block family matching the pattern in Fig. 2. Concretely, we approximate α Id 0d×K W ≈ . (6) 1 0K×d γ IK − K 1K 1⊤ K

Connection to mean shift. The resulting attention-driven recursion is reminiscent of the classical mean-shift procedure (Fukunaga & Hostetler, 1975). Mean-shift repeatedly moves each point toward a similarity-weighted local average, which causes groups of mutually similar points to collapse to shared representatives (modes). The same contraction appears here. Since each attention row is a probability distribution, Aℓ Xℓ is a weighted barycenter of the context, and the update X̃ℓ+1 = X̃ℓ +αℓ′ Aℓ Xℓ implements an averaging-style drift that merges clusters into prototypes. The crucial difference is that our similarity is supervised: Aℓ depends on both feature and label similarity via the Ỹ -term, so label-consistent points can attract and merge even when they are far in raw feature space. We therefore refer to this family as a coupled mean-shift: it imports the averaging-and-collapse geometry of classical mean shift into a supervised, label-aware classification setting.

Appendix B.2 shows that this abstraction closely matches the model’s implemented inference behavior. For layer ℓ, we denote the parameters of WQK,ℓ by (αℓ , γℓ ) and those of WVP,ℓ by (αℓ′ , γℓ′ ). Weights of the form (6) commute with the layerwise permutations in the symmetrized architecture, so the permutation sandwich leaves the computation unchanged. Plugging this weight abstraction into the transformer update (2) then yields the two-line dynamics (see Appx. B.2 for detail) X̃ℓ+1 = X̃ℓ + αℓ′ Aℓ Xℓ , Ỹℓ+1 = Ỹℓ + γℓ′ Aℓ Yℓc ,

(7)

Classwise separation. In the label-driven limit, attention over the context data becomes class-conditional: each point averages only with others sharing its label, giving a blockdiagonal attention matrix and partitioning the data into classwise clusters. This pushes class centroids apart in a way governed by the geometry of the class vectors {wk }, increasing the separability of the embeddings (Theorem C.1).

where the attention from all tokens to the training tokens is Aℓ = σ αℓ X̃ℓ Xℓ⊤ + γℓ Ỹℓc (Yℓc )⊤ ∈ R(n+1)×n (8) with σ the row-wise softmax. Here, X̃, Ỹ contain all features and labels as defined in (1) and we set Xℓ ∈ Rn×d and Yℓ ∈ Rn×K to denote the first n rows (in-context examples only). Labels are centered across classes via 1 Yℓc = Yℓ I − K 11⊤ )

Analysis: The Coupled Mean-Shift Mechanism

Test Point and Prediction Dynamics. While the training data undergoes the dynamics above, the initial movement of the test point xtest (ℓ), and the prediction ytest (ℓ) is not label-driven even in the above limit. This is because the initial ytest (0) is the zero vector, and so does not af-

1 and Ỹℓc = Ỹℓ I − K 11⊤ ). (9)

Thus Aℓ combines feature similarity and label agreement, and the same weights propagate both features and labels. 6

Layerwise Dynamics for In-Context Classification in Transformers

fect the “test-attention” (Atest ℓ )i ∝ exp(α⟨xtest (ℓ), xi (ℓ)⟩ + γ⟨ytest (ℓ), yi (ℓ)⟩)i . Instead, this dynamics is dominated by the interaction of xtest (0) with the training datapoints. As our main theoretical result, we argue that under explicit pointwise train/test separation conditions, the dynamics of the test point drive xtest (ℓ) towards the cluster of the correct label, and further drive its prediction ytest (ℓ) towards this label. Theorem 4.1. Let c∗ be the correct class prediction for xtest . Define the geometric alignments Rℓ := min ∗ ⟨xtest (ℓ), xi (ℓ)⟩, and

Figure 7. Mean-shift in Voronoi classification. Top: Simulation for (α, γ, α′ , γ ′ ) = (1, 5, 0.05, 0.2). Bottom: Trained transformer weights encode the same dynamics (7), matching Fig. 2.

i: ci =c

Lℓ := max∗ ⟨xtest (ℓ), xi (ℓ)⟩. i: ci ̸=c

this, in Appendix B.3 we recover the same extracted iteration even when each label is independently flipped to a uniformly random incorrect class with 30% probability.

Also define the test margin ∆ℓ = Rℓ − Lℓ , the training e ℓ := margin Γℓ as in Lemma 3, and the effective margin ∆ min(∆ℓ , Γℓ ). If the class-clusters are well-separated in e 0 > 0, and α∆0 ≥ the pointwise sense of Lemma 3, ∆ e 0 ), then log(K) + log(1 + 2/∆

Insights for Semi-Supervised Prediction. The analysis reveals a simple mechanism: a small set of labeled points steers the class prototypes, while unlabeled (and test) points move by geometric alignment, effectively “latching onto” the appropriate cluster. As a result, if unlabeled points concentrate around the class clusters and at least a few labels are present to anchor them, the same dynamics can propagate those labels through the geometry and correctly classify the test point whenever it aligns with the right cluster. In Section 5, we study this semi-supervised regime in detail.

1. The geometric margin diverges as ∆ℓ ≥ ∆0 (1 + α′ )ℓ . 2. Further, the label margin grows exponentially: For all ℓ ≥ ⌈(α′ )−1 log log(4K) α∆0 ⌉, ∀c ̸= c∗ , (ytest (ℓ))c∗ − (ytest (ℓ))c ≥ (1 + γ ′ )ℓ /2. In words, (R0 , L0 ) capture how well xtest aligns with points in the correct and incorrect classes, while Γ0 captures the corresponding pointwise training separation. The effective e 0 controls the assumption needed to keep the quantity ∆ train/test recursion favorable. Notice that if even if ∆0 > 0 is arbitrarily small, a large enough α enables the condition above: thus, as for large α, any initial bias in xtest towards the correct class-cluster is amplified exponentially through the geometry driven phase of this dynamics. Further, interpreting (ytest (ℓ))c∗ as logits of the final prediction, after an initial burn in period determined by α′ , these are driven exponentially as well. Figure 6 illustrates this dynamics.

Nonlinear Classification. Notice that in the label-driven regime, so long as the initial class-cluster means start out distinct, the dynamics will amplify their separation, even if the class distributions are anisotropic, or classes are separated by nonlinear boundaries. This in fact suggests that the same family of dynamics may be successful at classifying much richer problems than linear classification. We discuss the case of classifying into Voronoi cells in Section 4.3 below, and explore more settings in Appx. D. 4.3

Classification to Voronoi Cells

So far we have focused on the multiclass linear classification problem, where each class is defined by a fixed direction. We now consider nearest-centroid (Voronoi) classification where each class corresponds to a prototype, and the classwise decision boundaries are piecewise linear.

Figure 6. A simulation of the mean-shift trajectories with constant parameters (α, γ, α′ γ ′ ) = (1, 5, 0.08, 0.1). The test point ‘follows’ a trailing path towards its cluster.

Setup. We sample K centroids {ck }K k=1 and n feature vectors {xi }ni=1 independently from N (0, Id ) and assign labels by the nearest centroid (see Appx. B.4 for more detail),

Robustness to label noise. The analysis suggests robustness to label noise because the recursion updates points using averages over many context tokens: correctly labeled points in a class reinforce one another, while randomly flipped labels contribute mostly canceling noise and have limited influence on the resulting class prototypes. Consistent with

yi = arg min ∥xi − ck ∥2 , k∈[K]

(10)

which partitions the plane into K Voronoi cells. Results. Figure 7 (top) runs our extracted mean-shift itera7

Layerwise Dynamics for In-Context Classification in Transformers

tion on Voronoi data and shows that it correctly recovers the nearest-centroid clusters. This shows the mechanism is not specific to linear-score classifiers: the same mean-shift dynamics extends to prototype-based, piecewise-linear tasks. Consistent with this, Figure 7 (bottom) shows that a symmetrized transformer trained on the Voronoi classification task learns the same mean-shift classifier.

Transformers as Semi-Supervised In-Context Learners

We now move from fully supervised prompts to a semisupervised setting: each prompt contains a small set of labeled examples and many additional unlabeled points (their label token is set to a null value, e.g. 0). The question is whether a transformer can use these unlabeled points to improve classification, and why.

Figure 8. Transformers use unlabeled context to improve semisupervised in-context learning. Top: Each prompt has 8 labeled examples; longer contexts add unlabeled points only. We plot accuracy vs. context length (mean over 10,000 tasks; averaged over 3 runs; ±1 s.d. across runs is imperceptible). Bottom: Trained transformer weights encode the same dynamics (7), matching Fig. 2.

Our analysis of the emergent mean-shift classifier motivates this question and yields a simple prediction: in the early, geometry-driven phase, the test point’s behavior is controlled by how much it aligns with the correct versus incorrect clusters (the geometric margin ∆ℓ in Theorem 4.1). Adding unlabeled points from the same clusters sharpens this geometry, increasing the effective separation seen by the test point even without adding labels, and thereby makes the subsequent label-driven amplification more reliable.

racy gains as the context size n grows (rising from ∼ 72% to > 85%). This confirms that the model is not merely smoothing over the labeled examples, but actively aggregating geometric information from the unlabeled tokens to refine its decision boundary. Moreover, the weights of these transformers again encode the same mean-shift mechanism as in Sec. 4.1 (Fig. 8, bottom), showing that this motif extends to SSL. Appendix B.5 provides further details and an ablation showing that the gains come from the unlabeled tokens themselves, not from extra tokens increasing the transformer’s capacity or compute.

Setup. We start from the same multi-class linear classification instance with examples (xi , yi )ni=1 and a query xtest . We construct a semi-supervised variant in two steps. First, we make the class structure more pronounced by shifting each feature vector slightly in the direction of its class: xi ← xi + η wci , η > 0. This increases the geometric separation between classes by moving points toward their class direction, strengthening within-class coherence and between-class margins. Second, we keep remove the labels for all but nlab examples, yielding nunlab = n − nlab unlabeled inputs.

Transformer vs. classical SSL. Figure 8 highlights a sharp gap with standard graph-based SSL. Despite the geometric structure in the context, Label Spreading (purple/red) performs worse than simple supervised baselines. This phenomenon is consistent with known failure modes on highdimensional Gaussian data where Euclidean neighborhoods are fragile without careful kernel tuning. In contrast, the transformer reliably uses the unlabeled points and outperforms these SSL baselines as n grows, indicating that it has learned a more robust way to aggregate geometric signal.

Benefit of Unlabeled Context (Fixed Label Budget). To strictly isolate the model’s ability to leverage unlabeled geometry, we perform a controlled experiment where nlab = 8 is fixed to a scare budged, while nunlab varies in {0, 8, 24, 120}. We fix K = 3, d = 7, and L = 5, and train a single transformer with data shift parameter η = 0.5. With this choice of η, the classes are separable but the margin is small relative to the within-class spread, so the unlabeled geometry is only weakly informative and the SSL task is challenging.

Conclusion

In this work, we make in-context classification more identifiable by enforcing the task’s natural symmetries at every layer. This structure lets us extract an explicit attentiondriven recursion from the transformer. Strong behavioral alignment shows that the extracted dynamics closely matches the transformer’s actual test-time computation, giving a concrete characterization of its mechanism. We further

Figure 8 presents the results. Standard supervised baselines (Logistic Regression, SVM; dashed lines) cannot utilize the unlabeled points, and show flat performance curves. In stark contrast, the Transformer (Blue) achieves significant accu-

Layerwise Dynamics for In-Context Classification in Transformers

References

show that this family of dynamics can implement a labelaware, mean-shift-like update that pulls points toward class prototypes. The same mechanism also appears in more complex settings, including Voronoi-cell classification and semi-supervised classification, where we uncover a surprising ability of in-context learning transformers to benefit from unlabeled context examples.

Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36:45614–45650, 2023. Ahuja, K. and Lopez-Paz, D. A closer look at incontext learning under distribution shifts. arXiv preprint arXiv:2305.16704, 2023. URL https://arxiv. org/abs/2305.16704.

Our analysis focuses on synthetic tasks with clear, explicit symmetries. Many practical applications obey analogous symmetries; systematically identifying and cataloging them is left for future work. More broadly, our method relies on enforcing symmetries layer by layer, which makes it especially well suited to recovering permutation-commuting mechanisms. Extending this approach to in-context learning mechanisms that are not well captured by layer-wise symmetries is an important direction for future work.

Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. What learning algorithm is in-context learning? investigations with linear models. In International Conference on Learning Representations, 2023. URL https: //openreview.net/forum?id=0g0X4H8yN4I. Chan, S. C. Y., Santoro, A., Lampinen, A. K., Wang, J. X., Singh, A. K., Richemond, P. H., McClelland, J. L., and Hill, F. Data distributional properties drive emergent in-context learning in transformers. CoRR, abs/2205.05055, 2022. URL https://arxiv.org/ abs/2205.05055.

Impact Statement This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

Cohen, T. and Welling, M. Group equivariant convolutional networks. In Proceedings of the 33rd International Conference on Machine Learning, 2016. Comaniciu, D. and Meer, P. Mean shift: A robust approach toward feature space analysis. IEEE Transactions on Pattern Analysis and Machine Intelligence, 24(5):603– 619, 2002. Dai, D., Zhao, Y., Kang, L., Li, D., Zhou, T., and Tao, D. Why can GPT learn in-context? language models implicitly perform gradient descent as metaoptimizers. In Bouamor, H., Pino, J., and Bali, K. (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 729–741, Singapore, December 2023. Association for Computational Linguistics. doi: 10.18653/v1/2023.findings-emnlp. 54. URL https://aclanthology.org/2023. findings-emnlp.54. Deutch, D., Friedman, D., Zhou, D., and Bartal, Y. Incontext learning and gradient descent revisited. In Duh, K., Gomez, H., and Bethard, S. (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 913–928, Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10.18653/v1/ 2024.naacl-long.51. URL https://aclanthology. org/2024.naacl-long.51. Fukunaga, K. and Hostetler, L. The estimation of the gradient of a density function, with applications in pattern 9

Layerwise Dynamics for In-Context Classification in Transformers

recognition. IEEE Transactions on information theory, 21(1):32–40, 1975. Garg, S., Tsipras, D., Liang, P., and Valiant, G. What can transformers learn in-context? A case study of simple function classes. CoRR, abs/2208.01066, 2022. doi: 10.48550/ARXIV.2208.01066. URL https://doi. org/10.48550/arXiv.2208.01066.

Panwar, M., Ahuja, K., and Goyal, N. In-context learning through the bayesian prism. In International Conference on Learning Representations, 2024. URL https:// openreview.net/forum?id=HX5ujdsSon. Rigollet, P. The mean-field dynamics of transformers. arXiv preprint arXiv:2512.01868, 2025. URL https: //arxiv.org/abs/2512.01868.

Geshkovski, B., Letrouit, C., Polyanskiy, Y., and Satorras, V. G., Hoogeboom, E., and Welling, M. E(n) Rigollet, P. The emergence of clusters in selfequivariant graph neural networks. In Meila, M. and attention dynamics. In Advances in Neural InZhang, T. (eds.), Proceedings of the 38th International formation Processing Systems, volume 36, 2023. Conference on Machine Learning, volume 139 of ProURL https://proceedings.neurips. ceedings of Machine Learning Research, pp. 9323–9332. cc/paper_files/paper/2023/hash/ PMLR, 2021. URL https://proceedings.mlr. b2b3e1d9840eba17ad9bbf073e009afe-Abstract-Conference. press/v139/satorras21a.html. html. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, Han, C., Wang, Z., Zhao, H., and Ji, H. Understanding emerL., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Atgent in-context learning from a kernel regression perspectention is all you need. Advances in neural information tive. Transactions on Machine Learning Research, 2025, processing systems, 30, 2017. 2025. URL https://openreview.net/forum? von Oswald, J., Niklasson, E., Schlegel, M., Kobayashi, id=6rD50Q6yYz. S., Zucchet, N., Bach, F., and Flammarion, N. TransLee, J., Lee, Y., Kim, J., Kosiorek, A. R., Choi, S., and Teh, formers learn in-context by gradient descent. In ProY. W. Set transformer: A framework for attention-based ceedings of the 40th International Conference on Mapermutation-invariant neural networks. In Proceedings of chine Learning, volume 202 of Proceedings of Mathe 36th International Conference on Machine Learning, chine Learning Research, pp. 35151–35174. PMLR, 2019. 2023. URL https://proceedings.mlr.press/ v202/von-oswald23a.html. Li, X., Azar, B., Forsyth, D., and Dai, H. Transformers as algorithms: Generalization and stability in in-context Wies, N., Levine, Y., and Shashua, A. The learnlearning. In Proceedings of the 40th International Conferability of in-context learning. In Advances in ence on Machine Learning, volume 202 of Proceedings Neural Information Processing Systems, 2023. of Machine Learning Research, pp. 19565–19581. PMLR, URL https://proceedings.neurips. cc/paper_files/paper/2023/file/ 2023. URL https://proceedings.mlr.press/ v202/li23l.html. 73950f0eb4ac0925dc71ba2406893320-Paper-Conference. pdf. Lutz, P., Gangrade, A., Daneshmand, H., and Saligrama, V. Linear transformers implicitly discover unified numerical Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An exalgorithms. arXiv preprint arXiv:2509.19702, 2025. planation of in-context learning as implicit bayesian inference. In International Conference on Learning RepresenMaron, H., Ben-Hamu, H., Shamir, N., and Lipman, tations, 2022. URL https://openreview.net/ Y. Invariant and equivariant graph networks. In Inforum?id=RdJVFCHjUMI. ternational Conference on Learning Representations, 2019. URL https://openreview.net/forum? Xu, H., Xiang, L., Ye, H., Yao, D., Chu, P., and id=Syx72jC9tm. Li, B. Permutation equivariance of transformers and its applications. In Proceedings of the Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, IEEE/CVF Conference on Computer Vision and Pattern N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, Recognition (CVPR), pp. 5987–5996, June 2024. A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, URL https://openaccess.thecvf.com/ Z., Hernandez, D., Johnston, S., Jones, A., Kernion, content/CVPR2024/html/Xu_Permutation_ J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Equivariance_of_Transformers_and_Its_ Clark, J., Kaplan, J., McCandlish, S., and Olah, C. InApplications_CVPR_2024_paper.html. context learning and induction heads. arXiv preprint arXiv:2209.11895, 2022. URL https://arxiv. Yadlowsky, S., Doshi, L., and Tripuraneni, N. Preorg/abs/2209.11895. training data mixtures enable narrow model selection 10

Layerwise Dynamics for In-Context Classification in Transformers

capabilities in transformer models. arXiv preprint arXiv:2311.00871, 2023. URL https://arxiv. org/abs/2311.00871. Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R., and Smola, A. Deep sets. In Advances in Neural Information Processing Systems, 2017. Zhang, T., Frei, S., and Bartlett, P. L. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25:49:1–49:55, 2024. URL https: //jmlr.org/papers/v25/23-1042.html. Zhang, Y., Zhang, F., Yang, Z., and Wang, Z. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305.19420, 2023. URL https:// arxiv.org/abs/2305.19420. Zhou, D., Bousquet, O., Lal, T., Weston, J., and Schölkopf, B. Learning with local and global consistency. Advances in neural information processing systems, 16, 2003. Zhou, D., Zhou, L., Gu, J., Luo, T., Yu, Z., Dong, Y., Hu, J., Liu, Q., Li, X., Shan, Y., Wang, Z., and Zhao, D. The mystery of in-context learning: A comprehensive survey on interpretation and analysis. In Al-Onaizan, Y., Bansal, M., and Chen, Y.-N. (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 7045–7080, Miami, Florida, USA, November 2024. Association for Computational Linguistics. URL https://aclanthology.org/ 2024.emnlp-main.401. Zhu, X. and Ghahramani, Z. Learning from labeled and unlabeled data with label propagation. Technical report, Carnegie Mellon University, 2002.

Layerwise Dynamics for In-Context Classification in Transformers

Additional Related Work

Emergence of ICL from pretraining distributions. Beyond controlled “train-from-scratch” studies, several works emphasize that the pretraining data distribution can strongly influence whether and how ICL emerges. In particular, properties such as task-family mixtures, long-range coherence, or “burstiness” can shape the implicit meta-learning behavior that appears at inference time (Chan et al., 2022; Xie et al., 2022; Yadlowsky et al., 2023). This line of work is complementary to ours: we hold the task distribution fixed and instead use symmetry to make the learned inference procedure identifiable and comparable across perturbations. Bayesian and kernel perspectives on ICL. Several works interpret ICL as approximate Bayesian inference or Bayesian model averaging in latent-variable task models, providing statistical characterizations of the in-context predictor and its dependence on prompt length, depth, and pretraining distributions (Xie et al., 2022; Panwar et al., 2024; Zhang et al., 2023). A closely related view explains ICL behavior through kernel regression: attention weights act as similarity kernels, yielding predictors that resemble kernel smoothing and nearest-neighbor-like rules (Han et al., 2025). These perspectives are compatible with our findings at the level of similarity-based computation, but our extracted dynamics go beyond a single-shot kernel predictor: the model iteratively updates representations (and label-related state) across layers, producing a coupled diffusion/mean-shift process whose analysis yields concrete geometric predictions. Mechanistic interpretability and circuit analyses of ICL. Mechanistic interpretability work has identified specific transformer circuits implicated in ICL, notably the emergence of induction heads as a mechanism for copying and pattern completion that correlates with ICL capability (Olsson et al., 2022). Survey work synthesizes broader mechanistic and theoretical hypotheses for ICL across settings and model scales (Zhou et al., 2024). Our approach is complementary: rather than isolating a local circuit, we use symmetry to make the entire forward-pass computation identifiable in a controlled classification setting, enabling closed-form extraction and dynamical analysis.

Supplemental material on Empirical Protocols

B.1

Supplemental material on transformer training

Unless otherwise stated, all transformer experiments use K = 3 classes, feature dimension d = 7, and context length n = 64. We train with Adam at learning rate 10−3 , β1 = 0.9, β2 = 0.999, weight decay 0, batch size 8192, for 20,000 optimizer steps. We use no gradient clipping, dropout, label smoothing, or other regularization. The training loss is the cross-entropy between the ground-truth class of the query point and the transformer-predicted logits for that query. Training prompts are sampled online at each optimization step rather than from a fixed finite dataset. Accordingly, training is best viewed as stochastic optimization of the expected loss under the task-generating distribution, with each minibatch providing a Monte Carlo estimate of the population objective. 1 1 All weight matrices W ∈ R(d+K)×(d+K) are initialized entrywise independently as Wij ∼ Unif[− d+K , d+K ]. This uses a smaller initialization scale than standard Xavier-style schemes, which we found to improve training stability in our setup.

For the symmetrized model, during training and for each layer ℓ and batch element b, we sample independent random (x) (y) permutation matrices on the feature and label blocks, Pℓ,b ∈ Rd×d and Pℓ,b ∈ RK×K , and form (x) (y) Pℓ,b = diag Pℓ,b , Pℓ,b .

The layer update is then (b) (b) (b) ⊤ Zℓ+1 = Zℓ + Attn Zℓ Pℓ,b Pℓ,b . Thus, attention is computed on a block-permuted representation and the resulting update is mapped back by the inverse permutation before the residual addition. Empirically, the learned weights approximately commute with these permutations, so the sandwich becomes unnecessary after training; we therefore apply symmetrization only during training and omit it at evaluation time. Random seeds are sampled independently, and we average over 5 runs unless otherwise specified. B.2

Supplemental material on algorithm extraction

Influence and Function Alignment. We evaluate whether Unconstrained (U) and Symmetrized (S) transformers implement the same inference rule by probing (i) what in the prompt they rely on and (ii) what they output. Concretely, we track 12

Layerwise Dynamics for In-Context Classification in Transformers

three signals, each computed from the predicts query-point logit vector s(p) ∈ RK : (1) Query gradient (local decision rule at test time). We compute the gradient of each logit with respect to the query features, ∇xtest s(p) ∈ RK×d . (11) This probes the model’s local geometry: two models with aligned query gradients are locally sensitive to the same directions in input space, i.e., they induce a similar local decision boundary around xtest . (2) Context-gradient sensitivity (influence over demonstrations). For each context token i ≤ n−1 and class k, we measure how much the predicted logit changes under a small perturbation of that context example’s features, Sc,i = ∇xi s(p) c 2 . (12) This yields a map in S ∈ RK×n that functions like an influence / saliency map over the prompt: high values indicate which demonstrations the model is locally using for which test point class logits. (3) Prediction discrepancy (functional agreement). We compare outputs directly by looking at the probability assigned to the ground-truth test class, softmax(s(p)) ctest , (13) and report the squared differences in ptrue across models. This is the most direct measure of functional equivalence (up to calibration). For the two gradient-based probes (context sensitivity and query gradient), we flatten each per-prompt tensor and report agreement via Spearman and Pearson correlations. We report all three signals under three comparisons: (i) U–S (same prompt) is the main cross-family test of alignment; (ii) U–U (same prompt) measures training seed-to-seed variability within the unconstrained family; and (iii) A–B (same model) is a negative control using two independently sampled prompts, which should destroy prompt-specific influence structure (and increase prediction discrepancy), ruling out trivial correlations driven by architecture or preprocessing. Generalization across problem instances. While the main text demonstrates algorithmic equivalence via output correlation and internal attention dynamics, we also performed an extensive sweep to ensure this holds across problem settings. Figure 9 confirms that the unconstrained and symmetrized models achieve indistinguishable accuracy curves across varying the number of layers L ∈ [2, 16], context length n ∈ [8, 40], data dimension d ∈ [6, 18], and margin η ∈ [0, 1]. Both models track these task changes nearly identically, consistent with them implementing the same algorithmic behavior. Evidence that the weight abstraction is faithful. Step 3 argues that a low-parameter, diagonal weight abstraction preserves the transformer’s input–output behavior. We test this by progressively compressing the symmetrized model’s weights and measuring alignment with the original, unconstrained transformer: 1. Four clusters recover full accuracy. Figure 10 (left) shows that, in the symmetrized model, clustering each weight matrix into k = 4 groups of equal coefficients already matches the model’s accuracy. The clusters correspond to the four structural regions: the top-left block diagonal, the bottom-right block diagonal, the bottom-right background, and the zero background. Using this structure yields a three-parameter-per-layer abstraction, α Id 0d×K W ≈ , (14) 0K×d γ IK + δ1K 1⊤ K and Figure 10 (right) shows that its input–output alignment with the original transformer is essentially identical to that of the symmetrized model—i.e., the abstraction introduces no additional behavioral deviation beyond symmetrization. 2. Two parameters per layer suffice. Figure 11 shows that the transformer consistently learns a bottom-right background 1 coefficient of approximately δ ≈ − K . We therefore fix this parameter and reduce the parameterization to two per layer which gives the abstraction used in the main text. Figure 10 (left) shows that this two-parameter model still closely matches the original transformer’s outputs (R2 ≈ 0.9), with only a modest drop from R2 ≈ 0.96. This supports that the abstraction remains behaviorally faithful while being substantially simpler. 13

Layerwise Dynamics for In-Context Classification in Transformers

Figure 9. We compare predictive performance of unconstrained and symmetry-preserving transformers across problem instances (dimension, sample size, data margin) and model depths. Both architectures achieve the same performance throughout.

Figure 10. Weight abstraction preserves the symmetric transformer’s input–output behavior. Left: Clustering the learned weight matrices (Figure 2) reveals a simple, low-dimensional structure. In the symmetry-preserving model, summarizing each matrix with just four parameters retains essentially the full predictive performance. Right (three panels): Alignment with the original unconstrained transformer for (left to right) the symmetrized model, the weight-abstracted dynamics with three layer-wise hyperparameters (no explicit label centering), and the further-simplified dynamics with two layer-wise hyperparameters (with explicit label centering).

Developing transformer iteration under weight abstraction. To pass from the transformer update (2) to the recursion (7), define ⊤ WQ,ℓ WK,ℓ WQK,ℓ := √ , WVP,ℓ := WV,ℓ WP,ℓ . d+K Writing Zℓ = [ X̃ℓ Ỹℓ ], the full attention score matrix is Zℓ WQK,ℓ Zℓ⊤ + M . 1 Under the block approximation (6), WQK,ℓ acts as αℓ Id on features and as γℓ (IK − K 11⊤ ) on labels, so

Zℓ WQK,ℓ Zℓ⊤ + M ≈ αℓ X̃ℓ X̃ℓ⊤ + γℓ Ỹℓc (Ỹℓc )⊤ + M . Since M sets the last column (the query token) to −∞, the softmax assigns zero attention to that column. Thus attention is effectively taken only over the first n training tokens, giving Aℓ = σ αℓ X̃ℓ Xℓ⊤ + γℓ Ỹℓc (Yℓc )⊤ ∈ R(n+1)×n . Likewise, the value/output map reduces to Xℓ

Yℓ WVP,ℓ ≈ αℓ′ Xℓ 14

γℓ′ Yℓc .

Layerwise Dynamics for In-Context Classification in Transformers

Figure 11. Transformer learns label centering. Parameter δℓ as in intermediary weight abstraction (1) motivates fix1 ing δ = − K . Median and std. in symmetrized transformer over 5 training runs.

Figure 12. Accuracy on the noisy linear classification task as a function of context size, comparing the transformer (mean over 3 runs, L = 5, trained with n = 64) to logistic regression and SVM. Right: Ground-truth-class probability for the query point: symmetrized transformer vs. unconstrained transformer on the same noisy linear classification task, mean across 3 training seeds are shown.

Figure 13. Learned weight matrices for the unconstrained transformer (left) and the symmetry-preserving transformer (right) trained on noisy linear classification. The unconstrained model exhibits little visible structure, whereas enforcing the task symmetries produces a more regular pattern that is easier to interpret. The top and bottom rows show WQK,ℓ and WV P,ℓ , respectively.

Substituting into the residual update and separating feature and label blocks yields X̃ℓ+1 = X̃ℓ + αℓ′ Aℓ Xℓ ,

Ỹℓ+1 = Ỹℓ + γℓ′ Aℓ Yℓc ,

which is exactly (7). B.3

Supplemental material on the robustness to noise

In the main text, we extracted a mean-shift-style classifier from transformers trained on our synthetic tasks. Here we show that the same qualitative mechanism persists under substantial label noise: each labeled example is independently flipped to a uniformly random incorrect class with probability p = 0.3. Figure 12 (Left) reports the unconstrained transformer’s accuracy under noise across context sizes, showing performance comparable to standard baselines. Figure 12 (Right) compares the predicted probability of the ground-truth class for unconstrained versus symmetrized models; the close agreement indicates that the symmetry projection does not change the learned predictor. Finally, Figure 13 visualizes the learned weights under label noise. The dominant structure that emerges closely matches the noiseless case: weights again cluster into the same diagonal groups with the same functional roles in the update. This supports the interpretation from the main text that the transformer is still implementing the same latent algorithm—an attention-induced, coupled mean-shift classifier—even when supervision is substantially corrupted. B.4

Supplementary material on the Voronoi cell problem

In the main text, we showed that the extracted mean-shift classifier can solve the Voronoi cell classification problem. Here we provide additional experimental details confirming that (i) a transformer trained end-to-end also solves this task reliably, and (ii) its behavior is consistent with the same mean-shift-style iteration we identified in the linear classification setting. 15

Layerwise Dynamics for In-Context Classification in Transformers d Task construction. Each ICL instance is defined by K randomly sampled sites (centroids) {ck }K k=1 ⊂ R . For each batch element we draw ck ∼ N (0, Id ), xi ∼ N (0, Id ),

independently across k and i. Each point is labeled by the nearest site in Euclidean distance, yi = arg min ∥xi − ck ∥2 . k∈[K]

This induces a Voronoi partition of Rd with class regions determined jointly by the directions and magnitudes of the sites. For any pair of centroids, the tie set {x : ∥x − ca ∥2 = ∥x − cb ∥2 } is a hyperplane, so the classifier is a collection of linear boundaries. As before, the transformer input is a context of n feature–label pairs (xi , onehot(yi )) and a query point xtest ; the model is trained to predict the query label. Transformer performance and baselines. Figure 14 (Left) reports accuracy for K = 5 classes as a function of context length, comparing an unconstrained transformer to a nearest-neighbor baseline. The transformer consistently exceeds the nearest-neighbor baseline, confirming that it learns a nontrivial in-context strategy for the classification problem. Agreement between unconstrained and symmetrized models. Figure 14 (Right) compares the predicted class probabilities produced by the unconstrained and symmetrized transformers on the same test instances. The two models closely agree, indicating that enforcing the symmetry constraints does not change the learned solution, but instead exposes it more cleanly for analysis. Evidence for the same underlying algorithm. Finally, Figure 15 visualizes the learned weights with and without symmetrization. The qualitative patterns mirror those observed in our linear classification experiments: the same cluster structure and the same functional roles of weight groups reappear. This supports the interpretation from the main text that the transformer implements essentially the same coupled mean-shift iteration, now operating on the geometry induced by nearest-centroid (Voronoi) labels.

Figure 14. Left: Accuracy on the Voronoi (nearest-centroid) classification task as a function of context size, comparing the transformer (mean over 3 runs, L = 5, trained with n = 64) to nearest-neighbor baselines. Right: Ground-truth-class probability for the query point: symmetrized transformer vs. unconstrained transformer on the same Voronoi task, mean across 3 training seeds are shown.

Figure 15. Learned weight matrices for the unconstrained transformer (left) and the symmetry-preserving transformer (right) trained on in-context Voronoi cell classification. The unconstrained model exhibits little visible structure, whereas enforcing the task symmetries produces a more regular pattern that is easier to interpret. The top and bottom rows show WQK,ℓ and WV P,ℓ , respectively.

Layerwise Dynamics for In-Context Classification in Transformers

Figure 16. Left: Ground-truth query-class probability predicted by the symmetrized vs. unconstrained transformer on the same semisupervised classification task (mean over 3 training seeds). Right: Ablation with Unstructured Noise. When unlabeled context points are replaced with random noise N (0, Id ), the Transformer (Blue) ceases to outperform the supervised Linear SVM baseline (Green dashed). Compare this to Figure 8, where the Transformer leverages geometric structure to exceed the supervised baseline by > 10%. This confirms the model relies on the geometry of the unlabeled data to improve classification.

B.5

Supplementary material on the semi-supervised classification setting

In the main text, we showed that the transformer uses unlabeled context points to improve query prediction under a fixed label budget. Here we show that it does so via the same coupled mean-shift iteration identified in the fully supervised setting. Figure 16 (Left) compares the probability assigned to the ground-truth query class by the unconstrained and symmetrized transformers on the exact same inputs. The two predictions match closely across a wide range of problem instances, indicating that the models are functionally indistinguishable; symmetrization does not change the underlying computation, it simply makes it easier to expose and analyze. Figure 17 provides complementary evidence at the parameter level: the learned weights exhibit the same cluster structure as in our linear classification experiments. Together, these results support the interpretation that the transformer implements essentially the same coupled mean-shift iteration in the semi-supervised regime, now operating jointly on the feature geometry and the partially observed labels. Semi-Supervised Baselines. To evaluate geometric utilization (Figure 8), we compare the Transformer against Label Spreading (Zhou et√al., 2003), a graph-based semi-supervised algorithm. We utilize the k-nearest neighbor (k-NN) kernel with k = max(5, ⌊ n⌋). This kernel choice makes the baseline scale-invariant. The “Supervised” baseline models are a Logistic Regression and a linear SVM model trained only on the subset of labeled examples, ignoring the unlabeled tokens. Ablation: The Role of Geometric Structure. To verify that the performance gains in Figure 8 stem strictly from geometric aggregation (the mean-shift mechanism) rather than architectural capacity, we conduct an ablation where the unlabeled context tokens are replaced with unstructured Gaussian noise, zi ∼ N (0, Id ). Figure 16 (Right) shows the results. In this unstructured setting, the Transformer (Blue line) no longer surpasses the supervised ceiling. Instead, it merely converges to the performance of the Linear SVM (Green dashed line). Crucially, the ”super-supervised” lift observed Figure 8 (where accuracy exceeded 85% vs. the SVM’s 75%) completely vanishes. This confirms that the model is actively reading the manifold structure of the unlabeled data; when that structure is removed, the mean-shift mechanism defaults to a standard supervised solution, proving the gains are geometry-driven.

Analysis of the Mean-Shift Emergent Classifier

We have identified a family of attention-driven update rules that take the form: Xℓ+1 = Xℓ + αℓ′ Aℓ X̃ℓ , Yℓ+1 = Yℓ + γ ′ Aℓ Ỹℓ , 17

Layerwise Dynamics for In-Context Classification in Transformers

Figure 17. Learned weight matrices for the unconstrained transformer (left) and the symmetry-preserving transformer (right) trained on semi-supervised in-context linear classification. The unconstrained model exhibits little visible structure, whereas enforcing the task symmetries produces a more regular pattern that is easier to interpret. The top and bottom rows show WQK,ℓ and WV P,ℓ , respectively.

where X0 , Y0 are defined as in (1) and X̃ℓ ∈ Rn×d and Ỹℓ ∈ Rn×K are the first n rows of Xℓ and Yℓ , respectively. The matrix At ∈ R(n+1)×n is the row-stochastic softmax attention matrix: Aℓ = σ(αℓ Xℓ X̃ℓ⊤ + γℓ Yℓ Ỹℓ⊤ ). The transformer instantiates one member of this family during training by selecting the scalars {αℓ , γℓ , αℓ′ , γℓ′ }ℓ∈[L] . In the following, we are going to show that this family of updates is a novel approach to classification with surprising behavior and versatility. Our analysis exposes two key properties of the emergent mean-shift classifier: (1) it reinforces intra-class separation in the training data (Theorem C.1), and (2) it drives the test point towards the correct cluster, assuming it is initialized with some proximity towards that cluster (Theorem C.2) C.1

Class and Cone Invariances

For a point xi let us define as ci ∈ [K] the class in which it is originally placed, i.e. yi = eci . For each class c ∈ [K], we define its population as Sc := {i : yi = ec }. First, we focus on the evolution of the labels Yt . We show that in the regime γ ≫ α the labels of the training points all remain pointing to their original class. Lemma 1. Let κ = γ/α. For all t ≥ 0, the attention weight on any incorrect class index k ∈ / Sc(j) decays exponentially with κ: (t)

e (t) Ajk ≤ e−ακη ·

(t)

α⟨xj ,xk ⟩

(t)

Zj,in

where Zj,in is the partition function restricted to the correct class. Consequently, as κ → ∞, the ‘leakage’ of attention weights across class boundaries vanishes and the dynamics recover the exact label evolution yj = (1 + γ ′ )t · ec(j) . Proof. We first analyze the structure of the attention matrix A(t) in terms of the ratio κ = γ/α. Let Sc(j) = {k ∈ [n] : c(k) = c(j)} be the set of indices in the same class as j. We assume a label separation margin ⟨yj , yk ⟩ = η > 0 for k ∈ Sc(j) and 0 otherwise. Substituting γ = κα, the partition function is: X X Zj = eακη eα⟨xj ,xs ⟩ + eα⟨xj ,xs ⟩ . s∈Sc(j)

s∈S / c(j)

We analyze the probability mass Ajk in two cases: Case 1: Different Class (k ∈ / Sc(j) ). The attention weight is strictly bounded by the label margin κη. Dividing the numerator and denominator by eακη : Ajk = P

s∈Sc(j)

e−ακη eα⟨xj ,xk ⟩ P

eα⟨xj ,xs ⟩ + e−ακη

s∈S / c(j)

eα⟨xj ,xs ⟩

eα⟨xj ,xk ⟩ . α⟨xj ,xs ⟩ s∈Sc(j) e

≤ e−ακη · P

Layerwise Dynamics for In-Context Classification in Transformers

Thus, the leakage to incorrect classes decays exponentially with κ. Case 2: Same Class (k ∈ Sc(j) ). Factorizing eακη from the numerator and denominator yields the exact expression: eα⟨xj ,xk ⟩ P . α⟨xj ,xs ⟩ α⟨xj ,xs ⟩ + e−ακη s∈S / c(j) e s∈Sc(j) e

Ajk = P

Taking the limit κ → ∞ eliminates the cross-class term in the denominator, recovering the softmax over the class cluster: eα⟨xj ,xk ⟩ . α⟨xj ,xs ⟩ s∈Sc(j) e

lim Ajk = P

κ→∞

If we further assumed the feature scale α → 0 (relative to the label signal), the exponential terms approach unity, recovering a uniform distribution over the class-cluster: 1 lim lim Ajk = . α→0 κ→∞ |Sc(j) | Of course, in this limit, the test point observes no nontrivial evolution. Henceforth, we will exploit the block-diagonal structure of the training attention with (implicitly) nonzero α. Label Evolution. We now proceed by induction on the labels using the asymptotic limit derived above (Ajk = 1 |Sc(j) | I[c(k) = c(j)]). For t = 0, the statement holds by definition. Assuming the statement is true for some t > 0, the labels evolve as: (t+1)

(t)

= yj + γ ′

n X

(t) (t)

Ajk yk

k=1 (t) = yj + γ ′

(t) (t)

Ajk yk .

k:c(k)=c(j) (t)

(t)

By the inductive hypothesis, for all k ∈ Sc(j) , yk = yj . Since Ajk is a probability distribution concentrated on Sc(j) , the (t)

sum is equal to yj . Thus, (t+1)

(t)

= yj + γ ′ yj = (1 + γ ′ )yj .

This completes the proof. Next, we show that class membership is preserved geometrically as well. We define the polyhedral cone Cc capturing all the points classified as class c with positive margin against all other classes. Cc := x ∈ Rd : ⟨wc − wc′ , x⟩ > 0 ∀c′ ̸= c . (t)

We show that in the limit case γ ≫ α each point xi throughout time remains in its original cone. We show this by induction on the timestep t. Lemma 2 (Invariant Cones). Assume that at time t, (t)

xi ∈ Cyi

∀i.

In the regime κ → ∞, the same condition holds at time t + 1. Proof. By Lemma 1, in the high-γ regime, the attention matrix becomes block-diagonal with respect to the class partition: (t) (t) (t) Atrain = diag A(1) , . . . , A(K) , (t)

where each block A(c) is row-stochastic and has nonnegative entries. 19

Layerwise Dynamics for In-Context Classification in Transformers

Fix a class c and a point i ∈ Sc . The feature update satisfies (t+1)

(t)

= xi + α′

(t)

Aij xj .

j∈Sc ′

Let c ̸= c. Taking inner products with wc − wc′ gives (t+1)

⟨wc − wc′ , xi

(t)

⟩ = ⟨wc − wc′ , xi ⟩ X (t) (t) Aij ⟨wc − wc′ , xj ⟩. + α′ j∈Sc

All terms on the right-hand side are strictly positive by the induction hypothesis and nonnegativity of the weights. Thus (t+1) xi ∈ Cc . C.2

Monotonic Global Directional Margin

Consider the centroid of a class c ∈ [K] defined as µc :=

1 X xi . |Sc | i∈Sc

To quantify separation in the multiclass setting, we define a global directional margin. Intuitively, it is not enough to show (t) (t) that ||µc − µc′ ||2 grows to establish class separation because then the separation could be on a direction orthogonal to the boundary wc − wc′ . Hence we define Mt :=

K X K X

(t)

⟨wc − wc′ , µ(t) c − µc′ ⟩.

c=1 c′ =1 c′ ̸=c

We show that this margin increases in t: Theorem C.1. In the limit κ → ∞, the global directional margin Mt is strictly increasing in t. Proof. We analyze the evolution of a single centroid. Since κ → ∞, recall by Lemma 1 that in every timestep the class of (t) each token remains the same. Fix a class c. Using the block-diagonal structure of Atrain , the centroid update is 1 X (t+1) µ(t+1) = xi c |Sc | i∈Sc   X (t) (t) P (t) (t) 1 X  (t) (t+1) = xi + α ′ Aij xj  (xi = Aij xj when γ ≫ α) |Sc | j∈Sc j∈Sc i∈Sc ! X 1 X (t) (t) (t) ′ = µc + α Aij xj |Sc | j∈Sc i∈Sc X (t) (t) (t) ′ = µc + α πj xj , j∈Sc

where (t)

πj :=

1 X (t) Aij |Sc | i∈Sc

defines a probability distribution over Sc . Fix c′ ̸= c. Taking inner products with wc − wc′ yields ⟨wc − wc′ , µ(t+1) ⟩ = ⟨wc − wc′ , µ(t) c c ⟩ X (t) (t) + α′ πj ⟨wc − wc′ , xj ⟩. j∈Sc

Layerwise Dynamics for In-Context Classification in Transformers

By the invariant cone lemma, each inner product in the sum is strictly positive, implying ⟨wc − wc′ , µ(t+1) ⟩ > ⟨wc − wc′ , µ(t) c c ⟩. Summing this inequality over all c′ ̸= c and then over all classes c shows that Mt+1 > Mt . C.3

Test Point Dynamics

We now proceed to our main theoretical result - the analysis of the test point dynamics. In the high γ regime (γ ≫ α), the initial test label y (0) is uniform across classes. Consequently, the label-driven component of the attention mechanism is (t) initially uninformative, and the test point’s trajectory is governed primarily by the feature dot products ⟨x(t) , xi ⟩. This observation motivates our analysis strategy: we approximate the training data dynamics (X (t) , Y (t) ) using the simplified label-driven limit (where points cluster perfectly by class), while retaining the full, coupled dynamics for the test point xtest . This approximation captures the two-phase nature of the inference process. In the early phase, feature similarity breaks the symmetry, guiding the test point toward the correct cluster and reinforcing the correct class label. Once the label margin (t) ∆y becomes sufficiently large, the dynamics transition to a label-driven regime, resulting in the asymptotic geometric convergence of the test point and confidence. The following equations describe the coupled trajectories of the test point xtest and its label embedding. x

(t+1)

(t)

+α

n X

′

i=1 n X

y (t+1) = y (t) + α′

(t) (t)

(15)

(t) (t)

(16)

Ai x i Ai y i

i=1 (t)

(t)

where Ai denote the ith coordinate of the test-token attention, i.e., Ai = (Atest )i , with (t)

Atest = σ(αx(t) Xt⊤ + γy (t) Yt⊤ ). We define the class-level attention masses p(t) c :=

(t)

Aj ,

j∈Sc

and the class logits (t) s(t) c := log Zc + γλt yt,c ,

Zc(t) :=

(t) exp α⟨x(t) , xj ⟩ ,

X j∈Sc

where λt = (1 + α′ )t is the norm of the training label vectors after t steps. By observing that under a class-wise block diagonal attention dynamics, ∀t, Yit = λt ec(i) , we find that (t)

Aj ∝ exp(α⟨x(t) , xtj ⟩ + γ⟨yt , Yj ⟩) = exp(α⟨x(t) , xtj ⟩ + γλt yt,c ), we find that X j∈Sc

Since

(t) j Aj =

(t)

Aj ∝

(t)

eγyt,c eα⟨x

(t)

,xj ⟩

= exp(s(t) c ).

j∈Sc

(t) j∈Sc Aj , it follows thus that

P P c

(t)

esc p(t) . c = P (t) sr re Suppose the test point xtest belongs to class c⋆ . We begin by establishing structural bounds on the geometric evolution. Lemma 3. Assume the following initialization at time t = 0: 21

Layerwise Dynamics for In-Context Classification in Transformers

(i) (Test margin) Define (0)

(0)

L0 := max ⟨x(0) , xj ⟩,

R0 := min ⟨x(0) , xj ⟩, j∈Sc⋆

j ∈S / c⋆

and assume ∆0 := R0 − L0 > 0. (ii) (Training separation) Define (0)

(0)

ρ0 := min ⟨xi , xj ⟩,

Λ0 :=

i,j∈Sc⋆

(0)

max

i∈Sc⋆ , j ∈S / c⋆

(0)

⟨xi , xj ⟩,

and assume Γ0 := ρ0 − Λ0 > 0. (0)

(iii) (Norm bound) ∥x(0) ∥2 ≤ 1 and ∥xi ∥2 ≤ 1 for all i. For each t ≥ 0, define the test-point quantities (t)

(t)

Rt := min ⟨x(t) , xj ⟩,

Lt := max ⟨x(t) , xj ⟩,

j∈Sc⋆

∆t := Rt − Lt ,

j ∈S / c⋆

and the training-point quantities (t)

(t)

ρt := min ⟨xi , xj ⟩,

Λt :=

i,j∈Sc⋆

max

i∈Sc⋆ , j ∈S / c⋆

(t)

⟨xi , xj ⟩,

Γt := ρt − Λt .

e t := min(∆t , Γt ). Then, the following hold for all t ≥ 0: Finally, let ∆ (t)

(a) (Norm growth) For all i, maxj∈Sc(i) ∥xj ∥2 ≤ (1 + α′ )t . (b) (Training margin recursion) Γt+1 ≥ (1 + α′ )2 Γt . (c) (Test margin recursion) (t)

(t)

∆t+1 ≥ (1 + α′ )∆t + α′ (1 + α′ )pc⋆ Γt − 2α′ (1 − pc⋆ )(1 + α′ )2t+1 . e t satisfies (d) (Effective margin recursion) The quantity ∆ (t) ′ ′ 2t+1 e t+1 ≥ 1 + α′ + α′ (1 + α′ )p(t) e ∆ . c⋆ ∆t − 2α (1 − pc⋆ )(1 + α ) Proof. Part (a): Recall the update rule for a training point xj belonging to class c(j): X (t+1) (t) (t) (t) xj = xj + α′ Ajk xk . k∈Sc(j)

Taking the Euclidean norm of both sides and applying the triangle inequality yields: (t+1)

∥xj

(t)

(t) (t)

∥2 = xj + α′

(t)

≤ ∥xj ∥2 + α′

Ajk xk

k∈Sc(j)

(t)

Ajk ∥xk ∥2 .

k∈Sc(j)

(t)

Let Mt := maxk∈Sc(j) ∥xk ∥2 denote the maximum feature norm within the class at iteration t. We bound:   X X X (t) (t) (t) (t) Ajk ∥xk ∥2 ≤ Ajk Mt = Mt  Ajk  . k∈Sc(j)

k∈Sc(j)

By row stochasticity, we have: (t+1)

∥xj

(t)

∥2 ≤ ∥xj ∥2 + α′ Mt ≤ (1 + α′ )Mt 22

(17)

Layerwise Dynamics for In-Context Classification in Transformers

Taking the maximum over all j ∈ Sc(j) on the left-hand side establishes part (a). Part (b): For i, j ∈ Sc⋆ , write (t)

(t) (t)

∆xi := α′

(t)

∆xj := α′

Aik xk ,

k∈Sc⋆

(t) (t)

Ajk xk .

k∈Sc⋆ (t)

(t)

Since each sum is a convex combination of points in Sc⋆ , every inner product of the form ⟨∆xi /α′ , xj ⟩, ⟨xi , ∆xj /α′ ⟩, (t)

(t)

and ⟨∆xi /α′ , ∆xj /α′ ⟩ is bounded below by ρt . Hence (t+1)

⟨xi

(t+1)

, xj

(t)

⟩ ≥ ⟨xi , xj ⟩ + 2α′ ρt + α′2 ρt ≥ (1 + α′ )2 ρt .

Taking the minimum over i, j ∈ Sc⋆ gives

ρt+1 ≥ (1 + α′ )2 ρt . (t)

(t)

Now let i ∈ Sc⋆ and j ∈ / Sc⋆ . Since ∆xi /α′ is a convex combination of points in Sc⋆ while ∆xj /α′ is a convex (t)

(t)

combination of points in the incorrect class Sc(j) , every inner product of the form ⟨∆xi /α′ , xj ⟩, ⟨xi , ∆xj /α′ ⟩, and (t)

(t)

⟨∆xi /α′ , ∆xj /α′ ⟩ is bounded above by Λt . Therefore (t+1)

⟨xi

(t+1)

, xj

(t)

⟩ ≤ ⟨xi , xj ⟩ + 2α′ Λt + α′2 Λt ≤ (1 + α′ )2 Λt .

Taking the maximum over i ∈ Sc⋆ and j ∈ / Sc⋆ gives Λt+1 ≤ (1 + α′ )2 Λt . Subtracting yields Γt+1 = ρt+1 − Λt+1 ≥ (1 + α′ )2 (ρt − Λt ) = (1 + α′ )2 Γt . P (t) (t) P (t) (t) (t) Part (c): Let ∆x(t) = α′ i Ai xi denote the update to the test point, and ∆xj = α′ k∈Sc(j) Ajk xk denote the update to the training point j. Expanding the inner product by linearity: (t+1)

⟨x(t+1) , xj

(t)

⟩ = ⟨x(t) + ∆x(t) , xj + ∆xj ⟩ (t)

(t)

= ⟨x(t) , xj ⟩ +

⟨∆x(t) , xj ⟩ | {z }

Term 1 (Test Update Uj )

(t)

⟨x(t) , ∆xj ⟩ | {z }

Term 2 (Training Update Vj )

+ ⟨∆x(t) , ∆xj ⟩ . | {z } Quadratic term

(t)

We first explicitly bound the quadratic cross-term resulting from the coupled updates. Let Qj := ⟨∆x(t) , ∆xj ⟩. Recall P (t) (t) (t) (t) (t) that the training point update is ∆xj = α′ k∈Sc(j) Ajk xk . Since Ajk is row-stochastic, ∆xj /α′ represents a convex P (t) (t) (t) combination of feature vectors within the cluster c(j). We define this cluster-average vector as νc(j) := k∈Sc(j) Ajk xk . The quadratic term can thus be written as: (t)

Qj = ⟨α′

n X

(t) (t)

(t)

Ai xi , α′ νc(j) ⟩ = α′2

i=1

n X

(t)

Ai ⟨xi , νc(j) ⟩.

i=1

We derive bounds for this term depending on whether j belongs to the correct class c∗ . (t)

Lower Bound (j ∈ Sc∗ ): If j is in the correct class, then νc(j) is a convex combination of points in Sc∗ . We split the summation over the test attention masses into the correct class (i ∈ Sc∗ ) and incorrect classes (i ∈ / Sc∗ ). (t)

(t)

• For i ∈ Sc∗ , the inner product ⟨xi , νc(j) ⟩ is a weighted average of intra-class inner products. By the definition of ρt , (t)

(t)

⟨xi , xk ⟩ ≥ ρt for all k ∈ Sc∗ . Thus, ⟨xi , νc(j) ⟩ ≥ ρt . 23

Layerwise Dynamics for In-Context Classification in Transformers

• For i ∈ / Sc∗ , we apply the global norm bound. Since ∥x∥ ≤ (1 + α′ )t , the inner product is lower-bounded by −(1 + α′ )2t . Combining these, and noting that

(t) (t) i∈Sc∗ Ai = pc∗ :

(t) (t) (t) Qj ≥ α′2 pc∗ ρt − (1 − pc∗ )(1 + α′ )2t . (t)

Upper Bound (j ∈ / Sc∗ ): If j is in an incorrect class c ̸= c∗ , then νc(j) consists of points from Sc . (t)

(t)

• For i ∈ Sc∗ , the inner product ⟨xi , νc(j) ⟩ represents cross-class similarity between c∗ and c. By the definition of Λt , (t)

(t)

⟨xi , xk ⟩ ≤ Λt for every k ∈ Sc(j) . Thus, ⟨xi , νc(j) ⟩ ≤ Λt . • For i ∈ / Sc∗ , we again apply the worst-case norm bound (1 + α′ )2t . This yields the upper bound: (t) (t) (t) Qj ≤ α′2 pc∗ Λt + (1 − pc∗ )(1 + α′ )2t . We explicitly bound the contributions from the test update Uj and the training update Vj . We split the summations into indices belonging to the correct class Sc⋆ and the incorrect classes Scc⋆ . 1. Lower Bound for Rt+1 (Case j ∈ Sc⋆ ). For a point j in the correct class, we want to lower bound both update terms. • Test Update (Uj ): Uj = α′

(t)

i∈Sc⋆

(t)

Ai ⟨xi , xj ⟩ + α′

(t)

Ai ⟨xi , xj ⟩.

i∈S / c⋆

For i ∈ Sc⋆ we have ⟨xi , xj ⟩ ≥ ρt , while for i ∈ / Sc⋆ by Cauchy-Schwarz and the norm bounds from Part (a) we get a lower bound of −(1 + α′ )2t , so (t)

(t)

Uj ≥ α′ pc⋆ ρt − (1 − pc⋆ )(1 + α′ )2t .

• Training Update (Vj ): Vj = α ′

(t)

Ajk ⟨x(t) , xk ⟩.

k∈Sc(j) (t)

Since j ∈ Sc⋆ , the sum runs over k ∈ Sc⋆ . By definition, ⟨x(t) , xk ⟩ ≥ Rt . Thus: X (t) Vj ≥ α ′ Ajk Rt = α′ Rt . k∈Sc⋆

Combining these with the quadratic lower bound gives the total lower bound: (t) (t) Rt+1 ≥ Rt + α′ Rt + α′ (1 + α′ ) pc⋆ ρt − (1 − pc⋆ )(1 + α′ )2t 2. Upper Bound for Lt+1 (Case j ∈ / Sc⋆ ). For a point j in an incorrect class, we want to upper bound both update terms. • Test Update (Uj ): Uj = α′

(t)

Ai ⟨xi , xj ⟩ + α′

i∈Sc⋆

(t)

Ai ⟨xi , xj ⟩.

i∈S / c⋆

Using ⟨xi , xj ⟩ ≤ Λt for cross-class pairs and the Cauchy-Schwarz norm bound (1 + α′ )2t otherwise: (t) (t) Uj ≤ α′ pc⋆ Λt + (1 − pc⋆ )(1 + α′ )2t . 24

Layerwise Dynamics for In-Context Classification in Transformers

• Training Update (Vj ): Vj = α ′

(t)

Ajk ⟨x(t) , xk ⟩.

k∈Sc(j) (t)

Since j ∈ / Sc⋆ , the class c(j) ̸= c⋆ . Thus, any k ∈ Sc(j) is not in the correct class. By definition, ⟨x(t) , xk ⟩ ≤ Lt . Thus: X (t) Vj ≤ α′ Ajk Lt = α′ Lt . k∈Sc(j)

Combining these with the quadratic upper bound gives the total upper bound: (t) (t) Lt+1 ≤ Lt + α′ Lt + α′ (1 + α′ ) pc⋆ Λt + (1 − pc⋆ )(1 + α′ )2t Subtracting upper from lower bound yields (t)

(t)

∆t+1 ≥ (1 + α′ )∆t + α′ (1 + α′ )pc⋆ (ρt − Λt ) − 2α′ (1 − pc⋆ )(1 + α′ )2t+1 , which is exactly the claimed recursion for ∆t . (t)

Part (d): Since pc⋆ ≤ 1, part (b) implies (t) Γt+1 ≥ (1 + α′ )2 Γt ≥ 1 + α′ + α′ (1 + α′ )pc⋆ Γt . e t ≤ ∆t and ∆ e t ≤ Γt , part (c) gives Using ∆ (t) e (t) ′ ′ 2t+1 ∆t+1 ≥ 1 + α′ + α′ (1 + α′ )pc⋆ ∆ . t − 2α (1 − pc⋆ )(1 + α ) e t+1 . Combining the last two displays and taking the minimum proves the recursion for ∆ We finally prove that the margin ∆t grows indefinitely and that the label prediction remains correct, meaning that the test point is attracted to the correct class both geometrically and with its label. This requires a coupled induction on the geometric (t) (t) (t) margin ∆t and the label margin ∆y := yc⋆ − yc for any class c ̸= c⋆ . Theorem C.2. Assume α′ < 1 and the initialization satisfies: e 0 > 0, ∆

e 0) α∆0 > log(K) + log(1 + 2/∆

and

∆(0) y ≥0

(∀c ̸= c⋆ ).

(18)

Then, for all t ≥ 0: 1. The geometric margin grows geometrically: ∆t ≥ ∆0 (1 + α′ )t . (t)

2. The label margin is non-decreasing: ∆y ≥ 0. Proof. We proceed by induction on t. The base case t = 0 holds by assumption. Assume the hypotheses hold for step t. Bounding the Logit Gap: Consider the difference in logits between the correct class c⋆ and any incorrect class c: (t)

Zc⋆ (t) (t) sc⋆ − s(t) +γ (yc⋆ − yc(t) ) . c = log (t) | {z } Zc (t) | {z } Label Term ∆y Geometric Term

(t)

By the definition of Rt and Lt , we have Zc⋆ ≥ |Sc⋆ |eαRt and Zc ≤ |Sc |eαLt . Assuming balanced classes (|Sc⋆ | = |Sc |): (t)

log

Zc⋆

(t)

≥ α(Rt − Lt ) = α∆t . 25

Layerwise Dynamics for In-Context Classification in Transformers (t)

Using the inductive hypothesis, ∀c ̸= c∗ , yt,c∗ − yt,c ≥ ∆y ≥ 0. Thus: (t)

(t)

sc⋆ − s(t) c = log

Zc∗

(t)

+ γ(yt,c∗ − yt,c ) ≥ α∆t .

(t)

Bounding Attention Leakage: The probability mass on incorrect classes is 1 − pc⋆ = (t) 1 − pc⋆ =

(19)

(t) c̸=c⋆ pc . Using Eq. 19:

(t) X X (t) (t) esc e−(sc⋆ −sc ) ≤ e−α∆t ≤ Ke−α∆t . P s(t) ≤ r e ⋆ ⋆ ∗ c̸=c c̸=c c̸=c r

Using the inductive hypothesis ∆t ≥ ∆0 (1 + α′ )t , this leakage term decays super-exponentially. Inductive step on ∆t+1 From Lemma 3(b), the training margin obeys Γt ≥ Γ0 (1 + α′ )2t . e 0 = min(∆0 , Γ0 ), we in particular have Γ0 ≥ ∆ e 0. Since ∆ From Lemma 3(c), the test margin updates as (t)

(t)

∆t+1 ≥ (1 + α′ )∆t + α′ (1 + α′ )pc⋆ Γt − 2α′ (1 − pc⋆ )(1 + α′ )2t+1 . Using the lower bound on Γt gives (t) (t) ∆t+1 ≥ (1 + α′ )∆t + α′ (1 + α′ )(1 + α′ )2t pc⋆ Γ0 − 2(1 − pc⋆ ) . Therefore it suffices to show (t)

(t)

pc⋆ Γ0 ≥ 2(1 − pc⋆ ). This inequality is equivalent to (t)

(1 − pc⋆ )(Γ0 + 2) ≤ Γ0 . By the leakage estimate above and the inductive lower bound on ∆t , (t)

(1 − pc⋆ )(Γ0 + 2) ≤ Ke−α∆t (Γ0 + 2) ≤ Ke−α∆0 (Γ0 + 2). e 0 , we have Since Γ0 ≥ ∆ −α∆0

2 (Γ0 + 2) ≤ Ke Γ0 1 + . e0 ∆ 2 Ke−α∆0 1 + ≤1 e0 ∆ −α∆0

by Eq. (18), and hence (t)

(1 − pc⋆ )(Γ0 + 2) ≤ Γ0 . Therefore ∆t+1 ≥ (1 + α′ )∆t , which closes the induction. (t+1)

Label Margin Induction (∆y ). Finally, recall the update rule for the test label embedding: y (t+1) = y (t) + Pn (t) (t) ′ ⋆ α i=1 Ai yi . We track the margin between the correct class component c and an incorrect class component c. (t) From Lemma 1, the training labels maintain their class-alignment throughout the dynamics, satisfying yi = λt ec(i) for some scalar scaling factor λt > 0. Substituting this into the update rule, the k-th component evolves as: (t+1)

(t)

= yk + α′ λt

n X

(t)

Ai I(c(i) = k) = yk + α′ λt pk .

i=1

Layerwise Dynamics for In-Context Classification in Transformers (t)

(t)

Now, considering the margin ∆y = yc⋆ − yc : (t)

(t)

(t) ′ (t) ∆(t+1) = (yc⋆ + α′ λt pc⋆ ) − (yc(t) + α′ λt p(t) y c ) = ∆y + α λt (pc⋆ − pc ) . {z } | Drift Term

(t)

In Step 1, we established that the logit gap sc⋆ − sc ≥ α∆t is strictly positive. Since the softmax function is monotonic, (t) (t) (t) (t) a larger logit implies a larger probability mass: sc⋆ > sc =⇒ pc⋆ > pc . Therefore, the drift term is strictly positive. (t) (t+1) (t) Combined with the inductive hypothesis ∆y ≥ 0, it follows that the label margin strictly increases: ∆y > ∆y ≥ 0. This completes the induction. Growth of the Label Margin. Coupling the above with the growing size of λt in turn induces a growth in the label margin. Indeed, let t0 := inf{t : 1 − 2Ke−α∆t ≥ 1/2}. Then note that for all t ≥ t0 , (t)

(t)

−α∆t pc∗ − p(t) ≥ c ≥ pc∗ − (1 − pc∗ ) = 1 − 2(1 − pc∗ ) ≥ 1 − 2Ke

Thus, for all t ≥ t0 ,

1 . 2

′ ′ ∆(t+1) ≥ ∆(t) y y + α λt /2 ≥ α λt /2. (t)

Recalling from Lemma 1 that λt , the norm of the prediction vectors, is as (1+γ ′ )t , we conclude that for all t ≥ t0 +1, ∆y ≥ (1 + γ ′ )t . Finally, note that since ∆t ≥ (1 + α′ )t ∆0 , t0 ≤ inf{t : α∆t ≥ log(4K)} ≤ inf{t : (1 + α′ )t ≥ log(4K)/α∆0 } ≤

Additional insights into Mean-Shift Classification

D.1

The role of α

log 4K 1 log + 1. α′ α∆0

The hyperparameter α controls the influence data geometry exerts on the weights A. Recall that the data point updates take the form X (c) (t) (t+1) (t) = xi + α′ Aij xj . xi j∈Sc

When α = 0 the updates are solely driven by the relationships between labels, which are initialized to be intra-class orthogonal. At the same time however, this leads to the test point xtest to be pulled “equally” towards the direction of all classes, which means it does not move from its initial position. Thus, α > 0 plays an important role not only for the mean-shift dynamics of the train points but also for the trajectory of the probe point. D.2

Mean-shift classification on non-linear data distributions

In this section, we stress-test the extracted mean-shift classifier by applying it to a range of data distributions in two dimensions that go beyond the linear setting studied in the main text. Separated Ellipses We generate two classes of anisotropic Gaussian clusters. Each class k ∈ {0, 1} consists of points sampled from a multivariate normal distribution N (µk , Σ), where µ0 = (−s, −s) and µ1 = (s, s) provide linear separation. Voronoi Cells We simulate a multi-class problem by defining K random centroids {ck }K k=1 within a bounded region [−B, B]2 . Data points xi are sampled uniformly from this region. The label yi for each point is assigned according to the nearest centroid, effectively partitioning the space into Voronoi cells: yi = argmink ∥xi − ck ∥2 . 27

Layerwise Dynamics for In-Context Classification in Transformers

Figure 18. Setting α = 0 gives us label-driven mean-shift, and the probe point does not move.

Figure 19. Setting α > 0 leads to a probe point trajectory.

Concentric Circles This dataset consists of K classes arranged as nested rings. For the k-th class, points are generated using polar coordinates where the radius is fixed at rk = 1 + k and the angle θ is sampled uniformly from [0, 2π). Gaussian noise ϵ ∼ N (0, σ 2 I) is added to the Cartesian coordinates (rk cos θ, rk sin θ) to create diffuse ring structures that are not linearly separable. Surprisingly, the mean-shift classifier still separates these data well despite the inherently non-linear decision boundaries. Intertwined Spirals Two classes are generated as interleaving spiral arms. For each class k ∈ {0, 1}, we define a parametric curve where the radius r grows linearly with the angle θ ∈ [0, 4π]. The coordinates are given by x = (θ cos(θ + ϕk ), θ sin(θ + ϕk )), where ϕk is a phase offset (e.g., 0 and π). Gaussian noise is added to the points, creating a highly non-convex geometry that is difficult to separate with simple distance metrics. We did not find any parameter setting for which the mean-shift classifier reliably separates these class boundaries, and so this is a failure mode for these dynamics. Concretely, it appears that for this data, the strongly intertwined geometry overwhelms the label-separations, which are further hampered by the fact that the cluster-centroids for the two labels are essentially coincident.

Layerwise Dynamics for In-Context Classification in Transformers

Figure 20. Mean-Shift Effects on separated ellipses.

Figure 21. Voronoi Cells.

Figure 22. Concentric Circles

Layerwise Dynamics for In-Context Classification in Transformers

Figure 23. Spirals Mean-Shift

Record · ID 10348 · SHA-256 c21188785154e9cc
Conceptio Open Knowledge Archive — every document is proof-bundled with source, license, and retrieval metadata.