Distinct mechanisms underlying in-context learning in transformers
Authors: Cole Gibson, Wenping Cui, Gautam Reddy
Organizations: Lewis-Sigler Institute for Integrative Genomics, Princeton University, Princeton, NJ 08544, USA · Department of Physics, Princeton University, Princeton, NJ 08544, USA
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 S 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, K=∣S∣. The first (K1∗) is set by a kinetic competition between subcircuits and the second (K2∗) 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 K matrices in S sampled prior to training. Out-of-distribution sequences are generated from transition matrices drawn directly from DT . (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 ( K ) and training time ( t ) 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 M1 to M2 in contrast to the abrupt transition from G1 to G2 . Three quantities characterize the transition between the different phases: (i) τ2-Gen , which is the training time to transition from G1 to G2 given K>K1∗ , (ii) K1∗ , which together with τ2-Gen defines the ‘triple point’ between G1 , G2 and M1 , and (iii) K2∗ , which delineates G2 and M2 . The dashed phase boundary indicates an extrapolation as t→∞ .
Figure 2: (a) Training dynamics of the transformer training loss (top) and generalization loss (bottom) at low, medium, and high K . 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 K . 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 K 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 K (Appendix IV.2 for details). Here, nA(1) and nA(2) have been normalized by the sequence length, N=256 . The correspondence between phases and mechanistic readouts are as follows: G1 and M1 : high nA(1) and low ϕδ(1) ; G2 : high ϕδ(1) and ϕβ(2) ; M2 : high ϕδ(1) and nA(2) .
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 ≤0.25 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 K=64 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 log2K . 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 K→∞ display the same sharp transition from G1 (green line) to G2 (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 G2 (refer to Appendix VI for a detailed explanation). (i) Given the state embeddings in subspace Bstate , (ii) the first attention block selectively attends to the previous state (with specificity ∝δ ) and copies that state to an orthogonal subspace B1 , and (iii) the second attention block produces the empirical next-state distribution by comparing the current state with the contents of B1 (with specificity ∝β ) and copying the contents of B2 from matching positions.
Figure 6: (a) Loss landscape in the (β,δ) parameter space given by equation 3 and 4 (see equation 94 in Appendix) where wA−D are obtained by optimization over the full landscape. (b) Top: Scaling of the effective parameter β with training iterations t . Bottom: Scaling of the effective parameter δ with training iterations t . 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 δ∼t2 instead. (c) Scaling of the transition time τ2-Gen with sequence length N . The time τ2-Gen 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 ϕβ(2) through training is sigmoidal as the data diversity crosses over K1∗ . This indicates that the model either reaches the threshold and proceeds to G2 (high ϕβ(2) ) or does not reach the threshold and proceeds to M1 , 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 ϕβ(2) first crosses the threshold value ϕβ∗ at K=K1∗ ). See Appendix X for details. (b) Bimodality of the model dynamics near K1∗ , shown for K=94≈K1∗ and 64 seeds trained on the same DT . Different model seeds show either small τ2-Gen or never achieve G2 . All model seeds in the rightmost bin did not enter G2 after 5×104 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. K1∗ increases when the effective learning rate of the 2-Gen circuit γ2-Gen is reduced (LR perturbation). Increasing the rate of memorization by providing information about the ground truth task identity also increases K1∗ (TI perturbation). Kernel density estimated from 4 seeds.
Figure 8: (a) The time between the onset of G2 and its disappearance in favor of M2 diverges as (K2∗−K)−γ with γ≈2 and K2∗≈7000 . The model is assumed to enter M2 when it achieves a training loss Ltrain<L2-Gen . Median value over 4 different seeds trained on the same task distribution DT 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 K2∗ was used to compute the offset K2∗−K . Dashed line with slope −γ provided for reference. (b) The mean fraction of memorized tasks extrapolated to convergence decreases linearly past K≈28 . See Appendix XI for details. A linear extrapolation of these data suggests that the model will not memorize any chain for K≈5000 , which is consistent with the estimate for K2∗ 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 Ltrain−L2-Gen and Lgen−L2-Gen as functions of the task-vector embedding dimension Dφ and data diversity K . 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 K^2∗ over three independent trials. The labels (ℓ1,ℓ2) 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 ( ξMem ) and generalization ( ξGen ). (a) When the data diversity K is small relative to K1∗ , memorization is both faster than generalization and reaches a lower loss, so the model proceeds along ξMem alone. (b) For K between K1∗ and K2∗ , generalization develops before memorization and the model begins progressing along ξGen . However, the model eventually switches over to memorization, which can achieve a lower loss. (c) When K is well above K2∗ , constraints on model expressivity renders memorization unfavorable relative to generalization, and dynamics are constrained to ξGen 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 K and sequence length N computed over a large batch of sequences and averaged over 8 task distributions DT . The N=256 autoregressive loss is shown for K scaling. N scaling of the non-autoregressive loss is shown at data diversity K=128 .
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 x~n(ℓ−1):−xn(ℓ−1)−(xSRC,n−\expectationvaluexSRC,n) where xSRC,n is the source vector corresponding to SRC (Subtable ( 1(a) )) and \expectationvalue⋅ denotes the batch average. The target is then perturbed during the forward pass by injection of x~n(ℓ−1) 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) G1 and (b) M1 .
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 K=1024 ( K=128 ) 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 Fd with different d for C=10,α=1 .
Figure A7 : Loss landscapes near δ,β→0 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 M1 by replacing the output of Att1 and comparing to the 1-Mem predictor.
Figure A9 : The minimum achieved value of the generalization loss Lgen drops sharply across as the data diversity crosses over K1∗ . As in Figure 7 a, each seed has been shifted along the horizontal axis by the value of K1∗ determined for that seed by the ϕβ(2) 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.
Transformers perform computations through many successive attention and feedforward blocks, allowing them to learn complex correlations between a large collection of coupled variables. When does stacking successive attention-feedforward computations over many layers provide a computational advantage over a single transformer block? We address this question by examining in-context learning in generalized linear attention transformers. We first introduce a general theory of distributed inference in such transformers, subject to constraints on communication and depth. We show that such systems can exploit internal representations (`function vectors') to infer a latent context variable at increasingly finer scales over its layers. For an in-context linear regression task, the theory predicts that while one-layer transformers without feedforward blocks are optimal for Gaussian priors over the context variable, multi-layer transformers are superior for non-Gaussian, tree-like priors. Trained linear attention transformers reproduce quantitative predictions from the theory. Using causal key-patching experiments, we verify that function vectors in intermediate layers mediate adaptive routing of information. Our results suggest that depth and feedforward blocks enable transformers to implement adaptive inference, and this is advantageous when the distribution over latent variables has hierarchical structure.
Ravin Raj, Gautam Reddy
Joseph Henry Laboratories of Physics, Princeton University, Princeton, New Jersey 08544, USA
Transformer-based large models have demonstrated remarkable generalization abilities across different tasks by leveraging a context-aware attention module for in-context learning. With richer context, transformers adapt more effectively to the current use case without any parameter updates. However, the quadratic computational and memory complexity with respect to context length significantly slows data processing in softmax transformers. Linear transformers were proposed to address this issue by reducing the complexity to linear dependence on context length, but the design and understanding of the feature mapping in linear attention, from a theoretical viewpoint, remain unclear. In this paper, we investigate the approximation and generalization abilities of linear transformers under a two-staged sampling process from domain generalization. We show that linear transformers perform in-context learning as learning a mapping from context distributions to response functions. A dimension-independent convergence rate is obtained for our generalization analysis, which also exhibits the tradeoff between the regularities of data distributions and latent features. Guided by our theoretical framework, we propose a new perspective on activation and loss design for linearizing pretrained softmax large language models.
Peilin Liu, Ding-Xuan Zhou
School of Mathematics and Statistics, University of Sydney, NSW Australia
Induction heads are attention circuits believed to underlie in-context learning in transformers, yet a precise characterization of the estimators they implement remains elusive. We study transformers trained on order-k Markov chains and identify two complementary smoothing mechanisms. First, at finite attention-weight scale, the circuit implements a soft context-matching estimator: it aggregates contributions from exact and partial context matches, weighted exponentially by their overlap, and induces a data-dependent interpolation across context orders analogous to Jelinek-Mercer smoothing. Second, a beginning-of-sequence (BOS) token induces additive pseudo-counts, recovering Dirichlet-style smoothing. We construct a disentangled transformer implementing both mechanisms and show that trained transformers recover the predicted attention patterns. Across settings where pseudo-count smoothing is optimal or lower-order contexts provide structured evidence, trained transformers match or outperform classical count-based baselines. Our results bridge mechanistic interpretability of induction heads with classical statistical smoothing, revealing that transformers learn to regularize in-context estimation rather than simply count.
Francesco D'Angelo, Oguz Kaan Yuksel, Swathi Shree Narashiman +1
1TML Lab, EPFL, Switzerland. · 2IIT Madras, India. · †Work done during an internship at the TML Lab, EPFL.