What Can Transformers Learn In-Context? A Case Study of Simple Function Classes
Authors: Shivam Garg
Dimitris Tsipras
Percy Liang
Gregory Valiant
Venue: NeurIPS 2022 Accept (2022)
Keywords: in-context learning
transformers
linear regression

Abstract: In-context learning is the ability of a model to condition on a prompt sequence consisting of in-context examples (input-output pairs corresponding to some task) along with a new query input, and generate the corresponding output. Crucially, in-context learning happens only at inference time without any parameter updates to the model. While large language models such as GPT-3 exhibit some ability to perform in-context learning, it is unclear what the relationship is between tasks on which this succeeds and what is present in the training data. To investigate this, we consider the problem of training a model to in-context learn a function class (e.g., linear functions): given data derived from some functions in the class, can we train a model (e.g., a Transformer) to in-context learn most functions from that class? We show empirically that standard Transformers can be trained from scratch to perform in-context learning of linear functions---that is, the trained model is able to learn unseen linear functions from in-context examples with performance comparable to the optimal least squares estimator. In fact, in-context learning is possible even under two forms of distribution shift: (i) between the training data of the Transformer and inference-time prompts, and (ii) between the in-context examples and the query input during inference. We also show that we can train Transformers to in-context learn more complex function classes: sparse linear functions where the model outperforms least squares and nearly matches the performance of Lasso, and two-layer neural networks where the model performs comparably to neural networks trained on in-context examples using gradient descent.

