arxiv: 2604.11613 · v2 · submitted 2026-04-13 · 💻 cs.LG · cs.AI

Recognition: unknown

Layerwise Dynamics for In-Context Classification in Transformers

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

Authors on Pith no claims yet

Pith reviewed 2026-05-10 15:02 UTC · model grok-4.3

classification 💻 cs.LG cs.AI

keywords in-context classificationtransformer dynamicsequivarianceattention mechanismsclass separationlayerwise recursionin-context learningsoftmax attention

0 comments

The pith

By enforcing equivariance under feature and label permutations, transformer layers yield an explicit recursive update rule for in-context classification.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Transformers perform in-context classification from a handful of labeled examples, yet the step-by-step computation inside the model has remained difficult to read. The authors constrain every layer so that swapping features or labels leaves the overall function unchanged; this symmetry forces the weights into a highly structured form without reducing what the model can do. From the structured layers they then read off a precise recursion indexed by depth that tracks how attention updates the training points, their labels, and the test probe together. The recursion arises from attention matrices built on combined feature-label similarities and turns out to push classes farther apart in a geometry-driven way while aligning the test point with its correct class. A reader would care because the result replaces the usual black-box description of in-context learning with a concrete, layer-by-layer algorithm that can be written down and analyzed.

Core claim

Enforcing feature- and label-permutation equivariance at every layer of a softmax transformer trained for multi-class linear classification in the hard no-margin regime produces highly structured weights from which an explicit depth-indexed recursion can be extracted. This recursion is an end-to-end identified emergent update rule in which 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 that can provably amplify class separation and yields robust expected class alignment.

What carries the argument

The depth-indexed recursion, an emergent update rule extracted from the equivariant layers that couples updates via attention matrices on mixed feature-label Gram structure.

If this is right

Attention matrices couple the updates of training points, their labels, and the test probe at each layer.
The dynamics provably amplify separation between classes through the geometry of the Gram structure.
Expected alignment of the test probe to its true class becomes robust with increasing depth.
The model computes the same function as the unconstrained transformer yet exposes an interpretable algorithmic motif.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The explicit recursion could be used as a scaffold to initialize or constrain transformers on other few-shot tasks to encourage similar separation behavior.
The geometric amplification motif may link to classical iterative algorithms for linear separation such as margin-based updates.
Imposing analogous symmetry constraints on transformers for regression or generation tasks might surface comparable layerwise rules in those domains.

Load-bearing premise

Enforcing feature- and label-permutation equivariance at every layer preserves the functional behavior of the original transformer while producing structured weights that reveal the recursion.

What would settle it

Training the equivariant transformer and verifying whether its actual layer-by-layer updates match the predictions of the extracted recursion formula; mismatch at any depth would show the identification has failed.

Figures

Figures reproduced from arXiv: 2604.11613 by Aditya Gangrade, Arjun Chandra, Patrick Lutz, Themistoklis Haris, Venkatesh Saligrama.

**Figure 1.** 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 m… view at source ↗

**Figure 2.** 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 … view at source ↗

**Figure 3.** 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 r… view at source ↗

**Figure 5.** 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 [PITH_FULL_IMAGE:figures/full_fig_p006_5.png] view at source ↗

**Figure 7.** 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 [PITH_FULL_IMAGE:figures/full_fig_p007_7.png] view at source ↗

**Figure 6.** 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. 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 label… view at source ↗

**Figure 8.** 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 [PITH_FULL_IMAGE:figures/full_fig_p008_8.png] view at source ↗

**Figure 9.** 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 [PITH_FULL_IMAGE:figures/full_fig_p014_9.png] view at source ↗

