A Simple Generalisation of the Implicit Dynamics of In-Context Learning

1 Motivation and main result

• The linear-regression experiment shows an excellent match between theory and experiment.Figure 1 compares theoretical and empirical test losses for the last token in the final block, and prediction differences across blocks and training steps.Source: Figure 1 · p. 2, §1 · p. 2
• The generalized result equates any token's prediction in any contextual block, including Pre-LN blocks, with a query-only block using context-updated MLP parameters.For blocks beyond the first, the inputs are refined versions of the original context and query.Source: §1 · pp. 1–2
• The context-dependent update to the MLP's first weight matrix is rank one, with updates also specified for the final-layer biases.The rank-one update is an outer product of a column vector and a row vector.Source: §1 · p. 2
• The analysis covers one step of token generation; studying settings where the answer is a sequence of tokens remains an open direction.The authors also note potential validation on large-scale models and analysis of implicit updates in language-trained models.Source: §1 · p. 2

A.1.1 Extension to all sequence positions

• Indexing the contextual-layer output at position i extends the result to any token position.Different positions correspond to different implicit updates.Source: §A.1.1 · p. 4
• Stacking the updates for different positions into ΔB(C) also produces a rank-one update matrix.Source: §A.1.1 · p. 4
• Each position-specific weight update is rank one, formed as the outer product of a position-dependent vector and a shared vector.The position-dependent vector is u_i = WΔA(i), while v = A(x) is shared across positions.Source: §A.1.1 · p. 4

A.1.2 Extension to any contextual block

• Iterative application generalizes the result to any contextual block beyond the first.Source: §A.1.2 · p. 5
• For a later block, its context and query are refined versions of the original, unprocessed inputs.Source: §A.1.2 · p. 5
• The generalized equation reduces to Eq. 1 when the block is the first and the token is the last.Source: §A.1.2 · p. 5

A.1.3 Extension to any block with more accurate skip connections

• The more exact block structure feeds the input skip x into the MLP, unlike the preceding structure that omits it from the full contextual-layer output.Source: §A.1.3 · p. 5
• The proof extends the result to any token position, with Δz(i) defined as the difference between an input element and the query.Source: §A.1.3 · p. 5
• Direct computation establishes the equality for the skip-connection block without layer normalization.Source: §A.1.3 · p. 6

A.1.4 Generalisation to any block with skips and layer normalisations (Theorem 1)

• The general weight update uses the difference between the full MLP inputs with and without context, and has rank one.The inputs are denoted g_i and f, respectively.Source: §A.1.4 · p. 6
• With skip connections, the MLP input difference also includes the skip-connection change, and an implicit update to the last MLP layer's biases is needed.The contextual and context-free inputs differ by ΔA(i) + Δz(i).Source: §A.1.4 · pp. 6–7
• For Pre-LN blocks, nested layer normalisations prevent the full-input difference from simplifying, and the update accounts for both pre-attention and pre-MLP normalisation.The weight update is given in Eq. 29 and the last-layer bias update in Eq. 30.Source: §A.1.4 · p. 7
• For any token position and block, including skip connections and Pre-LN, the contextual block's prediction equals that of the block with context-specific MLP parameter updates.The equality is empirically verified for Pre-LN blocks.Source: §A.1.4 · pp. 6–7

A.2 Experimental details

• Training sequences contain input-output pairs from sampled linear functions, while inference tests on previously unseen functions using a query input.Each task samples a function coefficient vector w from a standard Gaussian; inputs are also sampled from a standard Gaussian.Source: §A.2 · p. 7
• Transformers are trained to minimize last-token prediction error over batches of tasks.The query target is the linear function's output for the query input.Source: §A.2 · p. 8
• Figure 1 uses B = 128, N = 51, d_x = 2, and L = 5 residual blocks, each with three-head causal attention and a two-layer GeLU MLP.Models are trained for 100 Adam steps with learning rate η = 5e−2.Source: §A.2 · p. 8
• The reported mean squared differences measure the deviation between theoretical and empirical predictions, averaged over batches, sequence positions, and input dimensions.Runs are repeated across random seeds; all experiments run on CPU.Source: §A.2 · p. 8

A.3 Alignment of implicit weight updates

• Directional alignment is measured as the normalized Frobenius inner product between two token-associated weight updates.The Frobenius inner product is defined as ⟨A, B⟩_F = Tr(A^T B).Source: §A.3 · p. 8
• For a given task sequence, the alignment structure among token-specific updates is qualitatively consistent across blocks.The reported pattern is shown for multiple example tasks and sequence inputs.Source: §A.3 · p. 8, Figure A.1 · p. 8
• The alignment of last-token updates across blocks shows no particular structure during training.Results were consistent across random seeds.Source: §A.3 · p. 9, Figure A.2 · p. 9

A.4 Supplementary figures

• Figure A.3 applies the Figure 1 metrics to Pre-LN blocks with other hyperparameters held constant; obtaining good generalization with layer normalization was more challenging across many hyperparameters.The caption notes that Pre-LN remains standard in most large language models.Source: Figure A.3 · p. 9
• Increasing the number of tasks B makes learning easier, with empirical and theoretical test losses plotted while other hyperparameters remain fixed..Source: Figure A.4 · p. 10
• Increasing input sequence length N facilitates learning, with empirical and theoretical test losses compared at different lengths.Other hyperparameters are held constant.Source: Figure A.5 · p. 10
• Increasing input dimensionality d_x makes learning more challenging, with empirical and theoretical test losses compared across dimensions.Other hyperparameters are held constant.Source: Figure A.6 · p. 10
• Figures A.7 and A.8 report the same results as Figure A.1, using a different example task or input sequence in A.7 and yet another example task in A.8.Source: Figure A.7 · p. 10, Figure A.8 · p. 11

• The checklist affirms that the abstract and introduction accurately reflect the paper’s contributions and scope, and that complete theoretical proofs are provided in Appendix A.The authors also state that the claims are verified by experiments.Source: NeurIPS Paper Checklist · pp. 12–13
• The checklist identifies one-step token generation as the work’s main limitation, noting that settings where the answer is a sequence of tokens remain to be studied.Source: NeurIPS Paper Checklist · p. 12
• The checklist says experimental details needed to understand and reproduce the results are provided in Appendix A.2, and code will be released upon publication.The authors report that experiments ran on a single CPU.Source: NeurIPS Paper Checklist · pp. 13–15
• The authors answer that error bars are not reported where results did not significantly vary across random seeds or runs.Source: NeurIPS Paper Checklist · p. 14
• The checklist marks societal-impact discussion as not applicable because the tested models are too simple for modern AI applications.Source: NeurIPS Paper Checklist · p. 15
• The checklist marks new assets, human-subject research, and non-standard core-method LLM use as not applicable.Source: NeurIPS Paper Checklist · pp. 17–18