Symmetry Reveals Layerwise Dynamics: How Transformers Perform In-Context Classification

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

cs.LGApr 13, 2026

Symmetry Reveals Layerwise Dynamics: How Transformers Perform In-Context Classification

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

Organizations: Boston University, Departments of Computer Science and Electrical & Computer Engineering.

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.

Figures & tables

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 $\leftrightarrow$ iterations), yielding a geometry-driven algorithmic motif in which features and labels co-evolve.

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 $\bm{W}_{QK,\ell}$ and $\bm{W}_{VP,\ell}$ , respectively.

Panel of Figure 3 — Figure 3: Transformer symmetrization preserves inference behavior. 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; $\pm$ 1 s.d. across runs). Right: After averaging predictions across 5 training runs, U and S are functionally nearly identical: ground-truth-class probabilities agree with $R^{2}>0.98$ over $10\,000$ tasks, yielding the same decision boundaries.

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

Figure 5: Label-driven mean-shift increases class separation. Simulated dynamics with fixed parameters $(\alpha,\gamma,\alpha^{\prime},\gamma^{\prime})=(1,5,0.08,0.1)$ , showing class centroids drifting apart.

Figure 6: A simulation of the mean-shift trajectories with constant parameters $(\alpha,\gamma,\alpha^{\prime},\gamma^{\prime})=(1,5,0.08,0.1)$ . The test point ‘follows’ a trailing path towards its cluster.

Panel of Figure 7 — Figure 7: Classification dynamics for Voronoi cells. Top: Simulation for $(\alpha,\gamma,\alpha^{\prime},\gamma^{\prime})=(1,5,0.05,0.2)$ . Bottom: Trained transformer weights encode the same dynamics ( 7 ), matching Fig. 2 .

Panel of Figure 8 — Figure 8: Transformers use unlabeled context to improve semi-supervised 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; $\pm$ 1 s.d. across runs). Bottom: Trained transformer weights encode the same dynamics ( 7 ), matching Fig. 2 .

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

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.

Panel of Figure 10 — 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).

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 $\bm{W}_{QK,\ell}$ and $\bm{W}_{VP,\ell}$ , respectively.

Panel of 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 5 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 $\bm{W}_{QK,\ell}$ and $\bm{W}_{VP,\ell}$ , respectively.

Panel of Figure 16 — Figure 16: Left: Ground-truth query-class probability predicted by the symmetrized vs. unconstrained transformer on the same semi-supervised classification task (mean over 5 training seeds). Right: Ablation with Unstructured Noise. When unlabeled context points are replaced with random noise $\mathcal{N}(0,I_{d})$ , 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.

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 $W_{QK,\ell}$ and $W_{VP,\ell}$ , respectively.

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

Figure 19: Setting $\alpha>0$ leads to a probe point trajectory.

Figure 20: Mean-Shift Effects on separated ellipses.

Explore similar work

Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent
May 7, 2026Chenyang Zhang, Yuan CaoLogistic Regression Self-Attention
School of Computing & Data Science, The University of Hong Kong
One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context
Sep 1, 2026Skanda Athreya, Yutong WangTransformer In-Context Learning
James B. Conant High School · Hoffman Estates, IL · Illinois Institute of Technology +1
Implicit Statistical Inference in Transformers: Approximating Likelihood-Ratio Tests In-Context
Mar 11, 2026Faris Chaudhry, Siddhant GadkariIn-Context Learning Transformer Interpretability
Department of Computer Science Imperial College London

May 7, 2026J/K move · Enter open · S save

Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent

Chenyang Zhang, Yuan Cao

School of Computing & Data Science, The University of Hong Kong

Logistic Regression Self-Attention Gradient Descent Transformer Attention In-Context Learning

Open paper AskPDF arXiv

Transformers have demonstrated remarkable in-context learning (ICL) capabilities. The strong ICL performance of transformers is commonly believed to arise from their ability to implicitly execute certain algorithms on the context, thereby enhancing prediction and generation. In this work, we investigate how transformers with softmax attention perform in-context learning on linear classification data. We first construct a class of multi-layer transformers that can perform in-context logistic regression, with each layer exactly performing one step of normalized gradient descent on an in-context loss. Then, we show that our constructed transformer can be obtained through (i) training a single self-attention layer supervised by one-step gradient descent, and (ii) recurrently applying the trained layer to obtain a looped model. Training convergence guarantees of the self-attention layer and out-of-distribution generalization guarantees of the looped model are provided. Our results advance the theoretical understanding of ICL mechanism by showcasing how softmax transformers can effectively act as in-context learners.

Abstract

Explore similar work

Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent

One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context

Implicit Statistical Inference in Transformers: Approximating Likelihood-Ratio Tests In-Context