cs.LGApr 14, 2026
SaveDistinct mechanisms underlying in-context learning in transformers
Organizations: Lewis-Sigler Institute for Integrative Genomics, Princeton University, Princeton, NJ 08544, USA · Department of Physics, Princeton University, Princeton, NJ 08544, USA
Abstract
Modern distributed networks, notably transformers, acquire a remarkable ability (termed `in-context learning') to adapt their computation to input statistics, such that a fixed network can be applied to data from a broad range of systems. Here, we provide a complete mechanistic characterization of this behavior in transformers trained on a finite set of discrete Markov chains. The transformer displays four algorithmic phases, characterized by whether the network memorizes and generalizes, and whether it uses 1-point or 2-point statistics. We show that the four phases are implemented by multi-layer subcircuits that exemplify two qualitatively distinct mechanisms for implementing context-adaptive computations. Minimal models isolate the key features of both motifs. Memorization and generalization phases are delineated by two boundaries that depend on data diversity, . The first () is set by a kinetic competition between subcircuits and the second () is set by a representational bottleneck. A symmetry-constrained theory of a transformer's training dynamics explains the sharp transition from 1-point to 2-point generalization and identifies key features of the loss landscape that allow the network to generalize. Put together, we show that transformers develop distinct subcircuits to implement in-context learning and identify conditions that favor certain mechanisms over others.
Figures & tables
Figure 1: (a) Data generating process for training and out-of-distribution sequences. Training sequences are generated from transition matrices drawn from the matrices in sampled prior to training. Out-of-distribution sequences are generated from transition matrices drawn directly from . (b) An illustration of 1-point (unigram) statistics and 2-point (bigram) directed, nearest-neighbor statistics. (c) The 2-layer transformer architecture, which contains one attention block and one multi-layer perceptron (MLP) block in each layer. After the operation of two layers, the linear block maps the embedded representation to a discrete distribution over states that is then normalized to form the logit. (d) Schematic of the four algorithmic phases as the number of chains ( ) and training time ( ) are varied. Phases are consistent with either memorizing or generalizing predictors, which use either 1-point or 2-point statistics, denoted by 1- and 2- prefixes respectively. Orange to red gradient in memorization phase indicates a gradual transition from to in contrast to the abrupt transition from to . Three quantities characterize the transition between the different phases: (i) , which is the training time to transition from to given , (ii) , which together with defines the ‘triple point’ between , and , and (iii) , which delineates and . The dashed phase boundary indicates an extrapolation as .
Figure 2: (a) Training dynamics of the transformer training loss (top) and generalization loss (bottom) at low, medium, and high . Notice the plateaus during training that occur near the predictor training and out-of-distribution losses. The model dynamics and predictor losses have been produced by evaluating model checkpoints and the predictors on a large batch of sequences. We have omitted for clarity the predictor losses near which each model does not exhibit plateaus. (b) Behavioral readouts of the transformer through training across a range of . Color intensity is inversely proportional to the divergence (in nats) of the transformer predictions from the indicated predictor (Appendix IV.1 ). Simultaneous similarity of the model to both 1-Gen and 1-Mem predictors at large results from the relatively high similarity of these predictors to the uniform predictor over states. (c) Mechanistic readouts and order parameters of the transformer through training across a range of (Appendix IV.2 for details). Here, and have been normalized by the sequence length, . The correspondence between phases and mechanistic readouts are as follows: and : high and low ; : high and ; : high and .
Figure 3: Circuit tracing. (a) The transformer may be unrolled into a directed graph with each block taking input from the previous blocks. (b-c) Left: Visualization of the complete circuit traced for the phase. See Appendix V for the definitions of each node in the circuit and an interpretation of the connections. Edge thickness is proportional to the importance factor for that edge, and dashed lines group the attention subcircuits for each block. Edges with importance are omitted, and any nodes with all incoming and outgoing edges omitted are depicted as desaturated. Right: Interpretative schematic of the circuit traced for the phase, with arrows indicating input-output relationships and layer function labeled. See the main text for more details. Isolated layers are omitted entirely. In panel (c), the asterisk denotes the task vector discussed in Section IV.3 .
Figure 4: (a) A 2-D t-SNE visualization of task vectors for the model demonstrates increased separability through training. (b) Results of task vector patching experiment. Replacing the task vector formed on a sequence from transition matrix A with that formed on a sequence from transition matrix B results in the transformer predicting according to the transition matrix B. See Appendix IX.3 for details. (c) The information content of the task vectors is optimized through training. Mutual information is presented as a fraction of the channel capacity . See Appendix IX.4 for details.
Figure 5: (a) Schematic of the SA-Transformer for generalization. (b) Loss dynamics for the SA-Transformer trained in the limit display the same sharp transition from (green line) to (blue line). Data is shown from 5 model seeds, with one darkened for clarity. (c) Model dynamics are dominated by the previous-state attention and the diagonal of one block of the key-query matrix. Top: Normalized weights from the relative positional bias in the first attention layer at different iterations. Bottom: Normalized token-wise attention derived from the bottom-left block of the key/query matrix in the second attention layer at different iterations. (d) A schematic of the induction-head mechanism underlying the model prediction in (refer to Appendix VI for a detailed explanation). (i) Given the state embeddings in subspace , (ii) the first attention block selectively attends to the previous state (with specificity ) and copies that state to an orthogonal subspace , and (iii) the second attention block produces the empirical next-state distribution by comparing the current state with the contents of (with specificity ) and copying the contents of from matching positions.
Figure 6: (a) Loss landscape in the parameter space given by equation 3 and 4 (see equation 94 in Appendix) where are obtained by optimization over the full landscape. (b) Top: Scaling of the effective parameter with training iterations . Bottom: Scaling of the effective parameter with training iterations . Different lines in each plot with the same color represent numerics with different seeds. In the “Control” setting, the direct- bias predicts a linear growth of both . In the “Bias-Ablation” setting, the second-to-last sequence state has been resampled to remove the direct- bias which results in instead. (c) Scaling of the transition time with sequence length . The time is measured as the number of iterations of SGD required until the sharp drop in loss indicating the development of the induction head. Top: Measurements for the SA-transformer (described in eq. 67 ) in the two settings considered in panel (b). Bottom: Measurements for a non-autoregressive 2-layer transformer with MLP1 removed and additive relative positional biases. Optimization uses SGD with momentum. Kernel density estimated from 8 seeds.
Figure 7: (a) The maximum value of the induction head parameter through training is sigmoidal as the data diversity crosses over . This indicates that the model either reaches the threshold and proceeds to (high ) or does not reach the threshold and proceeds to , with few intermediates. Evaluated over 16 different model seeds trained on the same task distribution. Each seed has been shifted along the horizontal axis such that first crosses the threshold value at ). See Appendix X for details. (b) Bimodality of the model dynamics near , shown for and 64 seeds trained on the same . Different model seeds show either small or never achieve . All model seeds in the rightmost bin did not enter after iterations. (c) Perturbations to the learning dynamics of the competing sub-circuits have effects on the memorization-generalization transition that are explainable by a kinetic competition. increases when the effective learning rate of the 2-Gen circuit is reduced (LR perturbation). Increasing the rate of memorization by providing information about the ground truth task identity also increases (TI perturbation). Kernel density estimated from 4 seeds.
Figure 8: (a) The time between the onset of and its disappearance in favor of diverges as with and . The model is assumed to enter when it achieves a training loss . Median value over 4 different seeds trained on the same task distribution shown, with errorbars indicating the 95% confidence interval. Dashed line indicates a divergence fit to the medians with a linear offset. Inset: Replot of the data highlighting asymptotic behavior. The fit value for was used to compute the offset . Dashed line with slope provided for reference. (b) The mean fraction of memorized tasks extrapolated to convergence decreases linearly past . See Appendix XI for details. A linear extrapolation of these data suggests that the model will not memorize any chain for , which is consistent with the estimate for from the data in panel (a). Errorbars show standard deviation over the 4 seeds. (c) Schematic of the minimal network inspired from the circuit-tracing results in Figure 3 , used to reproduce the 2-Mem behavior shown in the remaining panels. (d,e) Heat maps of and as functions of the task-vector embedding dimension and data diversity . Blue (red) indicates that the model is over (under) performing relative to the 2-Gen predictor; white denotes the crossover. (f) Average critical data diversity over three independent trials. The labels indicate the number of hidden layers in MLP1 and MLP2, respectively. The standard deviation error bar is estimated with three independent trials.
Figure 9: A schematic summarizing the factors that govern memorization-generalization transitions. Suppose the transformer’s training dynamics proceeds along a generalized coordinate for memorization ( ) and generalization ( ). (a) When the data diversity is small relative to , memorization is both faster than generalization and reaches a lower loss, so the model proceeds along alone. (b) For between and , generalization develops before memorization and the model begins progressing along . However, the model eventually switches over to memorization, which can achieve a lower loss. (c) When is well above , constraints on model expressivity renders memorization unfavorable relative to generalization, and dynamics are constrained to as a result.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure A1 : Scaling of the four predictor losses with data diversity and sequence length computed over a large batch of sequences and averaged over 8 task distributions . The autoregressive loss is shown for scaling. scaling of the non-autoregressive loss is shown at data diversity .
Table 1: Explicit definition of source vectors contributing to the residual stream, and the target perturbations. To ablate the edge from source SRC to a target TGT at layer , we compute the perturbed input as where is the source vector corresponding to SRC (Subtable ( 1(a) )) and denotes the batch average. The target is then perturbed during the forward pass by injection of according to the prescription for target TGT (Subtable ( 1(b) )). The layer-normalization each block applies to the perturbed input computes the statistics over the unperturbed input.
Figure A2 : Complete visualization and interpretive schematic of the traced circuits as in Figure 3 for the remaining two phases (a) and (b) .
Figure A3 : Connectivity matrices between the circuit components. See Table 1 for the definitions of sources and targets. Dark squares indicate connections that demonstrate a large change in model behavior when they are ablated. Generalizing (memorizing) circuits are traced from the ( ) models at the training checkpoint with lowest divergence to the corresponding predictor.
Figure A4 : Parameters of the attention circuits defined in equation 67 at different iterations. The corresponding loss dynamics is given in Figure 5 (b).
Figure A5 : Training dynamics of the mixture-of-experts weights in equation 67 .
Figure A6 : Scaling of with different for .
Figure A7 : Loss landscapes near for (a) original data and (b) data with the second-to-last token bias removed. The white dashed line is the corresponding dynamics of and .
Figure A8 : The task vector patching experiment performed for by replacing the output of Att1 and comparing to the 1-Mem predictor.
Figure A9 : The minimum achieved value of the generalization loss drops sharply across as the data diversity crosses over . As in Figure 7 a, each seed has been shifted along the horizontal axis by the value of determined for that seed by the threshold criteria. Results derived using the same models used for Figure 7 a.
Figure A10 : The fraction of memorized tasks displays sigmoidal training dynamics across data diversities. Individual results for 4 seeds are plotted as thin traces, and the mean is plotted by a thick line.