What Can Transformers Learn In-Context? A Case Study of Simple Function Classes
Shivam Garg"
Dimitris Tsipras"
Percy Liang
Gregory Valiant
Stanford University {shivang, tsipras, pliang, graliant)@essetaafordd. stanford. edu
Abstract
In-context learning refers to the ability of a model to condition on a prompt se- quence consisting of in-context examples (input-output pairs from some task) along with a new query input, and generate the corresponding output. Crucially, in-context learning happens only at inference time without any parameter updates to the model. While large language models such as GPT-3 exhibit some ability to perform in-context learning, it is unclear what the relationship is between tasks on which this succeeds and what is present in the training data. To make progress towards understanding in-context learning, training a model to in-context learn a function class (e.g., linear functions): given data derived from some functions the class, can we train model in- context learn "most" functions from this class? We show empirically that standard Transformers can be trained from scratch perform in-context learning of linear functions that is, the trained model is able to learn unseen linear functions from in-context examples with performance comparable to the optimal least squares esti- mator. In fact, in in-context learning is possible even under two forms of distribution shift: (i) between the training data of the model and inference-time prompts; (ii) be- tween the in-context examples and the query input during inference. We also show that we can train Transformers in-context learn more complex function classes- i.e., sparse linear functions, two-layer neural networks, and decision trees with performance that matches or exceeds task-specific learning algorithms/
consider the well-defined problem
Introduction
Large language models such GPT-3 [10 are able to perform in-context learning: given prompt containing examples from task (itput-output pairs) and new query input, the language model generate the corresponding output. For example, these models able to produce English translations of French words after being prompted on few such translations, e.g.:
maison + house, chat + cat, chien dog
prompt
completion
This capability quite intriguing as it allows models to adapt to a wide range of downstream tasks on-the-fly on-the-fy--wthhuu without the need to update the model after training [10 35 55 g However, it is unclear to what extent these models are able to learn new tasks from in-context examples alone as opposed to indexing into a vast set of known tasks from the training data (e.g., see Min et al. (1)).
"qqual contribution.
Our code and models available hhttps://giii "The term "in-context learning" may also refer a more general notion learning from a prompt [48. this work, we focus on the standard notion which refers to learning task/function given in-context examples [0]
wtttttttcottttceccctececct.t context-
leattleeaninng
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

To make progress towards understanding in-context learning, we consider the well-defined problem of learning a function class from in-context examples. That is, we say that a model can in-context learn function class F if, for "most" functions I E F, the model can approximate f (Tequry) ) for a new query input Fquery by conditioning on a prompt sequence (x1, f(11)...., ... Ikk ; (TT Tquery)
containing in-context examples and the query input.
Formally, let Dx be distribution over inputs and Dy be distribution over functions in J. A prompt P is a sequence (II, f(z1),... Uks I (2x2, (quer)) where inputs (i.e., and Equery) are drawn ii.d. from Dx and f is drawn from Dy. We say that a model M can in-context learn the function ) if it can predict f (Iquery) with an average error class F up to with respect to (Dx, Dx),
Ep (MM(P))) I (Iquery)) S 6, where {(; :) is some appropriate loss function, such as the squared error. Within this framework, we can now concretely ask:
Can we train model to in-context learn a certain function class?
Note that, here. being able to in-context learn a function class is a property of model M alone, independent of how it was trained. Training such a model can be viewed as an instance of meta- general paradigm for learning a model or method that can learn from data.
learning
We empirically study this question, focusing on Transformer models [69 the architecture behind recent large language models-trainel from scratch in-context learn simple, well-defined function classes (e.g. linear functions). Specifically, we sample prompts containing in-context examples (input-output pairs) generated using functions in given class and train models to predict the function value at the corresponding query inputs. (see illustration in Figure I. Our findings are as follows.
Training data
Inference
WI.
N(O. 1n) W/ ~
Wierr N(O. L.)
360
540
W
W
Ueg 3%)
auto-regressive model
60
10
(g. ?)
Oxp x)
I.....n n
Figure Can we train a model that in-context learns a function class (here linear functions)? We train Transformers by repeatedly sampling random function I from that class, as well as random inputs Z1...., Ty and training the model to predict each f'(xa) given the prompt an, [(1)...... ...))....... (i (xi-1), xy (wrt squared loss). Then, during inference, we evaluate the model's ability to predict accurately on new, unseen functions.
Transformers can in-context learn linear functions. We show empirically that we can train a standard Transformer from scratch to in-context learn the class of linear functions, with respect the input distribution Dx being an isotropic Gaussian in 20 dimensions, and D being the distribution over linear functions with weight vectors drawn from an isotropic Gaussian (the model was trained on prompts generated from the same distributions Dx and Dz). Specifically, the trained model achieves error comparable to the optimal least squares estimator, suggesting that it encodes an effective learning algorithm, at least for the distribution used to generate the training prompts.
Generalization -iittt-iittiittiinn prompts. To understand the extent to which the trained model encodes an algorithm that works beyond the training distribution, we consider in-context learning under two types of distribution shifts: (a) shift between the prompts encountered during training and inference (e.g., training on prompts without any noise in the in-context example outputs but testing with noisy outputs), (b) a shift between the in-context examples and the query input during inference (e.g., in-context examples lie in one orthant and the query input lies in another). We find that the performance of our model is quite robust to such shifts, indicating that has learned to perform linear regression with some generality.

---

More complex function classes. We also consider the function classes of 3-sparse linear functions, two-layer ReLU neural networks with 100 hidden units, and decision trees of depth 4, all with 20 dimensional inputs. We show that we can train Transformer models that can in-context learn these classes (with respect to isotropic Gaussian inputs and appropriately defined distributions over functions). For sparse linear functions, the trained model exploits sparsity, obtaining performance better than least squares and comparable to Lasso. For neural networks, the trained model performs comparably to neural networks of the same architecture trained using gradient descent on in-context examples. Moreover, it is also able to in-context learn linear functions. For decision trees, the trained model learns unseen trees with as few as 100 in-context examples, whereas greedy learning and tree boosting algorithms are unable to achieve competitive performance. Note that learning these function classes requires involved algorithms (e.g., gradient descent with the Lasso objective), and our results show that Transformers can encode algorithms with similar performance in a single forward pass.
Model capacity and problem dimension. Finally, we explore how the ability of Transformers to in-context learn linear functions scales with model capacity and problem dimensionality. We find that increasing model capacity improves performance significantly, and allows the model to in-context learn higher-dimensional functions. Also, increasing the capacity often improves performance under distribution shifts significantly, even when the absolute improvement in the standard error is small.
2 Training models for in-context learning
We now describe general methodology for training model that can in-context learn function class F with respect to a distribution Dy over functions, and Dx over inputs. We start by constructing random training prompts
follows. We sample random function f from the class according to Dr, then create a set of random inputs I Hi..... FK+1 drawn independently from Dx, and evaluate on these inputs to produce the prompt P = (Ir1, f(x (z))....., 1k++1 WR+11, f(=a++ 1)). For example, in the case of linear functions, inputs could be drawn from the isotropic Gaussian distribution (O, Ia), and a random function chosen by sampling weight vector w from N N(0, Ia) and setting fxx) = w T x.
Now, given such prompts, we train model to predict (xx) for given x; based on set of in-context examples. Concretely, let p denote the prompt prefix containing i in-context examples (the first i input-output pairs) and the (i + 1)th input: P = (I1, f( ) 22, I (12)..... ...,,,.................. Ti, f (z;)) z.+1). Then, we train a model Me parameterized by 0 to minimize the expected loss over all the prompt prefixes:
min Ep
Σ e (Mo (P') .f (E+41))
where ((.,)) .) is an appropriately chosen loss function. Below, we describe how this general methodol- ogy can be implemented for a concrete model family (see Appendix A for additional details).
Model structure. We use a decoder-only Transformer architecture (69 from the GPT-2 family [44 with 12 layers. 8 attention heads, and 256-dimensional embedding space (22.4M parameters). The model takes as input a sequence of vectors in its embedding space and predicts the next vector in the sequence within the same space (in language modeling, these vectors correspond to input tokens). We apply this model to our prompts (I1, fz1)..... Fk+1, f 1)) as follows. We map each prompt output fz;)) to the same dimension prompt inputs 2; by appending zeros, and map the prompt inputs and outputs into the latent embedding space of the Transformer through (learnable) linear transformation. We then another (learnable) linear transformation to map the vector predicted by the model to a scalar. Note that the Transformer architecture allows us to compute the prediction (Mo(Pi) for all prompt prefixes in single forward pass.
Training. We train the model according to the training objective in using squared error as the loss function. We sample batch of random prompts at each training step and update the model through a gradient update. We use a batch size of 64 and train for 500k steps. This training is done from scratch, that is, we do not fine-tune a pre-traincd language model, nor do we train actual text.
Curriculum learning. Many function classes contain functions of varying complexity. We exploit this by training our model using curriculum 20 60 (3), where we train on simpler distribution of functions in the beginning (e.g., linear functions with weight vectors restricted to a low-dimensional subspace) and gradually increase the function complexity. This speeds up training drastically, often allowing us to train models that would be significantly more expensive to train without a curriculum (see Section 6 for details).
3

---

3 In-context learning of linear functions
In the previous section, we describe general methodology for training Transformer models to in-context learn a class of functions. Here, we focus simple function class namely linear functions-and study how well models trained using our methodology can in-context learn this class.
Prompt distribution. We consider the class of linear functions F f(z) = w z, w E RR), in dimensions where 20. We sample T1,..... 2k, qquery. and independently from the isotropic Gaussian distribution N (0, Ia). We then compute each y; w x; and construct the prompt as P = (I1, V1, T2, 1/2,... Lk, yk: E.uery).
Baselines. To contextualize the performance of our trained model, compare it to other learning algorithms: (a) the least squares estimator, computing the minimum-norm linear fit to the in-context examples (Ii, yi). (b) n-Nearest Neighbors, averaging the y/ values for the n nearest neighbors of Equery- (c) averaging the values you; to estimate w and compute the inner product of this estimate with Fquurr- Least squares the optimal estimator for this problem and thus serves as a lower bound. The other basclines are consistent (but sub-optimal) estimators that are easier to compute and thus provide an estimate of the performance achieved by simple approaches. See Appendix A.3 for more details.
1.2
1.0
Least Squares
squared errer
Neighbors
0.8
Averaging
0.4
0.2
0.0
10
In-context examples
Figure 2: Evaluating the trained Transformer on in-context learning linear functions. We plot the normalized squared eтrоr of the Transformer ((M(P) - x=qury)² 2 /d) and the relevant baselines, as we vary the number of in-context examples. Transformer's error decreases at rate comparable to least squares. When the number of in-context examples reaches the problem dimension d (here 20), least squares achieves () error while the Transformer achieves an error of 0.02, improving to 0.0006 at 2d examples. The simple baselines perform better than the zero estimator (dashed), but still rather poorly. (Error averaged over 1280 prompts. 90% confidence intervals, 1000 bootstrap trials.)
3.1 Transformers can in-context learn linear functions
We show the in-context learning ability of the resulting model along with the relevant baselines in Figure 2] The trained Transformer is able to in-context learn the class of linear functions with respect the prompt distribution specified above, performing comparably the optimal least squares estimator for any number of in-context examples considered. While the simpler baselines achieve non-trivial error, they are far from optimal, indicating that the trained model encodes a more complex algorithm. Moreover, in Appendix [.7] we show that the model cannot be relying on memorization of training prompts or weight vectors, and thus encodes an algorithm capable of in-context learning linear functions that are very different from those seen during training.
3.2
What functions is the model learning in-context?
Recall that the goal of our model is: given the prompt P P= (#1, w Ti,..., Tk, w. Tk: Feuery). output w Equery- Thus, if we fix the prefix given by the in-context examples, we can view the output of the model function fw.mna (Equery), that approximates w T Tquery- When l < d (fewer in-context examples than dimensions), the ground truth cannot be recovered perfectly and the ideal model should approximate (proj (w)) Tyuarr- where proj (w) is the projection of onto the subspace spanned by Z1,..., Here, we will evaluate how accurately the model approximates this.
Visualizing along a random direction. For a randomly sampled fixed prefix, we visualize for ara (Iquery) we vary the query input along random direction (Figure 3ap. That is,

---

pick random unit vector æ, and evaluate fw (Az) as we vary A, the distance of the query input from origin. We observe that fu (Ax) and fu, (Az) closely match the ground truth and faw (Az) matches the projected ground truth, when the distance from origin is not too large compared to the norm of a typical randomly sampled input. In fact, in Appendix B.I we show that the model is quite robust to scaling the query input: the error doesn't t increase much as we scale up the query input by a factor of up to 2, or scale down by factor of up 16, and degrades slowly after.
ground truth -- ground truth projected
#dims n-corfext examples
#dins in-content examples
#dims in-cotext examples
20
20
20
frear trodk
runction value
1.0
10
0.0
-10
-10
-10
0.2
grodient and trac
gradient and projected
-20
-20
-20
-10
-30
10
-10
10
10
20
30
40
distance from origin
distance from origin
distance from origin
in-context examples
(a) function visualizations
(b) gradients
(a) Model prediction as we fix the in-context examples and vary the query input along random direction (three random prompts). The shaded regions denotes where the norm of a randomly training input lies with probability 0.99. When the scale of the query input is close this range, the model prediction is close the true linear function (or its projection to the space of in-context inputs when < d). (b) We compute the gradient of the model prediction with respect to the query input, and plot its (normalized) inner product with the true and projected w, averaged over 1280 prompts. The gradient aligns almost perfectly with w when k d, and with projected w for all indicating that the model locally aligns with the ground truth.
Figure 3: Understanding the wrefx-conditione
Local correctness. So far, have seen that the model is able to make predictions close to the ground truth for randomly drawn query inputs and in-context examples. We will now turn our attention to the local change of f around Equery by considering the gradient of the function fw.er. (Iquery) ) with respect to Zquery (our model is fully differentiable so we can compute it directly). Since f computed by the model should ideally approximate proj (w) T xx this gradient should lie in the direction of the projected ground truth projz,.. (w). In Figure 3b we show the inner product between the gradient and proj (w) (both normalized), averaged over 1280 random prompts, and observe that they align almost perfectly. Since proj.,. (w) = w almost surely when & 2 d. we observe that the gradient also aligns with perfectly in this regime. Thus the model is locally correct around the query input.
Extrapolating beyond the training distribution
In the previous section, we demonstrated that we can train model to in-context learn linear functions with respect to the distribution of prompts encountered during training. That is, we evaluate the in-context learning ability of the model with respect distributions Dx and D that were also used to train the model. Here, we evaluate the in-context learning performance of our model on prompt distributions different from the used for training. Our overarching goal is to better understand the learning algorithm encoded by our model by analysing how it performs on such prompts.
Formally, we refer to the distribution of functions used during training as Dirain and the corresponding distribution of prompt inputs as Dirain During inference, functions are sampled from a (potentially different) distribution D while prompt inputs from a distribution Diest Moreover, deviating again from our analysis so far, we also consider a separate distribution Diest query from which the query input is sampled, potentially dependent on the rest of the in-context inputs F ..,..., Ih (sampled from D
Within this framework, we consider the same model last section, and evaluate its performance on prompts that deviate from those encountered during training, either by
x/J or 2. introducing mismatch between in-context examples and query input, i.c., Diest # Disst.
1. sampling prompt inputs or functions from different distribution, that is Dtrain + Desst
5

---

Different orthants for in-context and query inputs. We fix the sign of each coordinate to be positive or negative for all in-context inputs &; (at random). As a result, all in-context inputs lie the same orthant, while the query input lies in another orthant with high probability. The model is not affected by the mismatch between in in-context and query inputs and closely match the performance of least squares. In this case, the model achieves errors 0.062 and 0.004 for 20 and 40 in-context examples respectively (see Figure 4c), whereas recall that it achieves errors 0.02 and 0.0006 on standard prompts. This indicates that the model is not relying on some variant of nearest neighbor We describe three such prompt distributions below, along with the corresponding results in Figure 4 and provide the full results in Appendix [3.] Overall, the model performs reasonably accurate in-context learning with respect to these prompt distributions, indicating that it has indeed learnt to perform linear regression to some generality.
Recall that we generate a training prompt / = (21, w Ii...., IR, w Tk. Equery) by drawing the prompt inputs (I, and Equory), and the weight vector (w) i..d. from (0, Ia), with d 20.
1.50
1.25
1.25
1.25
squared arron
1.00
1,00
1.00
Trandarmee
0.75
0.75
0.75
Least Squares
N-Naneett Neighbors
0.50
0.50
0.50
0.25
0.00
0.00
0.00
20
10
In-context examples
In-context examples
In-context examples
(a) Skewed covariance
(b) Noisy linear regression
(c) Different orthants
Figure 4: In-context learning on out-of-distribution prompts. We evaluate the model on prompts that deviate from those seen during training by: (a) sampling prompt inputs from a non-isotropic Gaussian, (b) adding label noise to in-context examples, (c) restricting in-context examples to single (random) orthant. Model error degrades gracefully and remains close to that of the least squares estimator, indicating that its in-context learning ability extrapolates beyond the training distribution.
Skewed covariance. We sample prompt inputs from (0, Σ) where E is a skewed covariance matrix with eigenbasis chosen uniformly random and 7hh eigenvalue proportional to 1/22. We normalize the inputs so that their expected squared norm is equal to that of inputs encountered during training, and study the effect of input scale separately in Appendix B.2 The model matches the performance of least squares until k 10, mimicking the sharp drop in the error in this regime, but its error plateaus afterwards (see Figure 4a). Thus, is not perfectly robust to this distribution mismatch but still does relatively well, achieving less than half the error of the nearest neighbor baseline in most cases.
Noisy linear regression. We add noise each prompt output: the hh output is equal to wT2; + e; where 6, ~ (0, The trained model closely tracks the performance of least squares for most prompt sizes (see Figure Interestingly, the model also exhibits the double descent error curve [2] that is known to manifest for the least squares estimator (46). Note that in this noisy setting, the optimal estimator corresponds to least squares with appropriate & rg-regularization, which we cannot
expect the model to learn since it was trained on noiseless data.
search as in that case, its error would have been higher (see the 3-nearest neighbor baseline).
More complex function classes
We now consider in-context learning for more complex function classes, namely sparse linear functions, decision trees, and two-layer ReLU neural networks. We are back in the setting where the prompt distribution during inference is same as that during training. The overall methodology remains the same: we sample random functions from these families and train a Transformer to approximate these functions given in-context examples. (See Appendix A.3 for details and baselines.) Note that, here, we are training a new Transformer model (from scratch) for each function class, independently of the model studied in the previous sections (which was trained in-context learn linear functions).
Sparse linear functions. First, we consider functions of the form f(x)) = w T x where w E IRd and has exactly non-zero coordinates. To sample prompt P f(1) ) ... Ik. f(xx), Equary), we

---

draw prompt inputs F; and Equery, and a weight vector w from (O,, Ia). and then zero out all but coordinates of uniformly at random. this setting, the least squares estimator is no longer optimal. One can perform better by leveraging sparsity, e.g., using Lasso (8), which involves solving the least squares objective with an & fi-nrrm weight regularizer. We plot the performance of our model trained for d = 20 and 8 = 3 in Figure 5a and observe that it is also able to leverage sparsity, nearly matching the performance of Lasso. Note that, unlike least squares, Lasso does not have a closed form expression and involves iterative minimization of the regularized objective, yet the Transformer able to achieve comparable performance in a single forward pass.
Transfomer
1.50
Transformer
1.0
Least Squares
3-Nearest Neighters
error
Averaging
Greedy Leamiin
Lasso
1.00
XGBocst
0.6
Greedy Thee Learning [w/ sign proproc.)
0.50
xCBBoost
0.2
tw/ sign preprec.)
0.00
10
20
30
40
0
20
40
60
80
100
in-contaxt examples
In- co-context axamples
(a) Sparse linear functions
(b) Decision trees
1.2
Transfomer
1.2
transformer
Least Squares
1.0
Squares
squared enor
3-Nearest Neighbbors
3-Nearest boighbors
0.8
2-layer NN, GD
0.8
2-layer NN. (D)
0.6
0.4
0.2
0.2
0.0
40
100
20
80
nn--ontext examples
in-context examples
(c) 2-layer NN
(d) 2-layer NN, eval on linear functions
Figure 5: Training a Transformer to in-context learn more complex function classes. (a) A Transformer trained on prompts generated using sparse linear functions can in-context learn this class, with error decreasing at a rate similar to Lasso, and significantly better than minimum norm least squares. (b) A Transformer trained on prompts generated using random decision trees can in-context learn this class, with much better performance than greedy tree learning or tree boosting. (c) A Transformer trained on prompts generated using random 2-layer ReLU neural networks can in-context learn this class. The error decreases at rate similar to the baseline which involves training neural network using a variant of gradient descent with in-context examples the training data. (d) The same model (from (c)) can in-context learn the class of linear functions. The error decreases at a rate slower than least squares, but comparable to neural network trained using variant of gradient descent. All errors are normalized so that the zero estimator achieves an error of 1 (dashed).
Decision trees. Next, we consider the class of depth 4 decision trees with 20 dimensional inputs. A function f in this class is represented by a full binary tree (16 leaves) where each non-leaf node is associated with a coordinate, and each leaf is associated with value. To evaluate f on an input we traverse the tree starting from the root node, going to the right if the coordinate associated with that node is positive and to the left otherwise. f(x) is given by the value associated with the leaf node reached. To sample a random prompt P = (xr J [xx) )..))................ Wg. [((xx), ), Jquery), we draw inputs æ;s and Viuuery from N Ta). and f corresponds to a tree where the coordinates associated with leaf nodes are drawn uniformly at random from .......d} and the values associated with the leaves are drawn from N (0, 1). In Figure (bb. we show that Transformers can be trained to in-context learn this class much better than greedy tree learning and boosting (via XGBoost [3)) Since the decision trees in our class of functions predict solely based on the sign pattern of æ,,,, we also consider baseline where we provide the greedy learning and XGBoost algorithms with the signs of each x; instead. This significantly improves their performance, but they still perform much worse than the Transformer (error 0.31 vs. 0.12 at 100 in-context examples).
Two-layer neural networks. Finally, we consider the class of two layer ReLU neural net- works containing functions of the form f(x) = Ei= 0s a a(w; 2). where as E IR, w; E R and o-) ma is the ReLU activation function. To draw a random prompt P
7