**Figure 10.** Figure 10: Weight abstraction preserves the symmetric transformer’s input–output behavior. Left: Clustering the learned weight matrices ( [PITH_FULL_IMAGE:figures/full_fig_p014_10.png] view at source ↗

**Figure 11.** Figure 11: Transformer learns label centering. Parameter δℓ as in intermediary weight abstraction (1) motivates fixing δ = − 1 K . Median and std. in symmetrized transformer over 5 training runs [PITH_FULL_IMAGE:figures/full_fig_p015_11.png] view at source ↗

**Figure 13.** 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 updat… view at source ↗

**Figure 14.** 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 [PITH_FULL_IMAGE:figures/… view at source ↗

**Figure 15.** 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. 16 [PITH_FULL_IMAGE:figu… view at source ↗

**Figure 16.** 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 [PI… view at source ↗

**Figure 17.** 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… view at source ↗

**Figure 18.** Figure 18: Setting α = 0 gives us label-driven mean-shift, and the probe point does not move [PITH_FULL_IMAGE:figures/full_fig_p028_18.png] view at source ↗

**Figure 19.** 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… view at source ↗

**Figure 20.** Figure 20: Mean-Shift Effects on separated ellipses [PITH_FULL_IMAGE:figures/full_fig_p029_20.png] view at source ↗

**Figure 21.** Figure 21: Voronoi Cells [PITH_FULL_IMAGE:figures/full_fig_p029_21.png] view at source ↗

**Figure 22.** Figure 22: Concentric Circles 29 [PITH_FULL_IMAGE:figures/full_fig_p029_22.png] view at source ↗

**Figure 23.** Figure 23: Spirals Mean-Shift 30 [PITH_FULL_IMAGE:figures/full_fig_p030_23.png] view at source ↗

read the original abstract

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 label-permutation 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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

They extract an explicit recursion for in-context classification by enforcing equivariance, which is new and traceable, but the constraint may narrow how much it explains standard transformers.

read the letter

They extract an explicit recursion for in-context classification by enforcing equivariance, which is new and traceable, but the constraint may narrow how much it explains standard transformers. The paper trains softmax transformers under feature- and label-permutation equivariance at every layer. This keeps functional equivalence while structuring the weights so attention matrices from mixed Gram features drive coupled updates to points, labels, and the test probe. From there they derive a depth-indexed recursion that implements a geometry-driven motif, with a claimed proof that it amplifies class separation in the hard no-margin linear case and produces robust alignment. That extraction is the concrete advance over prior empirical or partial ICL analyses. The work does well by making the internal dynamics identifiable and end-to-end traceable without changing the model's output behavior. The structured weights let them write down the update rule cleanly, which is useful for anyone wanting a mechanistic account rather than just accuracy curves. The main soft spot is generality. The stress-test concern holds: imposing the symmetry could alter the loss landscape or force the very structure they then extract, so the recursion might not appear in unconstrained transformers. The abstract gives no sign of direct comparisons, so the claim that this is the emergent rule inside a softmax transformer rests on the unverified assumption that the constrained class is representative. If the full text has those checks or a clear discussion of the restriction, the result strengthens; otherwise it stays narrower. The math and citation pattern look standard for this area, with no obvious circularity in the abstract. This is for researchers working on mechanistic interpretability or theoretical accounts of few-shot transformer behavior. A reader who wants explicit update rules to analyze or extend will get value if the derivations check out. It deserves a serious referee because the extraction of an identifiable recursion is worth detailed verification on the proofs and the generality question. I would send it for peer review with a request that reviewers focus on whether the equivariance preserves the essential dynamics of standard models.

Referee Report

1 major / 2 minor

Summary. The manuscript claims that enforcing feature- and label-permutation equivariance at every layer in transformers trained for multi-class linear classification in the hard no-margin regime renders the computation identifiable while preserving functional equivalence. This yields highly structured weights from which an explicit depth-indexed recursion can be extracted; the recursion is driven by attention matrices formed from mixed feature-label Gram matrices, implements a geometry-driven motif that provably amplifies class separation, and produces robust expected class alignment. The authors present this as the first end-to-end identified emergent update rule inside a softmax transformer.

Significance. If the extracted recursion is shown to be representative of unconstrained softmax transformers rather than an artifact of the imposed symmetries, the work would be significant for supplying the first explicit, interpretable algorithmic description of in-context classification dynamics together with a provable amplification property. Such a result would supply a concrete mechanistic motif that could be tested, extended, or used to design more interpretable architectures.

major comments (1)

[Abstract] Abstract: The load-bearing claim that the equivariant models 'maintain functional equivalence' to the original transformer while yielding an emergent recursion representative of general softmax transformers is not supported by any verification that the symmetry constraints leave attention patterns, loss landscapes, or optimization trajectories unchanged. The Gram-structured attention and coupled updates are derived only after the equivariance is imposed, raising the possibility that the structured weights and amplification motif are forced by the constraint rather than discovered.

minor comments (2)

[Abstract] The abstract asserts that the dynamics 'can provably amplify class separation' yet supplies no derivation outline, proof sketch, or reference to the relevant theorem; the full manuscript must include these steps with explicit assumptions.
No empirical comparison is mentioned between the equivariant models and standard (non-equivariant) transformers on the same task; such a check is needed to confirm that performance and class alignment are not degraded by the symmetry constraints.

Simulated Author's Rebuttal

1 responses · 0 unresolved

We thank the referee for their insightful comments on our work. We address the major concern regarding the verification of functional equivalence and the potential artifactual nature of the emergent recursion in the point-by-point response below. We have made revisions to strengthen the manuscript accordingly.

read point-by-point responses

Referee: [Abstract] Abstract: The load-bearing claim that the equivariant models 'maintain functional equivalence' to the original transformer while yielding an emergent recursion representative of general softmax transformers is not supported by any verification that the symmetry constraints leave attention patterns, loss landscapes, or optimization trajectories unchanged. The Gram-structured attention and coupled updates are derived only after the equivariance is imposed, raising the possibility that the structured weights and amplification motif are forced by the constraint rather than discovered.

Authors: We acknowledge that the original manuscript did not include explicit side-by-side comparisons of attention patterns, loss landscapes, or optimization trajectories between the equivariant and unconstrained models. In the revised version, we add experiments demonstrating that the equivariant models achieve nearly identical in-context classification performance (average accuracy difference < 2% across 10 seeds) and exhibit similar training loss trajectories to standard transformers. We also provide attention visualizations for unconstrained models, showing that their attention weights, when symmetrized over feature and label permutations, closely resemble the Gram-matrix structures observed in the equivariant case. A comprehensive analysis of the full loss landscape is computationally prohibitive given the model size; however, the comparable performance and trajectories suggest that the constraints do not fundamentally alter the optimization path or solution quality. We clarify in the revised abstract and introduction that the recursion is explicitly derived under the equivariance constraints, which were imposed to achieve identifiability, but that these constraints respect the inherent symmetries of the multi-class classification task. Thus, the motif is not artificially forced but rather made explicit by restricting to the equivariant subspace where the computation becomes interpretable. We do not assert that this exact recursion governs all unconstrained transformers, but it offers the first fully identified example of such dynamics and a testable mechanistic hypothesis for broader cases. The abstract has been updated to reflect this nuance. revision: partial

Circularity Check

0 steps flagged

No circularity: equivariance is an explicit modeling choice for identifiability, not a reduction by construction

full rationale

The paper's derivation begins by imposing feature- and label-permutation equivariance at every layer on a softmax transformer to render the weights identifiable and highly structured while asserting functional equivalence to the unconstrained model. From the resulting trained models an explicit depth-indexed recursion is extracted, with attention matrices formed from mixed feature-label Gram structure driving coupled updates. This recursion is presented as an emergent, geometry-driven motif that amplifies class separation. No equation or step in the abstract or described chain shows the extracted recursion reducing to the equivariance constraint by algebraic identity, a fitted parameter being relabeled as a prediction, or a load-bearing premise justified solely by self-citation. The equivariance enforcement is a deliberate architectural restriction chosen to enable extraction and interpretability; the resulting dynamics are derived from the trained equivariant models rather than being tautological with the inputs. The derivation chain therefore remains self-contained against external benchmarks and does not meet the criteria for any enumerated circularity pattern.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

The central claim rests on the domain assumption that equivariance can be imposed without changing functional behavior; no free parameters or invented entities are mentioned in the abstract.

axioms (1)

domain assumption Enforcing feature- and label-permutation equivariance at every layer maintains functional equivalence while making the internal computation identifiable.
This is the enabling step stated in the abstract that allows extraction of the recursion from structured weights.

pith-pipeline@v0.9.0 · 5432 in / 1430 out tokens · 65172 ms · 2026-05-10T15:02:13.402151+00:00 · methodology

discussion (0)

Reference graph

Works this paper leans on

10 extracted references · 4 canonical work pages · 1 internal anchor

[1]

Appworld: A controllable world of apps and people for benchmarking interactive coding agents

URL https://aclanthology.org/2023. findings-acl.247/. Deutch, G., Magar, N., Natan, T., and Dar, G. In-context 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 ...

work page doi:10.18653/v1/ 2023
[2]

Lee, J., Lee, Y ., Kim, J., Kosiorek, A

URL https://openreview.net/forum? id=6rD50Q6yYz. Lee, J., Lee, Y ., Kim, J., Kosiorek, A. R., Choi, S., and Teh, Y . W. Set transformer: A framework for attention-based permutation-invariant neural networks. InProceedings of the 36th International Conference on Machine Learning, 2019. Li, Y ., Ildiz, M. E., Papailiopoulos, D., and Oymak, S. Trans- formers...

work page arXiv 2019
[3]

URL https://openreview.net/forum? id=Syx72jC9tm. Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y ., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish...

work page internal anchor Pith review arXiv 2022
[4]

2023 , month = nov, journal =

URL https://proceedings.mlr.press/ v202/von-oswald23a.html. Wies, N., Levine, Y ., and Shashua, A. The learnability of in-context learning. InAdvances in Neural Information Processing Systems, 2023. Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An ex- planation of in-context learning as implicit bayesian infer- ence. InInternational Conference on Lea...

work page arXiv 2023
[5]

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

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...
[6]

Supervised

Two parameters per layer suffice.Figure 11 shows that the transformer consistently learns a bottom-right background coefficient of approximately δ≈ − 1 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 ma...

2003
[7]

•Test Update (U j): Uj =α ′ X i∈Sc⋆ A(t) i ⟨x(t) i , x(t) j ⟩+α ′ X i /∈Sc⋆ A(t) i ⟨x(t) i , x(t) j ⟩

Lower Bound forR t+1 (Casej∈S c⋆).For a pointjin the correct class, we want to lower bound both update terms. •Test Update (U j): Uj =α ′ X i∈Sc⋆ A(t) i ⟨x(t) i , x(t) j ⟩+α ′ X i /∈Sc⋆ A(t) i ⟨x(t) i , x(t) j ⟩. For i∈S c⋆ 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 U...
[8]

•Test Update (U j): Uj =α ′ X i∈Sc⋆ A(t) i ⟨x(t) i , x(t) j ⟩+α ′ X i /∈Sc⋆ A(t) i ⟨x(t) i , x(t) j ⟩

Upper Bound forL t+1 (Casej /∈S c⋆).For a pointjin an incorrect class, we want to upper bound both update terms. •Test Update (U j): Uj =α ′ X i∈Sc⋆ A(t) i ⟨x(t) i , x(t) j ⟩+α ′ X i /∈Sc⋆ A(t) i ⟨x(t) i , x(t) j ⟩. Using⟨x i, xj⟩ ≤Λ t for cross-class pairs and the Cauchy-Schwarz norm bound(1 +α ′)2t otherwise: Uj ≤α ′ p(t) c⋆ Λt + (1−p (t) c⋆ )(1 +α ′)2t...
[9]

The geometric margin grows geometrically:∆ t ≥∆ 0(1 +α ′)t
[10]

Proof.We proceed by induction ont

The label margin is non-decreasing:∆ (t) y ≥0. Proof.We proceed by induction ont. The base caset= 0holds by assumption. Assume the hypotheses hold for stept. Bounding the Logit Gap:Consider the difference in logits between the correct classc ⋆ and any incorrect classc: s(t) c⋆ −s (t) c = log Z(t) c⋆ Z(t) c| {z } Geometric Term +γ(y (t) c⋆ −y (t) c )| {z }...