Symmetry Reveals Layerwise Dynamics: How Transformers Perform In-Context Classification
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.
Lay Summary
Modern AI systems based on transformers can often learn a new task from a few examples in their input ("in-context learning"), but it is usually unclear how they make the final prediction. We study this question in controlled classification problems where the order of features, examples, or class names should not matter. By forcing every layer of the model to respect these natural symmetries, we keep its predictions essentially unchanged while making its internal computations easier to read. This reveals a repeated procedure: at each layer, the model compares examples using both their data values and labels, then moves related examples and the test point closer to the right class group. We provide an interpretation of this process as gradually sharpening clusters, so class groups become more separated as information passes through the layers. This explanation makes testable predictions, which we confirm with noisy labels, unlabeled examples, and more complicated (prototype-based) class boundaries. In particular, the model can use unlabeled examples to improve classification when they reveal useful geometric structure. Our results offer a concrete approach for revealing the algorithms transformers use at prediction time, especially in settings with clear symmetries, such as the classification task studied here.