---

(21, ff1), (ai))...,,z;; 2k. [xxx); Fqqury), we sample inputs 2;S and 2query from (0, Ia). along with net- work parameters a,s and w;s from (0. 2/r) and N (0, La) respectively. We set the input dimension d to 20 and the number of the hidden nodes r to 100. In Figure 5c] we show that Transformers can be fact, the Transformer performs comparably to
trained to in-context learn this class of functions. the baseline which trains two-layer neural network of the same architecture on in-context examples using Adam [31] a variant of gradient descent (see Appendix AA33 for details). Moreover, the model trained to in-context learn two-layer neural networks is also able to in-context learn linear functions (for which it not explicitly trained), albeit at a rate slower than least squares, but comparable to a neural network trained on in-context examples generated using a linear function. (see Figure 5d
Investigating what matters for in-context learning
We now return the setting of training models in-context learn linear functions and explore different factors that lead to successful in-context learning.
Problem Dimension and Capacity. In Section 3 and we saw that Transformer models can be trained to in-context learn 20-dimensional linear functions accurately and relatively robustly. To explore the interplay between problem dimensionality and capacity, we also consider models with fewer parameters (see Appendix I and train each architecture on (10, 30, 40, 50) -dimensiooa problems. In Figure [6] we plot the model error with 2d in-context examples we vary the problem dimension d and the model capacity. In the standard setting, i.e., when the training and inference time prompt distributions are the same, we observe that the error decreases as we increase the capacity or reduce the problem dimensionality (see Figure b model capacity helps. For out-of-distribution prompts, the settings with skewed covariance or with in-context example and query inputs lying in different orthants are challenging, especially for higher dimensional problems. However, the error decreases considerably (in most cases) as we increase the model capacity, even when absolute decrease in the standard error is small (see Figure 66b and (c). See Appendix B.3 for additional plots.
dimensions
104
10r
10y
10
20
20
10
10:1
20
10
10a
10
10
10a
3.4M
3.4M
7.6M
7.6M num of parameters
number parameters
number parameters
(a) Standard
(b) Different orthants
(c) Skewed covariance
Figure 6: Understanding the effect of model capacity and problem dimension on in-context learning (b,c) prompts. We train Transformers to in-context learn linear functions and plot the error with 2d in-context examples as we vary problem dimension d and model capacity. Capacity helps in most cases, especially on out-of diitt--iiittitution prompts, even when the absolute gains in the -n-distribution setting are small. We train 3 models in each case with different random seeds, and show the median eтror (solid lines), and the minimum and maximum errors (shaded region). (See Appendix B.4 for training variance analysis.)
performance for in-distribution (a) and out-of-distribun
Curriculum learning. When training our models, we initially draw the prompt inputs from a fixed dimensional subspace (by setting some of the coordinates to 0) with prompt length 11 (number of input-output pairs), and increase the subspace dimension by 1 and prompt length by 2 every 2, 000 training steps, until the subspace dimension reaches the ambient dimension d and prompt length reaches 2d + (see Appendix A.2 for details). This speeds up training drastically, especially for higher dimensional problems: for dimension 50, the loss barely decreases through the 500k training steps without curriculum but reaches close to the optimum with curriculum. For the 20 dimensional problem where we were able to train the model without curriculum within the training (step count) budget, we did not observe any qualitative difference in accuracy or robustness compared to the model trained with curriculum. We compare training with and without eurriculum in Appendix [..]
Notably, when training Transformers without curriculum, there an initial relatvlyyl long-period in training where the loss does not decrease, followed by a period of sharp decrease. The length of
8

---

this period varies with training randommess and seems to increase on average with problem dimension. Interestingly, Olsson al. [48] observe a similar jump in the in-context learning ability of language model which they attribute to the formation of "induction heads".
Number of distinct prompts functions seen during training. To estimate the amount of training data required for in-context learning, we perform two ablation studies. In the first study, we limit the number of distinct prompts seen during training. That is, we create a set of ny randomly generated prompts (as described in Section 2). and sample prompts from this set during training (here, we train without curriculum, as it would introduce additional prompts during the warmup phase). In the second study, we only limit the number of distinct functions used for training. That is we create set of fw randomly chosen vectors (corresponding to nww linear functions) and sample weight vectors uniformly from that set to generate the training prompts (the inputs are still sampled from (O, Ia) for each training prompt). We find that the amount of training data required is relatively small: non-trivial in-context learning is possible with ny = 100k or nhw = 1K, and the error drops close to that of the unrestricted model (discussed in Section (3) with " = 1M or tr 10k (details in Appendix B.6). For context, in Section 3
the model is trained on fresh prompts at each step and thus encounters 32M distinct linear functions and prompts (500k training steps with a batch size of 64).
7 Related work
In-context learning, Since Brown et al. [00] demonstrated the in-context learning ability of GPT-3,
there has been a significant interest in improving and understanding this capability [36 39 79 38 59
40
43 33 The works most relevant to ours are as follows. Xie et al. (74) propose a Bayesian inference framework explaining how in-context learning works despite formatting differences between training and inference distributions. Razeghi et al. [57] show that in-context learning for numerical reasoning tasks is better for instances whose terms are more prevalent in training data. Min al. 33 demonstrate tasks where in-context learning works even when the prompt outputs are chosen randomly, questioning to what extent these models are truly learning new tasks on-the-fly, while Rong [58] gives examples of novel tasks on which these models demonstrate on-the-fly flyy learning ability. Chan et al. [12 demonstrate that distributional properties such as long- tailedesss are crucial for in-context learning on an imae-basee few-shot dataset. Olsson et al. [48] and Elhage et al. [19] consider different framing of in-context learning, referring to any model behavior that utilizes information in a prompt to make predictions that improve with prompt size. They hypothesize the existence of special circuits inside Transformer models responsible for in-context learning, that copy similar patterns from the prompt sequence. Pesut [52] and Dinh et al. [16] Table 16] consider in-context learning for small tabular datasets and learning problems in one and two dimensions, and show that GPT-3 can obtain non-trivial accuracy. Our work contributes to this line of work, by posing in-context learning as a well-defined problem of learning function classes at inference time, and empirically investigating if we can train models that in-context learn simple function classes.
14
Transformers. There a long line of work investigating the capabilities 669
15
77
51
76 78)
limitations 6) applications 37.17 49, and internal workings 19 65
7 18
48
of Transformer
47 demonstrate the ability of Transformer models to solve prediction tasks using the input context, albeit in different settings. Müller al. 44 introduce -Paa fitted transformer network" that trained to approximate Bayesian inference with priors such as Gaussian processes and Bayesian neural networks, and use to perform downstream tasks such as tabular dataset classification and few-shot image classification. Nguyen and Grover 47 introduce Transformer neural processes, building on prior work on neural processes (24 23 30), and show that they achieve state-of-the art performance on tasks such as image completion and contextual multi-armed bandits. Our work complements these works, focusing on understanding the in-context learning ability of Transformers for various simple function classes and the extent to which this ability extrapolates beyond the training distribution.
models. Most similar to our work, Müller et al. 44 and Nguyen and Grover
Meta-learning. Training a model to perform in-context learning can be viewed as an instance of the more general te-aer or meta-learning paradigm [62 45 67]. Typical approaches from this extensive line of work (see (2 for a survey) include: training a meta-learner to update the learning parameter initializations which allow to quickly
parameters of downstream learner train for downstream tasks (2 56 learning latent embeddings for effective similarity search [66].
Most relevant to setting are approaches that directly take as input examples from a downstream

---

task and a query input and produce the corresponding output 226 42 61
24, 23, 32 Our work contributes this line of work, by investigating the learning-to-eaa abilities of Transformer models.
Data-driven algorithm design. Another line of work aims to discover algorithms that perform (as opposed algorithms with guarantees on their worst-case performance). See Balcan [1] for a survey on advancements on the theoretical foundations of such algorithms. Our work can be viewed as part of this line of work, as we train Transformer models to discover algorithms for different learning problems.
well on distribution of inputs 27
75 70
29
64 [5]
8 Discussion
In this work, we formalize and study the question: can we train models that learn different classes of functions in-contex(? We show that Transformer models trained from scratch can in-context learn the class of linear functions, with performance comparable to the optimal least squares estimator, even under distribution shifts. Moreover, we show that in-context learning is also possible for sparse linear functions, decision trees, and two-layer neural networks; learning problems which are solved in practice with involved iterative algorithms such as gradient descent.
At the same time, understanding the implications of our results for language models requires further investigation. A pertinent question regarding the in-context learning capabilities of language models how they leverage in-context examples [41] Our results demonstrate that Transformers can encode complex learning algorithms that utilize in-context examples in a far-fro far-from-ur -rivial manner. In fact, this the case for standard Transformer architectures trained with standard optimization procedures. The extent to which such non-trivial in-context learning behavior exists in large language models is still open, but we believe that our work takes step towards formalizing and understanding this question.
Our work lays the groundwork for several future directions.
Complexity of in-context learning. We empirically show that model capacity helps in performing
in-context learning accurately and robustly. How does the in-context learning loss (T) depend on the complexity of the function class F, the capacity of model M, and data used to train M. Understanding this question for models explicitly trained to perform in-context learning may suggest an upper bound for the in-context learning performance of models such as GPT-3 that are not explicitly trained for it.
Curriculum learning. Within our framework, there natural notion of curriculum learning where, during training, we gradually increase the complexity of the function class learned in-context. This leads to drastic training speed-ups. What is the reason behind such a speedup? Are similar speedups also possible for training large language models (thus significantly reducing training time and energy)?
Inductive bias of model families. Our framework presents an opportunity to understand and compare the inductive biases of different model families (e.g., Transformers vs. LSTMs) in well-defined setting. For instance, a concrete question is: Are there function classes that are easier to in-context learn using Transformers but harder for LSTMs and vice-versa?
Understanding the learning algorithms encoded in Transformers. The models we train can perform in-context learning, and are thus themselves encoding learning algorithms. However, we do not really understand the encoded algorithms. It would thus be worth investigating the internal workings of these models to get better understanding of these algorithms. Moreover, for settings such as decision trees, we do not have good understanding of what the optimal learning algorithms are or when known heuristics work [I] Nevertheless, in Section we found that Transformers are able to discover sample efficient algorithms, thus suggesting an intriguing possibility where we might be able to discover better learning algorithms by reverse engineering these models.
Acknowledgements
We thank Niladri Chatterji, Micah Goldblum, Rohith Kuditipudi, Shibani Santurkar, Carmen Strassle, Mirac Sugzun, Li-Yang Tan, and anonymous reviewers for helpful comments.
SG was funded by Stanford Interdisciplinary Graduate Fellowship. DT was funded by Open Philanthropy, and partially supported by NSF Award CCF-1813044 GV was supported by NSF Awards CCF-1704417, CCF-18 3049, Frontier Award 1804222 and DOE award DE-SC0019205. We performed our experiments on the Stanford NLP cluster.
10

---

References
[1] Maria-Florina Balcan. Data-driven algorithm design. In Tim Roughgarden, editor, Beyond Worst Case Analysis of Algorithms. Cambridge University Press, 2020.
[2] Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine- learning practice and the classical bias-variance trade-off. Proceedings of the National Academy of Sciences, 2019.
[3] Irwan Bello, Hieu Pham, Quoe Le, Mohammad Norouzi, and Samy Bengio. Neural combina- torial optimization with reinforcement learning. arXiv preprint arXiv:1611.0999 2016.
[4] Samy Bengio, Yoshua Bengio, Jocelyn Cloutier, and Jan Gecsei. On the optimization of a synaptic learning rule. In Preprints Conf. Optimality in Artificial and Biological Neural Networks, 1995.
[5] Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In International Conference on Machine Learning (ICML), pages 41-48, 2009.
[6] Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of trans- formers to recognize formal languages. arXiv preprint arXiv:2009. 1/644, 2020.
[7] Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of transformers and its implications in sequence modeling. arXiv preprint arXiv:2006.09288. 2020.
[8] Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony. Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, al. Gpt-neox-20h: An open-source autoregressive language model. arXiv preprint arXiv:2204.0674: 2022.
[9] Guy Blanc, Jane Lange, Mingda Qiao, and Li- Yang Tan. Decision tree heuristics can fail, even in the smoothed setting. arXiv preprint arXiv:2107.008/9 2021.
[10] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Neural Information Processing Systems (NeurIPS), 2020.
[11] Alon Brutzkus, Amit Daniely, and Eran Malach. Id3 learns juntas for smoothed product distributions. In Conference on Learning Theory, pages 902-912 202-915. PMLR, 2020.
[12] 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 arXiv:2205.05055 2022.
[13] Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In conference on knowledge discovery and data mining (KDD), 2016.
[14] Yanda Chen, Ruiqi Zhong, Sheng Zha, George Karypis, and He He. Meta-learning via language model in-context tuning. arXiv preprint arXiv:2110.07814 2021.
[15] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Uni- versal transformers. arXiv preprint arXiv: 0tt.:/1877.001... 2018.
[16] Tuan Dinh, Yuchen Zeng, Ruisu Zhang, Ziqian Lin, Shashank Rajput, Michael Gira, Jy-yong Sohn, Dimitris Papailiopoulos, and Kangwook Lee. Lift; Language-interfaced fine-tuning for non-language machine learning tasks. arXiv preprint (.XXv:2206006665 2022.
[17] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x words: Transformers for image recognition scale. arXiv preprint D2Xiv:2010.11920 2020.
[18] Benjamin L. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and vari- able creation in self-attention mechanisms, In International Conference on Machine Learning (ICML), 2022.
11

---

[19] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Hemighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, 7ac Hatfield-Dodee Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. tttps:/ttansformeeriiiiuisssssp pwtpiiss.ppp/2222111 /tt.iuulsssppt//////aa/aeeo/nn/ httt//////wwwwwineeee.iieo.
[20] Jeffrey L Elman. Learning and development in neural networks: The importance of starting small. Cognition, 1993.
[21] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adapta- tion of deep networks. In International conference on machine learning (ICML), 2017.
[22] Jerome H Friedman. Greedy function approximation: a gradient boosting machine. Annals of statistics, 2001.
[23] Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shamahan, Yee Whye Tch. Danilo Rezende, and SM Ali Eslami. Conditional neural processes. In International Conference on Machine Learning, pages 1704-1713, PMLR, 2018.
[24] Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint urXiv:/807.0/660 2018
[25] Michael Hahn. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 2020.
[26] Sepp Hochreiter, A Steven Younger, and Peter R Conwell. Learning to learn using gradient descent. In International conference on artificial neural networks (ICANN), 2001.
[27] Eric Horvitz, Yongshao Ruan, Carla Gomes, Henry Kautz, Bart Selman, and Max Chickering. A bayesian approach to tackling hard computational problems (preliminary report). Electronic Notes in Discrete Mathematics, 2001.
[28] Timothy Hospedales, Antreas Antoniou, Paul Micaelli, and Amos Storkey. Meta-learning in neural networks: A survey. arXiv preprint arXiv:2004. 05439, 2020.
[29] Elias Khalil, Hanjun Dai, Yuyu Zhang, Bistra Dilkina, and Le Song. Learning combinatorial optimization algorithms over graphs. Neural Information Processing Systems (NeurIPS), 2017.
[30] Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. arXiv preprint (001:09900.07611 2019.
[31] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14Xiv:1412.6980, 2014.
[32] Louis Kirsch and Jürgen Schmidhuber. Meta learning backpropagation and improving it. Neural Information Processing Systems (NeurIPS), 2021.
[33] Andrew K Lampinen, Ishita Dasgupta, Stephanie CY Chan, Kory Matthewson, Michael Henry Tessler, Antonia Creswell, James L McClelland, Jane X Wang, and Felix Hill. Can language models learn from explanations in context? arXiv preprint 22X.0:200..02329, 2022.
[34] Ke Li and Jitendra Malik. Learning to optimize. arXiv preprint (0.iv:6666.01/00 2016.
[35] Opher Lieber, Or Sharir, Barak Lenz, and Yoav Shoham. Jurassic-1: Technical details and evaluation. White Paper. AI21 Labs, 2021.
[36] Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for gpt-3? arXiv preprint 2rX1v:2101.06004, 2021.
[37] Kevin Lu, Aditya Grover, Pieter Abbeel, and Igor Mordatch. Pretrained transformers as universal computation engines. arXiv preprint 2rXiv:2103.05244 2021.
12

---

[38] Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where find them: Overcoming few-shot prompt order sensitivity. arXiv preprini arXiv:2/04.08780, 2021.
[39] Sewon Min, Mike Lewis, Hannaneh Hajishirzi, and Luke Zey..emoyer Noisy channel language model prompting for few-shot text classification. arXiv preprint arXiv:2108.44/00 2021.
[40] Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. arXiv preprini arXiv:21 /0. 15943, 2021.
[41] Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202.1282 2022.
[42] Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta-learner. In International Conference on Learning Representations (ICLR), 2018.
[43] Swaroop Mishra, Daniel Khashabi, Chitta Baral, Yejin Choi, and Hannanch Hajishirzi. Refram- ing instructional prompts to gptk's 's language. arXiv preprint arXiv22109.08330 2021.
[44] Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. arXiv preprint arXiv:2112.10510, 2021.
[45] Devang K Naik and Richard J Mammone. Meta-neural networks that learn by learning. In International Joint Conference on Neural Networks (IJCNN), 1992.
[46] Preetum Nakkiran. More data can hurt for linear regression: Sample-wise double descent. arXiv preprint arXiv:1912.072423 2019.
[47] Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. arXiv preprint arXiv:2207.04179 2022.
[48]
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfeeld-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. heads/index.htre
[49] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning (ICML), 2018.
[50] Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and Duchesnay. Scikit-learn: Machine learning Python, Journal of Machine Learning Research, 2011.
[51] Jorge Pérez, Javier Marinković, and Pablo Barceló. On the turing completeness of modern neural network architectures. arXiv preprint arXiv:1901.03422 2019.
[52] Lovre Pesut. Who models the models that model models?
am ex- URL
ploration of gpt-3's in-context model fitting ability, 2022.
https: //ww: www.:aalignmeetorru..cmomm sho-modeels-the-modeee wwwps/oorwonnoooooosoaooososoose
[53] Alec Radford, Karthik Narasimhan, Tim Salimans, and llya Sutskever. Improving language understanding by generative pre-training. OpenAl blog, 2018.
[54] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 2019.
13

---

[55] Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican. Jordan Hoffmann, Francis Song. John Aslanides, Sarah Henderson. Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint 2r 1146, 2021.
[56] Sachin Ravi and Hugo Larochelle. Optimization model for few-shot learning. International Conference for Learning Representations (ICLR), 2017.
[57] Yasaman Razeghi, Robert L Logan IV, Matt Gardner, and Sameer Singh. Impact of pretraining term frequencies on few-shot reasoning. arXiv preprint arXiv:2202.07000 2022.
[58] Frieda Rong. Extrapolating to unnatural language processing with gpt-3's in-context learning: The good, the bad, and the mysterious), 2021. URL http://ai. stanf wew.oduddddddeoeeooeooooouuueuue blog/ in- in- www.eett---eeeinnnngeet----
[59] Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for in-context 2021.
learning. arXiv preprint
[60] Terence D Sanger. Neural network learning control of robot manipulators using gradually increasing task difficulty, IEEE transactions on Roboties and Automation, 1994.
[61] Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta-leariig with neemory-aggeentee neural networks. In International conference on machine learming (ICML), 2016.
[62] Jürgen Schmidhuber. Evolationary principles in self-referential learning, or on learning how to learn: the meta-meta-.. hook. PhD thesis, Technische Universität München, 1987.
[63] Avi Schwarzschild, Eitan Borgnia, Arjun Gupta, Furong Huang, Uzi Vishkin, Micah Goldblum, and Tom Goldstein. Can you learn an algorithm? generalizing from easy to hard problems with recurrent networks. Neural Information Processing Systems (NeurIPS), 2021.
[64] Daniel Selsam, Matthew Lamm, Benedikt, Percy Liang. Leonardo de Moura, David I. Dill, al. Learning sat solver from single-bit supervision. In International Conference Learning Representations (ICLR), 2018
[65] Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprini (xXiv:2103.0760) 2021.
[66] Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. Neural Information Processing Systems (NeurIPS), 2017.
[67] Sebastian Thrun and Lorien Pratt. Learning to learn. Springer Science & Business Media, 2012.
[68] Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society: Series B (Methodological,, 1996.
[69] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Neural Information Processing Systems (NeurIPS), 2017.
[70] Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. Pointer networks. In C. Cortes, N. Lawrence D. Lee, M. Sugiyama, and R. Garnett, editors, Neural Information Processing Systems (NeurlPS), 2015.
[71] Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In International Conference Machine Learning, 2021.
[72] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: Ofate-of-the- art natural language processing. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing: System Demonstrations. Association for Computational Linguistics (ACL), 2020.
2020.
14

---

[73] Xiaoxia Wu, Ethan Dyer, and Behnam Neyshabur. When do curricula work? arXiv preprint arX1v:2012.03107. 2020.
[74] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference. In International Conference on Learning Representations (ICLR), 2022.
[75] Lin Xu, Frank Hutter, Holger H Hoos, and Kevin Peyton-Brown Satzilla: portfolio-based algorithm selection for sat. Journal of artificial intelligence research, 2008
[76] Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2105. 2021.
[77] Chulhee Yun, Srinadh Bhojanapalli,, Ankit Singh Rawat, Sashank J Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequeene functions? arXiv preprint arXiv: 10077, 2019.
[78] Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.0433 0430,, 2022.
[79] Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In International Conference on Machine Learning (ICML), 2021.
15

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A] We don't t
perceive any negative impacts beyond those of building large-scat scale models general.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes] Appendix [
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Appendix
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes] Appendix A
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any assets either in the supplemental material or as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
16


Reviews:
[{"review_id": "wgL7RcBuvLU", "reviewer": "Reviewer_VCf1", "summary": "This paper proposes to study the ability of transformers model to predict values of a function given an input via sample prompting. Overall, this is a nice study that sheds some light on the capability of transformers model to deduce what function the context refers to and execute/compute the value of that function given an input in the context. ", "questions": "- For the input format x_1, f(x_1), x_2, f(x_2), .. if x_1 is a 20-dim vector, how is it presented to the model. This is a sequence of 20 numbers followed by 1 number for f(x_1, and so on?\n- any comparison regarding the overall loss function proposed?", "limitations": "I think the paper can easily be expanded to include complex functions. 2-layer ReLU is more complex than linear but also not as interpretable. Extending to something like degree-m polynomial and to observe trends with respect to capacity / how much it takes to train / etc can shed more light into the ability of transformers model for this type of sequence modeling.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n\n- Quite a thought-provoking paper studying transformers capabilities to learn explicit mathematical functions\n- Nice analysis / experiments. Everything seems well designed given the scope\n\nCons:\n- There's a lot of mentioning for 'in-context'. However, I do fear that this could be a bit misleading. All the model is doing is doing sequential modeling that given a series of inputs and the corresponding function values, the model predicts the function value of  the previous input. So in a way this can also be seen as sequence modeling. ", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "p-d50mApUVV", "reviewer": "Reviewer_aHn3", "summary": "**Goal:** Investigate to what extent decoder-only transformer models are capable to in-context learn functions from simple functions classes, like linear functions, sparse linear functions and shallow neural networks. If the answer is positive, understand the training conditions that lead to such in-context learning capabilities. \n\n**Main results:**\n* Transformers are capable of in-context learning all of the simple function classes tested in the paper to a surprising extend. The mean squared error values computed using the transformers often match the best-case known algorithms (such as min-norm linear regression and LASSO).  (The extent to which this is true is quite striking) \n* Carefully designed experiments show that this capability isn't explained by simple memorization of training examples. Rather, the transformers seem to learn in-context learning strategies that generalize to nontrivial distribution shifts. \n  * Distribution shifts in the in-context exemplars, as well as between in-context and query examples mostly leave the performance unchanged. \n  * In the case of learning linear functions, even the gradient of the learned transformer matches that of the min-norm linear regression solution. This is also a very strong and surprising finding. \n\n====================\n\n**Post rebuttal update**\n\nI thank the authors for their response. I still think that this is a strong submission that deserves a place in NeurIPS2022, hence maintain my score. ", "questions": "* Regarding curriculum learning: Do you think one can do away with curriculum learning by using a fixed dataset that contains a mixture of examples of different difficulties? Establishing this could be useful for facilitating future research in this topic, as one does not need to deal with the additional subtleties regarding a curriculum learning setup. \n* Inputs to the transformer: I don't quite get the description between lines 107 and 109. How exactly are the exemplars processed before being fed to the first transformer layer? \n* Consistency of results: Does training work all the time, or do you need to train a number of times before you come across something that displays the impressive in-context learning results? \n* Figure 3 (b): Is this plot averaged over multiple prompts? (This is a very interesting finding!) \n\n\nNitpick: Typo at line 335: \"when being trained\" instead of \"when being training\" \nAnother nitpick: In line 307, you write that you study \"what classes of functions can a model learn in-context\". The current set of results don't seem to address this exact framing, as answering that question would require experiments with many more types of functions. Perhaps rephrasing this would be helpful? ", "limitations": "I don't wish to report any glaring limitations beyond what has been listed as weaknesses above. ", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "STRENGTHS\n* **Comprehensive experiments:** The paper is full of clever experiments that strongly support the claims made in the paper. A number of these experiments yield nontrivial and (in my opinion) unexpected and surprising results. \n* **Sound and convincing:** The claims made in the paper are sound and convincing (assuming the experiments were run correctly).\n  * Some of the results look \"too good to be true\" (i.e. to the extent transformers match the performance of optimal algorithms like LASSO and min-norm linear regression), though the experiments look carefully done and I don't currently have any suspicions as to their validity. \n* **Well scoped and crisp problem formulation and analysis** The paper is very well scoped, and the goal and the contributions are very clear. \n\n\n\nWEAKNESSES\n* **Lack of discussion on types of in-context learning** There are multiple forms of in-context learning that can perhaps be studied separately (Olsson et. al. [31] discuss this briefly). This paper tackles only one of those - it would be better if the authors disambiguated this earlier in the paper. \n* **More distribution shifts:** Perhaps one or two more distributions shifts in addition to the skewed Gaussian distribution could be helpful. \n(The paper is very well-scoped, and (in my opinion) doesn't have that many glaring limitations within that scope. Some of what's below can be regarded as recommendations to make the paper better instead of weaknesses)\n* **Visualizing Sampled 2-layer neural network functions:** It'd be useful to provide some kind of visualization to make sure that the approximated neural network functions are non-trivial in the domain they're evaluated at. For example, all of the sampled neural nets could be roughly linear around the domain most of the samples fall, which'd make the result a bit less interesting (though still cool). \n* **Lack of mechanistic understanding:** After the authors make a very compelling case that transformers do indeed learn generic in-context learners, perhaps the authors could have spent some time analyzing how a GPT-2 like architecture could be achieving this mechanistically. ", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "KfYGp2F1872", "reviewer": "Reviewer_NrcW", "summary": "The paper studies in-context learning in transformers by training models to predict the output of a linear function given a few input/output examples presented in context. The main findings are that transformers trained in this way are able to match the performance of the optimal estimator in various settings, and show some ability to generalize out of distribution. Surprisingly, models trained on sparse linear functions have an in-context performance that matches Lasso. Also surprisingly, models trained on the output of 2-layer ReLU networks perform reasonably well on linear functions. Finally, the authors find that curriculum learning (starting with easier examples and progressing to harder ones) greatly speeds up training.", "questions": "It's interesting to explore how models are able to achieve in-context learning of linear and sparse linear functions. For example, can one recover the linear function's weights from the transformer's activations? And does the transformer implement some version of Lasso, or does it work in a different way? I realize this is out of scope for the present work. I'm only proposing these questions in case the authors have more to share on this.\n\nIn Figure 2 \"trivial error\" is introduced but is not defined. I believe it is defined to be the zero estimator only in Figure 5.", "limitations": "This work has no obvious societal impact.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper studies in-context learning in a well-scoped and well controlled setting. The setup and results are presented clearly, and the experiments are designed to produce sharp results with clear baselines and conclusions. I found several of the results surprising and inspiring for future investigation. The experiments on OOD generalization were especially interesting. Understanding in-context learning in transformers is an important problem, and this paper is a good step in that direction.\n\nIn terms of weaknesses, the settings studied in this work are relatively simple, and it is still hard to see how they connect with examples of in-context learning using natural data.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EZBy9imkWUY", "reviewer": "Reviewer_oarv", "summary": "This paper tries to study the in-context learnability of function classes in transformers. While the existing research is actively exploring this phenomenon in language models, this work presents a novel analysis of the phenomenon using function classes (in lieu of tasks) as the objects to be learnt. The paper shows that transformers can learn linear functions with high precision (obtaining almost the same loss as the least squares estimator) and that learning more complex  function classes (e.g. 2 layer neural nets) entails the in-context learning of simpler function classes (e.g. linear). The paper also shows that (similar to many language tasks in large-scale language models), learning function classes could also be classified as an \"emergent\" behavior during training. Further, the paper also analyze model behavior under distribution shift and how that in-context learning of function classes is surprisingly robust to two common forms of distribution shift. Altogether, this analysis is a solid contribution to the literature, with simple, easy to follow experiments that can shed more light into the mechanics of in-context learning.", "questions": "1. In section 5 (lines 229-232), does the model trained to learn the 2-layer neural net also learn sparse linear functions? \n2. For a model with a given capacity, how does the data size change in-context learning performance?\n3. The paper shies away from analyzing the interaction of training data and model capacity. Does the in-context learning performance depend on effective model complexity and not just on model capacity (as characterized by number of parameters)? Any comments on incorporating training data size into the analysis?\n4. Unless I missed something, is there an experiment where model training is done a mixture of function classes? How sensitively do the the results depend on the mixture ratios?\n", "limitations": "The authors do not include a section on limitations and potential negative societal impacts of the work. I urge the authors to add a limitations section, explaining any factors (e.g. training data size, effective model complexity, etc.) that could further complicate the analysis.\n", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The strengths of the paper are:\n\n1. The paper is well written, with a clearly defined (and well motivated) experimental protocol. This simple setup could encourage further experiments.\n2. The empirical results are intriguing and will be of interest to the broader ML research community, given transformer has become such an indispensable architecture. \n3. To my knowledge this is the first exposition of \"emergent\" in-context learning in function classes as well as curriculum learning of function classes.\n\nThe paper has the following weaknesses:\n\n1. There is no analysis of the size of training data that is required for successful in-context learning. A surprising omission in the experiments is the lack of a plot of data size against in-context learning performance. I think this is an important component that is missing.\n2. The paper does not shed any light on the in-context learning phenomenon for language tasks, and one of the missing pieces is it's lack of analysis on training data size, along with training on a mixture of function classes.\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept