​​What learning algorithm is in-context learning? Investigations with linear models
Authors: Ekin Akyürek
Dale Schuurmans
Jacob Andreas
Tengyu Ma
Denny Zhou
Venue: ICLR 2023 notable top 5% (2023)
Keywords: in-context learning
transformers
sequence models
deep learning
meta learning

Abstract: Neural sequence models, especially transformers, exhibit a remarkable capacity for in-context learning. They can construct new predictors from sequences of labeled examples $(x, f(x))$ presented in the input without further parameter updates. We investigate the hypothesis that transformer-based in-context learners implement standard learning algorithms implicitly, by encoding context-specific parametric models in their hidden representations, and updating these implicit models as new examples appear in the context. Using linear regression as a model problem, we offer three sources of evidence for this hypothesis. First, we prove by construction that transformers can implement learning algorithms for linear models based on gradient descent and closed-form computation of regression parameters. Second, we show that trained in-context learners closely match the predictors computed by gradient descent, ridge regression, and exact least-squares regression, transitioning between different predictors as transformer depth and dataset noise vary. Third, we present preliminary evidence that in-context learners share algorithmic features with these predictors: learners' late layers encode weight vectors and moment matrices.  These results suggest that in-context learning is understandable in algorithmic terms, and that (at least in the linear case) learners may work by rediscovering standard estimation algorithms.

Published as a conference paper at ICLR 2023
WHAT LEARNING ALGORITHM IS IN-CONTEXT LEARN- ING? INVESTIGATIONS WITH LINEAR MODELS
Ekin Akyürek 12..
Dale Schuurmans'
Jacob Andreas
Tengyu Ma
Denny Zhou"
Google Research
MIT CSAIL
3 Stanford University
collaborattie advising
ABSTRACT
Neural sequence models, especially transformers, exhibit a remarkable capacity for in-context learning. They can construct new predictors from sequences of labeled examples (x, f()) presented in the input without further parameter up- in-context learners implement standard learning algorithms implicilly, by encoding smaller models in their activations, and updating these implicit models as new examples appear in the context. Using linear regression prototypical problem, we offer three sources of evidence for this hypothesis. First, we prove by construction that trans- formers can implement learning algorithms for lincar models based gradient descent and closed-fore form ridge regression. Second, we show that trained in-context learners closely match the predictors computed by gradient descent, ridge regres- sion, and exact least-squares regression, transitioning between different predictors transformer depth and dataset noise vary, and converging Bayesian estima- tors for large widths and depths. Third, we present preliminary evidence that in-context learners share algorithmic features with these predictors: learners' late layers non- non-linearly encode weight vectors and moment matrices. These results suggest that in-context learning understandable in algorithmic terms, and that (at least in the linear case) learners may rediscover standard estimation algorithms.
dates. We investigate the hypothesis that transformer-bis
INTRODUCTION
One of the most surprising behaviors observed in large neural sequence models is in-context learn- ing (ICL; Brown et al., 2020). When trained appropriately, models can map from sequences of (I, fz)) pairs to accurate predictions f(x') on novel inputs x'. This behavior occurs both in mod- els trained on collections of few-shot learning problems (Chen et al., 2022; Min et al., 2022) and surprisingly in large language models trained open-domain text (Brown et al., 2020; Zhang al., 2022; Chowdhery et al., 2022). ICL requires model to implicitly construct map from in-context examples a predictor without any updates to the model's parameters themselves. How can a neural network with fixed parameters to learn new function from new dataset on the fly?
This paper investigates the hypothesis that some instances of ICL can be understood as implicir implementation of known learning algorithms: in-context learners encode an implicit, context- dependent model their hidden activations, and train this model on in-context examples in the course of computing these internal activations. As in recent investigations of empirical properties of ICL (Garg et al., 2022; Xie et al., 2022), we study the behavior of transformer- (Vaswani et al., 2017) in past work, our goal is not to understand what functions ICL can learn, but how it learns these functions: the specific inductive biases and algorithmic properties of wansformer-basee ICL.
predictors
a restricted class of learning problems, here linear regression. Unlike
In Section 3, we investigate theoretically what learning algorithms transformer decoders can imple- ment. We prove by construction that they require only modest number of layers and hidden units to train linear models: for d-dimensional regression problems, with (d) hidden size and constant depth, transformer can implement single step of gradient descent; and with Odd2) hidden size
Correspondence to akyurek cdu. Ekin a student MIT, and began this work while he was intern at Google Research. Code and reference implementations are released at this web page 'The work done when Tengyu Ma works a visiting researcher Google Research.

---

Published as a conference paper at ICLR 2023
and constant depth, a transformer can update a ridge regression solution to include single new observation. Intuitively, n steps of these algorithms can be implemented with times more layers.
In Section 4, we investigate empirical properties of trained in-context learners. We begin by con- structing linear regression problems in which learner behavior is under-determined by training data (so different valid learning rules will give different predictions on held-out data). We show that model predictions are closely matched by existing predictors (including those studied in Section 3), and that they transition between different predictors as model depth and training set noise vary, be- having like Bayesian predictors large hidden sizes and depths. Finally, in Section 5, we present preliminary experiments showing how model predictions are computed algorithnically. We show that important intermediate quantities computed by learning algorithms for linear models, including parameter vectors and moment matrices, can be decoded from in-context learners' hidden activa- tions.
A complete characterization of which learning algorithms are (or could be) implemented by deep networks has the potential to improve both our theoretical understanding of their capabilittess and limitations, and our empirical understanding of how best to train them. This paper offers first steps toward such characterization: some in-context learning appears involve familiar algorithms, discovered and implemented by transformers from sequence modeling tasks alone.
2 PRELIMINARIES
Training a machine learning model involves many decisions, including the choice of model archi- tecture, loss function and learning rule. Since the carliest days of the field, research has sought to understand whether these modeling decisions can be automated using the tools of machine learning itself. Such "meta-learning" approaches typically treat learning (Schmidhuber et al., 1996; Andrychowicz et al., 2016; Finn et al., 2017): they define "inner" and "outer" models and learning procedures, then train an outer model set parameters for an inner procedure (e.g. initializer or step size) to maximize inner model performance across tasks.
bi-level optimization problem
Recently, a more flexible family of approaches has gained popularity. In in-context learning (ICL), meta-Icarning is reduced to ordinary supervised learning: a large sequence model (typically imple- mented as a transformer network) is trained map from sequences f 22 fx2)), (.........3.. zn,] to predictions (xx) (Brown et al., 2020; Olsson et al., 2022; Laskin et al., 2022; Kirsch & Schmidhu- ber, 2021). ICL does not specify an explicit inner learning procedure; instead, this procedure exists only implicitly through the parameters of the sequence model. ICL has shown impressive results on synthetic tasks and naturalistic language and vision problems (Garg et al., 2022; Min et al., 2022; Zhou et al., 2022).
Past work has characterized what kinds of functions ICL can learn (Garg et al., 2022; Laskin al., 2022) and the distributional properties of pretraining that can elicit in-context learning (Xie et al., 2021; Chan et al., 2022). But how ICL learns these functions has remained unclear. What learning algorithms (if any) discovered in the course of training? This paper takes first steps toward answering these questions, focusing on widely used model architecture (the transformer) and an extremely well-understood class of learning problems (linear regression).
implementable by deep network models? Which algorithms are actually
2.1 THE TRANSFORMER ARCHITECTURE
Transformers (Vaswani al., 2017) are neural network models that map a sequence of input vectors x = [rI,..., Fn. to a sequence of output vectors y = "vi,...,yn], Each layer in a transformer maps a matrix ( (interpreted as a sequence of vectors) to sequence H To do so, a transformer interested in autoregressive (or
layer processes each column h. (!) of H in parallel. Here, "decoder-only") transformer models in which each layer first computes a self-attention:
a: Attention(h" ( W.,, W?, W W.)
W. bb...... 6m]
(2)
where each b is the response of an "attention head" defined by: bj softmax (I.)h.) T(w3" H.)) (Wy H.).
(3)
2

---

Published as a conference paper at ICLR 2023
then applies feed-forward transformation:
h ((+1) FF(a;; W1, W2) W1o(W22a + () + a; h!)
(4)
(5)
Here denotes nonlinearity, e.g. Gaussian error linear unit (GeLU; Hendrycks & Gimpel, 2016):
< 2
1 + erf
x
(6)
V2
and À denotes layer normalization (Ba et al., 2016):
x E a
(())
Var [
where the expectation and variance are computed across the entries of æ. To map from æ to y, transformer applies sequence of such layers, each with its own parameters. We use 0 to denote model's 's full set of parameters (the complete collection of W matrices across layers). The three main factors governing the computational capacity of transformer are its depth (the number of layers), its hidden size (the dimension of the vectors h), and the number of heads (denoted m above).
2.2 TRAINING FOR IN IN-CONTEXT LEARNING
We study transformer models directly trained on an ICL objective. (Some past work has found that ICI. also "emerges" in models trained on general text datasets; Brown al., 2020.) train transformer 7 with parameters 0 to perform ICI., we first define a class of functions F, distribution p(f) supported on F, a distribution p(æ:) over the domain of functions in F, and loss function L. We then choose 0 to optimize the auto-regressive objective, where the resulting To is an in-context learner:
arg min
E
(s(s),, To ((mi)... mx))))
(8)
Rn (p(() f~p(f)
2.3 LINEAR REGRESSION
Our experiments focus linear regression problems. In these problems, F is the space of linear functions f(x) = w x where w, a E Rd and the loss function the squared error C(y v') = model problem in machine learning and statistical estimation, with diverse algorithmic solutions. thus offers an ideal test-bed for understanding ICL. Given a dataset with inputs X ........ nn and y ......., . Vnl. the (regularized) linear regression objective:
(y- 2²² Linear regression
E Iz 3., Al 21|20|1
(9)
is minimized by: w = (X X A
(10)
With À = 0, this objective is known as ordinary least squares regression (OLS); with À > 0, it known ridge regression (Hoerl Kennard, 1970). (As discussed further in Section ridge
regression can also be assigned a Bayesian interpretation.) To present a linear regression problem to a transformer, we encode both x and fax) as d 1-dimensional vectors: z; [0, zz). D where 0,, denotes the d-dimensional zero vector.
[3t. 0²],
WHAT LEARNING ALGORITHMS CAN TRANSFORMER IMPLEMENT?
For a transformer-based model to solve Eq. (9) by implementing an explicit learning algorithm, that learning algorithm must be implementable via Eq. (1) and Eq. (4) with some fixed choice of transformer parameters& 6. In this section, we prove constructively that such parameterizations exist, giving concrete implementations of two standard learning algorithms. These proofs yield upper bounds on how many layers and hidden units suffice to implement (though not necessarily learn) each algorithm. Proofs given in Appendices A and B.
3

---

Published as a conference paper at ICLR 2023
3.1 PRELIMINARIES
It will be useful to first establish a few computational primitives with simple transformer implemen- tations. Consider the following four functions from R - -RR" >RR"
m! s,t, j').)"")): selects the entries of the Thh column of H between rows and and copies them into the tl column (t 2 8) of between rows i' and j'. yielding the matrix:
H Hegg's Hja
High
HI.--
; in each column h of H, interprets the entries between and j as an a x matrix A1. and the entries between i' and j' as a x c matrix A2. multiplies these matrices together, and stores the result between rows i" and j", yielding matrix in which each column has the form [fag"-1, A2, tegr."".
mul e, (1,3), ('))',, (1", ")):
div(H; (i. i,, (t" ))")) in each column h of H, divides the entries between and j by the abso- lute value of the entry at i', and stores the result between rows i" and j". yielding a matrix in which every column has the form [h hag [yy hyrd""
aff (1,5), (t',"), (i", .""), W1, W2, b): in each column h of H, applics an affine transformation to the entries between and j and i' and j', then stores the result between rows 1" and j" yielding a matrix in which every column has the form 11, Wiha; + W2h + b, hgrr.".
Lemma 1. Each of mov, mul, div and aff can be implemented by single transformer decoder layer: in Eq. (1) and Eq. (4), there exist matrices WQ, W.K, WY, W., W1 and W2 such that, given a matrix H as input, the layer's output has the form of the corresponding function output above.
With these operations, we can implement building blocks of two important learning algorithms.
3.2 GRADIENT DESCENT
Rather than directly solving linear regression problems by evaluating Eq. (10), standard approach to learning exploits a generic loss minimization framework, and optimizes the ridge-regression ob- jective in Eq. (9) via gradient descent on parameters u. This involves repeatedly computing updates:
w' w dw (C(w" x,
(C(wTa,,,) v.) + A||20| w 22x(w x yæ + Xuv)
(11)
for different examples (æ, yi), and finally predicting w x,, new input Fa. A step of this gradient descent procedure can be implemented by transformer:
Theorem 1. A transformer can compute Eq. (11) (i.e. the prediction resulting from single step of gradient descent on an in-context example) with constant number of layers and O(d) hidden space, where d is the problem dimension the imput x. Specifically, there exist transformer parameters such that, given an input matrix of the form:
(0))
0
11
(12)
0
the transformer's output matrix (L)) contains an entry equal to wT xn (Eg. (11)) at the column index where In is input.
3.3 CLOSED-FORM REGRESSION
Another way to solve the linear regression problem is to directly compute the closed-form solution Eq. (10). This is somewhat challenging computationally, as requires inverting the regularized covariance matrix X X XI. However, one can exploit the Sherman-Morrioo formula (Sherman & Morrison, 1950) to reduce the inverse to a sequence of rank-one updates performed example-by- example. For any invertible square A,
1 A- (A + uo T) A-1 -1
(13)
+ v A-l
We omit the trivial size preconditions, e.g. mul: (i j a b, - - b, - c* c*
4

---

Published as a conference paper at ICLR 2023
Because the covariance matrix X X in Eq. (10) can be expressed as a sum of rank-one terms each involving a single training example æ;, this can be used to construct an iterative algorithm for computing the closed-form ridge-regression solution.
Theorem 2. A transformer can predict according single
update:
w' (MI + +2,0,) Difl (
t0,0 1 1 L
Ztt/
(14)
with constant layers and (dd) hidden space. More precisely, there exists set of transformer parameters 9 such that, given an input matrix of the form in Eq. (12), the transformer's 's output matrix (LL) contains entry equal to w' In (Eq. (14)) at the column index where nn is input.
Discussion. There are various existing universality results for transformers (Yun et al., 2020; Wei al., 2021), and for neural networks more generally (Hornik al., 1989). These generally require very high precision, very deep models, or the use of an external "tape", none of which appear to be important for in-context learning the real world. Results in this section establish sharper upper bounds on the necessary capacity required to implement learning algorithms specifically, bringing theory closer to the range where it can explain existing empirical findings. Different theoretical constructions, in the context of meta-learning, have been shown for linear self-attention models (Schlag et al., 2021), or for other neural architectures such as recurrent neural networks (Kirsch & Schmidhuber, 2021). We emphasize that Theorem I and Theorem each show the implementation of a single step of an iterative algorithm; these results can be straightforwardly generalized to the multi-step case by "stacking" groups of transformer layers. As described next, it is these iterative algorithms that capture the behavior of real learners.
WHAT COMPUTATION DOES AN IN-CONTEXT LEARNER PERFORM?
The previous section showed that the building blocks for two specific procedures gradieen descent on the least-squares objective and closed-form computation of its minimizer are implementable by transformer networks. These constructions show that, in principle, fixed transformer parameteriza- tions are expressive enough to simulate these learning algorithms. When trained on real datasets, however, in-context learners might implement other learning algorithms. In this section, we in- vestigate the empirical properties of trained in-context learners in terms of their behavior. In the framework of Marr's (2010) "levels of analysis", we aim to explain ICL at the computational level by identifying the kind of algorithms to regression problems that transformer- based ICL implements.
4.1 BEHAVIORAL METRICS
Determining which learning algorithms best characterize ICL predictions requires first quantifying the degree which two predictors agree. We use two metrics to do so:
Squared prediction difference. Given any learning algorithm A that maps from a set of input-
output pairs D [ V1. ... Tn, ynl to predictor f(x) (z)) we define the squared prediction difference (SPD):
SPD(A1 Az) D (pp() x (ppp(
(A,(D)(z)) Az())=²))
(15)
where D is sampled as in Eq. (8). SPD measures agreement at the output level, regardless of the algorithm used to compute this output.
Implicit linear weight difference. When ground-truth predictors all belong to a known, paramet- ric function class (as with the linear functions here), may also investigate the extent to which different learners agree on the parameters themselves. Given an algorithm A, we sample a con- text dataset D above, and an additional collection of unlabeled test inputs Dx then compute A's prediction on each x/ yielding a predictor-specific dataset D = {(x// }: = ((D)((()) } encapsulating the function learned by A. Next we compute the implied param- eters:
{x/). We
+ = arg min Σ
(0: w x2)22
(16)
5

---

Published as a conference paper at ICLR 2023
08
0.30
200
0.25
U.N.A
0.4
0.00
0.55
00
0.00
#esemplon
(a) Predictor-ICI. fit w.r.t. prediction differences.
(b) Predictor-ICI. fit w.r.t implicit weights.
Figure 1: Fit between ICL and standard learning algorithms: We plot (dimension normalized) SPD and ILWD values between textbook algorithms and ICL noiseless linear regression with d = 8. GD(a) denotes one step of batch gradient descent and SGD(a) denotes one pass of stochas- tic gradient descent with learning rate a. Ridge(A) denotes Ridge regression with regularization parameter A. Under both evaluations, in-context learners agree closely with ordinary least squares, and are significantly less well approximated by other solutions to the linear regression problem.
We can then quantify agreement between two predictors A1 and A by computing the distance between their implied weights in expectation over datasets:
ILWD(A1, A2) EDED,, [ A. 112
(17)
When the predictors are not linear, ILWD measures the difference between the closest linear predic- tors (in the sense of Eq. (16)) each algorithm. For algorithms that have linear hypothesis space (e.g. Ridge regression), we will use the actual value of V instead of the estimated value.
4.2 EXPERIMENTAL SETUP
We train Transformer decoder autoregresively on the objective in Eq. (8). For all experi- ments, we perform a hyperparameter search over depth L E 21, 2, 4, 8, 12, 16), hidden size W {16, 32, 64, 256, 512, 1024) and heads M E 21, 4, 4,84, 8). Other hyper- eyper-parameters are noted in Appendix D. For main experiments, we found that L 16, H 512, M 4 minimized loss on validation set. We follow the training guidelines els for 500, 000 iterations, with each in-contex, dataset consisting of 40 (æ. y) pairs. For the main experiments we generate data according to p(w)= = (0, and p(x) = (0,
Garg et al. (2022), and trained mod-
4.3 RESULTS
ICL matches ordinary least squares predictions noiseless datasets. We begin by comparing a (L = 16, H = 512, M = transformer against variety of reference predictors:
k-nearest neighbors: the uniform variant, models predict 9; = E, My where j the top-3 closest data point to x; where j i. In the weighted variant, weighted average yi x 3 E, [r. z is calculated, normalized by the total weights of the yys.
One-pass stochastic gradient descent: ús w x, where w, is obtained by stochastic gradient descent on the previous examples with batch-size equals to 1: w; w;-1 -7-19/-1 Aw;-1).
One-step batch gradient descent: Й = w.J where w; is obtained by one of step gradi- ent descent on the batch of previous examples: w; wo - 2ax(X T "ww " x Awo).
Ridge regression: We compute Ú; wn where wT (XTx x - We denote the case of X 0 as OLS.
The agreement between the transformer-basee ICL and these predictors is shown in Fig. 1. As can be seen, there are clear differences in fit to predictors: for almost any number of examples, normalized SPD and ILWD are small between the transformer and OL.S predictor (with squared error less than 0.01), while other predictors (especially nearest neighbors) agree considerably less well.
When the number of examples less than the input dimension = 8, the linear regression problem under-determined, in the sense that multiple linear models exactly fit the in-context training

---

Published as a conference paper at ICLR 2023
datig XCL)
125
1-4-0M
1.5wdh
1:30-04
3/304
112M4
528-04
25344
(19, XC3
2 148-4
WA8c4d
3
21866
19
6-4-0M
17
(50/L0/ +0
(8251 U1B
(0 50.5
(0.5) 0/
0.50055
Figure 2: ICL under uncertainty: With problem dimension d= 8, and for different values of prior MSPD values for each predictor
variance 72 and data noise o², display ciimension-normaliine pair, where MSPD is the average SPD value over underdetermined region of the linear problem. Brightness proportional with MSPD ICL most closely follows the minimum- minimum-Bayes-Tiin Bayes-risk Ridge regression output for all x²
7r values.
dataset. In these cases, OLS regression selects the minimum-norm weight vector, and (as shown in Fig. 1), the in-context learner's predictions are reliably consistent with this minimum-nome predictor. Why, when presented with ambiguous dataset, should ICL behave like this particular predictor? One possibility is that, because the weights used to generate the training data are sampled from a Gaussian centered zero, ICL. learns to output the minimum Bayes risk solution when predicting under uncertainty. Building on these initial findings, our next set of experiments investigates whether ICL behaviorally equivalent Bayesian inference more generally.
ICL matches the minimum Bayes risk predictor on noisy datasets. To more closely exam- ine the behavior of ICL algorithms under uncertainty. we add noise to the training data: now we present the in-context dataset as a sequence: (x1)) + 61...., En. [(x,) + 6a] where each & (O, 2) Recall that ground-truth weight vectors are themselves sampled from a Gaussian distribution; together, this choice of prior and noise mean that the learner cannot be certain about the target function with any number of examples.
Standard Bayesian statistics gives that the optimal predictor for minimizing the loss in Eq. (8)
is: (18)
Ely xr DJ.
This is because, conditioned on r and D, the scalar ((z, Elyl>. x D) is the minimizer of the loss E[(y- Iz, D). and thus the estimator ÿ is the minimzier of [[(y - ú) E,, g) D]]. For linear regression with Gaussian priors and Gaussian noise, Bayesian estimator Eq. (18) has closed-form expression:
= (xTx+ -21)''XTY;
i x
(19)
Note that this predictor has the same form as the ridge predictor from Section 2.3, with the reg- ularization parameter set to F In the presence of noisy labels, does ICI. match this Bayesian predictor? We explore this by varying both the dataset noise G² and the prior variance 72 (sampling w (0, 7))). For these experiments, the SPD values between the in-context learner and various regularized linear models is shown in Fig. 2. As predicted, as variance increases, the value of the ridge parameter that best explains ICL behavior also increases. For all values of o², the ridge parameter that gives the hest fit to the transformer behavior is also the that minimizes Bayes risk. These experiments clarify the finding above, showing that ICL in this setting behaviorally matches minimum- Bnnimum-Bayes-riss predictor. We also note that when the noise level a + 0+, the Bayes predictor converges to the ordinary least square predictor. Therefore, the results on noiseless datasets studied in the beginning paragraph of this subsection can be viewed corroborating the finding here in the setting with o + +
ICI exhibits algorithmic phase transitions as model depth increases. The two experiments above evaluated extremely high-capacity models in which (given findings in Section 3) computa- tional constraints are not likely to play role in the choice of algorithm implemented by ICL. But what about smaller models-dos the size of in-context learner play role in determining the learning algorithm it implements? To answer this question, we run two final behavioral experi- ments: one in which vary the hidden size (while optimizing the depth and number of heads in Section 4.2), then vary the depth of the transformer (while optimizing the hidden size and number of heads). These experiments are conducted without dataset noise.
7

---

Published as a conference paper at ICLR 2023
0.4
(DLA B0B5
(Ritpeth St,
0.3
(ONNO anilann
MSPDLA
0.2
7SC2118.04 KIL)
(CENU ICLJ
0.1
(orthe/01)) KTT)
(CENC (CL)
0.0
20
2'
2
2a
2
2 H (hidden size)
(depth)
(a) Linear regression problem with d
0.4
(Rudan XCL)
0.3
MSPD(A,.
OUNNLL wosteeds KCL
(KCNN/ vntoon) (C3)
0.2
(SGDX085), ICC)
v0bebot: KCt.
(ODABOTI
0.0
2
2
(depth)
(hidden
(b) Lincar regression problem with 16
We show SPD averaged over underdetermined re- gion of the linear regression problem. In-context learners behaviorally match ordinary least squares predictors there is enough number of layers and hidden sizes. When varying model depth (left background), algorithmic "phases" emerge: models transition between being closer to gradient de- scent, (red background), ridge regression (green background), and OLS regression (blue).
Figure 3: Computational constraints on ICL:
Results are shown in Fig. 3. When we vary the depth, learners occupy three distinct regimes: very shallow models (IL) are best approximated by single step of gradient descent (though not well- approximated in an absolute sense). Slightly deeper models (2L-4L) are best approximated by ridge regression, while the deepest (+8L) models match OLS as observed Fig. 3. Similar phase shift occurs when we vary hidden size in 16D problem. Interestingly, we can read hidden size re- quirements to be close to ridge-regression-lik solutions as 16 and H 2 32 for 8D and 16D) problems respectively, suggesting that ICL. discovers more efficient ways to use available hidden state than our theoretical constructions requiring (dd)) Together, these results show that ICL does not necessarily involve minimum-risk prediction. However, even in models too computationally constrained to perform Bayesian inference, alternative interpretable computations can emerge.
DOES ICL ENCODE MEANINGFUL INTERMEDIATE QUANTITIES?
Section 4 showed that transformers are a good fit to standard learning algorithms (including those constructed in Section 3) at the computational level. But these experiments leave open the question of how these computations are implemented at the algorithmic level. How do transformers arrive at the solutions in Section 4, and what quantities they compute along the way? Research on extract- ing precise algorithmic descriptions of learned models is still in its infancy (Cammarata et al., 2020; Mu & Andreas, 2020). However, we can gain insight into ICL by inspecting learners' intermediate states: asking what information is encoded in these states, and where.
To do so, we identify two intermediate quantities that we expect to be computed by gradient descent and ridge-regression variants: the moment vector X Y and the (min-norm) least-square estimated weight vector WOLs. each calculated after feeding n exemplars. We take a trained in-context learner, freeze its weights, then train auxiliary probing model (Alain & Bengio, 2016) attempt to recover the target quantities from the learner's 's hidden representations. Specifically, the probe model
takes hidden states at a layer (1)) as input, then outputs the prediction for target variable. We define probe with position-attentinn that computes (Appendix E):
= softmax(3, 8.))
(20)
i FF, W.H(!)
(21)
8

---

Published as a conference paper at ICLR 2023
0.06
X gole kyes
Lagen Ei's adtemtinn buntrap CaM
Fale
0075
MSE
0.04
0.50
Tageel
0.02
0.25
10
14
16
10
15 20
25 95 40
Loper
positon
You probe onco lesr
Layse I2tammitionn hormap courii
Pare
M6E
-0.5
0.00
25
Figure 4: Probing results d 4 problem: Both moments Y (top) and least-square solu- tion WOLs (middle) are recoverable from learner representations. Plots in the left column show the accuracy of the probe for each target in different model layers. Dashed lines show the best probe accuracies obtained on a control lask featuring a fixed weight vector w = 1. Plots in the right col- umn show the attention heatmap for the best layer's probe, with the number of input examples on the x-axis. The value of the target after n exemplars is decoded primarily from the representation of yr. or, after n examplars, uniformly from 8n24-
We train this probe to minimize the squared error between the predictions and targets L(v, ö) (v - -/. The probe performs two functions simultancously: its prediction error on held-out rep- resentations determines the extent to which the target quantity is encoded, while its attention mask, a identifies the location in which the target quantity is encoded. For the FF term, we can insert the function approximator of our choosing; by changing this term which the targel quantity is encoded-e.g if FF is a linear model and the probe achieves low error, then we may infer that the target is encoded linearly.
determine the manner in
For each target, we train a separate probe for the value of the target on each prefix of the dataset: i.c. one probe decode the value of w computed from a single training example, second probe to decode the value for two examples, etc. Results are shown in Fig. 4. For both targets, a 2- layer MLP probe outperforms a linear probe, meaning that these targets are encoded nonlinearly (unlike the constructions in Section 3). However, probing also reveals similarities. Both targets are decoded accurately deep in the network (but inaccurately in the input layer, indicating that probe success is non-trivial.) Probes attend to the correct timestamps when decoding them. As in both constructions, TTy appears to computed first, becoming predictable by the probe relatively early in the computation (layer 7); while becomes predictable later (around layer 12). For comparison, we additionally report results on a control task in which the transformer predicts yss generated with fixed weight vector significantly worse at recovering moment matrices (see Appendix E for details).
1 (so no ICI. is required). Probes applied to these models perform
6 CONCLUSION
We have presented a set of experiments characterizing the computations underlying in-context learn- ing of linear functions in transformer sequence models. We showed that these models are capable in theory of implementing multiple linear regression algorithms, that they empirically implement this range of algorithms (transitioning between algorithms depending on model capacity and dataset noise), and finally that they can be probed for intermediate quantities computed by these algorithms.
While our experiments have focused on the linear case, they can be extended to many learning problems over richer function classes e.g. to a network whose initial layers perform non-linear feature computation. Even more generally, the experimental methodology here could be applied to larger-scale examples of ICL, especially language models, to determine whether their behaviors are also described by interpretable learning algorithms. While much work remains to be done, our results offer initial evidence that the apparently mysterious phenomenon of in-context learning can be understood with the standard ML toolkit, and that the solutions to learning problems discovered by machine learning researchers may be discovered by gradient descent well.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
We thank Evan Hernandez, Andrew Drozdov, Ed Chi for their feedback on the early drafts of this paper. At MIT, Ekin Akyürek supported by an MIT-Amazon ScienceHub fellowship and by the MIT-IBM Watson Al Lab.
REFERENCES
Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes. ArXiv preprint, abs/1610.01644, 2016. URL https::s///aii
Marcin Andrychowicz, Misha Denil, Sergio Gomez Colmenarejo, Matthew W. Hoffman, David Pfau, Tom Schaul, and Nando de Freitas. Learning to learn by gradient descent by gradi- ent descent. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 29: Annual Con- ference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 3981 -3999 2016. URL https: http::/:d///ooeeenngsss/. neurips httpppaapr//////////// 0a1.8855555222444445555558333333 Abstrrraatbdd-tbbaatccccaaaaaaae Abtt
Jimmy Lei Ba. Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. ArXiv preprint, 0bs//607.00m 167.004500 2016. URL http:/::/wwwxvv.////eeee
06450.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert-Voss Gretchen Krueger. Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse. Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: An- nual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, Decem- ber 6-12, 2020, virtual, 2020. URL httpss:/prooceeding..oo ieuruii wwteurips.ccppppeeeeeeeeeee p20ppppeer/220/haaaa 1457c0d6bf 1457:00666 BB433644ff47797bbbb555
html.
Nick Cammarata, Shan Carter, Gabriel Goh, Chris Olah, Michael Petrov, Ludwig Schubert, Chelsea Voss, Ben Egan, and Swee Kiat Lim. Thread: circuits. Distill, 5(3)=224, 2020.
Stephanie CY Chan, Adam Santoro, Andrew K Lampinen, Jane X Wang, Aaditya Singh, Pierre H Richemond, Jay McClelland, and Felix Hill. Data distributional properties drive emergent few- shot learning in transformers. ArXiv preprint, fals2205.050055 2022. URL httttttss///xiitt org/ abs/ 05055.
Chi Chen, Maosong Sun, and Yang Liu. Mask-align: Self-supervised neural word alignment. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 4781-4791, Online, 2021. Association for Computational Linguistics. doi: 10.18653/v1/ 2021.acl-long.33 URL https::t///.. httt::://aalaanhologggo. wrwph///oooyoooo22211111
wwg/ acl- llng..363
Yanda Chen, Ruiqi Zhong, Sheng 7ha, George Karypis, and He He. Meta-learning via lan- guage model in-context tuning. In Proceedings of the 60th Annual Meeting of the Associa- fion for Computational Lingwistics (Volume 1: Long Papers). pp. 719-730, Dublin, Ireland, 2022. Association for Computational Linguistics. doi: 10.118655/112022.aal-ooogg http::///aclanttoogy.oo/////////
URL
acl aol.long.55.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, al. Palm: Scaling language modeling with pathways. ArXiv preprint, abs/2204.02311, 2022. URL 02311.
https://arxiv.. wtt.:/orgggb/bb/b/eeo/oo
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation
of deep networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th Interna- tional Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-1 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 1126-1135. PMLR, 2017. URL http://proceediim mlr. wttppp///s////nni//// httm...mmmm..
10

---

Published as a conference paper at ICLR 2023
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. ArXiv, abs/2208.01066, 2022.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). ArXiv preprint, abs/1606.08415,, 2016. URL hww.hhtps////tt http://aaavv.../abbbeeesse orggabs//6000 08415.
Arthur E Hoerl and Robert W Kennard. Ridge regression: Biased estimation for nonorthogonal problems. Technometrics, 12(1):55-67, 1970.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are uni- versal approximators. Neural networks, 2(5):359-366, 1989.
Louis Kirsch and Jürgen Schmidhuber. Meta learning backpropagation and improving Advances in Neural Information Processing Systems, 34:14122-141344 2021.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. ArXiv preprint, abs/2210.14215, 2022. URL https://aaxtis .rgg/ abs/2210. 14215.
David Marr. Vision: A computational investigation into the human representation and processing of visual information. MIT press, 2010.
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetalCI.: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chapter of the As-
sociation for Computational Linguistics: Human Language Technologies, pp. 2791-2809, Seat- tle. United States, 2022. Association for Computational Linguisties. doi: 10.18653/w1/2022. URL htttp::/acclanthol.e org/ 2022..2..222 maacl-main.20
Jesse Mu and Jacob Andreas. Compositional explanations of neurons. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Confer- ence on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtual, 2020. URL /ttps://proceeding www.gi.....is...mmmm pwtpppppe/2//////eeeeeeeep c74956f 0w1.998f833bb/8ccc67777776222222 -bsp:222t--bttttctttttt. html.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, T. J. Henighan, Benjamin Mann, Amanda Askell, Yushi Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Gan- guli, Zac HatfieldDoodds Danny Hermandez, Scott Johnston, Andy Jones, John Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom B. Brown, Jack Clark, Jared Kaplan, Sam McCan- dlish, and Christopher Olah. In-context learning and induction heads. 2022.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th Interna- fional Conference on Machine Learning, ICML, 2021, 18-24 July 2021, Virtual Event, vol- ume 139 of Proceedings of Machine Learning Research, pp. 9355-9366. PMLR, 2021. URL http:://proceeings.. www. wttpppsees/////eeeee/ httml.
Juergen Schmidhuber, Jieyu Zhao, and Marco A Wiering, Simple principles of metalearning. 1996.
Jack Sherman and Winifred J Morrison. Adjustment of an inverse matrix corresponding to change in one element of given matrix. The Annals of Mathematical Statisticy, 21(1):124-121 1950.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan. and Ro- Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Confer- ence on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998-6008. 2017. URL https: http:::::///preeeenneee 3f5ee243547dee91 fbd053c1 c4a845aa http:/orrreddnnn..eriir..ppppssp newrips....... cc/ cc/pappr/200 heteer/211//haa ///ash// html.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: study approximating turing machines with transformers. ArXiv preprint, abs/2107.13163. 2021. URL
http:://aaxxv.ooo/abbb/////
1316..
11

---

Published as a conference paper at ICLR 2023
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. ArXiv preprint, 0bs/2111.02066 2021. URL https: //arxiv.org/abs/2111 02080..
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference. ArXiv, abs/21 1.02080, 2022.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequene functions? In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https: wwtp:wwwwwwwwwwweeeeeeeeeeeeeeee httt//ooummidByyMMoNEE
net
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models, 2022.
Kaiyang Zhou, Jingkang Yang, Chen Change Loy, and Ziwei Liu. Learning to prompt for vision- language models. International Journal of Computer Vision, 130(9):2337-2348, 2022.
12

---

Published as a conference paper at ICLR 2023
A THEOREM 1
The operations for 1-step SGD with single exemplar can be expressed following chain (please
proof's for the Transformer implementation of these operations (Lemma 1) in Appendix C):
(4,11
(move x:)
aff((,11++000000444-000 w)
(w x:)
4F((1--22-0000022444--------111- =-1
(w æ y)
m
(x x y))
af((0.00.3-223333))==..)))
(write w)
aF W1 1,2 = - (x x y) Xw)
38 (3+ 34, 3d), W1
2d,3
I, 2, -2x,
(w')
3d)
(move w')
(; (3, 2d, 3d), (1, d), 3d, 4 4+ 3d))
(w' 222)
This will map:
5/1
0
0
1/
21
0
.72
w 211
w I1
w 222
21
Ø1 y/
W 222
21
(w [11
y)
H
10
16
W
TIw T Ø1 Xw
21 (w
11
y)
Nw
T2
Ø Aw
w 2a(x Ø1 Aw)
w'
w
2 x22
21
Aw)
2a(r1w) w Aw)
w'
w/
(w 20(x1 x Aw)) x1
w 1
w' X2
We can verify the chain of operator sep--by-step. In each step, we show only the non-zero rows.
mov(; 1, 0, (1, d), (1, + d))
(move a:)
91
0
91 0
0
22
21
222
aff(; (1, + d).
0. (1
d, + d), W1 w)
(w x))
0
11
0
V1
21
301
271
22
21
x2
aff(; (1 + 2,2 + d), (0, 1), (2 + d, 3 + d), W1 = I, W2
II)
(w T x - y)
1/1
0
1/1
T1
21
31
16 Σ1
1 T
T1
W T
D22
w
W
W.
T2
W Σ1
W 37
!1
W
12
mul(: (: d. 1, (1, + d), (2 + d,3 + d), (3 + d,3 2d))
(x(w T x y))
8/1
0
V1
F2
Æ1
H
w T D1
W I1
w
.2
W 201
W J1
w 12
T x1
w
w
œ -
w
1 F1
W
I'1
91
W
12
I w
21
31
-yy
22
0. 2d.3 +333+3d).4 w.
(write
13

---

Published as a conference paper at ICLR 2023
0
y1
0
51
0
I1
I
222
22
w 21
W II
w 2/1
W 222
W T 21
W 21
W 21
W
21
W 22
21
31 (w)
21
y)
T2W
21
31 10 21
I1
(w)
8/1
2'2 W
01
aff(; (3+d 3++dd)+ (3+2,, (3 (3+3d, 3+ 4d), Wi = 1, W2 = 2X) (x x ---y y) 2Xw)
0
1/1
0
1/1
21
I1
22
31
.22
W 21
w I1
w 22
W 21
w
W
22
W T1
W T J1 - y
w 2/2
W
21
W
21
W
22
I w T I
3 (w' 21 - y)
I T2 x1
2 10
11
11 (w)
21
y)
F2
21
W
W
21 w
21
Aa
[w T1 y) Xw
22 Aw
aff(; (3
2d, + 3d), (3 3d, + 4d), (3 + 2d,3 3 + 3d), W I, = -2a,)
(w))
0
211
I1
222 222
21
w T1
21
I
w
T1 w
x1
T1 (w T 21 y)
T2W X1
w
W
37
Xw
(w I1
y)
Aw
F2W 31
Aw
V1
0
#1
T1
22
x1
22
w
w y
222
w x1
I1 (w) y
w - 20r [ w 11 Xw)
w
2a 222 w σ Xw)
W T1 Aw
D (w T1 - 3) Aw
22W
I1 - Aw
m, 2d. 3d), + -22, 3 3d)
(move wr)
1/1
0
21
01
872
w
W
w. T .222
w
w I1
w
I1 w
D11
D1 (w" Ø1 - y)
TzW 41
20( ( w
.1
Xw)
w/
20 22 w
41
Xw)
21 w
Xw
21 (w)
D1 - v)
Xu
32 w
R1
Xw
1/1
21
D1
22
w
w T
T
w x1
to" T I1 y
w T
T2
4 w T F1
I (w T x - 11)
T2W x1
20( w 41 Xw)
w/
2a (. w
4r Xw)
Z1t T1 Aw
21 (w
x1 3) Aw
x2 I1 Aw w/
20(x1 x1 Aw)
w/
mul (1 (3 2d. 34), (1,1 d), (3 3d, 3d))
(wo x2)
14

---

Published as a conference paper at ICLR 2023
1/1
0
I
222
w
W T
w T
w
w I1
22
I1 w 21
T1 (w) 21 - y)
I2" I1
w
(x u 21 Aw)
w'
20 22 w
21 Xw)
30 31 Aw
21 (w x1 v)
Aw
322 w 21
Xw
20(2 1 w
11 Xw)
w
w
0
21
21
D2
w .n
D22
w: x
w
w T
21 (w T I y)
I œ1
to I1
w 22( w 21 - Aw)
w 2a T2 w
20 - Xw)
21W T1 Au
21("") 21 y)
Aw
x2 21 Au
20(x1 w x1 Xw)
w w 21
w
(w 20 T 21 Aw)) æ1
X2
We obtain the updated prediction in the last hidden unit of the third time-step.
Generalizing to multiple steps of SGD. Since w' written in the hidden states, we may repeat this iteration to obtain ús = u=" æy where w" the one step update w'' 20(22 222 y2 + Aw, requiring total of (n) layers for single pass through the dataset where is the number of examplers.
As empirical demonstration of this procedure, the accompanying code release contains ref- erence implementation of SGD defined in terms of the base primitive provided in an anymous links https: //ttps://icl.... s3. us-east-2 2. amazonaws wwt./ooooort/eeec (primi tives, sgd, ridge] py (to preserve the anonymity we did not provide the library dependencies). This implementation predicts consecutive SGD updates on
Un un, Fn where w,, the weight vector resulting from previous examples. It can be verified there that the procedure requires O(n + d) hidden space. Note that, is not (nd)) because we can reuse spaces for the next iteration for the intermediate variables, an example of this performed in (w') step above highlighted with blue color.
B THEOREM 2
We provide a similar construction to Theorem 1 (please see proofs for the Transformer implemen- tation of these operations in Appendix C, specifically for div see Appendix C.6)
mov(; d1, d)) mL1:4+11..11--0000)00--444+++
(move xi)
(x1
(A5 = )
f0=1)
m11::4441164032222222-0000444443 (VA = x: +)
mu(:2:111.2222"13,,3,33334444444 4d 2d²))
412,1 x1) )
oFF[44444|2.22221144444332421113 AA-'u=
(+2
21f6f(4444422110de33322222244443 (right term)
6Ff11122|222220,,444323444433333 W1 = 1, W₂ = -1)
(A
(Aitxr)
(A
15

---

Published as a conference paper at ICLR 2023
(2 4d 3d 5d +dd)) 4 mul(;d, 1, (2 + 4d 3d 5d 3d (1, + +d,, (2 + 3d, 33 5d + 3d)))
3d2))
(move w²)
(w' x2)
Note that, in contrast to Appendix A, we need ((d)) space to implement matrix multiplications. Therefore over-all required hidden size is ((d))
As Theorem 1, generalizing to multiple iterations will least require ((n) layers, as we repeat the process for the next examplar.
C LEMMA 1
All of the operators mentioned in this lemma share a common computational structure, and can in fact be implemented as special cases "base primitive" we call RAW (for Grotd-Arttmeti--Wee This operator may also be useful for future work aimed at implementing other algorithms.
The structure of our proof of Lemma is follows:
1. Motivation of the base primitive RAW.
2. Formal definition of RAW.
3. Definition of dot, aff, mov in terms of RAW.
4. Implementation RAW terms of transformer parameters.
5. Brief discussion of how to parallelize RAW, making it possible to implement mul.
6. Seperate proof for div by utilizing layer norm.
C.I RAW OPERATOR: INTUITION
At a high level,
primitives in Lemma involve a similar sequence operations:
1) Operators read some hidden units from the current or previous timestep: dot and aff read from two subsets of indices in the current hidden state h&", while mov reads from previous hidden state he. This selection is straightforwardly implemented using the attention component of transformer layer.
We may notate this reading operation as follows:
Σ
① h [r]
(22)
W
KEK()
Read with Attention
Here denotes list of indice to read from, and K denotes a map from current timesteps to target timesteps. For convenience, we use Numpy-like notation to denote indexing into vector with another vector:
Definition C.1 (Bracket), a[.] is Python index notation where the resulting vector, x' = æ(r]:
yj Br, j 1,
The first step of our proof below shows that the attention output ( can compute the expression above.
2) Operators perform element-wise arithmetic between the quantity read in step and another set of entries from the current timestep: This step takes different forms for aff and mul (mov ignores values at the current timestep altogether).
For notational convenience, will use h to refer to sequence of hidden states (instead of H in Eq. (1).), by will be the hidden state time step t!
16

---

Published as a conference paper at ICLR 2023
Wa [K(]]
Σ
+ [r]
Wh
(multiplicative form)
(23)
Read with Atlention
W.
Σ h [r]
+ Wh [s]
(additive form)
(24)
K(O)
(E
Read with Attention
The second step of the proof below computes these operations inside the MLP component of the transformer layer.
3) Operators reduce, then write to the current hidden state Once the underlying element-wiss operation calculated, the operator needs to write these values to the some indices current hidden state, defined by list of indices w. Writing might be preceded by reduction state (e.g. for computing dot products), which can be expressed generically as a linear operator W.. The final form of the computation is thus:
Elementwise operation
h
[1+1)
[v]
t
Wo
W
Σ
(2) h
( [r]
+ Wh;
(25)
(i)
KEK()
Read with Attention
Here, + means that the other indices $ are copied from h/-1
C.2 RAW OPERATOR DEFINITION
We denote this "master operator" as RAW:
Definition C.2. RAW(h; &, r,w, Wo, Wa W, K) is function R# + R It is parameter- ized by elementwise operator E [+,@). three matrices W Rdx[ss, W. Rdxirl, W. glw[xd three index sets and and a timestep map K Z + (Z+*). Given an input matrix h, it outputs a matrix with entries:
(1-1)
W.
((TT))
We
Σ
[r]
+ [n]
1)
i= 1,..,, T;
(26)
[(O]]
kEK())
h (1-1)
h
(2)
I...,T::
(27)
We additionally require that j K(1)
= j < é (since self-attention is causal.)
(For simplicity, did not include possible bias term in linear projections W., W.. W, we always assume the accompanying bias parameters by ba, when needed)
C.3 REDUCING LEMMA 1 OPERATORS TO RAW OPERATOR
Given this operator, the mul and div. Instead of the matrix multiplication operator mul, we will first show the dot product dot (a special of mul), then later in the proof, will argue that parallelize these dot products in Appendix C.5 to obtain mul. We will show how to implement div separately in Appendix C.6.
define each primitive in Lemma using single RAW operator, except
17

---

Published as a conference paper at ICLR 2023
Lemma 2.
We can define mov, aff operator, and the dot product case of mul Lemma by using
single RAW operator
dot(h; (1,5), ",00::/0)000=0m.....//0000000000
mu
=RAW(h; O = I,W, = W. = (1 -,,
aff(h: (1, (f.5), (""",,
-RAW(h; +. W W., = W2, W. =b [,
(0.()))))
mov(h; (1,5).(( (('j'))
= =RAW(h; +. W = 0, W. = I, W. (,
((t,
Proof. Follows immediately by substituting parameters into Eq. (26).
C.4 IMPLEMENTING RAW
It remains only to show:
Lemma 3. A single transformer layer can implement the RAW operator: there exist settings of trans- former parameters such that, given arbitrary hidden matrix h input, the transformer computes h' (Eq. (26)) as oulpul.
Our proof proceeds in stages. We begin by providing specifying initial embedding and positional embedding layers, constructing inputs to the main transformer layer with necessary positional in- formation and scratch space. Next, we prove three useful procedures for bypassing (or exploiting) non-linearities in the feed-forward component of the transformer. Finally, we provide values for maining parameters, showing that we can implement the Elementwise and Reduction steps described above.
C.4.1 EMBEDDING LAYERS
Embedding Layer for Initialization: Rather than inserting the input matrix h directly into the transformer layer, we assume (as is standard) the existence linear embedding layer. We can set this layer to pad the input, providing extra scratch space that will be used by later steps of our implementation.
We define the embedding matrix We as:
Wa ((a+I)x(at1)
0
(28)
Then, the embedded inputs will
i W. = [0, 2úg 0 On-d-1
(29)
9. W.y: = = 1
(30)
Implementing RAW ultimately requires con- trolling which position attends to which position in each layer. For example, may wish to have layers in which each position attends to the previous position only, or in which even positions at- tends other even positions. We can utilize position embeddings, p;, control attention weights. In a standard transformer, the position embedding matrix a constant matrix that is added to the inputs of the transformer after embedding layer (before the first layer), the actual input to the transformer is:
Position Embeddings for Attention Manipulation:
h ã. P:
(31)
We will use these position embeddings to encode the timestep map K. To do this, we will use 2p units per layer (p will be defined momentarily). p units will be used to encode attention keys k,, and the other p will be used encode queries q
We define the position embedding matrix as follows:
Pt (0a41; 4% 20 ....... k 4
(L)
(32)
18

---

-N it1 Published as a conference paper at ICLR 2023
With K' encoded in positional embeddings, the transformer matrices We and Wx are easy to define: they just need to retrieve the corresponding embedding values:
0
W =
dpxp 0exp
W
Opxp
Tpxp
(33)
The constructions used in this paper rely on two specific timestep maps K, each of which can be implemented compactly in terms of & and q:
Case 1: Attend to previous token.
This can be constructed by setting:
k; E; qi = Ne:-1
where N is sufficiently large number. In this case, the output of the attention mechanism will be:
a softmax
(W])..) (W h::))
softmax
9: [ki.... k:l)
softmax
0 N...,,N.0)
0,
1
0]
(1-1)
Case 2: Attend to a single token. For simpler patterns, such as attention to a specific token:
(t) i2t
K(t)
(34)
10 i<t
only hidden unit is required. We set:
k
N i=t
= N
from which it can be verified (using the same procedure as in Case 1) that the desired attention pattern is produced.
Intricacy: How K be empty? We can also cause K (() attend an empty set by assuming the softmax has extra ("imaginary")) timestep obtained by prepending a 0 to attention vector pot-hoc (Chen al., 2021).
Cumulatively, the parameter matrices defined in this subsection implement the Read with Attention component of the RAW operator.
C.4.2 HANDLING & UTILIZING NONLINEARITIES
The mul operator requires elementwise multiplication of quantities stored in hidden states. While transformers are often thought of as only straightforwardly implementing affine transformations on hidden vectors, their nonlinearities in fact allow elementwise multiplication a high degree of approximation. We begin by observing the following property of the GeLU activation function in the MLP layers the Transformer network:
Lemma 4. The GeLU nonlinearity can be used to perform multiplication: specifically,
V (eLLU(x + GeLU((z) - GeLU(y)) = ry + (z:3 +y²)
(35)
19

---

Published as a conference paper at ICLR 2023
Proof. A standard implementation of the GeI.U nonlinearity defined as follows:
1 + tanh (v2 (x+ (=+000004477)))))
I GeLU(x)= 2
(36)
Thus
2 2 -x² + 0(x2)
GeLU(x)= = + 2
(37)
2 GeLU(r + v) GeLU((z) GeLU(U) = V + O +
=
(38)
क / + v) - GeLU(() - GeLU
zy ~
(39)
For small x and y, the third-order term vanishes. By scaling inputs down by constant before the GeLU layer, and scaling them up afterwards, models may use the GeLU operator to perform elementwise multiplication.
We can generalize this proof other smooth functions discussed further in [TODO REF]. Previous work also shows, in practice, Transformers with ReLU activation utilize non-linearities to get the multiplication in other settings.
When implementing the aff operator, we have the opposite problem: we would like the output of addition to be transmitted without nonlinearities the output of the transformer layer. Fortunately, for large inputs, the GeL.U nonlinearity is very close to linear; to bypass it suffices to add to inputs a large N:
Lemma 5. The GeLU nonlinearity can be bypassed: specifically,
GeL.U( + xx) N x N N>1
(40)
Proof.
1 tanh (v2 (N 0011715N7)) )) N
GeLU(N - = N
(41)
2
N
+ 1) - N
(42)
(43)
For all verions of the RAW operator, it is additionally necessary to bypass the LayerNorm operation. The following formula will be helpful for this:
Lemma 6. Let N be a large number and X the LayerNorm function, Then the following approxi- mation holds:
/2 L
NA Σ x. 0)) ~ 2N, -2N
Σ x. 0] N
(44)
Proof.
E[z] =
(45)
Var[æ]
- (N2 + N2 +x²)
1
2N2
(46)
L
(47)
Then,
/ 2 NNN
0) -N
V
2
NIV 2N2 x, V2L, -VLL
L
L
Σ x
2NN2
= [x, 2N, -2N Σ x
20

---

Published as a conference paper at ICLR 2023
By adding a large number N two padding locations and sum the part of the hidden state that we are interested to pass through LayerNorm, we make to the output of LayerNorm pass through. This addition can be done in the transformer's feed-forward computation (with parameter W prior to layer norm. This multiplication of v Σ N can be done first layer of ML. back, then linear layer can output/use æ. For convenience, we will henceforth omit the LayerNorm operation when it is not needed.
We may make each of these operations precise as desired (or allowed by system precision). With them defined, we are ready to specify the final components the RAW operator.
C.4.3 PARAMETERIZING RAW
We want show a layer of Transformer defined in above, hence parameterized by (W, W1. W2, (W2, W Wm)), can well-approximate the RAW operator defined in Eq. (25).
We will provide step by step constructions and define the parameters in Begin by recalling the transformer layer definition:
a = softmax (W]ha) ho) (
(48)
b a(WyH. H.))
(49)
a, W [bi,..., bm)
(50)
h (++) = FF((; W1, W2)
(51)
Wio h)) + h
(52)
Attention Output We will only use m = 2 attention heads for this construction. We show in Eq. (32) that we can control attentions to uniformly attend with a pattern by setting key and query matrices. Assume that the first head parameters W.,, W have been set in the described way to obtain the pattern function K.
Now we will set remaining attention parameters WY. W.. W wy and show hat we can make the a, + h such that:
(1) term in Eq. (4) contain the corresponding term in Eq. (25), some unused indices
Wa Σ
(a 1 + h +
Whirl [K(()] kEK(
(53)
(a h Deee = (n()))eet
(54)
Then the term on the RAW operator can be obtained by the first head's output. In order achieve that, we will set Wa as a part of actual attention value network such that WY is sparse matrix 0 everywhere expect:
(WY mm))rrl (Wam,, m I,..., |t), n E 1,...//r)
(55)
Now our first heads stores the right term in Eg. (53) in the indicies t. However, when we add the residual term h." this will change. remove the residual term, we will use another head to output h (1) by setting W w such that K(i) and wy (similar to Eq. (34)):
(WY) r(m) m 1,..../t/..... It.,,,//,, 1,.... |r|
(56)
Then, W E H222H is zero otherwise:
(W")e(m):s(m) =1 1,..., (W")gmi.emm)a m 1,... [t]
(57)
(58)
(59)
We already defined (WQ, WK VV), and V and obtained the first term in the Eq. (25) in (a, +
h
(
Dreet
21

---

Published as a conference paper at ICLR 2023
Arithmetic term Now we want to calculate the term inside the parenthesis Eq. (25). We will calculate it through the MLP layer and store in m. and substract the first term. Let's denote the input to the MLP æ; = (a, + h ) the output of the first layer u4 the output of the non- con-ineerrtt as ag and the final output as m, The entries of m. will be:
(m,)rew = W.
Wa [ kEK())
Σ
h [r)
@ Wh [s]
2r,(w)
(60)
(mi)ret = -x;(t)
(61)
(m.)e = 0
(62)
We will define the MI.P layer to operate the attention term calculated above with a part of the current hidden state by defining W1 and W2. Let's assume we bypass the LayerNorm by using Lemma 6.
Let's show this seperately for + and @ operators.
RAW -+)-) If the operator, Eq. (25) and overwrite the space where the attention output term Eq. (53) is written, and add large positive bias term N by pass GeLU as explained Lemma 4. We will use an available space in the I2 same size as t.
+, first layer of the MLP will calculate the second term in
(ui)eet = Wh [ss] + +x([t] +
(63)
(u)VEE [x,[t] + N
(64)
(ui)eww = -xx,[w] + N ( =-N
(65)
(66)
This can be done by setting W1 (weight term of the first layer of the MLP) zero except the below indices:
(Wi [W)ea(])a(n) = (W) mm m. E 1....
[t], n EEE E....,,
(67)
(W1 [mpi.tfn = +1 E I,...,
[t], E ......
(68)
(W) Depm].t[m]
m E l,... [t]
(69)
(W5 [(mm)
-1 E 1, 1,..., [w]
(70)
(71)
(72)
and the bias vector 61
(bi)eet N
(73)
(bi) N
(74)
(bi)ret = N
(75)
(bi) = -
(76)
(77)
Note the second term is added to make unused indices t U Uf become zero after the gelu, which outputs zero for large negative values. Since we added a large positive term, we make sure gelu behaved like linear layer. Thus we have,
(ui)eet Wh Wh([s] + xr[t] + N
(78)
(vi)ret [x, + N
(79)
(Di)e'ew [w](w] + N
(80)
(D:) =0 Now, we need to set Wz, simulate Wo E R|m1||11
(81)
(W2) (W2)wjm).ffri m,,m E 1, /w], 1,..../t
(82)
(W₂ fmm]./mm +1 m I,..., [t]
(83)
(W₂) Wmm].mmm +1 m 1,..., [w]
(84)
(85)
22

---

Published as a conference paper at ICLR 2023
(b2)m/m) = - Σ (Wo)mi) - N m E /w...//w
(86)
(b2) -
(87)
(b2) =0
(88)
(89)
Therefore, mw[w] Wo + WoWh >r[w] equals what we promised in Eq. (60) for
+ case. If we sum this with the residual z; term back Eq. (53), so the output of this layer can be written as:
(b)) = Wo ((NYC)
We kEK())
(K((i)]
Σ
+ Wh:
(o [s]
(90)
(ha)ttt (h)))ee
(91)
RAW -AW. If the operator, O need to three extra hidden units the same size as |t], let's name the extra indices as La to. te, and output w space. The (u,) will get below entries to be able to use where N is large number:
(u.)EEt (Wh((s) + xx(t))/N
(92)
(u:)rEt, xi[t]/N
(93)
(u:)ret. WWh|ss|NN
(94)
(ui)e'et [x + N
(95)
(u:)rre -x, [w] +N
(96)
uu ((tU, -N
(97)
(98)
All of this operations are linear, can be done Wi zero except the below entries:
(W1) (W2)e.[m].sl = (W) m. / m E 1.... |tal. n E 1, ..., [s]
(99)
(W1) e,, [m]] 1/N m I,... tal, E 1,... |t|
(100)
(Wi i) vs [m]] = 1/N m I,...
to n2 I.... [t]
(101)
(W1) +. [m].s[mm] 1/N m I,....
ttlln 1....,
(102)
(Wi [mmmi
m l,...
(103)
(W, (1(.)t(mm)
1,... |t]
(104)
(105)
and b, to:
(b (bi)eeuuruunee))
(106)
(bi)eertu)) = N
(107)
(bi)egrunnntoonnnu) N
(108)
(109)
The resulting with the approximations become:
(u.)et,. = gelu + (t)
(110)
(M))et = gelu(a;(E//N))
(111)
(Ds)rre. = gelu(Wh, gelu(Wh((ss/NN)
(112)
(N:)ret ax.(t] + N
(113)
(0:) = x.[[]] + N
(114)
(w))e(ue.uuwe.. 0
(115)
(116)
23

---

Published as a conference paper at ICLR 2023
Now, we can use the GeLUU trick Lemma 4, by setting W2
(W2) (12).(m)) (Wo) N2
E 1, I...., w n I..... ttl
(117)
(W2)umm).tunl -WW m2 N2
m L,... w 1,..,, tt
(118)
(W2))(m).tul)) (W mm, N2
m L..... w 1, 1.....,.tt |tel
(119)
V 2
(W2) [mmm..mmm = 1 m E ..... [w]
(120)
(W2) cfmm/mm.t = m
[t]
(121)
(122)
We then set bg:
(b2) (b2) N
(123)
(b2) (b2)) =0
(124)
(125)
With this, m.(w] Womn * WoWhh-1ss - a.sww) [.] and
(ha)ttt) W.
W.
Σ
n
(r)
(126)
[K(i]]
kE
(h))t = (h!)rew
(++1)
(127)
We have used 4/4] space for internal computation of this operation, and finally used [w| space to write the final result. We show RAW operator is implementable by setting the parameters of Transformer.
C.5 PARALLELIZING THE RAW OPERATOR
Lemma 7. With the conditions that K is constant, the operators are independent (i.e (ri U 84 w:) n Wiyt = 0), and there is Ex(4tt + |Wxl) available space in the hidden state, then Transformer layer can apply k such RAW operation in parallel by setting different regions of W1. W2, y and (WV), }E matrices.
Proof. From the construction above, is straightforward to modify the definition of the RAW operator to perform & operations as all the indices of matrices that we use in Appendix C.4.3 do not overlap with the given conditions in the lemma.
This makes it possible to construct Transformer layer not only to implement vector-vector dot products, but general matrix-matrix products, as required by mul. With this, we show that we can implement mul by using single layer of Transformer.
C.6 LAYERNORM FOR DIVISION
Let say we have the input [c, y, OJ calculated before the attention output in Eq. (53), and we want to divide y This trick is very similar to the on in Lemma 6. We can use the following formula: Lemma 8. using LayerNorm for division. Let N, M to be large numbers, LayerNorm the following approximation holds:
1 MNA(NNG, M
-Nc - Σ y
0)) ~ [MN, y -MN y
(128)
M
Proof.
F(x0] 0
(129)
222² 1 Σ 2² (Ne+ Σ y 2)
))
2N
Var[æ] I
(130)
M
M
L
(131)
24

---

Published as a conference paper at ICLR 2023
Then,
2 MNA((NG, M' M'
y -Nc- Σ y
M 0))
M [MN, y
-22
/2
M 2
L
L y
2 MN
-V
L
2
L
Σ
MNC'
MMN 0]
wamted resull
To get the input to the format used in this Lemma, we can easily use W to convert the head outputs. Then, after the layer norm, this way, we can approximate scalar division in one layer.
W to pull the back and write it to the attention output. By
Lemma 1
By Lemmas 2, 3, 3, and 8; constructed the operators in Lemma using single layer of a Transformer, thus proved Lemma 1n
D DETAILS OF TRANSFORMER ARHITECTURE AND TRAINING
We perform these experiments using the Jax framework on P100 GPUs. The major hyperparameters used these experiments are presented in Table 1. The code repository used for reproducing these experiments will be open sourced at the time of publication. Most of the hyperparameters adapted from previous work Garg et al. (2022) to be compatible, and we adapted the Transformer architecture details. We use Adam optimizer with cosine learning rate scheduler with warmup where number of warmup steps set to be 1/5 of total iterations. We use larned absolute position embeddings.
Parameter
Search Range
Number of heads
1, 4, 8 1,2, 16, 32, 64, 256, 512, 1024
Number layers
Hidden size
Batch size
64
Maximum number of epochs
500.000
Initial Learning rate (Ir;)
1c-4, 2.5c-4
Weight decay
0, le-5
Bias initialization
uniform scaling, normal(1.0)
Weight initialization
uniform scaling, normal(1.0)
Position embedding initialization
uniform scaling, normal(1.0)
Table I: Hyperparameters used in the ICL. The best parameter for each hyperparameter high- lighted.
In the phase shift plots in Fig. we keep the value in the x-axis constant and used the best setting over the parameters: (number of layers, hidden size, number of heads and learning rate).
E DETAILS OF PROBE
We will use the terms probe model and task model to distinguish probe from ICL. Our probe is defined as:
softmax(s..)
(132)
FF.(a (W,h)
(133)
The position scores s, E R are learned parameters where T the max input sequence length (T' = 80 in our experiments). The softmax of position scores attention weights for each position and for each target variable. This enables us to learn input-independent, optimal target locations for each target (displayed on the right side of Fig. 4). We then average hidden states by using these attention weights. A linear projection, W, E RT''H', is applied before averaging. FF is either linear layer or 2-layer MLP (hidden size=512) with GeLU activation function. For each layer, train different probe with different parameters using stochastic gradient descent. H' equals to the 5 512. The probe is trained using an Adam optimizer with a learning rate of 0.001 (chosen from among {0.01, 0.001, 0.0001) validation data).
25

---

Published as a conference paper at ICLR 2023
X'Y grobe eroor layer
Layer antationn Testitap over time
Liaaa Podee
3GF Tick
Targes Index
MSE
0.04
0.4
0.02
0.2
0.00
layer
postlon
Winor probe Layee
Layer-12's atention hestmup over
0.04
Tirget Index
0.4
MSE
-02
0
2
X
10
11
12
13
14
15
16
10 15
20 25 30 35 40 45 50 55 60 65 70 75
Figure 5: Detailed error values of the control probe displayed in Fig. 4.
1.0
0.8
Acsemplan
Figure 6: R² of linear weight estimation on d = 8 problem
Control Experiments In Fig. 4, dashed lines show probing results with task model trained on control task, in which always the all-ones 1. This problem structurally resembles our main experiment setup, but does not require in-context learning. During probing, we feed this model data generated by w sampled form normal distribution as the original task model. We observe that control probe has significantly higher error rate, showing that the probing accuracy obtained with actual task model is non-trivial We present detailed error values the control probe in Fig. 5.
F LINEARITY ICL
In Fig. 1b, we compare implicit linear weight of the ICL against the linear algorithms using ILWD measure. Note that this measure do not assume predictors be linear: when the predictors are not linear, ILD measures the difference between closest linear predictors (in Eq. (16) sense) to each algorithm.
To gain more insight to ICL's algorithm, we can measure how linear ICL in different regimes of the linear problem (underdetermined, determined) by using R² (coefficient of determination) measure. So, instead of asking what's the best linear fit in Eq. (16), can ask how good is the linear fit, which the R² the estimator. Interestingly, even though our model matches min-norm least square solution in both metrics in Section 4.3, show that ICL is becoming gradually linear in the
regime Fig. 6. This is an important result, enables us to say the in-context learner's hypothesis class is not purely linear.
26

---

Published as a conference paper at ICLR 2023
0.010
approcimation
0.010
apposimaiin
0.008
0.008
0.006
0.006
0.004
0.002
0.002
0.000
0.000
-0.10
-0.05
0.00
0.05
0.10
-0.10
-0.05
0.00
0.05
0.10
(a) Approximating x² using GeLU Eq. (136),
(b) Approximating x? using tanh, Eq. (140), where
8 le
0.010
Rel.U spproximation
0.008
0.006
0.004
0.002
0.000
-0.10
-0.05
0.00
0,05
0.10
X
picce-wise linear approximation x by using
ReLU, Eq. (141).
Figure 7: Approximations of multiplication via various non- non-lincarities.
G MULTIPLICATIVE INTERACTIONS WITH OTHER NON-LINEARITIES
We can show that for a real-valued and smooth non-linearity f(x), we can apply the same trick in in the paper body. In particular, write Taylor expansion
f(x) Σ aax ao + ai az²
(134)
which converges for some sufficiently small neighborhood: X E [-c, First, assume that the second order term d2 dominates higher-order terms in this domain such that:
azz²
0:s2x
wherer E X
It's easy to verify that the following is true:
(f f(=+t))f(s)) f(x)= f(y) + ao) ry + O(x + y²)
(135)
20g
So, given the expansion for GeLU in Eq. (37), we can use this generic formula obtain the multi- plication approximation:
GeLU + v) GeLU(x) GeLU(w) x zy
(136)
We plot this approximation against 2 for (0.1, 0011 range in Fig. 7a.
In the case of a2 is zero, cannot get any second order term, and in the case of d2 is negligible (x + )² will dominate the Eq. (135), so we cannot obtain a good approximation of zy. this case, resort to numerical derivatives and utilize the ag term:
('(x) a1
202x + 30zzz +
(137)
27

---

Published as a conference paper at ICLR 2023
(3 not negligible, azæ² & ai>zr* in the same domain, we can use numerical derivatives to get a multiplication term:
f(x + y +0 +(x f(x +8) f(x) fly + 8) f(v) + a1
xx++ +v²")
64g
8
5
(138)
For example, tanh has no second order term in its Taylor expansion:
x² 2x
tanh =
(139)
3 15
Using above formula we can obtain the following expression:
tanh(2 + + -
tanh(æ y) tanh(x x 00 tanh(x)
8
(140)
tanh(y + 6) tanh(u) + +1) xy
Similar to our construction in Eq. (110), we can construct a Transformer layer that calculates these quantities (noting that is small, input-independent scalar).
We plot this approximation against x² for [0.1, -011] range in Fig. 7b. Note that, if we use this approximation our constructions we will need more hidden space as there are 6 different tanh term as opposed to 3 GeLU term in Eq. (1))).
Non-smooth non-linearities ReLU is another commonly used non-linearity that is not differen- tiable. With ReLU, we can only hope get piece-wise linear approximations. For example, we can try to approximate z² with the following function:
0.0375 * ReLU(() + 0.0375 ReLU(-z)) + R0.0 (x 0.05))+
ReLU(-0.05 (x 0.05)) ReLU (x 0.025)) ReLU(-00025 (x + 0.025) x²
(141)
We plot this approximation against 2.2 for [0.1, -0.1] range in Fig. 7c.
H EMPIRICAL SCALING ANALYSIS WITH DIMENSIONALITY
In Figs. 3a and 3b, showed that ICL needs different hidden sizes to enter the "Ridge regression phase" (orange background) or "OLS phase" (green background) depending on the dimensionality d of inputs However, cannot reliably read the actual relations between size requirements and the dimension of the problem from only two dimensions. To better understand size requirements, ask the following empirical question for each dimension: how many layer/hidden size/heads needed to better the least-squares solution than the Ridge(A = <) regression solution (the green phase in Figs. 3a and 3b)?
To answer this important question, we experimented with d = 2,, 8, 12, 16, 20} and run an experiment sweep for each dimension over:
number of layers (L): {1, 4,8, 12, 16),
hidden size (H): {16, 32, 64, 256, 512, 1024).
number of heads (M): (1,2,4,8)
learning rate: {1c-4. 2.5e-4).
For each feature that affects computational capacity of transformer (L, H. M), optimize other features and find the minimum value for the feature that satisfies SPD(OLS, ICL) < SPD(Ridgo(A c), ICI.). We plot experiment with = 0.1 in Appendix H. We find that single head is enough for all problem dimensions, while other parameters exhibit step-function- like dependence input size.
Please note that other hyperparameters discussed in Appendix D (c.g weight initialization) were not optimized for each dimension independently.
28

---

Published as a conference paper at ICLR 2023
30
hiddlen size (H)
6
Layers (L)
25
4
20
2
4
8
16
20
4
8
16
20
Problem dimension (d)
Problem dimension (4)
(a) hidden size requirements.
(b) layer requirements.
1.04
1.02
Heads (M)
1.00
0.98
0.96
20
Problem dimension (d)
(c) number of head requirements.
Figure 8: Empirical requirements model parameters satisfy SPD(Ridge() O.T), > SPL(OLS, ICL) when other parameters optimized.
29


Reviews:
[{"review_id": "r6OiDLrugE", "reviewer": "Reviewer_NUsG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Previous work has shown that transformers can be trained to be \"in-context learners (ICL) of linear functions”. That is, given a couple of (x, y) pairs sampled from a line (that is also sampled randomly at test time), and a new x' from the same line in its context, a transformer can be trained to output the correct y'. \n\nThis work is focused on understanding *how* this happens. \n\nThis paper provides three pieces of evidence for the hypothesis that transformers in fact learn standard learning algorithms to acquire the skill described above. \n1. Construction: The authors provide a bottom-up construction demonstrating how a transformer can represent both a single step of gradient descent and a single step of Sherman-Morrison update in O(1) layers. This implies these steps can be composed in  a single model to “run optimization” during the forward pass. \n2. What “computation” an in-context learning transformer displays by testing how similar its outputs on held-out data is to well-known algorithms: The authors systematically show that 1) the best performing model behaves very similar to ordinary least squares when no noise is present, 2) when noise is added, the best model learns a solution that minimizes Bayes risk (very interesting). 3) Transformers that have their capacity restricted interpolate between different predictors. \n3. How learned in-context-learning transformers actually implement ICL: The authors present preliminary evidence, using auxiliary probes, that the activations of ICL transformers might extract the moment matrix and the parameter vector during the forward pass. \n\nIn addition to these pieces of evidence, the authors also propose a neat way of thinking about how a transformer might represent algorithmic operations (named the \"raw\" operation in the appendix), which might be useful in analyses in other domains. ", "summary_of_the_review": "The authors provide a constructive, behavioural and mechanistic account of how transformer networks can in theory and practice learn to in-context-learn linear functions. \n\nThe paper is well-written with insightful questions and clear answers. I believe this is a good paper that deserves to be highlighted in NeurIPS. ", "strength_and_weaknesses": "(The requests from the authors during the rebuttal phase is tagged below.)\n\n**STRENGTHS:**\n* **Scope and relevance:** Given the surge of interest in in-context learning, this paper is very timely in taking on a very important topic.\n* **Significance of contributions:** The paper does a great job illustrating how a transformer might implement an iterative optimizer in its weights. While I expect this work not to be the final word in how a transformer architecture might learn to be a linear function approximator, the questions asked by the authors seem to be the right ones and the provided answers seem novel and meaningful. \n* **Clarity:** The main body of the paper is written very well.\n* **\"RAW\" operator is useful:** While not discussed in the main body of the paper, the raw operator defined and used in the appendix for the constructions seem like a quite neat way to think about constructive proofs for transformers. \n\n**WEAKNESSES:**\n* **Dependence on the GELU activations:** If I'm not mistaken, the authors make use of the numerics of the GELU activation to show how a the feed-forward layers can implement the dot product and matrix multiplication operations. I have two concerns about this: \n  * It's not clear how small the activations need to be for the approximation to hold. (REQUEST) Would it be possible to add a plot in the appendix that shows how well $x * x = x^2$ is approximated with the GELU approximation? The x-axis could be the values between $-0.1$ to $+0.1$, and the y-axis would be the error in the approximation, or the proportion of the error to the x value.  \n  * Transformers that don't use this activation also display ICL capabilities, including the ability to be linear function approximators (happy to be corrected about this - perhaps GELU is much more important than I anticipate). \n  * In light of this, I'm not sure if the construction of $mul$ truly represent what the transformers might be learning (happy to be convinced otherwise during the rebuttal) . (REQUEST) Perhaps it'd be good to at least briefly discuss the dependence on GELU to get the $mul$ construction in the discussion section following Theorem 2? \n* **Mechanistic interpretability results lack control groups:** As the authors acknowledge, interpreting probing results is often tough. Hence, the results presented in Section 5 appear to be speculative as of now. (This limitation is acknowledged by the authors). (REQUEST) One approach to strengthen the results would be to produce Figure 4 (left and right) with randomly initialized, slightly trained and fully trained (but poorly performing in ICL) transformers. If these also yield comparable loss values, then discussion of the results of Section 5 should be reconsidered. \n* **Making Figure 2 denser to confirm a claim:** The authors claim that for all values of $\\sigma^2$ and $\\tau^2$, the right parameter that gives the best fit also is the one that minimizes Bayes risk. Figure 2, which seems to be what this sentence is based on, is a bit too sparse to really reach this conclusion. If it's possible, it would be great to make the plots a bit denser to see the trends a bit more clearly. That is, sample tau and sigma a bit denser to add more measurement points? Marking the MSPD between the transformer and the predictor using the optimal ridge parameter with a different marker for each value $\\tau$ and $\\sigma$ on the plots could also bring extra clarity perhaps? (nitpick) Lastly, perhaps a colormap that has a gradient determined by the value of the ridge parameter could also make the plot easier to read.\n* **Ambiguity in Figure 4 (right):** Figure 4 (right) isn't very clear to me. What are the x and y axes showing? \n\n\n**QUESTIONS TO AUTHORS:** (a bunch of REQUESTS): \n* **Position embeddings:** What kind of position embeddings did you use? (REQUEST) Space permitting, could you briefly discuss how you expect the choice of position embeddings to affect ICL result? If I'm not mistaken, the derivations in the appendix assume absolute position embeddings a'la the original transformer paper. \n* **Gradient descent with n steps as another benchmark predictor:** I was wondering if you considered observing whether a few steps of gradient descent (not just a single step) matches ICL performance of medium-shallow transformers? If the initial weights are initialized with zeros, gradient descent will presumably consistently lead to an increase in the norm of the implied weight vector. Since having an L2 penalty on the weight vector also has a similar effect, perhaps this is why ICL performance of medium-shallow transformers are matched well with the predictors with a ridge term. (Analogously, perhaps gradient descent with a couple of steps has low MSPD with ridge regression with an appropriately chosen lambda?)\n* **MLP used in the probe experiments:** Could you give further details regarding the MLP probe used in Section 5? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "DMu7zBenKI", "reviewer": "Reviewer_p6X8", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "\nThis paper presents an investigation into large-language models' (LLMs) abilities to learn algorithms for in-context linear regression given only sample problems consisting of inputs and linearly related outputs, but no information about the true (linear) hypothesis class underlying these relationships.\n\nThe authors prove that transformer LLMs are capable of representing and performing the matrix operations needed to support 2 commonly algorithms for linear regression -- gradient descent and regression via rank-1 updates.\n\nAnd they also show that for a set of underdetermined regression problems generated by a particular (idealized) process, LLMs are capable of learning an algorithm for learning the minimum Bayes risk predictor (relative to that process). Further experiments are conducted to ellicit the effect of LLM model structure and capacity on the learning algorithms learned. The authors show that LLM capacity does influence which learning algorithms are derived, but how and to what extent is left unexplored.\n\nFinally, the authors use a set of \"probe\" experiments to understand how in-context regression fitting works inside an LLM. The experiments suggest that the LLMs are computing fundamental quantities like the moment and weight vectors, but exactly how these are being computed (or if the hypothesis class is even truly linear) is not determined and is left for future work.\n", "summary_of_the_review": "This paper is novel in that it provides the first examination of how in-context learning works for the simple but non-trivial learning problem of linear regression. Because in-context learning is a phenomenon that is just starting to be understood, and because this paper provides positive (though not conclusive) results, I advocate for acceptance to ICLR.\n", "strength_and_weaknesses": "\nStrengths:\nThis paper appears to be the first to investigate the neural mechansisms underpinning in-context learning for the very simple but non-trivial hypothesis class of linear regression models. The authors' finding that LLMs seem to mimic sensible and even Bayes-optimal learning algorithms is surprising and will surely motivate further research into understanding the power and limits of in-context learning.\n\nWeaknesses:\nAs the authors admit, their findings are suggestive but the exact mechanisms of in-context learning for linear regression are still not understood.\n\nAdditionally, there were some experiments I was surprised the authors did not perform:\n- examining how the dimensionality of the regression problem affects the learning algorithm derived for an LLM with a given capacity.\n- determining whether the LLM derives Bayes optimal learning rule for other choices of p(w) and p(x). In particular, ridge regression is the correct choice for the p(w) used. But would the LLM mimic Lasso regression if p(w) was a Laplace prior over the regression coefficients?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "9zPapCTXVt", "reviewer": "Reviewer_5hCi", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper investigates the capability of transformer-based in-context learners (ICL) to implicitly implement learning algorithms. The authors consider a study-case of linear regression models to validate their hypothesis that ICL encode context-specific parametric models: 1) ICL implement linear models by encoding gradient descent 2) ICL implements closed-form computation of regression parameters with Sherman–Morrison update. These results are summarized in Theorem 1 and 2 for a single-step update. The paper also investigates how the learned algorithm is implemented by investigating what intermediate quantities are being computed. The authors conducted probing experiments to assess if moment and weight vectors are computed for training linear regression ICL task. ", "summary_of_the_review": "The paper gives important insights on how ICL encode its learning mechanism. The theoretical results of the paper are derived in the case of linear regression. It is shown that ICL is cable of encoding gradient descent and regression parameters via closed-form. This contribution is useful for better understanding how transformers work in general and specifically in the case of In-Context-Learning.", "strength_and_weaknesses": "Strength\n- The paper provides theoretical insights on how ICL implements learning algorithms. This can be used to introduce new paradigm for training ICL\n- The paper is supported by experimental analysis on simulated data related to linear regression to validate the theoretical results \n- The paper investigated different scenario such as the case of noisy data\n\nWeakness\n- The paper derives the results for a single step of larger iterative algorithm\n- The presented results are derived for the case of linear regression and it is unclear how it extends to non-linear problems. \n- The experiments are conducted on toy datasets with normal distributions.\n- The probing experiments are not very conclusive\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-5%