In-Context Learning through the Bayesian Prism
Authors: Madhur Panwar
Kabir Ahuja
Navin Goyal
Venue: ICLR 2024 poster (2024)
Keywords: In-context Learning
Transformers
Inductive Biases
Meta Learning
Language Modelling
Bayesian Inference
Primary area: transfer learning, meta learning, and lifelong learning

Abstract: In-context learning (ICL) is one of the surprising and useful features of large language models and subject of intense research. Recently, stylized meta-learning-like ICL setups have been devised that train transformers on sequences of input-output pairs $(x, f(x))$. The function $f$ comes from a function class and generalization is checked by evaluating on sequences generated from unseen functions from the same class. One of the main discoveries in this line of research has been that for several function classes, such as linear regression, transformers successfully generalize to new functions in the class. However, the inductive biases of these models resulting in this behavior are not clearly understood. A model with unlimited training data and compute is a Bayesian predictor: it learns the pretraining distribution.
In this paper we empirically examine how far this Bayesian perspective can help us understand ICL. To this end, we generalize the previous meta-ICL setup to hierarchical meta-ICL setup which involve unions of multiple task families. We instantiate this setup on a diverse range of linear and nonlinear function families and find that transformers can do ICL in this setting as well. Where Bayesian inference is tractable, we find evidence that high-capacity transformers mimic the Bayesian predictor. The Bayesian perspective provides insights into the inductive bias of ICL and how transformers perform a particular task when they are trained on multiple tasks. We also find that transformers can learn to generalize to new function classes that were not seen during pretraining. This involves deviation from the Bayesian predictor. We examine these deviations in more depth offering new insights and hypotheses.

Published as a conference paper at ICLR 2024
IN-CONTEXT LEARNING THROUGH THE BAYESIAN PRISM
Madhur Panwar* Kabir A
Navin Goyal
V Microsoft Research India
(t (t-mpanwar, htw.l/gg//m.rrofo..ooom
University of Washington
kahujalcs www.i..sssinngtnnedd..
ABSTRACT
In-context learning (ICL) is one of the surprising and useful features of large lan- guage models and subject of intense research. Recently, stylized meta-learning- like ICL setups have been devised that train transformers on sequences of input- output pairs (z. f (a)). The function comes from function class and generaliza- tion is checked by evaluating
sequences generated from unseen functions from the same class. One of the main discoveries in this line of research has been that for several function classes, such as linear regression, transformers successfully generalize to new functions in the class. However, the inductive biases of these models resulting in this behavior are not clearly understood. A model with un- limited training data and compute is a Bayesian predictor: it learns the pretraining distribution. In this paper empirically examine how far this Bayesian perspec- tive can help us understand ICL. To this end, we generalize the previous meta-ICL setup to hierarchical meta-ICL setup which involve unions of multiple task fami- lies. We instantiate this setup on a diverse range of linear and nonlinear function families and find that transformers can do ICL in this setting as well. Where Bayesian inference is tractable, we find evidence that high-capacity transformers mimic the Bayesian predictor. The Bayesian perspective provides insights into the inductive bias of ICI. and how transformers perform a particular task when they are trained on multiple tasks. We also find that transformers can learn to gener- alize to new function classes that were not seen during pretraining. This involves deviation from the Bayesian predictor. We examine these deviations in more depth offering new insights and hypotheses
INTRODUCTION
In-context learning (ICL) is one of the major ingredients behind the astounding performance of large language models (LLMs) (Brown et al., (020). Unlike traditional supervised learning. ICL is the ability to learn new functions f without weight updates from input-output examples (x. fx)) provided as input test time. For instance, given the prompt up down, low -> high, small ->, a pretrained LLM will likely produce output bi it apparently infers that the function in the two examples is the antonym of the input and applies on the new input. This behavior often extends to more sophisticated and novel functions unlikely to have been seen during training and has been the subject of intense study, e.g., Min et al. (2022b); Webssn & Pavlick (2022); Min al. (202aa); Liu et al. (2023); Dong et al. (2023). More broadly than its applications in NLP, ICL can also be viewed as providing method for meta-learning (Hospedales al. (022) where the model learns to learn a class of functions.
Theoretical understanding of ICI. is an active area of research. Since the real-world datasets used for LLM training are difficult to model theoretically and are very large, ICL has also been studied in stylized setups, e.g., Xie et al. (2022): Chan et al. (2022b); Garg et al. 2022); Wang et al. (2023);
"Equal Contribution
1 A part of this work was done while Kabir was a Research Fellow at Microsoft Research India. http://tttuuccoc
We release our code at /utp:pp://gggttbb..ooo///rrrrrrr -ayesi an-prrannniiiiiio

---

Published as a conference paper at ICLR 2024
Hahn & Goyal (023). These setups study different facets of ICL, In this paper. we focus on the meta-learning- like framework of Garg et al. (2022). Unlike in NLP where training is done on doc- uments for the next-token prediction task, here the training and test data look similar in the sense that the training data consists of input of the form is ( where æ, IRd and are chosen i.i.d. from a distribution, and f IRd - IR is func- tion from a family of functions, for example, linear functions or shallow neural networks. We call this setup MICL. A striking discovery in Garg et al. (2022) was that for several function families, transformer- based language models during pretraining learn to implicitly implement well- al- gorithms for learning those functions in context. For example, when shown 20 examples of the form (æ., w" T x) where x, w E IR?20, the model correctly outputs w Trst 2 on test input test Apart from linear regression, they show that for sparse linear regression and shallow neural networks the trained model appears to implement well-known algorithms; and for decision trees, the trained model does better than baselines. Two follow-up works AAküürek et al. (2022) and von Oswald et al. (2022) largely focused on the case of linear regression. Among other things, they showed that transtormers with one attention layer learn to implement one step of gradient descent on the linear regression objective with further characterization of the higher number of layers.
f(x)),,.... (Ixx, [(xx)) #k+11 ) and output
Bayesian predictor. An ideal language model (LM) with unlimited training data and compute would learn the pretraining distribution as that results in the smallest loss. Such LM produces the output by simply sampling from the pretraining distribution conditioned on the input prompt. Such an ideal model often called Bayesian predictor. Many works make the assumption that trained LMs are Bayesian predictors, c.g. Saannshi et al. (2021); Xie et al. (2022); Wang et al. (002)). Most relevant to the present paper, Akküürk al. (2022) show that in the MICL setup for linear regression, in the underdetermined setting, namely when the number of examples is smaller than the dimension of the input, the model learns to output the least L2-norm solution which is the Bayes-optimal prediction. In this paper we empirically examine how general this behavior is across choices of tasks.
Prior work has investigated related questions but we are not aware of any extensive empirical verifi- cation. F.g., Xie et al. (2022) study a synthetic setup where the pretraining distribution is given by a mixture of hidden Markov models and show that the prediction error of ICL approaches Bayes- optimality as the number of in-context examples approach infinity. In contrast, we test the Bayesian hypothesis for ICL over a wide class of function families and show evidence for equivalence with Bayesian predictor at all prompt lengths. Also closely related, Müller et al. (2022); Hoolmann et al.) (003) train transformer models by sampling data from a prior distribution (Prior Fitted Networks), so it could approximate the posterior predictive distribution inference time. While these works focus on training models to approximate posterior distributions for solving practical tasks (tabular data), our objective to understand how in-context learning works in transformers and to what extent we can explain it as performing Bayesian Inference on the pre-training distribution.
Simplicity bias. Simplicity bias, the tendency of machine learning algorithms prefer simpler hypotheses among those consistent with the data, has been suggested as the basis of the success of neural networks. There are many notions of simplicity (Mingard et al. 2023 Goldblum et al., 2023)
Does in-context learning also enjoy simplicity bias like pretraining?
Our contributions. brief, our contributions are
1. A setup for studying ICL for multiple function families: First, we extend the MICL setup from Garg et al.] (2022) to include multiple families of functions. For example, the prompts could be generated from mixture of tasks where the function f is chosen to be either a linear function or decision tree with equal probability. We call this extended setup HMICL. We experimentally study HMICL and find that high-capacity transformer models can learn in context when given such task mixtures. (We use the term "high-capacity informally; more precisely, it means that for the task at hand there is sufficiently large model with the desired property.)
2. high-capacity transformers perform Bayesian inference during ICL: To understand how this abil- ity arises we investigate in depth whether high-capacity transformers simulate the Bayesian predic- tor. This motivates us to choose a diverse set of linear and nonlinear function families as well as prior distributions in HMICL and MICL setups. Function families we consider were chosen because ei- ther they permit efficient and explicit Bayesian inference or have strong baselines. We provide direct and indirect evidence that indeed high-capacity transformers often mimic the Bayesian predictor.
2

---

Published as a conference paper at ICLR 2024
The ability to solve task mixtures arises naturally as consequence of Bayesian prediction. In con- current work, Bai et al. (2023) also study the multiple function classes setup for ICL like us, showing cn-context learn individual function classes from the mixture. However, there
that transformers are three main differences between our works. |ai et al. (2023) interpret the multi-task ICL case as algorithm selection, where transformer based on the in-context examples selects the appropriate al- gorithm for the prompt and then executes it. We provide an alternate explanation that there is no need for algorithm selection and follows naturally from the Bayesian perspective. Further, while they show that their constructions approach bayes-optimality at large prompt lengths, we show through experiments that this actually holds true for all prompt lengths. Finally, we test this phenomenon on much larger set of mixtures. For Gaussian Mixture Models, we also compare the performance with the exact Bayesian predictor and such a comparison missing in Bai et al. (2023).
3. Link between ICI. inductive bias with the pretraining data distribution: We also investigate the inductive bias in a simple setting for learning functions given by Fourier series. If ICI. is biased towards fitting functions of lower maximum frequency. this would suggest that it has bias for lower frequencies like the spectral bias for pretraining. We find that the model mimics the Bayesian predictor; the ICL inductive bias of the model is determined by the pretraining data distribution: during pretraining all frequencies are equally represented, then during ICL the LM shows no preference for any frequency. On the other hand, if lower frequencies are predominantly present in the pretraining data distribution then during ICL the LM prefers lower frequencies. Chan et al.) 20222abb studies the inductive biases of transformers for ICL and the effect of pretraining data distribution on them. However, the problem setting in these papers is very different from ours and they do not consider simplicity bias.
4. Generalization to new tasks not seen during training in HMICL: In HMICL setup, we study generalization to new tasks that were not seen during pretraining. We find that when there's 's sufficient diversity of tasks in pretraining, transformers generalize to new tasks. Similar study was made in the concurrent work of Raventós et al. (2023) for the noisy linear regression problem within MICL.
5. Study of deviations from Bayesian prediction: Finally, we study deviations from the Bayesian predictor. These can arise either in multitask generalization problems or when the transformer is not of sufficiently high capacity for the problem at hand. For the former, we study the pretraining inductive bias and find surprising behavior of transformers where they prefer to generalize to a large set of tasks early in the pretraining which then they forget. For the latter, drawing on recent work connecting Bayesian inference with gradient- based optimization, we hypothesize that in fact transformers may be attempting do Bayesian inference.
BACKGROUND
We first discuss the in-context learning setup for learning function classes as introduced in Garg et al. (2022). which we call Meta-ICL or MICL. Let Dx be a probability distribution on Rd. Let J be a family of functions f Rd + R and let D be distribution on F. For simplicity, we often use f ~ F to mean f ~ Dr. We overload the term function class to encompass both function definition as well as priors its parameters. Hence, linear regression with standard gaussian prior and sparse prior will be considered different function classes based on our notation.
To construct a prompt P (æ1, f(x),.... p: (Tp), 1)) of length p. we sample f ~ F and inputs I, Dx i..d. for i (1....p. A transformer- based language model Ma is trained to predict f( given P. using the objective: ming 1 p+1 E (Mo(Pi), If ((x:11)) where pi denotes the sub-prompt containing the first input-output examples as well as the (i (i+ 1)- 2.. ..,, f(xx), #++1) and Ti.p = (#I,..., zp). While other choices of the loss function E(:.)) are possible, since we study regression problems we use the squared-error loss (i.c., f(y, v') (y- 22)²) in accordance with Garg et al. (2022).
th input, i.c. (I1, fxx1,,
At test time, we present the model with prompts Ptest that were unseen during training with high probability and compute the error when provided i in-context examples: lossek: mophabiiityaaa
EyP.oo (Ma ((x++1)]] for l E [1,....p],
PME. Our work uses basic Bayesian probability described, e.g., in Murphy (022). We mentioned earlier that an ideal model would learn the pretraining distribution. This happens when using the cross-entropy loss. Since the square loss in the objective definition, the predictions of this
3

---

Published as a conference paper at ICLR 2024
ideal model can be computed using the posterior mean estimator (PME) from Bayesian statisties.
For each prompt length i, and any prompt Q = (1 g(æ1). .......... Tp. g (p)) Tp+1) where g is a function in the support of Dr, we compute the PME by taking the corresponding summand
in objective definition above, which will be given by (Q') = E [(#++1) P = Q"1 for all
i S p. This is the optimal solution for prompt Q, which we refer to as PME. Please refer AA.1
technical details behind this computation.
2.1 HIERARCHICAL. META-ICI.
We generalize the MICL setup, where instead of training transformers from functions sampled from single function class, we sample them from mixture of function classes. Formally, we define a mixture of function classes using a set of m function classes FF = {F1. FmF J and sampling probabilities a = [01,-.. aml 1d with Σ a; 1. We use a to sample a function class for constructing the training prompt P. We assume the input distribution Dx to be same for each class FT. More concretely, the sampling process for P defined as: F;; ~ F s.t. P(F F;) = a;; ii) ~ Fú iii) 2, ~ Dx. Vj E (1.....2) and finally, iv) P (z1, f ... Zp.
We call this setup Hierarchical Meta-ICL or HMICL, as there is an additional first step for sam- pling the function class in the sampling procedure. Note that the MICL setup can be viewed as special case of HMICL where m = 1. The HMICL setting presents a more advanced scenario to validate whether the Bayesian inference can be used to explain the behavior of in-context learn- ing in transformers. Further, our HMICL setup is also arguably closer to the in-context learning in practical LLMs which can realize different classes of tasks (sentiment analysis, QA, summarization etc.) depending upon the inputs provided. (For additional discussion on HMICL and MICL, refer to Appendix IC.I) The PME for the hierarchical case is given by:
(P) = 31 Mo.F, (P) +... Mo.F..
(1)
where B; a.P; (P)ppp for i m. Probability density p.(-) is induced by the function class F; on the prompts in a natural way, and pP)( = a:p;: +... (mmPm(P) (P). Please refer to 1 in the Appendix for the derivation. The models are trained with the squared error loss mentioned above and at test time we evaluate loss@k for each task individually.
2.2 MODEL AND TRAINING DETAILS
We use the decoder-only transformer architecture Vaswani et al. (017)) as used in the GPT models Radford et al.) (2019). Unless specified otherwise, we use 12 layers, 8 heads, and a hidden size (dx) of 256 in the architecture for all of our experiments, We use a batch size of 64 and train the model for 500k steps. For encoding the inputs x;'s 's and f=:;)ss we use the same scheme as Garg et al. )022) which uses a linear map E E R xd to embed the inputs x; 's as Ex, and f(æ,)' )s as pad (..). where fpaa( (z,) = [z,) DT E Rd. In all of our experiments except the ones concerning the Fourier series, we choose Dx as the standard normal distribution i.e. N/O0, 1). unless specified otherwise. To accelerate training, we also use curriculum learning like Garg et al. (2022)
for all our experiments where start with simpler function distributions (lower values of and p) at the beginning of training and increase the complexity as we train the model.
3 TRANSFORMERS CAN IN-CONTEXT LEARN TASK MIXTURES
In this section, we provide evidence that transformers' ability to solve mixture of tasks arises nat- urally from the Bayesian perspective. We start with a Gaussian Mixture Models (GMMs) example where the exact Bayesian solution is tractable and later discuss results for more complex mixtures.
3.1 GAUSSIAN MIXTURE MODELS (GMMs)
We define mixture of dense-linear regression classes FGMM- = [FDR,,..... FDR... ). where FDR; {f : æ + w x w; E Rd and w; ~ NN(u,, E,). In other words, each function class in the mixture corresponds to dense regression with Gaussian prior on weights (but different means or covariance matrices). We report experiments with m = 2 here, and the mean vectors are given by μi = (3, .., 0) and μ2 = (-3, ..., 0) for the two classes The covariance matrices are equal ( = Σe = E*), where E is the identity matrix La with the top-left entry replaced by 0. Note that we can equivalently view this setup by considering the prior on weights as a mixture of Gaussians Le. PM (M))
NaNa((1 En) + azVa(He. Ez). For brevity, we call the two function classes Ti and Tz. We train the transformer on a uniform mixture ie. ag 0g are 7. We use d = 10 and the prompt length p E (10, 20).
4

---

Published as a conference paper at ICLR 2024
Duluatio TI
Maw R.O
Ted malariie people to N2t 1,1
Evd pooeds A11
preass tn
22000
FME 2R0
FME
coamph))
Incombes
Figure 1: Transformers simulate PME when trained on dense regression task-mixture with weights having a mixture of Gaussian prior (GMM). (left): Comparing the performance of the Transformer with PME of individual Gaussian components (PME (Ti) and PME (T2)) and of the mixture PME (GMM). (right): MSE between the probed weights of the Transformer and PMEs.
Recovering implied weights. To provide a stronger evidence for the Bayesian hypothesis, apart from the loss curves, we also extract the weights implied by transformers for solving the regression task in-context. Following Akyürek et al. (2022), we do this by generating model's predictions {u{} on the test imputs { ~ Dx and then solving the system of equations to recover Probe, We then compare the implied weights wrobe with the ground truth weights w as well as the weights extracted from different baselines by computing the their MSE.
Results. Figure T (left), we note that Transformer's errors almost exactly align with those of the PME of the mixture, PME (GMM), when prompts come from either Tj T2. (For details on computation of PME. please refer to $2.2 in Appendix). For each plot, let T prompt and T Tother denote the component from which prompts are provided and the other component respectively. When = 10 examples from prompt have been provided, the Transformer, PME prpommpe), and PME (GMM) all converge to the same minimum error of O. This shows that Transformer is simulating PME (GMM), which converges to PME (pprompt) ) at k d. PME (Tohher)'s errors keep increasing as more examples are provided. These observations are in line with Eq. 3 As more examples from the prompt are observed, the weights of individual PMEs used to compute the PME (GMM) (i.e., the f's) evolve such that the contribution of T promps increases in the mixture with is (Fig. 22 in the Appendix). In Figure (right), MSE between weights from different predictors are plotted. Transformer's ' implied weights are almost exactly identical to PME (GMM) for all k. Please refer to S..2 for additional details and results.
More complex mixtures. We test the generality of the phenomenon discussed above for more complex mixtures, involving mixtures of two three different linear inverse problems (e.g. dense regression, sparse regression, sign vector regression) as well as some mixtures involving non-linear function classes like neural networks and decision trees. In all of these cases we observe that trans- formers trained on the mixtures are able to generalize on the new functions from the mixture of function classes and match the the performance of single-function class transformer models depend- ing upon the distribution of input prompt. Please refer to $ for details.
Implications. Our GMM experiments challenge the existing explanations for the multi-task case, e.g. the models first recognizes the task and then solves it. When viewed through the Bayesian lens, transformers do not need recognize the task separately and recognition and solution are intertwined as we show in Equation [
SIMPLICITY BIAS IN ICL?
In this section, we explore if transformers exhibit simplicity bias in ICL. In other words, when given a prompt containing input output examples, do they prefer to fit simpler functions among those that fit the prompt? To study this behavior we consider the Fourier Series function class, where the output is a linear function of sine and cosine functions of different frequencies. By training transformers on this class, during ICL we can study if transformers prefer fitting freeueyyyy functions the prompt over higher frequencies, which can help us study the presence of a simplicity bias.
More formally, we can define Fourier series by the following expansion: f(x) dg + ΣN On cos (nTz/L) L) Σ bn sin (nzx/L) L) where, x E [-L, L). and ap, a 's and b are known Fourier coefficients and cos no/ and sin un/ define the frequency components.
5

---

Published as a conference paper at ICLR 2024
Fourrer Series MICL
Fourier
M 4
0.25
0.2
2/
55 0.50
0:25
32
0.0
0000
4.344 5
5 101122
101112
121112
522450 6 0
Figure 2: Measuring the frequencies of the simulated function during ICL by transformer.
MICL Setup. In the MICL setup we train transformer on a single function class de- fined as Ffourier
{f(; (=; b) w (2))), Rd) with standard gaussian prior Note that here &N as the Fourier feature map i.c. Øx(x) [1, cos L), ................................ cos (NTx/ L), sin (zæ/L),
on weights
sin (Nxz/L)J². For training transformers to in-context-learn Frouierr we fix value of N and sample functions f Floureer We consider the inputs be scalars, i.e. 24 E [-L, L] and sample them i.i.d. from the uniform distribution on the domain: I2 ~ L(-L,L). In all of our experiments, we consider N = 10 and L = 5. At test time we evaluate on Frourier for M 10), i.e. during evaluation we also prompt the model with functions with different maximum frequency as seen during training.
HMICL Setup. We also consider mixture of Fourier series function classes with different max- imum frequencies, i.c. Flouieer = {Fluurier ;... Flourier. We consider N = 10 in our experiments and train the models using a uniform mixture with normalization. During evaluation, we test indi- vidually on each Floueer where M E [1.
Measuring inductive biases. To study simplicity bias during ICL, we propose method re- cover implied frequency from the transformer model. We start by sampling in-context examples (Z)) Zgg f(xx), ), and given the context obtain the model's predictions on a set of test in- puts {a (= i.c. 3/ = MM ((z1, ((I1,((z) f(z1), ... Ik, (axx), We can then perform Discrete Fourier Transform (DFT) on f1/1... .... Vmm to obtain the Fourier coefficients of the function output by M, which we can analyze to understand the dominant frequencies.
Results. In both MICL and HMICL setups discussed above observe that transformers are able to in-context learn these function classes and match the performance of the Bayesian predictor or strong baselines. Since, in this section we are primarily interested in studying the simplicity bias, here we only report the plots for frequencies recovered from transformers at different prompt lengths in Figure 2 (more details in Figures 12 and 32 of Appendix). As can be seen in Figure 2 (left), in the single function class case, transformers exhibit no bias towards any particular frequency. For small prompt lengths (A: = 2), all N frequencies receive similar absolute value of coefficients as implied by the transformer. As more examples are provided (k 21), transformer is able to recognize the gold maximum frequency (M = 4) from the in-context examples, and hence coefficients are near zero for > M, but as such there is no bias towards any particular frequencies. However, when we consider the mixture case in Figure 2 (right), the situation is different. We see clear bias for lower frequencies at small prompt lengths; however, when given sufficiently many examples they are able to recover the gold frequencies. This simplicity bias can be traced to the training dataset present in most of the functions of the mixture while higher frequencies will be more rare: Frequency 1 will be present in all the function classes whereas frequency N will be present only in Flounerr We perform additional experiments biasing pre-training distribution to high frequencies and observe complexity bias during ICL (Appendix CC...T).
for the mixture since lower frequencies
Implications. These results suggest that the simplicity bias (or lack thereof) during ICL can be attributed to the pre-training distribution which follows naturally from the Bayesian perspective i.c. the biases in the prior are reflected in the posterior. Transformers do not add any extra inductive bias of their own as they emulate the Bayesian predictor.
MULTI-TASK GENERALIZATION
In this section we test the HMICL problems on out-of -istrribbtt (OOD) function classes to check generalization. We work with the d monomials regression problem, F- which is given by function class where the basis is formed by feature set S, subset of degree-2 monomials

---

Published as a conference paper at ICLR 2024
1.1
100
Figure 3: Multi-task generalization results for Monomials problem. ID and OOD evaluation for K 10, 100 is presented. As task diversity (K) increases, the model starts behaving like Lasso and BP pracy and its ID and OOD losses become almost identical, i.e. it generalizes to OOD.
d}. We can then define the feature map Ls(2) = (Izz²) (o.j)s and f(x:) = ((x) is function of this class, where w ~ Nis [ (0, I). We compare the performance of on this class with OL.S performed on the feature set S (OLSs) which is the Bayesian predictor in this case. We find that the error curves of the TF trained and evaluated on this class follow OLSs baseline closely for all prompt lengths, on both in- and out-of-distribution evaluation. (Refer to BB2.3 for a detailed description of setup and results.)
S
Extending to HMICL setting. For HMICL, we use multiple feature sets S₁. to define the mixture. Each S& defines function class JS, )m The pretraining distribution is induced by the uniform F distribution ((F) over collection of such function classes, IF (F ... Sx J, where Sy C. M. K feature sets Sy 's, each of size D, are chosen at the start of the training and remain fixed. K is the task diversity of the pretraining distribution. To sample training function for the TF, we first sample a function class Fimmont)) with replacement from ((FF) and then sample a function from the chosen class; f(x) w" Sy()). where w (D I). Our aim is to check if TF trained on (FF) can generalize to the full distribution of all function classes (for feature sets of size D) by evaluating its performance on function classes corresponding to feature sets S' & (S (.k].
Experimental Setup. We choose D d 10,p 124. There is no curriculum learning and p remain fixed throughout training. Note that the total number of degree-2 monomials = Mtst = ( + (1) = 45 + 10 = 55; and the total number of distinct feature sets S 's (and hence function classes)
(Meat) (19) s 310 We train various models for different task diversities; K E {10, 20, 40, 100. 500, 1000, 5000). We evaluate on batch of B 1280 functions in two settings: (a) In-Distribution (ID) test functions formed using randomly chosen function classes from the pretraining distribution; (b) Out-of-Distribution (OOD) Test functions formed using randomly chosen function classes not in the pretraining distribution.
Baselines. We compare the performance of multi-task transformer models with the following base- lines: 1. OLSs: Here, we perform OLS on the basis formed by the gold feature set S, which was used to define the function in the prompt that we wish to evaluate. This will correspond to upper bound on the performance as at test time the transformer model has no information about the correct basis. 2. OLS.,. Here, OLS is performed on the basis formed by all degree-2 monomials &as (æ)) for an input æ. Hence, this baseline can generalize to any of the feature set S. However, since all degree-2 monomial features are considered by this baseline, it would require higher number of input-output examples (equal to M1ot) for the problem to be fully determined. 3. Lasso Similar to OLS we operate on all degree-2 monomial features, but instead of OLS we perform Lasso with = 0.1. It should also generalize to arbitrary feature sets S, however, Lasso can take advantage of the fact that |S| = < Mtoti hence should be more efficient than OLS
) we use the transformer trained on the full distribution of function families, since the computation of the exact predictor is expensive. From the plots in Figure 3 we observe that while for small values of K, the OOD generalization is poor but as we move to higher values of K, the models start to approach the performance of OLS& y and eventually Lasso provement in OOD performance comes at the cost of ID performance as task diversity (K) increases. Eventually, at larger K, both ID and OOD performances ticularly interesting since the models learn to generalize to function classes out of the pre-training distribution and hence deviate from the Bayesian behavior which would lack such generalization
Results. As proxy for the Bayesian predictor (BP poxy)))),
unseen S's. Further, they also start behaving like BP Prox.. However, this im-
identical. These observation
par-
7

---

Published as a conference paper at ICLR 2024
tasks
tasks
16
tasks
tasks
128 tasks
256 tasks
2048
tasks
-1048576 tasks
000 Evaluation
3.4
MiC.
1.2
Ridge
OMMSE
3.0
0.5
0.2
22
22
2
2
2
100k 2004 300K 400k 500k
Phetrainiing Taskss
Figure 4: Left: Evolution of ID (solid lines) and OOD (dashed lines) losses during pretraining for representative task diversities. Task diversities {27 ... 2112) represent the Gaussian forgetting region. The moving average (over 10 training steps) of the losses is plotted for smoothing. Right: OOD loss given the full prompt length of 15 for the final checkpoint of models trained various task diversities. Task diversities {2 .2"...2" ..1] represent the transition region.
and the pre-tre pre-tranning distribution instead. We observe similar results for another family of func- tion classes coming from Fourier series (details of these in Appendix $D.2).
In concurrent work Raventós al. (2023) also present a multi-task setting within MICL where set of weight vectors define the pretraining distribution for the Noisy Linear Regression problem. Since we work with HMICL, our setting more general; moreover, generalization new function classes in our setting happens in similar way generalization new tasks in Raventós et al. (2023). They emphasized deviation from the Bayesian predictor. What leads to these deviations? To understand this, in the next section we study pretraining inductive bias of transformers.
DEVIATIONS FROM BAYESIAN INFERENCE?
In the previous section observed deviations from the Bayesian predictor in multitask generaliza- tion. To investigate this we study the pretraining dynamics of transformers in the first subsection. Another of apparent deviations from Bayesian prediction observed in literature when the prob- lem is too hard or the transformer has limited capacity. We discuss these in the second subsection.
6.1 ICL TRANSFORMER FIRST GENERALIZES THEN MEMORIZES DURING PRETRAINING
We observe very interesting phenomenon (which term "forgetting") from multi-task experi- ments: For certain task diversities, during pretraining, HMICL Transformer first generalizes (fits the full distribution) and later forgets il and memorizes (fits the pretraining distribution).
The 'orrgetting' phenomenon is general and occurs in our HMICL experiments in However, here focus the the Noisy Linear Regression problem from Raventós al. (2023) since forgetting is the cleanest in this setting. We briefly mention the problem setup and display the evidence for forgetting during pretraining, followed by its relation the agreement of HMICL Transformer with the Bayesian predictors on the pretraining and full distributions.
Problem Setup. We follow the Noisy Linear Regression (NLR) setup from Raventós et (2033) d = 8,p = 15. (For details, see (D.3 The pretraining distribution (PT ais is induced by the uniform distribution on fixed set of tasks (weight vectors). Several models are trained, one per task diversity K E {22 22. ... 220). The full distribution of weight vectors is standard normal. (Hence we use the term "Gaussian distribution" to refer to the full distribution (FGaist). ).) To form a function for training,
randomly choose weight vector w from the pretraining distribution and define f(x) = w" x + E, where e Na o² 0.25).
Evidence of forgetting and agreement with Bayesian predictors. As we did in $ we evaluate TF on tasks from both in- and out-of nut-oofprettraiing distribution, where the tasks used to construct the test function come from pretraining distribution the standard Gaussian distribution respectively; corresponding losses are called ID (Pretrain test) loss and OOD (Gaussian) loss. We also plot the Bayesian predictors for both pretraining (dMMSE) and full (Gaussian) distribution (Ridge regres-
8

---

Published as a conference paper at ICLR 2024
sion) as defined in Raventós al. 2023 In Figure 4 (left) we plot the evolution during pretraining of ID and OOD losses for representative task diversities (more details in $D3) ID loss se 0 for all task diversities. We group them into the following 4 categories based on OOD loss and describe the most interesting one in detail (full classification in D.33): (1) 2 to 2 no generalization; no forgetting; (2) 2 to 2 some generalization; no forgetting: (3) 2 to 211 full generalization and forgetting - OOD loss improves, reaches a minima Emin. at which it is same as ID loss, then it worsens. At Emin. OOD loss agrees with Ridge, then gradually deviates from it and at lend (end of pretraining). it is in between dMMSE and Ridge. We refer to this group of task diversities as the "Gaussian forgetting region" since the model generalizes the full (Gaussian) distribution over tasks at miin but forgets it by lend; (4) 2 to 2 full generalization; no forgetling.
The agreement of TF in OOD evaluation with Ridge or dMMSE as mentioned above shown in $D.3 Figure (right) plots the OOD loss given the full prompt length of 15 for the final checkpoint of models trained for various task diversities. As can be seen, smaller task diversities (up to 29) agree with dMMSE (Bayesian predictor on PT and larger task diversities (from 2 onwards) agree with Ridge regression (Bayesian predictor on FG (This observation was originally made by |Raventós et al. (2023) and we present it for completeness.) Intermediate task diversities (2? to 21 agree with neither of the two and we term them collectively as the transition region. We note that both the Gaussian forgetting region and the transition region consist of the same set of task diversities viz. (2?, ... 221). The phenomenon of forgetting provides an interesting contrast to grokking literature (e.g. Nanda et al. (2023) and can possibly be explained via the perspective of simplicity bias. The extent of forgetting is directly proportional to the input dimension (d) and is robust changes in hyperparameters (details, in section SD.3.
6.2 GRADIENT DESCENT AS A TRACTABLE APPROXIMATION OF BAYESIAN INFERENCE
Some recent results in the literature within MICL suggest that transformers compute their answer by gradient descent on in-context examples. Could this be related to Bayesian inference? We provide preliminary evidence for this in Appendix $
7 SUMMARY OF FURTHER RESULTS
In this section we summarize further results from the Appendix that verify the generality of Bayesian hypothesis. We test the hypothesis on a variety of linear and non-linear inverse problems both MICL and HMICL setups and find that transformers are able in-context learn and generalize to unseen functions from these function classes. In the cases where PME computation is tractable, we compare transformers with the exact Bayesian predictor (PME) and establish the agreement between the two. Where PME is intractable, we compare transformers with numerical solutions obtained using a Markov Chain Monte Carlo (MCMC) sampling algorithm Homan & Gelman (2014) (Figure 10 in Appendix). When even sampling based solutions do not converge, we compare with strong baselines that are known to be near optimal from prior work. For linear problems, we test on Dense, Sparse, Sign Vector, Low Rank and Skewed Covariance Regression. For these problems, we show that not only do transformers' errors agree with the Bayesian predictor (or the strong baselines), but also the weights of the function implied by the transformer. For the non-linear case, we explore regression problems for Fourier Series, Degree-2 Monomials, Random Fourier Features, and Haar Wavelets. For Bayesian inference the order of demonstrations does not matter for the class of problems used in our setup. Figure TT experimentalyy verifies it for Dense Regression, where transformer's performance is independent of the permutation of in-context examples across different prompt lengths. Further, we note that in the HMICL setup, generalization to functions from the mixture might depend on different factors such as normalizing the outputs from each function class. We provide complete details for each of these function families and corresponding results in Appendix B and $C
8 CONCLUSION
In this paper we provided empirical evidence that the Bayesian perspective could serve as uni- fying explanation for ICL. In particular, it can explain how the inductive bias of ICL comes from the pretraining distribution and how transformers solve mixtures of tasks. We also identified how transformers generalize to new tasks and this involves apparent deviation from Bayesian inference. There are many interesting directions for future work which we discuss in Appendix F

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
K.A. was supported in part by the National Science Foundation under Grant No. IIS2125201. We would like to thank all the anonymous reviewers for their constructive feedback.
REFERENCES
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm is in-contexi learning? investigations with linear models. CoRR, abs/221 2022. doi: BtRLhhttss:///oii... /rg/ 48550/ 2211 15661
10.48850/arx...211...555
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection, 2023.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the Ability and Limitations of Trans- formers to Recognize Formal Languages. In Proceedings of the 2020 Conference on Empir- ical Methods in Natural Language Processing (EMNLP), pp. 7096-7116 16, Online, November 2020. Association for Computational Linguistics. doi: 10. 10.1166/1122220.emm www.emmmmllmmm..m.... URL
https: http///aaaaathoooo....o.o org/ 2020. emnle-meee in. 576
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 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: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020. URL httpttsss////occeeeee Abstract.hir 1457c0d6bf cb www.196666111100111111111111 Abstaataaaa.aatc.tttttt......... www.ee/wwis..cppaaeee/////// wwtpapper/222//aaatttee
E.J. Candes and T. Tao. Decoding by linear programming. IEEE Transactions Information Theory, 51(12):400--4215, 2005. doi: 10.1 10.110991TT2220588888777
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang. Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. In S. Koyejo, S. Mohamed, A. Agar- wal. D. Belgrave, K. Cho. and A. Oh (eds.), Advances in Neural Information Pro- cessing Systems, volume 35, pp. 18878-18891, Curran Associates, Inc., 2022a. URL 201///2222////1//////
https://proceeiingg...e... wtw.iiign............e oorp::/000......0..ooeeeeoo-eooe
Stephanie C. Y. Chan, Ishita Dasgupta, Junkyung Kim, Dharshan Kumaran, Andrew K. Lampinen, and Felix Hill. Transformers generalize differently from information stored in context vs in weights. CoRR, abs/2210.05675. 2022b. doi: 10448550/axx1.221005577777 URL. https:
//doi.org/1v 48550/ 2210. 05675
Venkat Chandrasekaran, Benjamin Recht, Pablo A. Parrilo, and Alan S. Willsky. The convex geometry of linear inverse problems. Foundations of Computational Mathematics, 12(6))805- 849, oct 2012. doi: 10. 10110077712208--01-91333.. URL ftttss://dci.oor/11.... 2Fs10208 012 9135-7
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros, Marie Pel- lat, Kevin Robinson, Dasha Valter, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei. Scaling instruction finetuned language models, 2022.
10

---

Published as a conference paper at ICLR 2024
Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, Edouard Grave, Myle Ott, Luke Zettlemoyer, and Veselin Stoyanov. Un- supervised cross-lingual representation learning scale. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault (eds.), Proceedings of the 58th Annual Meeting of the Association for Computational Linguisties, pp. 8440-8451, Online, July 2020. Association for Computational org/
Linguistics. doi:
012653//1/2220.... aa1.000.cl-miinn main.747. URL 2020. main.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, Lei Li, and Zhifang Sui. A survey on in-context learning, 2023.
D.L. Donoho. Compressed sensing. IEEE Transactions on Information Theory, 52(4):1289-1306, 2006. doi: 00.1109/TT:2000.8875555
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can trans- formers learn in-context? a case study of simple function classes, In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems, volume 35, pp. 30583-30598. Curran Associates, Inc., nwwrips...... wttppper.cpppppppppppppp..
2022. URL httpps:/ httte:///proeeeeeesssrrrrrrrrree httppppperr.illeeaapee/2222222 Porpaaaeeeeappeee--oeerrrrrrrae
fi
Micah Goldblum, Mare Finzi, Keefer Rowan, and Andrew Gordon Wilson. The no free lunch theorem, kolmogorov complexity, and the role of inductive biases machine learning. CoRR. abs/2304.05366, 2023, doi: D0448500axXiv.2344.0555 205550/axiiv.304...
URL https:://o// htttp:://doi.org/I... org/10.
Michael Hahn and Navin Goyal. A theory emergent in-context learning as implicit structure induction. CoRR, abs/2303.07971, 2023. doi: 10,48550/rxiv.2233.07999 URL https://
doi.. 2ttp://1...550///xi..223333 07971
Adi Haviv, Oni Ram, Ofir Press, Peter Izsak, and Omer Levy. Transformer language models without positional encodings still learn positional information. In Findings of the Association for Compu- tational Linguistics: EMNLP 2022, pp. 1382-1390, Abu Dhabi, United Arab Emirates, Decem- ber 2022. Association for Computational Linguistics. URL/Itttss//aacaacchll..oo 2022. 1222.fine findings-ennsnnn www.nngsseennl.......mmmmmm..m le..lppp9... 99
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models. 2022.
Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. TabPFN: A transformer that solves small tabular classification problems second. In The Eleventh International Con- net ference on Learning Representations, 2023. URL htttss://operreeee nwt./frrumf forum?
id=cp5PvcI 1688
Matthew D. Homan and Andrew Gelman. The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo. J. Mach. Learn. Res., 15(1):1933-1633, jan 2014. ISSN 1532-4435.
T. Hospedales, A. Antoniou, P. Micaelli, and A. Storkey. Meta-learning in neural networks: A survey. IEEE Transactions on Pottern Analysis and Machine Intelligence, 44(09))5149-5116 sep 2022. ISSN 1939-3539, doi: 10.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA. May 7-9, 2015, Conference Track Proceedings. 2015. URL Rtthtt: / //arxiv..o www.vorg/aaae hwtp:/ba/111e /1al. 6980
Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing, ACM Compus. Surv., -5:9):155 2023. doi: 10.1145/3560815. URL

---

Published as a conference paper at ICLR 2024
O.L. Mangasarian and Benjamin Recht. Probability of unique integer solution to system of linear equations. European Journal of Operational Research, 214(1):27-30, 2011. ISSN 0377-2217. doi: http://bbbooor///110000/////2212 ccm/sccincc///aiaecccc http:///nnn//aataaaaa//////71112
URL http:///ww..eieencccrrrtt.... sciencedirectt
003511
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetalCL; Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chapter of the Asso- ciation for Computational Linguistics: Human Language Technologies, pp. 2791-2809, Seattle, United States, July 2022a. Association for Computational Linguistics. doi: 10.18653/v1/200 naacl- mmacl-main.200 201 URL
org/202222 (tli22.nacc--ann manll-mmain... 201
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? In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp- 11048-11064, Abu Dhabi, United Arab Emirates, December 2022b. Association for Computa- tional Linguistics. URL https: : hatas///olcaathooogg....o org.yy..rg//222 emnlp-main. 759
C Mingard, G Valle-Perez, Skalse, and AA Louis. Is sgd bayesian sampler? well, almost. Journal of Machine Learning Research, 22:1 2021.
Chris Mingard, Henry Rees, Guillermo Valle Pérez, and Ard A. Louis. Do deep neural networks have an inbuilt occam's razor? CoRR, abs/2304.06670, 2023. doi: 10445550/aXXi:223440677...
URL https::://// //to:://ddo.ooo org.i...g//o.
www.oo.........oooooo..oo....... 100 4 1188555//Xiiii www.aawwwwv.......... 06670
Aaron Mueller and Tal Linzen. How to plant trees in language models: Data and architectural effects on the emergence of syntactic inductive biases. 2023.
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Trans- formers can do bayesian inference. In International Conference on Learning Representations, 2022. URL/Bttts::/oppeeee/ee
net/ forum? httee///oorrummgGSSuuuuCCEEE
Kevin P. Murphy. Probabilistic Machine Learning: An introduction. probml. .
MIT Press, 2022. URL
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability, In The Eleventh International Conference on Learn- ing Representations, 2023. /ttpttts::// http:////eeereeiieeemmm net/ htte://rmmmmiippppeeeee-
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and R. Garnett (eds.). Advances in Neural Information Processing Systems, volume 32, Cur- Associates, Inc., 2019. http::/22.......ooeoo2eeeeeeee22
www.eeewwww..n...e...ee..e..oo.e cc/pspsrrppppppppeppp files/
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL
wwwpr//wwwwweeeeeeeeeemmmmmma net forumm saweet/ffoummmdddffme foo
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language cloudfront. pdf 188)9, 2019.
models are unsupervised multitask learners.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernll machines. In J. Platt, D. Koller, Y. Singer, and S. Roweis (eds.), Advances in Neural In- formation Processing Systems, volume 20. Curran Associates, Inc., 2007. URL
pttpp:://proceddingg wew.roccediinng..uri.........sss neurips...eu......e...e. ccc.spp..cc/aae cc/pppr /ttpeeeeeeeee/////0000011111 01 0133006603acc333994ffa8888888883
12

---

Published as a conference paper at ICLR 2024
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of nonnnonyaa bon-bayesian in-context learning for regression, 2023.
Yasaman Razeghi, Robert L. Logan IV au2, Matt Gardner, and Sameer Singh. Impact of pretraining term frequencies few-shot reasoning. 2022.
Nikunj Saunshi, Sadhika Malladi, and Sanjeev Arora. A mathematical exploration of why lan- guage models help solve downstream tasks. In 9th International Conference on Learning Rep- resentations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL htttps://ppeneeiiie
net Det/oorum?dddv IW3sEcls
Leslie N. Smith and Nicholay Topin. Super-convergence: Very fast training of neural networks using large learning rates, 2018.
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer
with rotary position embedding. CoRR, abs/2104,09864 2021. URL http:://axxxi.....////// org/
abs/ abs/2104 22012114....0 09864
Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Sta- fistical Society: Series B (Methodological) 58(1):267-288 1996. doi: https://doi.org/100 1/1/ 1111/ j.21776161. 1b00000........ URL https: //tptss::/ss....e onliinnibbarryy wi wwwey. com/ com/doi/
abs 1t55/00.1111//1 2517..251111 6161 1996 tb00080.x
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL file/ (el.eeeeeeeeeeeeleleleaaaaaaaaaa
files/paperr2011/
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. 2022.
Xinyi Wang, Wanrong Zhu, and William Yang Wang. Large language models implicitly topic models: Explaining and finding good demonstrations for in-context learning. CoRR, abs/2301.11916, 2023. doi: 10448500aaXi...330.1144 11:50//aXiv..33111 11916
URL lhttps://dii. wttp/....org/....
Albert Webson and Ellie Pavlick. Do prompt-based models really understand the meaning of their prompts? In Proceedings of the 2022 Conference of the North American Chapter of the Asso- ciation for Computational Linguistics: Human Language Technologies, pp. 2300-2344, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 11216653/1/222... M
org. org/2022 2022.0222222 nwacl-mainnnnnaaa 167
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Píerric Cistac, Tim Rault, Remi Louf, Morgan Funtowiez, 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 Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference
Empirical Methods in Natu- ral Language Processing: System Demonstrations, pp. 38-45, Online, October 2020. Associ- ation for Computational Linguistics. doi: 10.18653/v 1/tp23//ww/222.oom/m 10.18653/l//000.mmmmpppp. //aclanthology.... wrtpllogy.rr///220
URL https:
demos.6
Sang Michael Xie, Aditi Raghunathan, Percy Liang. and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https://cpenreviie net ott.: http:///nnneeeeeeeeeeeooooooooaa
13

---

Published as a conference paper at ICLR 2024
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in- context, 2023.
14

---

Published as a conference paper at ICLR 2024
CONTENTS
Introduction
Background
Hierarchical
2.2
Model and training details
3
Transformers can in-context learn task mixtures
3.1
Gaussian Mixture Models (GMMs)
Simplicity blas ICL?
Multi-task generalization
6
Deviations from Bayesian inference?
6.1
ICL Translormer first generalizes then memorizes during pretraining
6.2
Gradient Descent tractable approximation of Bayesian inference
Summary of further results
Conclusion
Technical Details
16
A.1 PME Theoretical Details
16
A.2
The curious case positional encodings.
17
A.3
Experimental Setup
17
Linear and Non-linear inverse problems
18
B.
Linear inverse problems
20
B.I.I
Function classes and baselines
20
B.1.2
Results
22
B.2
Non-linear functions
23
B.2.
Fourier Series
23
B.2.2
Random Fourier Features
24
B.2.3
Degree-2 Monomial Basis Regression
27
B.2.4
Haar Wavelet Basis Regression
30
Detailed Experiments for HMICL setup
30
Why HMICL?
30
C.2
Gaussian Mixture Models (GMMs)
32
More complex mixtures)
37
Fourier series mixture detailed results
40
15

---

Published as a conference paper at ICLR 2024
C.4.1 Complexity Biased Pre-training
42
Conditions necessary for multi-task ICL
44
D
Details regarding Multi-task generalization experiments
48
D.I Monomials Multi-task
48
D.2 Fourier Series Multi-task
48
D.3 Details on the phenomenon of forgetting
49
Gradient Descent tractable approximation of Bayesian inference
55
Further Concluding Remarks
55
A TECHNICAL DETAILS
A.1 PME THEORETICAL DETAILS
We mentioned earlier that an ideal LM would learn the pretraining distribution. This happens when using the cross-entropy loss. Since we use the square loss in the ICL training objective, the pre- dictions of the model can be computed using the posterior mean estimator (PME) from Bayesian statistics. For each prompt length i we can compute PME by taking the corresponding summand in the ICL training objective
min Byæ.. E (Ma(P"), f(x++1) mjn E, p pi é (Ma(P") 0
min Ep E 1 (Mo f(x++1)) | [Pi]
EEp: min B, [E (M((P)), f(x;+1)) IP*].
The inner minimization is seen to be achieved by MP = E, [f(#+11 Pi] as we use the squared-error loss. This is the optimal solution for prompt P and what we refer to as PME.
PME for a task mixture. We describe the PME for mixture of function classes. For simplicity we confine ourselves to mixtures of two function classes; extension to more function classes is analogous. Let F1 and F3 be two function classes specified by probability distributions Dyr and Dy., resp. As in the single function class case, the inputs are chosen i.i.d. from a common distribution Dx- For 01,02 E [0, 1] with an + Q2 1, an (0),0')) (2),xx))mixturr F of Fi and F2 is the meta-task in which the prompt P (f1,f(z)) f(z)). Tp. f(zp), x Ip+1) is constructed by first DD,,- Thus p=(f) a1Pr, (f) + a2p5,, (f), where p)(-) is the probability density under function class F which defines Dr. For conciseness in the following we use pi (-) for p. (-) etc. Now recall that PME for function class F is given by
picking task Fi with probability a; for i (1,2) and then picking
etc.
(o,((P)=ETDDA [1(x+1) - P=(1)))(()
Here df is volume element in F; this makes sense as all our function families F are continuously parametrized. We would like compute MP.=(P) in terms of PMEs for Fi and F2. To this end, we first compute
P=(f|P) (P(S)P2())) p(P(f)=F() p(Plf) PF(P)
(aIp((f) + (2)
PF(P) PP(P)
(IPi(P) (P))(P(())(()))
(P((P) (P( + PP(P) (P(
(P))(P)
PP(P) (f(f)))+ Q2P2(P)
PT(P) p=(f)P)
=BiPi(fP)) (FP))
16

---

Published as a conference paper at ICLR 2024
Dense Regression
Sparse Regression ICL
With Position Encodings
With Position Encodings
2.5
Without Position Encodings
Without Position Encodings
2.0
Max Training Lena
10280k
0.6
Max Trieing Leneth
1ozs8k
1.5
0.4
0.2
0.0
0.0
20
60
100
0
400
80
100
(#f in-context examples)
(#f in-context examples)
Figure 5: Impact of positional encodings on length generalization during in-context learning for dense and sparse linear regression tasks. For both tasks, the model was trained with p = 40 i.e. the maximum number of in-context examples provided.
where 81
(PP PF( and B2
Pop (P)
Plugging this in equation
we get
(P) = (i(lIP) ( + 3, pe(l)P) ( B2Mo
(3)
A.2 THE CURIOUS CASE OF POSITIONAL ENCODINGS.
Positional encodings both learnable or sinusoidal in transformer architectures have been shown result in poor length generalization Bhattamishra et al. (2020); Press et al. (2022), i.e. when tested on sequences lengths greater than those seen during training the performance tends drop drasti- cally. In our initial experiments, we observed this issue with length generalization in our in-context- learning setup as well (Figure 5). While there are now alternatives to the originally proposed posi- tion encodings like Rotary Embeddings et al. 2 and ALiBi Press et al. (2022) which perform better on length generalization, we find that something much simpler works surprisingly well in our setup. We found that removing position encodings significantly improved the length generalization both dense and sparse linear regression while maintaining virtually the same performance the training regime can be seen in Figure 5 These observations are in line with Bhattamishra et al. (2220) which shows that decoder-only transformers without positional encodings fare much better recognizing formal languages models without explicit position encodings can still learn positional information. Both works at- tribute this phenomenon the presence of the causal mask in decoder-only models which implicitly provides positional information to these models. Hence by default in all our experiments, unless specified, do not use any positional encodings while training our models.
well as Haviv et al. (0022) that shows transformers language
A.3 EXPERIMENTAL SETUP
We use Adam optimizer Kingma & Ba 2015) train models. We train all of models with curriculum and observe that curriculum helps in faster convergence, i.e.. the same optima can also be achieved by training the model for more training steps as also noted by Garg et 20022). Table states the curriculum used for each experiment, where the syntax followed for each column specify- ing curriculum [starts rt, end, increment, interval) The value of the said attribute goes from start to end, increasing by increment every interval train steps. Our exper- iments were conducted on a system comprising 32 NVIDIA V100 16GB GPUs. The cumulative training time of all models for this project was 30,000 GPU hours. While reporting the results, the error averaged over 1280 prompts and shaded regions denote a 90% confidence interval over 1000 bootstrap trials.
We adapt Garg et al. (2022) code-base for our experiments. We use PytorekPaszke et al. (2019) and Huggingface Woll al. (2020) libraries to implement the model architecture
17

---

Published as a conference paper at ICLR 2024
Table 1: The values of curriculum attributes used for each experiment. Ca. Cp and Crte denote
the curriculum on number of input dimensions (d), number of points (p) and any other experiment- specific attribute respectively. For Fourier Series Cneyy refers to maximum frequency N)
Experiment
Section
Ca
Cp
Cire n/a
Dense, Sparse and Sign-Vector Vector Regression
B.1.1
6, 1, 2000]
[10, 40, 2. 2000
Low-Rank Regression
B.1.1
Fixed (d 100)
Fixed (p 114)
n/a
Fourier Series
B.2.1
Fixed 1)
17. 43. 4, 2000
11 2000
Fourier Series Mixture
Fixed (d 1)
Fixed (p 40)
Fixed (NV 10)
GMM Regression (d 10,p = 10)
3.1.
C22
[5, 10, 20001
[5. 10, 1, 2000]
n/a
GMM Regression (d 10,p = 20)
5.1
[S. 10, 1, 2000]
[TO, 20, 2. 2000]
n/a
Degree-2 Monomial Basis Regression
B
Fixed (d = 20)
Fixed (p = 290)
n/a
Haar Wavelet Basis Regression
B.2.4
Fixed (d 1)
Fixed (p 32)
n/a
Dense Regression ICL
Dense Regression ICL
1.0
1.0
Transformer
(witte w/)
OLS
(wetode wOLL)
0.8
0.8
Ridge (0.01)
mean squared error
10s366
0.6
Bound
0.6
Bound
0.4
0.2
0.2
0.0
0.0
10
20
40
20
A
(# in context examples)
(# in con-omeex examples)
(a)
(b)
Figure 6: Results on the Dense Regression tasks mentioned in section BB1.1
and training procedure. For the baselines against which we compare transformers, we use implementation of OLS, Ridge and Lasso, and for T and L., norm mini-
seikit- rn't 2 mization given the linear constraints we use C
B LINEAR AND NON-LINEAR INVERSE PROBLEMS
Here, we discuss the results omitted from the I.21 for conciseness. Figure 6 shows the results on the Dense Regression task and our experiments corroborate the findings of Akyürek et al. (2022). where transformers not only obtain errors close to OL.S and Ridge regression for the dense regression task (Figure 6a) but the extracted weights also very closely align with weights obtained by the two algorithms (Figure 6b). This does indicate that the model is able to simulate the PME behavior for the dense regression class.
For sparse and sign-vector regression, we also visualize the weights recovered from the transformer for of the functions for each family. As can be observed in Figure g for sparse regression at sufficiently high prompt lengths (k > 10), the model able to recognize the sparse structure of the problem and detect the non-zero elements of the weight vector. Similarly, the recovered weights for sign-vector regression beyond k > 10, start exhibiting the sign-vector nature of the weights (i.e. each component either being +1 or -1).
We evaluate transformers on a family of linear and non-linear regression tasks. On the tasks where it is possible to compute the Bayesian predictor, we study how close the solutions obtained by the transformer and Bayesian predictor are. In this section, we focus only on single task ICL setting
hhans://sskii-aam..mmssssleleiie html https:
18

---

Published as a conference paper at ICLR 2024
Evaluation on Dense Regression Prompts
Transformer (Fox) On Fin Prompts
Trurolom (Pa)
Transformer (Fivn) On Fioa Promots
Tranffrm (Zox)
1.2
(7x)
CLS
1.0
1oss0k
0.8
loss@k
0.6
0.4
0.5
0.2
0.0
0.0
10
20
30
40
10
20
30
40
k
(#f in-context examples)
(# in-context examples)
(a)
(b)
Figure 7: Evaluating transformer model trained for regression on task FoR with w E {2: 2 2 E
{-2, -,, 2), satisfying the convex geometry conditions of Chandrasckaran et 2012).
Left: Comparing the performance of this model, i.e. Transformer (Fzx), with Transformer (Fsv
when both are tested on their respective prompts. Right: Comparing the performance of Trans-
former (F2R) with Transformer (FSvR) on Dense Regression (FDR) prompts. Transformer (F2x)
provides better performance than Transformer (FsvR) on in-distribution prompts but on OOD
prompts (from FDR), Transformer (FsvR) performs better.
trck
15
2
in-contest ecamples)
1.0
69 in contert exmppes))
1.0
0.5
0.5
8
0.0
10
10
0.0
122
0.5
12
14
1.0
-0.5
16
18
1.5
16
18
-20
2
10 Dim
12
14
16
18
8
10 Dim
12
14
16
18
(a)
(b)
Figure 8: Visualizing recovered weights for sparse and sign vector regression for of the examples the test set.
19

---

Published as a conference paper at ICLR 2024
(i.e. the model is trained to predict functions from a single family), while the mixture of tasks is discussed $
B.1 LINEAR INVERSE PROBLEMS
In this section, the class of functions is fixed to the class of linear functions across all problems, i.e. F == {f
x + w² x w RR what varies across the problems is the distribution of w. Problems in this section are instances of linear inverse problems. Linear inverse problems are classic problems arising in diverse applications in engineering, science, and medicine. In these problems, one wants to estimate model parameters from few linear measurements. Often these measurements are expensive and can be fewer in number than the number of parameters (p d). Such seemingly ill-posed problems can still be solved if there are structural constraints satisfied by the parameters. These constraints can take many forms from being sparse to having a low-rank structure. The sparse case was addressed by a famous convex programming approach Candes & Tao (2005): Donoho (2006) also known compressed sensing. This was greatly generalized in later work to apply to many more types of inverse problems; see Chandrasekaran et al. (2012). In this section, will show that transformers can solve many inverse problems context in fact all problems that we tried. The problem-specine structural constraints are encoded in the prior for w.
B.1.1 FUNCTION CLASSES AND BASELINES
Dense Regression (FDR). This represents the simplest case of linear regression as studied in Garg al. (2022); Akyürek al. 20222; yon Oswald et al. (2022). where the prior
is the standard Gaussian í.e. w ~ Wo.,, f). We are particularly interested in the underdetermined region i.e. < d. Gaussian prior enables explicit PME computation: both PME and maximum a posteriori (MAP) solution agree and are equal to the minimum L2-norm solution of the equations forming the training examples, i.e. min || /2//22 s.t. w" a, fa;), Vi Squares (OLS) solvers return the minimum L2-norm solution, and thus PME and MAP too, in the underdetermined region, i.e. < d.
Standard Ordinary Least
Ekewed-Covariance Regression (DS) This setup is similar dense dense-regression, except that we assume the following prior on weight vector: w [O, E), where E E I is the covariance matrix with eigenvalues proportional to 1/7², where E [1, d]. For this prior on w, we can use the same (but more general) argument for dense regression above to obtain the PME and MAP which will be equal and can be obtained by minimizing w " - w w.r.t to the constraints w x; = f(xx). This setup motivated by Garg et al. (2022), where it was used to sample æ, values for out-of- distribution (OOD) evaluation, but not as a prior on w.
be an s-sparse vector in Rd Le. out of its d components only s are non-zero. Following Garg et al. (2022). to sample w for constructing prompts P. we first sample w ~ NO.,,, I) and then randomly set its d components as O. We consider s = throughout our experiments. While computing the PME appears to be intractable here, the MAP solution be estimated using Lasso by assuming Laplacian prior w Tibshirani (199).. We tune the Lasso coefficient following Garg et al. (2022), i.e., by using a separate batch of data (1280 samples) and choose the single value that achieves the smallest loss.
Sparse Regression (FsR). In sparse regression,
assume
Sign- Sign-Vector Vector Regression (FsvR). Here, we assume w to be a sign vector in [11,+1]d. For con- structing prompts P, sample independent Bernoulli random variables by with mean of 0.5 and obtain = [2b, - ..... 2b 1]. While computing the exact PME remains intractable in this case as well, the optimal solution for A > d/2 can be obtained by minimizing the -nom
[ww] oo w.r.t. the constraints specified by the input-output examples (wo æ; = f(x:)) Mangasarian & Recht (201T).
Low-Rank Regression (FLgwRank-DR). In this case, is assumed to be flattened version of matrix W E Rixa (d = q²) with a rank r, where r & q. A strong baseline, in this case, is to minimize the nuclear norm L., of W, i.e. || W ||. subject to constraints w x; f(x:). To sample
20

---

Published as a conference paper at ICLR 2024
Stoart
Neew
Loa bark Peguraoo
Trandoon
0LS
Moli is
0:00
Cf inootaat compla!
d-cotaxt sampat
(< Incotan somp at
(c)
Loe Fan, Regrssca 121
(amntr +) (urnts
Lome of
18 (dinoweed maapli))
[A (Ale mample))
insedde mempii))
munele))
(c)
g
(h)
Figure 9: Comparing ICL in transformers for different linear functions with the relevant baselines. Top: lossek k values for transformers and baselines on skewed covariance, sparse, sign-vector, and low-rank regression tasks. Bottom: Comparing the errors between the implicit weights recovered from transformers waube with the ground truth weights w and weights computed by different base- lines. PME: denotes the weights obtained by minimizing w -- -ww for the skewed covariance regression task.
Low- Rank Regression ICL
Sign- Vector Regression ICL
2.0
Transformer
1.50
Transformer OLS
OLS
1.25-
NUTS
NUTS
1.5
1.00
loms8k
lose&k
0.75
0.50
0.5
0.25
0.0
0.00
10
15
20
0
A:
in-context examples)
(# in-context examples)
(a)
(b)
Figure 10: Computing the PME by Markov Chain Monte Carlo sampling method (NUTS) for (a) Low-Rank Regression, and (b) Sign-Vess Vector Regression. The problem dimension d for Low-Rank Regression and Sign-Vector Vector Regression is 16 (4 x 4 matrix) and 8 respectively. As can be seen, the Transformer is close to the respective JUTS-approximatte PME in both cases.
the rank-r matrix W. we sample A .N(0 1), s.t. A E Raxr and independently matrix B of the same shape and distribution, and set W ABT..
An artificial task. Techniques in prior work such as Chandrasekaran et al. (012)) require that for the exact recovery of vector w, the set of all these vectors must satisfy specific convexity conditions. However this requirement seems to be specific to these techniques, and in particular it's not clear if Bayesian approach would need such condition. To test this define a task FzR where the convexity conditions are not met and train transformers for regression on this task. Here, w E {2z x E (-2, 1, 1, 2) d/2), ). where zz denotes z concatenated to itself. Note that the size of this set is 2d, the same as the size of (-1, 13d, and many elements, such as zz with 2 E 1///1 lie strictly inside the convex hull.
21

---

Published as a conference paper at ICLR 2024
15
Error
10
5
2
4 6
8
10 11
12 Num Prompts
13
14
15
16
17
18
19
20
21
Figure 11: Box plot showing the effect the order of prompts on the Transformer's ICL perfor- mance for Dense Regression. On x-axis plot the number of in-context examples and on y-axis we have the quartiles for the errors obtained for different permutations of the examples. At each prompt length consider 20 permutations. As be seen, all the boxes nearly flat, meaning that the errors have nearly zero variance with the permutations, and hence the model in this case is robust the order of prompts.
Recovery bounds. For each function class above, there is a bound on the minimum number of in-context examples needed for the exact recovery of the solution vector w. The bounds for sparse, sign-vecton and low-rank regression are 2s log( d 5s/4, d/2, and 3r(2q -r) respectively Chan drasekaran et al. (0012..
B.1.2 RESULTS
We train transformer-basee models the five tasks following 22. Each model is trained with d = 20 and = 40, excluding Low-Rank Regression where we train with d = 100, = 114, and r Figures 9b compare the loss@k values these tasks with different baselines. Additionally, we also extract the implied weights wrobb from the trained models when given prompt following Akyürek et al. (2022) by generating model's predictions {uf} on the test inputs {=/)²== Dx and then solving the system of equations to recover probe, We then compare the implied weights probe with the ground truth weights well the weights extracted from different baselines to better understand the inductive biases exhibited by these models during in- context learning (Figures 91 9h).
Comparison with exact PME. Since results for dense regression have been already covered in Akyürek al. (2022), we do not repeat them here, but for completeness provide them Figure regression, we observe that the transformer follows the PME solution very
For skewed-covariane closely both in terms of the loss @k values (Figure 9ap as well the recovered weights for which the error between probe and w PME PME-Skew (weights obtained by minimizing w" Σ --'w) is close to zero at all prompt lengths (Figure De).
Comparison with numerical solutions. For Low-Rank Regression and Sign- Vector Regression, we provide comparisons with the numerical PME solutions obtained using No-U-Turn Sampling (NUTS) in Figure 10 and find that emror by transformers strongly agree with those of the numerical solution, and in some instances transformers actually perform slightly better (which we attribute to the numerical solutions also being an approximation).
Comparison with strong baselines from Chandrasekaran et al. (012) As can be seen in Figure 9 on all the tasks, transformers perform better than OLS and are able to solve the problem with samples i.e. underdetermined region meaning that they are able to understand the structure of the problem. The error curves of transformers for the tasks align closely with the errors of Lasso (Figure 9b), Loo minimization (Figure 9c). and L. mininization (Figure 9d baselines for the respective tasks. Interestingly for low-rank regression transformer actually performs better. Though, due to the larger problem dimension, (d 100) in this. it requires bigger model: 24 layers, 16 heads,
and 512 hidden size. In Figures 9f 9g and 9h we observe that at small prompt lengths poobe and
OLS are close. We conjecture that this might be attributed to both probe and OLS being close
22

---

Published as a conference paper at ICLR 2024
for small prompt lengths (Figure 8). Prior distributions for all three tasks hence, at small prompt lengths when the posterior is likely to be close to the prior, the PME is close to the mean of the prior which is O. At larger prompt lengths transformers start agree with w Lasso w wr and wh. This is consistent with the transformer following PME, assuming w Laseo, w whos, and w close to PME- we leave it to future work determine whether this is true (note that for sparse regression Lasso approximates the MAP estimate which should approach the PME solution as more data is observed). The recovered weights wrobe also agree with Lissoo w foo and wt. for their respective tasks after sufficient in-context examples are provided.
Finally, refer to Figure 7 for the results on task FZR- We observe that Transformers trained on this task (F2R) provide better performance than those trained on Sign-Vector Vector Regression (FsvR). Therefore, we can conclude that Transformers do not require any convexity conditions on weight vectors.
B.2 NON-LINEAR FUNCTIONS
Moving beyond linear functions, we now study how well transformers can in-context learn fune- tion classes with more complex relationships between the input and output, and if their behavior resembles the ideal learner i.e. the PME. Particularly, we consider the function classes of the form Fg = (J(:; (=; +) w (zz). t E RAJ. where + IRd + RA maps the input vector æ to an alternate feature representation. This corresponds to learning the mapping ((x) and then per- forming linear regression on top of it. Under the assumption of a standard Gaussian prior on w, the PME for the dense regression can be easily extended for Fa: min.. || [|,, s.t. w wT (T(;)) f(x,) for E (1,... .P).
B.2.1 FOURIER SERIES
A Fourier series is an expansion of periodic function into a sum of trigonometric functions. One can represent the Fourier series using the sine-cosine form given by:
N
f(x) ao + Σ an cos LtTx/L) Σ b, sin (nzx/L)
where, x E I. and an. an's and hn's are known as Fourier coefficients and cos BonnT/L. T. and sin nt/l. define the frequency components. We can define the func- tion class Flourier by considering & as the Fourier feature map i.c. 2x(z) [1, cos (Tar L),... ,os cos (Nxx/L), L), sin (xx/L),... sin (Nxx/L)) L)]T" and as Fourier coefficients: w = [a00 a1,'.' ,a, b1, "b,y" Hence, (x(x) E Rd and w E Rd. where d = 2N +
fix value of N and sample functions E F by sampling the Fourier coefficients from the standard normal distribution i.e. w (OO,, I). We consider the inputs to be scalars, i.c. I4 E [-L, L] and we sample them i.idd from the uniform distribution on the domain: ~ ((-L, L). In all of our experiments, we consider N 10 and L also prompt the model with functions with different maximum frequency as seen during training. As a baseline, we use OLS on the Fourier features (denoted as OLS Fourier Basis) which will be equivalent to the PME.
For training transformers to in-comtext-learn Flourier
At test time we evaluate on Fiounier for M E [0], i.e. during evaluation we
Measuring inductive biases. Once we train transformer-based model to in-context Iearn Fouuier how can we investigate the inductive biases that the model learns to solve the problem? "we would like to answer questions such as, when prompted with input-output examples what are the prominent frequencies in the function simulated by the model, or, how do these exhibited frequencies change as we change the value of k?? We start by sampling in-context examples (f11) dis I' (æx)). and given the context obtain the model's predictions on a set of m test in- puts {-]}= ". i.c. 8/ = M, ((z1, f(z1), (- ... Zx. (xx), x.)). We can then perform Discrete Fourier Transform (DFT) on {v1- ... Vnn to obtain the Fourier coefficients of the function output by M, which we can analyze to understand the dominant frequencies.
The results of our experiments concerning the Fourier series are provided in Figure 12 Transformers obtain 1 loss@k k values close the OLSS Fourier Basis baseline (Figure T2a indicating
Results.
23

---

Published as a conference paper at ICLR 2024
Function Predicted M 1D
080
0.00
(a)
(b)
Inductive Biases M
4
Transformrr k
01S
Transformer
01. Fourier Rasis
h- 21
75
0.2
Ogr 0.00
2
0.50
0.L
% 0.25
0.3
0.0
0.00
0.20
1 171.5 6 19100112
5
1011121
101112
123456 7 201122
Figure 12: Effectiveness of ICL. in transformers for Fourier series family of functions. Top left: lcss@k values for transformer and OLS Fourier Basis baseline. Top Right: Visualizing the func- tions simulated by the transformer and the OLS Fourier Basis. Bottom: Measuring the frequencies of the simulated function by the transformer and the baseline.
at least for the smaller prompt lengths the model is able simulate the behavior of the ideal predictor (PME). These plots use 12-layer transformers to obtain results, but we also investigate if bigger models help. Figure 13 plots bigger models with 18 and 21 layers where the agreement with PME is much better. Further, in Figures 12b and [2c we could only discuss results for subset of values of M and ks The function visualizations for the transformer and Fourier OLS baseline for different combinations of M and i are provided in Figure 15] We have observations consistent with Figure 12b] where the function outputs of the transformer and the baseline align closely. Similarly, in Figure [4] we present the distribution of frequencies in the predicted functions for the two methods and again observe consistent findings.
Since the inputs T4, in this case, are scalars, we can visualize the functions learned in context by transformers. We show one such example for a randomly selected function f ~ Ffounee for prompt- ing the model in Figure T2b As can be observed, the functions predicted by both the transformer and baseline have a close alignment, and both approach the ground truth function f as more examples provided. Finally, we visualize the distribution of the frequencies for the predicted functions in Figure 12c For a value of M, we sample 10 different functions and provide k in-context examples to the model to extract the frequencies of the predicted functions using the DFT method. As can be observed, when provided with fewer in-context examples (k = 2) both Transformer and the baseline predict functions with all the 10 frequencies (indicated by the values of a² + b2 in a similar range for E [L. 10)) but as more examples are provided they begin to recognize the gold maximum frequency (i.e. M- 4). The function visualizations for the transformer and Fourier OLS baseline for different combinations of M and & are provided in Figure 15] We have observations consistent with Figure 12b where the function outputs of the transformer and the baseline align closely. Sim- ilarly, in Figure 14 we present the distribution of frequencies in the predicted functions for the two methods and again observe consistent findings. This suggests that the transformers are following the Bayesian predictor and are not biased towards smaller frequencies.
B.2.2 RANDOM FOURIER FEATURES
Mapping input data to random low -iw-dimeniionl features has been shown to be effective to approx- imate lcrge-scale kernels Rahimi & Recht (2007). In this section, are particularly interested in Random Fourier Features (RFF) which can be shown to approximate the Radial Basis Function kernel and are given as:
24

---

Published as a conference paper at ICLR 2024
Fourier ICL
M = 10
0.12
TT: 04 31, 513)
0.10
OLS Fouior Budis
lose0k
0.08
0.06
0.04
0.02
10
30
(#f in-context examples)
Figure 13: Bigger models achieve better results on the Fourier Series task. Plotting the squared error (averaged over 1280 prompts) for bigger transformer (TF) models trained for 500k steps on the Fourier Series task. Training setup is the same as used for the model plotted in Figure 12a (Section B.2.1). which is also plotted here (blue color) for comparison. L and E denote the number of layers and embedding size for TF models respectively.
Figure 14: Measuring the frequencies of the simulated function by the transformer and the baseline for different values of M (maximum frequeney) and (number of in-context examples)
25

---

Published as a conference paper at ICLR 2024
11
10
10
M
Ma
M
Figure 15: Visualizing the functions simulated by the transformer and the OLS Fourier Basis, for different values of (maximum frequency) and & (number of in-context examples)
26

---

Published as a conference paper at ICLR 2024
de 1 10
Je Dwc
4-
Tankaon 191-015
1.0
Toadken 01-015
191-015
65
Hoad
Boand
8-4
[6 Excomrent esarglis)
(A ixcerien cxargles)
(A Iscareen csarpho)
(a)
(b)
(c)
D 100
d- 10, D
d 32, D
0.5
031
(9 invorted cxampho)
(+ incammer courgfo))
(6 (kinccmen curppo))
(f)
Figure 16: Comparing transformers performance RFF function family (F (FBP) with the RFF-OLS baseline for different values and D.
ΦD(x) = coos (wf x + 81), ... coss (WDE + óD))²
where w, E Rdd and & E RRIEE D). such that DD Rd + RD, Both w; and 8 are sampled randomly, such that w, E (N and 8, E (0,2 2T). We can then define the function family JRFF as linear functions over the random fourier features i.c. f = w" " (zx such that f ~ FRFF While training the transformer on this function class, we sample w; 's and 6 once and keep them fixed throughout the training. As a baseline, we use OLS over (x), y) pairs which will give the PME for the problem (denote this RIF-OLS).
4p
Results. For this particular family, we observed mixed results for transformers, i.e. they fail to generalize to functions of the family when the complexity of the problem is high. The complexity of this function class is dictated by the length of the w; vectors (and the inputs x) i.e. d and the number of random features D. We plot the lossêk values for transformer models trained on JRH for different values of and D in Figure [6] As can be observed, the complexity of the problem for the transformers is primarily governed by d, where they are able to solve the tasks for even large values of D, however, while they perform well for smaller values of d (d = 1 and d = 4), for d = 10, they perform much worse compared to the RFF-OLS baseline and the loss @ doesn't 't improve much once 15 in-context examples are provided.
B.2.3 DEGREE-2 MONOMIAL. BASIS REGRESSION
As stated in I.2.1 the Fourier Series function class can be viewed as linear regression over the Fourier basis consisting of sinusoidal functions. Similarly, we define a function class F with the basis formed by degree-2 monomials for any d-dimensional input vector x.
Using the notation introduced in B.1.1 the basis for F#x is defined as P {2ąx 11 i < d). Each function / ox is linear combination of basis and w i.e. f(x) = w D (x).
where is | -dimensional vector sampled from standard normal distribution.
For experimentation, we define sub-family FS Fm under F.moo(2) by choosing a proper subset C 2M and linearly combining the terms in to form f. This is equivalent to explicitly setting coefficients w; of terms in &x - S to O. We experiment with d 20, with the prompt length
27

---

Published as a conference paper at ICLR 2024
p 290 and |S] 20. We do not use curriculum (d, p. are fixed for the entire duration of the training run).
Baselines. We OLS fitted to the following bases as baselines: S basis (OLSs), all degree- 2 monomials i.c., LM basis (OLS and to a basis of all polynomial features up to degree-2 (OLSpuyy:2). We also compare Lasso (a 0.01) fitted to all degree-2 monomials i.e., Pas basis (Lassop as baseline.
1.50
Transformer
OLSe,
OLSs
1.00
OLSpoly
loss@k
Lassoe
0.75
0.50
0.25
0.00
50
100
150
200
250
300
k
(# in-context examples)
Figure 17: In-Distribution evaluation results on F=mom(2) sub-family degree-2 monomial basis regression. Evaluation of transformer prompts generated using the same S used during training.
Results. In Figure 17 we show the In-Distribution (ID) evaluation results for the F. experi- ments. Here, the test prompts contain functions formed by S (the same basis used during training)- We observe that Transformers closely follow OLSs. The increasing order of performance (decreas- ing loss@. @ for & 2 (S) of different solvers is: OLSpoly.2) S OLS < Lasso < Transformers < OLSs. Transformer's squared error takes little longer than OLSs to converge. Lasso4 is able to take the advantage of sparsity of the problem and is hence better than both OLS@ and OLSpoly.2). which respectively converge at & 210 and k 23 We also conduct an Out-of- Distribution (OOD) evaluation for Fmoon(2) whose results are shown in Figure 18| Here, we generate prompts from basis S' C 4M of the same size as S but differing from S in degree-2 terms, i.e. SS' S| = n. We show the results for different values n. Figure [8a shows the OLSg undergoes steep rise in error region Figure 18a where we notice that Transformer mimics OLSs, while OLSg is the best-performing baseline (since it fits the S' basis used construct the prompts). Transformer does not undergo double descent (for n = 1) and hence momentarily better than OLSs at = (S]. Similar plots are shown for n {2, 4, 5, 10, 15, 20). As n increases, the height of OLSs peak increases and the Transformer also starts to have a rise in errors at k = |S]. For
momentarily at k [S] (double descent). Figure 18b zooms into the lower
20, S and S have nothing in common, and Transformer still follows OLSg (OLS fitted the training basis S). As mentioned under ŞB.2, when the prior on weights w is Gaussian, the PME is the minimum T---norm solution. For F that solution is given by OLSs. Therefore, the results suggest that the transformer computing PME. In summary, transformers closely follow OLSg in this set-up. and more so on the OOD data, where they even surpass OLSs's performance when it experiences double descent.
4210 and 231 are the sizes of the bases to which OLS and OL fitted. Hence, they converge right when the problem hecomes determined in their respective hases.
28

---

Published as a conference paper at ICLR 2024
-
10-8
8--8
Tradarn 0.36
0..00 0Cdy
0.30 0ady
ocdy
303
Seeartt
(et ioveetl cample)
Of m-carted cumples)
(9 incorttdd mumpll))
(a)
(b)
0C3y
0.A6
(6 la socgle))
(6 moceple))
(d)
0.S5 0C.y
0.S. OC.y
Loanet
Leauet
mumplla))
mamdlo))
sumple))
Figure 18: Outto--Distribution evaluation results on F F basis regression. Evaluation of transformer trained on prompts generated using S', where S' con- tains n degree-2 monomials not present in S that was used during training. We show results for different values of n.
sub-family of degree-2 monomial
1.0
Transformer
OLSh
0.8
1oss@k
0.6
0.4
0.2
0
10
15
20
25
30
(# in-context examples)
Figure 19: Evaluating Transformer trained on Haar Wavelet Basis Regression task (F
29

---

Published as a conference paper at ICLR 2024
B.2.4 HAAR WAVELET BASIS REGRESSION
Similar to Fourier Series and 2 Monomial Basis Regression, we also define another non- linear regression function family (Fit"") using a different basis, $ PH, called the Haar wavelet basis.
& h is defined on the interval [0, 1] and is given by:
(z) f E + : (O) ())(())
22/2 - k), E
|0,1)
(xx) -1
2/ ,
0
otherwise,
where 1 is the constant function which is 1 everywhere on [0. 1l. To define f. we sample w from No,, 1) and compute its dot product with the basis, i.. w (z We construct the prompt P by evaluating f at different values of x ((0, 1). The Transformer model is then trained on these prompts P.
We use d = 1 and p = 32, both of which are fixed throughout the training run, i.e. we do not use curriculum. We only consider the basis terms corresponding to n 1, 2, 3). The baseline used is OLS on Haar Wavelet Basis features (OLSH). Note that for the model used throughout the paper 2.2). at k 32 the lossêk value is 0.18, while for bigger model and OLSH it is 0.07. Therefore, for this task we report the results for the bigger model which has 24 layers, 16 heads and 512 hidden size.
Results. In Figure 19 we observe that Transformer very closely mimics the errors of OLSH (i.c. OLS fitted the Haar Wavelet Basis) and converged OLSH at k 32. Since the prior on the weights w is Gaussian, OLSH is the PME. Hence, Transformer's performance on this task also suggests that simulating PME.
C DETAILED EXPERIMENTS FOR HMICL SETUP
C.I WHY HMICL?
The distinction between MICL and HMICL is reminiscent of the distinction between the usual supervised learning and meta-learning. Consider linear regression as an example the MICL setup, where all tasks are instances of linear regression and each weight vector defines task. Further, consider a mixture of two "meta-tasks" or "function classes", say linear regression and decision trees. It is true that in an abstract sense, one could potentially consider each task to be defined by its parameters and thus ignore which type of "meta-task" it instantiates (linear regression or decision tree). Therefore, under this interpretation, MICL is equivalent to HMICL
However, this view is too poarse-grained for our purposes. What is of interest, both from the appli- cation and theory perspectives, is more fine-grained view about whether and how models learn to perform these different "meta-tasks". The hierarchical structure is central to this discussion. Previ- ous work, e.g. Garg et al. (2022); Zhang et al.] 2023). considers MICL with only single "meta- task" and is thus not suitable for the type of analyses we perform. Compared to MICL, HMICL arguably closer to the ICL for LLMs which can perform a vast variety of "meta-tasks." Is the training data of real-world LLMs hierarchical? Due to the complex nature of real-world training data distributions, it is hard to find concrete evidence of them being hierarchical, but we believe multi-task training in LMs like T5 Raffel et al. (2023) and FLAN-TS Chung et al.) (2022) (with caveat for the latter being true for fine-tuning and not pre-training) or training of multilingual models Conneau et al. (2020) which involve pre-training corpora in different languages are some examples of the hierarchical nature of the training distribution weal-world LMs. To sum up. HMICL allows for better terminology for investigating
models, with potential for being related to real-worth LLMs more closely than the vanilla MICL setting. Hence, we make the distinction between the two and treat them as separate settings in work.
30

---

Published as a conference paper at ICLR 2024
1st dim. of tioe On Ti prompts
1st dim. of zu probe
On Ta prompts
( in-context examples les)
8
20
9
10
31
12
11
10
-11
28
28
-1
33
10
10
Samples of
Samples of
dim. Pploe
special
prompt
ss dim.
PME (GMM) special prompt
0025
0022
0.02 0.244
AP
432
114
3:
0.284
822
l: (# in-context examples)
2
31T
80
79
1.
332
0:44
0.953
3
4
RSM
2
8:07
5
3.28
314
4.14
0.013
5
0111
7
112
8
224
8.7
9
3318
10
10
-3
4
7
10
1
2
4
7 8
10
Samples of
Samples of
Figure 20: Transformers simulate PME when trained on dense regression task-mixture (d
10, p = 10, 01 = Q2 = 1/) with weights having a mixture Gaussian prior (GMM), (top): 1st dimension of Transformer's probed weights across the prompt length. (bottom): 1 dimension of Transformer's probed weights and PME (GMM) across the prompt length for specially constructed prompt.
31

---

Published as a conference paper at ICLR 2024
Evaluation prompts (w Na(mn, En)
Evaluation on Ta prompts (w NNlau,, Ez))
Transformer (GMM)
5
5
PME (7i)
PME (7a)
4
4
PME (GMM)
loss@k
los8@k
OLS
3-
3-
2
2
0
10
15
20
0
10
15
20
k
(#) in-context examples)
(# in-context examples)
Figure 21: Transformers simulate PME when trained on dense regression task-mixture (d 10, p = 10, a1 = a2 = with weights having a mixture of Gaussian prior (GMM). Comparing the performance of the Transformer, PMEs, and OLS in under- and over-determined regions. For all context lengths, the transformer follows PME(GMM) and is far from OLS in the under-determined region.
C.2 GAUSSIAN MIXTURE MODELS (GMMS)
Here we discuss some details regarding $3.1 and more results on GMMs. We start with a description of how we calculate PMEs for this setup.
Computation of PMEs. As mentioned in A and B.2, we can compute the individual PMEs for components T3 and T2 by minimizing the L2 distance between the hyperplane induced by the prompt constraints and the mean of the Gaussian distribution. In particular, compute PME for each Gaussian component of the prior, we solve system of linear equations defined by the prompt constraints (w x; yi, E {1, 2, .., p)) conjunction with an additional constraint for first coordinate, i.e. (w)) +3 (for (a (m, En) or w -3 (for. Na E2)). Given these individual PMEs, we calculate the PME of the mixture using Eq.
Now discuss
results for GMMs. First, see the evolution of I's (from Eq.
3 PME Next, we
(GMM), and Transformer's s probed weights across the prompt length (Figures 22 and 23
the results for the Transformer models trained on the mixture with unequal weights, i.e. 01 A 02 (Figure 24 and for the 20 model (Figure 25)
Agreement of weights between Transformer and PME(GMM). Figure 20 (top) shows the
lution of the first dimension of the Transformer weights, i.e. Probe, )1 with prompt length k. We that Transformer is simulating PME (GMM), which approaches PME (T (ppompps) with increasing prompt length (k). Note that regardless of k., the first dimension of PME (T,) is (J4)1. the first dimension of the mean of the prior distribution 73 since the Gaussian has a fixed value in the first dimension. Note that PмЕ (GMM) approaches PME (I prompt)) with increasing k (Eq. 3). Also note that our setting. regardless of the first dimension of PME (T,) (pt)1, the first dimension of the mean of the prior distribution T, since T; has fixed value (i.e. zero variance) in the first dimension. Hence, if Transformer simulating PME (GMM), the first dimension of Transformer's weights ( (wob)))), must approach (HT)1 (when T prompt Ti) and (J22) (when T prrompe T2). This is exactly what we observe as (wobe))) ) approaches +3 and -3 on Ti and T2 prompts respec- tively. At prompt length 0, in the absence of any information about the prompt, (wope) ) Rs This agrees with Eq. since 0 (11)1.3) .3, + (122)..32 where (μ1)) = +3, (122) -
0i = 0.5 and B2 = 02 = 0.5 when prompt P is empty. The figure shows that with the increasing evidence from the prompt, the transformer shifts its weights Tprompt' 's weights evidenced by the first coordinate changing from 0 to +3 or -3 based on the prompt. In (
we check the behavior of Transformer and PME (GMM) specially constructed prompts P. where (æ:)1
32

---

Published as a conference paper at ICLR 2024
On prompts
On 73
prompts
1.0
6:47
9:44
CR
[4 in-context osamples)
2
0.2
10
0.0
2
3
10
2
8
9
10
Samples
oz
Samples
(a)
On Ti
prompts
Bg
On T3
prompts
1.0
J0W
053
8: (4 in-context examples
6
0.2
0.0
2
9
10
2
8
10
Samples oz
Samples m
(b)
dim.
PME (GMM) On Ti prompts
dim.
PME (GMM) On T3 prompts
& (# in context examples
6
10
10
2 3
9
10
1
2
3
8
9 10
Samples W
Samples 10
(c)
Figure 22:
Evolution (as heatmaps) with prompt length (k) of H's and PME (GMM) appearing in Eq. 3 for the model trained with d 10, p 10, ag a2 J.2 We show 10 different samples of w for each plot.
33

---

Published as a conference paper at ICLR 2024
On Ti prompts
On T. prompts
1.0
&
0.8
0.8
Betas
0.6
Betas
0.6
0.0
0.0
4
10
2
10
(# in-context examples)
(# in-context examples)
(a)
On Ti prompts
On Ta prompts
3-
3-
wilmbe
1s dim. of weight vectors
1 dim. of weight vectors
2
PME (GMM)
0-
0-
33
-3
2
10
0
(# in-context examples)
(# in-context examples)
Figure 23: Evolution (as line plots) with prompt length (k) of B's, PME (GMM), and Probe for the model trained with d 10, = 10, aq 02 = We show the values averaged over 1280 samples.
34

---

Published as a conference paper at ICLR 2024
adtin medt Alay fill
Edtatioo paople Nat
(c)
PAe PME 1306
Sunded
Sandde
(b)
(d)
Figure 24: Transformers simulate PME when frained on dense regression task-mixture (d 10, = a0 = 3. 02 = 3) with weights having mixture of Gaussian prior (GMM). (a): Comparing the performance of the Transformer with Posterior Mean Estimator (PME) of individual Gaussian components (PME (Ti) and PME (T2)) and of the mixture PME (GMM). (b): MSE be- tween the probed weights of the Transformer and PMEs. (c): 1st dimension of Transformer's probed weights across the prompt length. (d): 1st dimension of Transformer's s probed weights and PME (GMM) across the prompt length for specially constructed prompt.
Exalution en pursti Kil
Enlation er 3) perptt (er Kxll
Enalastion ce 2, perpla (s Aii 33
Eallatic on parpp [o Wine Sal
Tranforrir (GHM)
(CMM)
PME
2013)
PAT (6007)
PME 2013
(9) invonttrnt ccampio))
(+) In-ceetod carpla)
(8 Incoettot cample)
(e Incontee: cargla)
(a)
(b)
Figure 25: Transformers simulate PME when trained dense regression task-mixture (d 10, 20, a1 = a2 } with weights having a mixture of Gaussian prior (GMM), Left: Comparing the performance of the Transformer with Posterior Mean Estimator (PME) of individual Gaussian components (PME (T;) and PME (T2) and of the mixture PME (GMM). Right: MSE between the probed weights of the Transformer and PMEs.
35

---

Published as a conference paper at ICLR 2024
and (zx)2:d (0, 1), Vi E {1,... .p). For our setup, choosing such æ;'s is guarantees that no in- formation about the distribution of w becomes known by observing P (since the only distinguishing dimension between Ti and T2 is the 1st dimension and that does not influence the prompt in this case as (æ:)) = 0). We note that Transformer's weights are all R 0 regardless of the prompt length, agreeing with the PME (GMM). Observing more examples from the prompt does not reveal any information about the underlying distribution of w in this case. Moreover, in Figure 21 we plot the errors for the Transformer model, PMEs, and OLS for the over-determined region. In the over- determined case (d > 10), the solution is unique, hence all the predictors including Transformer give the same solution and have errors R O. Also, as shown, Transformer's errors are smaller than OLS errors and agree with the PME (GMM) errors for all context lengths. This shows that Trans- former is indeed simulating the mixture PME and not OLS. All of this evidence strongly supports our hypothesis that Transformer behaves like the ideal learner and computes the Posterior Mean Estimate (PME).
Evolution of f''s, PME (GMM), and wrobbe Figure 22 plots the evolution of 8 and 1st dimension of PME (GMM) for 10 different w's. The 3"s (Figures 22a and (2b) are 0.5 (equal to a's) at = (when information is observed from the prompt). Gradually, as more examples are observed from the prompt, 31 promps approaches 1, while Byareer approaches 0. This is responsible for PME (GMM) converging to PME (T promptt) as seen in $3.1 The 1st dimension of PME (GMM) (Figure 22cb starts at 0 and converges to +3 -3 depending on whether T prompt is Ti or T2- Figure 23 shows the same evolution in the form of line plots where we see the average across 1280 samples of w. In Figure 23a. ST @ppromps approaches 1, while &Tetter approaches 0 as noted earlier. Consequently, in Figure 23b 1 dimension of PME (GMM) approaches +-3 or -3 based on the prompt. The 1st dimension of Transformer's ' probed weights, i.c. (worrb)) almost exactly mimics PME (GMM).
Unequal weight mixture with a1 2 & ag = 3. Figure 24 shows the results for another model
where a's are unequal (d 10. p- = 10, a. = 3.02 z) The observations made for Figure in $3.1 still hold true, with some notable aspects: (1) The difference between prediction errors, i.e. lass@k l (2a)), of PME (GMM) and PME (Ti) smaller than that of the uniform mixture (a) = a2 () case, while the difference between prediction errors and weights of PME (GMM) and PME (f2) is larger. This is because, at prompt length = 0, PME (GMM) is a weighted combination of component PMEs with a's coefficients (Eq. 3 Since 01 > 022 PME (GMM) starts as being closer to Tj than Tz. Also, since the Transformer follows PME (GMM) throughout, its prediction errors also have similar differences (as PME (GMM)'s) s) with PMEs of both components Ti and T2. (2) Transformer's s probed weights (wobe,,,, which used to have the same MSE with PME (T₂) and PME (T2) k O, now give smaller MSE with PME (T₂) than PME (T2) on prompts from both Ti and T2 (Figure 24b). This is a consequence of PME (GMM) starting out as being closer Ti than T3 due unequal mixture weights as discussed above. Since Transformer is simulating PME (GMM), probe is also closer to PME (Ti) than PME (72) at k 0 regardless of which component (T₁ or T₂) the prompts come from. Due to probe mimicking Tj more than T2 we also observe in Figure 24b that Probe gives smaller MSE with w (ground truth) when Tprompt = Ti compared to when T pprompt = =T2. (3) The 1st dimension of Transformer's weights ((ob)))))) ))) and PME (GMM) is instead of 0 when the prompt is either empty (24c) or lacks
information regarding the distribution of 24db. It happens because (wobe))), ] x T dimension
of PME (GMM) = (31)1.31 + (32211.32 = (+3) --3)(5)) = 1. Note that 31 = a1 = 3 and
32 = 02 3 when prompt P is empty at k = 0 (Eq. 3 When P is inconclusive of w. 31 ai
and 32 02 Vk (1,2,-...,2) .p).
Transformer model trained with longer prompt length (p = 20). Figure 25 depicts similar evidence as Figure of Transformer simulating PME (GMM) for a model trained with d 10. p = 20, a1 = Q2 = 1 We see that all the observations discussed in 1 also hold true for this model. Transformer converges to PME (GMM) and PME proppe)) w.r.t. both lossek (Figure 25a and weights (Figure 25b at k- = 10 and keeps following them for larger i as well.
In summary, all the evidence strongly suggests that Transformer performs Bayesian Inference and computes PME corresponding to the task at hand. If the task is a mixture, Transformer simulates the PME of the task mixture as given by 3
36

---

Published as a conference paper at ICLR 2024
Evoluation on Denze Regression Prompts
Evaluation on Sporse Regression Prompts
1:25
Trando ner (Poona
1.00
1,00
Toudsi (Fix)
[Pix]
loss0k
0.70
0.70
Laue
Sound
0.50
Hound
0.25
0.35
10
20
(# no-otte examples)
Figure 26: Comparing the performance of Transformer model trained on dense and sparse re- gression mixture F(DR. SR} with baselines, as well as single task models, trained on FDR and FSR individually.
Fualistion on Donse Regression Prompts
Evalustion Sparse Regression Prompts
(winnde
(se Joube (DEESE: SI
Jedar
[BB,58y Min
men aqured erter
0.15
(urinde
mecon squared emay
0.30
0.32
(st Imilee (DA,RR
Poile (sc [DRSKR
Bound
0.10
Ungand
0.06
0.00
0.00
0.00
0.00
10
201
(# in contert examples)
(+ incconteat ecample))
(a)
(b)
Figure 27: Comparing the errors between the weights recovered from the mixture model trained on F(DR. SR} mixture and different single task models and baselines while evaluating on FDR and FsR prompts
C.3 MORE COMPLEX MIXTURES
We start by training transformer models on the mixture of dense linear regression (FDR) and sparse linear regression (FsR) function classes. The function definition remains the same for both these classes i.c. f x + w æ, but for FDR we consider a standard gaussian prior on w and a sparse prior for FSR. We use the sparse prior from Garg et al. 20022), where we first sample ~ I) ~ and then randomly set its d components as We consider s = 3 throughout our experiments. Unless specified we consider the mixtures to be uniform i.e. a; 0.5 and use these values to sample batches during training.
During the evaluation, test the mixture model (denoted as Transformer F(DR. SR)) on the prompts sampled from each of the function classes in the mixture. We consider the model to have in-context learned the mixture of tasks if it obtains similar performance the single-task models specific to these function classes. For example, transformer model trained on the dense and sparse regression mixture (Transformer F(DRR should obtain performance similar to the single-task model trained on dense regression function class (Transformer FUR), when prompted with a function f ~ FDR and vice-versa.
The results for the binary mixtures of linear functions are given in Figure 26 As can be ob- served, the transformer model trained on FiDR, SR) obtains performance close to the OLS baseline as well as the transformer model specifically trained on the dense regression function class FDR when evaluated with dense regression prompts. On the other hand, when evaluated with sparse regression prompts the same model follows Lasso and single-task sparse regression model (Transformer (FsR)) closely. As a check, note that the single-task models when prompted with functions from a family different from what they were trained on, observe much higher errors, confirming that the transform- ers learn to solve individual tasks based on the in-context examples provided. Similar to GMMs in $3.1 here also compare the implied weights from multi-task models under prompts for both FoR
Results.
37

---

Published as a conference paper at ICLR 2024
Evaluation Dense Regression Prompts
Evaluation on VignoVector Regression Prompts
Transformer (F (FinssvR))
1.0
1.0
Transformer (Fua)
Transformer (Fsva)
0.8
0.8
OLS
locs0l:
1oss0k
Minimize fou
0.6
0.6
Bound
Bound
0.1
0.1
0.2
0.2
0.0
0.0
10
20
30
40
(# in-context examples)
in-context examples)
(8)
Regression Promots
Evalustion on Sign Regrasion Promats
0.31
(pnavan p
(w [DN soma P10v
0.235
0.32
mnee squered errer
(nn.svyy 01.)
0.10
(M(DR sony
01.)
0.08
Bound
0.13
0.06
Bound
6.00
(# in-contect examples)
[# in-context exampies)
(b)
(c)
Figure 28: Comparing the performance of Transformer model trained dense and sign-vector regression mixture F(DR. SVRj. with baselines, as well as single task models, trained on JDR and FsvR individually. Top: Comparing lossëk values of the mixture model with single-task models with different prompt distributions. Bottom: Comparing the errors between the weights recovered from the mixture model and different single task models and baselines while evaluating on FDR and Fsvr prompts.
38

---

Published as a conference paper at ICLR 2024
Evaluation Dense Regression Prompts
Evaluation on Neewed-Covarianne Regression Prompts
1.2
Transformer (F Faaastre-ns))
1.0
Transformer (Foe)
Transformer (F Faees DR)
OLS
leeaok
loasōk
Minimize w 'u
0.6
Bound
0.4
0.5
Bound
0.2
0.0
0.0
10
20
40
A
(# in context examples)
(#
context examples)
(a)
Evolustion on Dense Regresiion Prompts
Evalustion on Skewed- Covariance Regression Prompts
(M (DRSbey DR
0.10
Ludee (bel MMa Da!
meean sqpared areo
0.08
0.44
SDR DRC
uooke
Skon)
0.00
[BBLStew Day
PME Suny
Bound
0.2-
Bound
6.0
0.00
10
15
25
30
35
10
20
(# in context examples)
[+] context exemp es)
(b)
(c)
Comparing the performance of a Transformer model trained on dense and skewed- covariance regression mixture F(DRR SR)w with baselines, well single task models, trained on FDr and Fskew.DR DR individually. Top: Comparing loss@k @k values of the mixture model with single- task models with different prompt distributions. Red (OLS) and orange (Transformer (FDR)) curves overlap very closely, so are a bit hard to distinguish in the plots. Similarly in the top right plot, pur- ple (Minimize w" E h-' "w) and green (Transformer FSkew-DR) curves overlap. Bottom: Comparing the errors between the weights recovered from the mixture model and different single task models and baselines while evaluating on FDR and Fskew-DR prompts.
Figure 29:
Danie Regresiion Promets
Spurse Repession Prompts
Exaluation on VicteV Remession Promots
1.2
Transfomer (Farstsvg)
10-
3,0
Trandommer (Faa)
Transfomer (Fil
0.8
0.5
Trandhemer (Fax)
100038
0.6
6
Bound
Bound
0.4
0.2
0.2
00
0.0
0.0
examples)
1+ examlees)
(4 evamples)
Figure 30: Comparing the performance of transformer model trained to in-context learn F(DR. SR. SVR) mixture family with the corresponding single task models.
39

---

Published as a conference paper at ICLR 2024
and FsR and show that here again they agree with the weights recovered from single-task models as well as the strong baselines in this case (OLS and Lasso). We provide the plots for the weight agreement in this case in Figure 27
Next, we describe the results for other homogeneous mixtures F(DR, sVR). F(DR, Skew.DR) and FIDR, SR. SVR), as well as heterogeneous mixtures ((DRRDD DT) and F(DT, NNj. As can be seen in Figure 28 the transformer model trained on F(DR. SVR) mixture, behaves close to OLS when prompted with
Te FOR and close to the Lo minimization baseline when provided sign-vector regression prompts ( E FSVR). We also have similar observations for the F(DR. Skew-DR) mixture case in Figure 29) where the multi-task ICL model follows the PME of both tasks when sufficient examples are pro- vided from the respective task. Similarly, for the model trained on the tertiary mixture F (DRR, sR, svR) (as can be seen in Figure 30). the multi-task model can simulate the behavior of the three single- task models depending on the distribution of in-context examples. On FER and JsvR prompts the multi-task model performs slightly worse compared to the single-task models trained on FsR and FSVR respectively, however once sufficient examples are provided (still < 20), they do obtain close errors. This observation is consistent with the PME hypothesis i.c. once more evidence is observed the 3 values PME of the mixture should converge the PME of the task from which prompt P is sampled. The results on heterogeneous mixtures we discuss in detail below:
Heterogeneous Mixtures: Up until now, experiments for the multi-task case have been focused on task mixtures where all function families have the same parameterized form i.e w x for linear mixtures and w² (xx) for Fourier mixtures. We now move to more complex mixtures where this no longer holds true. In particular, we consider dense regression and decision tree mixture F(DR, DT) and decision tree and neural network mixture F(DI,,N).
We follow Garg et al. 2022)> 's setup for decision trees and neural networks. We consider decision trees of depth 4 and 20-dimensional input vectors æ. A decision tree is sampled by choosing the split node randomly from the features at each depth, and the output of the function is given by the values stored in the Ieaf nodes which are sampled from ((, 1). For neural networks, we consider 2-layer (1 hidden + I output) multi-layer perceptrons (MLP) with ReLU non-linearity i.e. f()) = E ae Re (w: æ²), where a E R and w; E Rdd The network parameters a,s and w,s are sampled from (0, 2/r) and ( 1) respectively. The input vectors æ,s are sampled from (( for both tasks. We consider greedy tree learning and stochastic gradient descent over a 2-layer MLP as baselines for decision trees and neural networks respectively. The values of hyperparameters for baselines such as the number of gradient descent steps, initial learning rate for Adam, etc. the same as Garg et al. (2022).
The results for the two mixtures are provided in Figure 31
The mixture model Transformer (F(DR, DTY) follows the single task model Transformer (FDR) when provided in-context examples from f ~ FDR and agrees with Transformer (FDr) when prompted with f ~ For (Figure[1aa) Sim- ilarly, we have consistent findings for F(DT, NN) mixture well, where the model learns solve both tasks depending upon the input prompt (Figure 3lb).
C.4 FOURIER SERIES MIXTURE DETAILED RESULTS
We consider a mixture of Fourier series function classes with different maximum frequencies, i.e. Ffourier = { ..... Ffourier). We consider N 10 in our experiments and train the models using a uniform mixture with normalization. During evaluation, we test individually on each each Flourier where M E []. We compare against consider two baselines: i) OLS Fourier Basis F i... performing OLS on the basis corresponding to the number of frequencies M in the ground truth function, and ii) / Ffourier which performs OLS on the basis corresponding to the maximum number of frequencies in the mixture i.e. N.
Figure 32a plots the lossêk metric aggregated over all the M E [1 for the model and the baselines. The performance of the transformer lies somewhere in between the gold-frequency base- line (OLS Fourier Basis F and the maximum frequency baseline (Flourier), with the model performing much better compared to the latter for short prompt lengths (A < 20) while the former baseline performs better. We also measure the frequencies exhibited by the functions predicted by the transformer in Figure 32b
We observe that transformers have a bias towards lower frequen-
'In practice, we use Adam just like Garg al. 2022
40

---

Published as a conference paper at ICLR 2024
Evaluation on Dense Regression Prompts
Evaluation on Decision Tree Prompts
1.2
Transformer (F(DR.DT))
2.0
Transformer (F(DRDDT))
WW
1.0
Transformer (FiR)
Transformer (Fivr)
OLS
1.5
Greedy Tree Learning
0.8
loss0k
loss@k
0.6
1.0
0.4
0.5
0.2
0.0
0
10
20
30
40
20
40
60
80
100
k
(# in-context examples)
(# in-context examples)
(a)
Evaluation on Decision Tree Prompts
Evaluation on Neural Network Prompts
2.0
Transformer (F(DTNN))
Transformer (F(DTNNN)
1.0
Transformer (Fir)
Transformer (FNN)
Greedy Tree Learning
2-layer NN, GD
1.5
0.8
loss0k
loss0k
1.0
0.6
0.4
0.5
0.2
0.0
20
40
60
80
100
20
40
60
80
100
(# in-context examples)
(# in-context examples)
Figure 31: Multi-task in-context learning for heterogeneous mixtures.
41

---

Published as a conference paper at ICLR 2024
Transformer Inductive Bizzes
Tankere
Ru kod
0L5 Foeee Eata E
133117
2011
(a)
(b)
Triining sissed high frequencies Ma
9303122
Figure 32: In-context learning on the Fourier series mixture class. Top Left: Comparing transform- with the baselines. Errors computed batches of 128 for M E [1, 10j and aggregated in the plot. Top Right: Visualizing the frequencies of the simulated function by transformers. Bottom: Training transformer on high-frequency biased Fourier mixture F fauuecirr and visualizing the simu- lated frequencies of the trained model.
cies when prompted with few examples; however, when given sufficiently many examples they are able to recover the gold frequencies. This simplicity bias can be traced to the training dataset for the mixture since lower frequencies are present in most of the functions of the mixture while higher frequencies will be more rare: Frequency 1 will be present in all the function classes whereas frequency N will be present only in Fiuuiier Our results indicate that the simplicity bias in these models during in-context learning arises from the training data distribution. We confirm the above observations by detailing results for different combinations of M and l in Figure 33
C.4.1 COMPLEXITY BIASED PRE-TRAINING
To further verify this observation, we also consider the case where the train- ing data is biased towards high frequencies and check if transformers trained with such data exhibit bias towards high frequencies (complexity bias). To mo- tivate such a mixture, we first define an alternate fourier basis: & (z) Joos (No L)), sin (ngT/I), cos ((no + 1)) ((ng + 1)m/L), cos (Nv/L), LT/LL,,iin sin (Nx/L)),
where 710 > 0 is the minimum frequency in the basis. 2no.N defines the function family Frourier and equivalently we can define the mixture of such function classes as Falsener [Floyrier FF }. One can see such a mixture will be biased towards high frequency; frequency N is present in each function class of the mixture, while frequency 1 is only present in Ffourier We train a transformer model on such a mixture for N = 5 and at test time, we $1.x"
T:N.N
Ffourier Figure 32c shows the inductive biases measure from @mo.. this trained model and we can clearly observe case of complexity bias, where at small prompt lengths, the model exhibited a strong bias towards the higher end of the frequencies that it was trained on i.e. close to 5.
evaluate the model on functions
We also trained models for higher values of the maximum frequency i.e. N = 10 for the high- frequency bias case, but interestingly observed the model failed learn this task mixture. Even for N = 5, we noticed that the convergence was much slower compared to training on the simplicity bias mixture Fóguner, This indicates, while in this case, the origin of simplicity bias comes from the training data, it is harder for the model to learn capture more complex training distributions, and simplicity bias in the pre- pre-training data distribution might lead to more efficient training Mueller & Linzen 22033).
42

---

Published as a conference paper at ICLR 2024
Figure 33: In-context learning of Fourier series mixture class. Measuring the frequeneies of the simulated function by the transformer for different values of M (maximum frequency) and & (num- ber of in-context examples). Showcases the simplicity bias behavior exhibited by the model low frequencies.
43

---

Published as a conference paper at ICLR 2024
C.5 CONDITIONS NECESSARY FOR MULTI-TASE ICL
We observed that the training setup can also influence the ability of transformers to simulate the Bayesian predictor during ICL. Particularly, in our initial experiments with F(DR. SR) mixture ($C.3) transformers failed to learn to solve the individual tasks of the mixture and were following OLS for both FDR and FSR prompts. To probe this, we first noted that the variance of the function outputs varied greatly for the two tasks, where for dense regression it equals d and equals the sparsity parameter s for sparse regression. We hypothesized that the model learning to solve just dense regression might be attributed to the disproportionnaely high signal from dense regression compared to sparse. To resolve this, we experimented with increasing the sampling rate for the FsR task family during training. Particularly on training the model with osR = 0.87, observed that the resulting model did learn to solve both tasks. Alternatively, normalizing the outputs of the two tasks such that they have the same variance and using a uniform mixture (osR 0.5) also resulted in multi-task in- context learning capabilities (also the setting of our experiments in Figure 26 Hence, the training distribution can have significant role to play in the model acquiring abilities to solve different tasks as has been also observed in other works on in-context learning in LLMs Razeghi et al. (2022) Chan jt al. 2022a).
We also studied if the curriculum had any role to play in the models acquiring multi-task in-context learning capabilities. In our initial experiments without normalization and non- oon-uniform mixtures, we observed that the model only learned to solve both tasks when the curriculum was enabled. However, training the model without curriculum for a longer duration ( more training data), we did observe it to eventually learn to solve both of the tasks indicated by a sharp dip in the evaluation loss for the sparse regression task during training. This is also in line with recent works Hoffmmann
ctt al. (2022) Touvron et al. (2023), which show that the capabilities of LLMs can be drastically improved by scaling up the number of tokens the models are trained on. Detailed results concerning these findings are in Figure 34 of the Appendix.
Figure comopaes transformer models trained on FIDR, SR) mixture with different setups i.e. train- ing without task-normalization and uniform mixture weights ag's (Figure (34a), training without task-normalization and non-uniform mixture weights (Figure (34b), and training with task normal- ization and uniform mixture weights (Figure 34c). As described above, we perform task normal- ization by ensuring that the outputs f()) for all the tasks have the same variance, which results in all the tasks providing a similar training signal to the model. To perform normalization, we simply divide the weights w sampled for the tasks by normalization constant, which is decided according to the nature of the task. With this, we make sure that the output y = wTa x has a unit variance. The normalization constants for different tasks are provided in Table 2
Normalization constants used for different tasks to define normalized mixtures for multi- task ICL experiments. Here d denotes the size of the weight vectors used in linear-inverse problems as well as the last layer of the neural network. s refers to the sparsity of sparse regression problems, r is the hidden size of the neural network and N refers to the maximum frequency for Fourier series.
Table 2:
Function Family
Normalization Constant
Dense Regression
d
Sparse Regression
Sign- Vector Regression
/ d
Fourier-Serri
N
Degree-2 Monomial Basis Regression
S
Decision Trees
Neural Networks
2
All the experiments discussed above (like most others in the main paper) were performed using curriculum learning. As discussed above, investigated if the curriculum has any effect on multi-
task ICL capabilities. The results for the same are provided in Figure 35
We also explore the effect of normalization multi-task ICL in Figure 36 for F(DRR sVR) task. As can be seen in Figure [3aa| for this particular mixture even while training the model without normalization, the model exhibited multi-task ICL, which can be explained by both tasks having
44

---

Published as a conference paper at ICLR 2024
Unnormalized Mixture
ODD
0.5, OsR
Evaluation Dense Regression Prompts
Evaluation Sparse Regression Prompts
1,25
1.25
Tranaformer (F(nns)))
Transformer (Fur)
Transformer (Jax)
10ss8k
0.75
0.75
01.5
0.50
Bound
Bound
0.25
0.25
0.00
0.00
10
211
30
301
(# in-contest examples)
in-contect examples)
(a)
Unnormalized Mocture
= 0.13,
0.87
Evaluation on Dense Regression Prompts
Evaluation on Sparze Regression Prompts
1,25
Transformer (7 pp
1,00
1.00
Transformer (7un)
Transformir (Par)
10n063
0.75
losp0t
01.5
0.50
Bound
0.50
Bound
Lacso
0.25
0.25
(# in-contest examples)
(# in in-content examples)
(b)
Normalized Mixture
opn 0.3, osn
Evaluation Dense Regression Prompts
Evoluation Sparse Regression Prompts
1.25
1.25
Transformer (Fipasn:)
1,00
Transformer (Fun)
1oss
0.75
0c
0.75
OLS
Transformer (Jar)
Bound
0.50
Bound
0.25
0.00
0.00
201
300
40
20
2001
300
(# in-context examples)
(# in-context examples)
(c)
Figure 34: Conditions affecting multi task ICL in transformers. Top: Evaluating lossek for trans- former model trained on F(DR, SR) task family without normalization and considering uniform mix- tures (i.e. ODR osR = 0.5), and comparing with single-task models and baselines. While the blue curve (Transformer (DR. SR)) ) is hard to see here, it is because it overlaps almost perfectly with the red curve corresponding to OLS in both cases. Center: Similar plots as above but for the model trained on the mixture FŞDR, SR\ with non-uniform weights i.c. aDR 0.13, asR 0.87. Bottom: Training the model with the normalized (and uniform) mixture such that outputs for the two tasks have the same variance. All the models are trained with the curriculum. The discussion continues in Figure 35 for the models trained without curriculum.
45

---

Published as a conference paper at ICLR 2024
Unnormalized Moxture
= 0.13, 0..7
Training Step: 500k
Evaluation Dense Regression Prompts
Evaluation Sparse Regression Promets
Transformer (FioR sR:)
Transformer (7on)
1,00
1.00
Transformer (Fan)
Jossak
0.75
0.75
01.5
0.50
Bound
0.50
Bound
Lasso
0.25
0.25
(# in-contest examples)
(#t in-content examples)
Unnormalized Mixture
OpR 0.13, osn
Training Step: 800k
Evaluation Dense Regression Prompts
Evaluation Spurse Regression Promots
1.25
Transformnr (Fipoten:
1.25
Transformer (Fun)
1:00
Transformar (Fas)
300006
0.75
0.75
OLS
0.50
Bound
Bound
0.25
0.25
0.00
0.00
20
30
40
10
30
in-context examples)
examples)
(b)
Training Dynamics
1 1oas010 for SR prompts
3.5
2.5
10s9810
2.0
Only solves DR
1.5
1.0
Solves both DR and SR
0.3
0.2
0.4
Training Steps
x10"
Figure 35: Evaluating transformer model trained without curriculum on F(DR. SRY task family without normalization and non-uniform weights (DR 0.13, aSR = 0.87 (similar to Figure 34b). Top: Evaluating the checkpoint corresponding to the 500k training step of the aforementioned model. Again, the blue curve (Transformer F(DR, SR) is hard to see here, but it is because it overlaps almost perfectly with the red curve corresponding to OLS in both cases. Center: Evaluating the same model but a much later checkpoint i.e. at 800k training step. Bottom: Evolution of losse 10 on FsR prompts while training the above model.
46

---

Published as a conference paper at ICLR 2024
Codeccionn on
1.5
0LS
bo comple
oomple
(a)
(c)
Puturier
Maken
mm
Ie sampin
C1 inconaa ausple
C6 womoeo corgle)
12 aesceedaaagga
(b)
(d)
Figure 36: Effect of output normalization on multi-task ICL in transformers. Top Left (a): A transformer model is trained on a uniform mixture of F(DR, SvR; task family (i.e. &DR = asVR = 0.5) without normalization. Evaluating 1oss@k for this model outputs f(x) are not normalized to have unit variance i.e. same as training). Note that for the F(DR, SVR) task family even without normalization the outputs fx) have the same mean and vari- ance (μ+ 02 20) for both the tasks. Bottom Left (b): Evaluating loss@k for the model in (a) on normalized prompts (where outputs f(r) for both tasks are normalized to have unit vari- ance). Top Right (c): A transformer model is trained on uniform mixture of F(DR, SVR; task family (i.e. ODR OsvR 0.5) with normalization. Evaluating loss @ for this model on unnormalized prompts. Bottom Right (d) Evaluating lossek for the model in (c) on normalized prompts. All the models are trained with the curriculum.
unnormalized prompts (where
the same output variance (i.e. d). Interestingly, when we evaluate this model (i.e. the one trained on unnormalized mixture) on in-context examples which have the outputs f(x;)'s normalized, the model fails to solve FSvR and follows OLS baseline for both the tasks. We hypothesize that since this situation represents Out of Distribution (OOD) evaluation and the model might not be robust towards performing multi-task ICL on prompts that come from a different distribution than those seen during training. Exploring OOD generalization in the multi-task case is compelling direction that we leave for future work.
Table 3: Multi-task generalization results for Monomials problem. The first row is ID evaluation, second row is OOD evaluation. As task diversity (K) increases, the model starts behaving like Lasso and BP BP proxy, and its ID and OOD losses become almost identical, i.c. it generalizes to OOD.
25. CoMaati
40. Dialationn
K-008
4.00
222
3.0
532
880
8:00
500
6.0
20
10
60
60
C-10 000 Sesaatis
<-46 000 Ivstater
C-39 202 EVBLUtOO
-30
1.20
1.22
312
1.33
1.31
1.20
3.50
3.50
2.23
2.33
533
033
47

---

Published as a conference paper at ICLR 2024
Table 4: Multi-task generalization results for Fourier Series problem. The first row is ID eval- uation, second row is OOD evaluation. As task diversity (K) increases, the model starts behaving like x and BP prosys and its ID and OOD losses become almost identical, i.e. it generalizes to OOD.
E-1345 OI Ditriaater
18 Inconder suandle)
Incontar ssarple))
D DETAILS REGARDING MULTI-TASK GENERALIZATION EXPERIMENTS
Like section I.I.I we tune the Lasso coefficient on a separate batch of data (1280 samples) and choose the single value that achieves the smallest loss.
D.1 MONOMIALS MULTI-TASK
Plots for various task diversities that we experiment with are shown in Table
D.2 FOURIER SERIES MULTI-TASK
Fourier Series problem MICL setting, Please refer to the setup defined in BB2.1 which comes under the MICL setting as it corresponds to single function class, Frourier
multi- Extending Fourier Series problem to HMICL setting. For extension to HMICL, ple subsets of frequencies S to define the mixture. Each S₃ defines a function class F The pretraining distribution is induced by the uniform distribution (FF over a collection of such func- tion classes, F [Jrourier
Ffouriery, }, where Sg (x)) the full basis. For example, Sg could be [1, cos (2Tx/ L), cos (6Tx/ L), cos (9Tx/ L), sin (2mzz/ L), sin (Grx/L) sin (9mx/L)J² consist- ing of sine and cosine frequencies corresponding term, is a part of every Sz). K feature sets Sg ss, each of size D, are chosen at the start the train- ing and remain fixed. K is the task diversity of the pretraining distribution. To sample training function for the TF, we first sample a function class Flourier with replacement from (FF and then sample a function from the chosen class; f(x) w" Sxxx), where w (D O. I). Similar to the Monomials problem, our aim to check if TF trained on (FF) can generalize to the full distribution of all function classes (for feature sets of size D) by evaluating its performance on function classes corresponding to feature sets S' & (.
integers 2, 6 and 9. (Note that 1, the intercept
Training Setup. d 1=1,p = 82, N 20. So, the full basis, (x),,, had 20 frequencies. D 3. We experiment with K < {1, 10, 100, 200, 400, 800, 1140).
Evaluation Setup. The baselines we consider are OLSs. OLS and Lasso For Lasso, we use a = 0.1. We again evaluate in two settings: (a) ID: On functions from the pretraining distribution. (b) OOD: On functions not in the pretraining distribution by sampling from function family not corresponding to any of the Sks used to define the pretraining distribution.
Results. Plots for various task diversities that we experiment with are Table 4 The trend is
the same it was for Monomials problem, f.e. ID performance degrades and OOD performance
48

---

Published as a conference paper at ICLR 2024
improves K increases. As K increases, TF's performance on ID and OOD becomes identical (from K = 100 onwards) and similar to the Lasso and Procy baselines.
D.3 DETAILS ON THE PIENOMENON OF FORGETTING
Problem Setup. We follow the Noisy Linear Regression (NLR) setup from Raventós et al. (2023) :
= 8. 15 (without curriculum learning). The noise variance x² 0.25. For this problem, the
transformer has 8 layers, 128-dimensional embeddings, and 2 attention heads, and is trained with batch size of 256 for 500k steps. One-cycle triangle learning rate schedule /Smihh & Topin (2018) is used with 50% warmup. Detailed plots for the four groups of task diversities mentioned in $6.1 OOD loss curves with Bayesian predictors for various checkpoints for task div 2
are in Figure 37
are in Figure 38 For other representative task diversities, the OOD loss curves of TF and Bayesian
predictors are in 5 Plots showing mean squared errors of implied weights of TF with Bayesian
predictors are in Table
Classification of task diversities. ID loss Rs 0 for all task diversities during pretraining. We group them into the following 4 categories based OOD loss:
1. 21 to 2 (no generalization; no forgetting) - OOD loss never decreases, converges to a value worse than same at the start of the training (fo), agrees with dMMSE at the end of the training (tena). [Figure 37a
2. 2 to 2 (some generalization and forgetting) OOD loss improves, reaches minima bmin, then worsens. OOD loss is worse than ID loss throughout pretraining and agrees with dMMSE at tend (i.e.. any generalization to Gaussian distribution is forgotten by tend). [Figure 37b)
3. 27 to 2 (full generalization and forgetting) OOD loss improves, reaches a minima Emi at which it is same as ID loss, then it worsens. At Emin OOD loss agrees with Ridge (Figure 38a then
gradually deviates from it and at fend, it is in between dMMSE and Ridge (e.g., Figure 38c We
refer this group of task diversities as the "Gaussian forgetting region" since the model generalizes to the full (Gaussian) distribution over tasks at tmin but forgets it by tend- [Figures 37c
38
22 to 20 (full generalization; no forgetting) - Throughout pretraining, OOD and losses are identical and OOD loss agrees with Ridge. [Figure 37d
Relation Simplicity bias? The phenomenon of forgetting (displayed by task diversity groups 2 and 3 above) is an interesting contrast to the grokking literature and in particular to Nanda et al. (2023), where they find that the model first memorizes and then generalizes (which, on the surface, is the opposite of what we observe). We can explain forgetting from the perspective of simplicity bias. Since PT is discrete and perhaps contains lots unnecessary details, the model instead finds easier to generalize to the 'simpler' Gaussian distribution which is continuous and much more nicely behaved. Hence, we speculate that the simplicity of the PT is inversely proportional to the number of tasks it contains. Very small task diversities (group 1) are exceptions to this rule since their P' is arguably much simpler than FG So, we do not see forgetting in those cases as the model prefers to only learn PTaia. Thus, we hypothesize that the simplicities of the distributions have the following order (G, denotes group i): PT (G2) ~ PT (Gs) < PT PTaiit. (Ga) < FG < PTaiit (G1).
Robustness and effect of the number of dimensions. The phenomenon of forgetting is robust to model sizes (Figure (39), to changes learning rate and schedule (Figure 40), and position encodings (Monomials and Fourier Series multi-task setups use 12-layer transformer that does not have position encodings). We also experimented with NLR problems having dimensions = 3 and d 16 (Figure 41) and found that the extent of forgetting (denoted by the disagreement of TF's and Ridge's 's loss on OOD evaluation) is directly proportional to the input dimension (d). Note that following Raventós al. (2023) we keep the signal-to-noiss ratio (d/o²) constant across these experiments by adjusting the noise scale to ensure that noise has a proportional effect and the observations due to change in dimension alone.
49

---

Published as a conference paper at ICLR 2024
ATTHE MPTEL WITH
100c
2004
2004
4008
7306
(a) No generalization; no forgetting)
144
sccee 2011 taco Pats talk Inseee
3.5
52
100x
400k
(b) Some generalization and forgetting
arran
and and Daee
0.5
0.4
(c) Full generalization and forgetting
sehett a011 2011
10m00: a010 S0i00 cate look stoee oed
(d) Full generalization; no forgetting
Figure 37: Evolution of ID and OOD losses during pretraining for different task diversity groups for the Noisy Linear Regression problem. The forgetting phenomenon depicted by groups in Figures (b) and (c). The moving average (over 10 train steps) of ID (*eval) and OOD (*eval_ood) losses are plotted, with the original (non-averaged) curves shown in a lighter shade. A checkpoint towards the end of the training is highlighted. We see that we increase task diversity (i.e. go from group (a) towards (d)). the difference between ID and OOD losses decreases. Groups (b) and (c) are noteworthy as they display the phenomenon of forgetting, where the models' OOD loss at an earlier checkpoint is the same as ID loss, but it increases later.
50

---

Published as a conference paper at ICLR 2024
Table 5: OOD loss curves of TF and Bayesian predictors for various checkpoints of models cor- responding to task diversities (K) 23, 29, 28, 2 2 respectively in rows. Each plot presents the loss across different prompt lengths. For task diversities 2 and 2 plots in the first column represent the point of minima (lmin). For task diversities 2 and 216 plots in the first column represent an earlier checkpoint.
K minima (fmin) or an earlier checkpoint
checkpoint after 100k train steps
checkpoint after 500k train steps
3.0
1.0
1520
0.5
04
0.4
00355
09/cc
23
10
22
24
0
23
12
34
210 nicentot
5
10
14
Marn moeer 32; Cape. 22006
tide:rde::20 32: Clipe 100003
Wim Pscce 12; Cage. 500350
Pedge
2.0
crance
0.6
044
0.4
1430
25
8.06 11001-001 (OTTT(E)))
210 1oco soumicis)
2/0 (14060ll crams
3.0
0.0
0.6
0.2
0.2
28
2
10
32
14
6
33
12
34
0
5
10
32
34
Nom balle: (3535: Oct: 33000
Nam bacer 51536 Dlee: 100000
Num Eaile: 63532 Oct: 500000
2.5
3.5
0.6
0.4
216
KC0 OOOORE ororpee
216 novrtes comilees
1 Dicoroot 0A3
Num tasks: 256, Cloc: 31300
Num tasks. 256 Ckit: 100000
Num tacks, 256, Ckst: 500000
TF
hidg:
migg:
0.6
0.4
10
12
K14 n.contoxcoammllo
K14 in nxcontaatooo oxo mxlis)))
K14 n- 20-E0E OKo milis)
(a) after 31k train steps (tmen)
(b) after 100k train steps
(c) after 500k train steps
Figure 38: Plotting the OOD loss for various checkpoints during training for task diversity 28, along with the Bayesian predictors. At tmin the model agrees with Ridge regression for all prompt lengths but later deviates and converges to somewhere in the middle of two Bayesian predictors.
51

---

Published as a conference paper at ICLR 2024
Table 6: Mean squared errors of implied weights of TTF with Bayesian predictors during OOD evalu- ation for various checkpoints of models corresponding task diversities (K) 2 29, 2 2 respec- tively in rows. Each plot presents the implied difference of weights across different prompt lengths. For task diversities 2 and 2 plots the first column represent the point of minima (Imin). For task diversities 2 and 2 plots in the first column represent an earlier checkpoint.
K
minima (Emin) or an earlier checkpoint after 100k train steps
checkpoint after 500k train steps
checkpoint
Clpt- 37200
K-m Cropt 100020
Chpt 502000
maY carriquee mr
54
ITC, eM60
60
2
iii.croood sxameto
Inconta ourppes
s.i.omicnttt moloa
32: Dipt 23050
K2 Clis 100030
Ke Diple 503200
1.4
25
Ckpt
(1pt=520003
(7. 9.5
ITs, Knyi
caIn 22.00d eror
o1
116, cm
0.3
84
1.4
28
1.5
K=01133:Cogt 33003
X- 05536 10:000
X= Caps 300020
2.8
6.6
2
52

---

Published as a conference paper at ICLR 2024
LISA, EXi44
D7123
2.2
0.4
(a) 2
(b) 2
XB4
(c) 2
(d) 2
(e) 2
(
0.4
0.2
XBC
15%
(g) 2 10
(h) 2
Figure 39: Effect of model size on forgetting. The moving average (over 10 train steps) of OOD losses for various models of different sizes trained on pretraining data with task diversities 2 to 2 are plotted. the legend, I. and F. denote the number of layers and embedding size of the trained models respectively. As can be observed, all the models show forgetting across task diversities. Additionally, for larger task diversities (rows 3, 4), viz. 2 to 2dd, bigger models exhibit higher OOD loss over the course of training, i.e. show more forgetting! For smaller task diversities (rows 1, 2), viz. 2 to 27, there is no clear trend the extent of forgetting across model sizes.
53

---

Published as a conference paper at ICLR 2024
00D
128 tasks learning rate 1e-3 w/ (triangle) rate schedule
00D ID 128 tasks
learning rate = 1e-3 w/o rate schedule
00D
ID 128 tasks learning rate = 1e-4 w/o rate schedule
1.2
1
0.8
0.6
0.4
0.2
Step
0
100k
200k
300k
400k
500k
Figure 40: The moving average (over 10 train steps) of ID (solid lines) and OOD (dashed lines) losses for task diversity 2 for different learning rates, with & without learning rate schedule are plotted. While the nature and extent of forgetting changes, the phenomenon itself robust and observed across all settings.
1.2
1
0.8
0.6
0.4
0.2
Step
0
100k
200k
300k
500k
Figure 41: The moving average (over 10 train steps) of ID (solid lines) and OOD (dashed lines) losses for task diversity 27 for different input dimensions (3 (green), (yellow), 16 (red)) are plotted. The extent of forgetting is directly proportional to the input dimension (d).
54

---

Published as a conference paper at ICLR 2024
E GRADIENT DESCENT AS A TRACTABLE APPROXIMATION OF BAYESIAN INFERENCE
We describe TWO recent results showing apparent deviation from Bayesian prediction: Garg et al. (022) showed that training transformers on functions sampled from the class of 2-layer neural networks, resulted in transformers performing very close GD during ICL. For linear regression, Akyürek et al. (202)) and yon Oswald et al. (2022) showed that low-capacity transformers (i.e., those with one few layers) agree with GD on the squared error objective. Why does GD arise in ICL when Bayesian prediction is the optimal thing to do? Towards answering this, we propose the following hypothesis: The reuson transformers appear to do GD is because GD may providing the best upproximation to Bayesian inference within the capacity constraints One line of evidence for this comes from Mingard et al. (202T) (and the references therein) who showed that for a range of deep neural network architectures GD solutions correlate remarkably well with the Bayesian predictor. Further, it's well known that for convex problems, GD provably reaches the global optima. hence approaching the Bayes-optimal predictor with increasing number of steps. von Oswald et al. (2022) shows that multiple layers of transformers can simulate multiple steps of gradient descent for linear regression and Akküüek al. (2022) highlights that as the capacity of transformers is increased they converge towards Bayesian inference.
F FURTHER CONCLUDING REMARKS
Much more remains to be done to determine how extensively transformers mimic the Bayesian predictor. Relation between the pretraining distribution and ICL inductive bias and its relation to real-world LLMs needs to be further fleshed out. The intriguing forgetting phenomenon needs to be better understood. How is it related to pretraining simplicity bias? Further progress on the re- lation between gradient-based optimization and Bayesian inference would be insightful. The case of decision trees studied in Garg et al.) (2022) is an interesting specific problem where the relation- ship between Bayesian inference and gradient descent remains unclear. While we studied out-of- distribution performance on new function families, ouu-t--distributtin performance on new input distributions is also of interest.
The present work focused on continuous functions setting which is easier to study from the Bayesian perspective. What happens for the real-world LLMs? How strongly does the Bayesian view hold there and what kind of deviations exist? The order of demonstrations is known to have significant influence the output of the LLMs. Thus more nuance would be necessary for LLMs. Assuming that the Bayesian view does hold for LLMs in some form potential practical implication is that it can help choose demonstrations in an informed manner make the conditional distribution con- verge faster to the intended output. Moreover, it can also potentially help us understand real-world LLM phenomena like hallucination and jailbreaking assuming the nature of the implied posterior distribution characterizes them. Finally, we treated transformers as black boxes: opening the box and uncovering the underlying mechanisms transformers use to Bayesian prediction would be very interesting.
55


Reviews:
[{"review_id": "dVrD6eJKl3", "reviewer": "Reviewer_P1i1", "summary": "This paper considers in-context learning for large language models. This paper extends the previous work \"What can transformers learn in-context? A case study of simple function classes.\" by considering multiple function families, out-of-distribution detection, and Bayesian prediction tasks. This paper conducts extensive experiments and large-scale analysis in the experiment section.", "questions": "- Considering the paper's main topic is about prompting engineering and large language models, the author might consider trying for a compatible conference in NLP.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The author conducts extensive experiments on prompting engineering for large language models over multiple tasks.", "weaknesses": "- The paper is not well organized which makes it very hard to read through. \n    - The Figures are very hard to understand, making it challenging to summarize the conclusion from the figures.\n    - It's not a good idea to keep stating \"refer to Appendix\" for the main methodology part of the paper. \n    - It is also not a good idea to keep stating that some steps in the proposed method are similar to a specific paper and cite that paper. At least, the authors need to give a clear reason for doing so.\n- The paper mentioned a lot of technique terms without giving any rational and proper definitions. For example, \n    - the \"Bayesian predictor\" is not properly explained in the context of \"in-context learning\". \n    - I'm not familiar with The \"deviations from Bayesian prediction\" . The author should explain that clearly.\n    - The \"Simplicity bias\" should be better replaced with \"Occam’s razor\" or \"no free lunch theorem\". Because the suggested terms are used very often in the literature. \n    - If you mean \"generalization is out of distribution (OOD) detection\". Then just using OOD detection would reduce a lot of confusion.\n- The correctness of Equation 2 is questionable. What is the exact definition of \"$df$\", the differentiation with respect to a function $f$?\n- As the main component of this work, the paper lacks the reason for extending ICL to hierarchical ICL. Why do you want to sample from a mixture of the function family?", "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": "", "extra_scores": {}}, {"review_id": "mSlWhLkD2F", "reviewer": "Reviewer_C7Ly", "summary": "The paper focuses on the phenomenon of in-context learning (ICL) in large language models (LLMs) and aims to understand the underlying inductive biases that enable successful generalization to new functions. The authors propose a hierarchical meta-ICL setup that encompasses multiple task families and investigate the empirical performance of LLMs in this expanded setting. They find evidence that high-capacity transformers mimic the behavior of Bayesian predictors in cases where Bayesian inference is tractable. This Bayesian perspective sheds light on the inductive biases of ICL and how transformers perform tasks when trained on multiple tasks. Additionally, the paper explores situations where transformers deviate from the Bayesian predictor, leading to new insights and hypotheses.", "questions": "It is well-known that the order of demonstrations has a significant impact on the final results of in-context learning. However, it appears that the current experimental settings do not account for this influential factor. Please correct me if I have misunderstood any aspects. Moreover, I am curious if increasing the model size of transformers can alleviate the forgetting phenomenon.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The paper addresses a significant problem by conducting extensive experiments that aim to provide an explanation for the in-context learning capabilities of large language models (LLMs) from a Bayesian inference perspective. This research direction holds promise for advancing our understanding of LLMs.\n2. The paper is well-organized and effectively communicates its ideas. The authors present a natural and logically motivated experimental setting, which not only facilitates comprehension but also encourages further exploration and experimentation in the field.", "weaknesses": "1. The settings examined in this study predominantly rely on synthetic data, which creates a notable disparity between the experiments and real-world data. This limits the generalizability and applicability of the findings to realistic scenarios.\n2. Although the experiments are comprehensive and thought-provoking, there are certain definitions that would benefit from further clarification. Specifically, many explanations in the current version are contingent upon the capacity constraints of transformers. To enhance the illustration of the impact of transformer capacity, it is suggested that the authors include plots depicting the size of transformers in relation to the observed deviation phenomenon. This would provide a clearer understanding of the relationship between capacity and performance.\n3. As an empirical investigation paper, I would suggest that the authors consider providing the code to facilitate further investigation and replication of their findings.", "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": "", "extra_scores": {}}, {"review_id": "4KssggeauI", "reviewer": "Reviewer_Rt1y", "summary": "This paper empirically showed the relationship between trained transformers and Bayesian optimal estimators. They showed that in Gaussian Mixture Models(GMM) with equal rates and two components, the prediction of trained transformers follows the PME over the two components, which is kind of different to the PME over prior distribution being the distribution of either. They also showed that the prediction follows the PME of a strong baseline (when the PME is intractible) for linear regression problems like dense linear regression, sparse linear regression, signed vector linear regression or low-rank linear regression. \n\nThen they investigated the simplicity bias of ICL of transformers. They showed with a single task distribution, the trained TF (transformers) do not give any bias on the frequency of training data, but with multiple sources of tasks  distribution, they will tend to give more weight to the low-frequency --- which is natural, since the low-frequency data appears in more tasks than high-frequency ones. They then show the in-distribution and OOD generalization capacity of IC-trained TFs on linear regression on some polynomial features, and showed that they mimics the OLS on the subset of features that appears in the test prompts. \n\nThey also showed some interesting phenomenon which are novel and  worth investigating in the future. For example, they showed the deviation from the IC-trained TFs and the Bayesian optimal estimators, which well supplements the experiments of Garg;s paper and Akyureck's paper who both showed that TFs mimic the PME. Additionally, they showed a very interesting 'forgetting' phenomenon.\n\n I will vote for acceptance for this paper. Based on the authors' reply, I raised my score to eight.", "questions": "Here, I have some questions, and as well as some personal suggestions. I will appreciate it if you can address my confusion.\n\n1. You have proposed a setting called HMICL, which you claim is different from the MICL setting proposed in [1]. I am wondering what is the difference between two settings (more formally, from the definitions in [2])? In HMICL, the task is sampled from a hierarchical distribution, which can still be viewed as 'a specific distribution', right? So to my understanding, you did not propose a 'new setting' but instead, you are considering the exact MICL setting where the task distribution is hierarchical, right? I am also wondering why you think this hierarchical structure of task distributions matters in reality? Is there any evidence showing that the linguistic data or some other real-world data follows this hierarchical distribution?\n\n2. In PME paragraph in page 3, you claim that 'the predictions of the model can be computed using the posterior mean estimator (PME) from Bayesian statistics.' I think this is not that rigorous, since this happens only when the function class you are considering (here, it's the TF class) are expressive enough so that the PME is included in this function class, right? (Namely, this only happens in the realizable case). When the model class is not large enough, the TFs can not express the PME and hence, can only 'approximate' it.\n\n3. In GMM experiments, it's good to see that TFs follows the PME over the prior of two mixed components. Two questions I have: (1): In the figure, you only show k <= 10, which is a under-determined case for a linear model, since the dimension you use is d=10. In the over-determined case, what is the PME (is it OLS) and how does the TFs behave?  (2). Can you also plot the curve for OLS for all context length? I am wondering what the difference between a TFs and OLS here. \n\n4. For the experiments on multiple linear models (DR, SR, SVR, Skew-DR,etc), even though the PME is not analytically tractible, it is still possible to numerically compute the PME. Have you tried numerical computation for these PME and compare the TF to them? I think this will be a strong evidence to say whether the TFs are really mimicking the Bayesian optimal estimators. Although this could be hard in high-dimension case (I am not an expert in Bayesian but I guess so), it should be doable in a d=5 or d=10 case?\n\nAnother question is, how do you determine the regularization coefficient of LASSO? Does this require more data (if so, that mean you are actually using a longer context to do LASSO than to do ICL using TFs.)\n\n5. For the experiments on Frourier series and second-degree polynomials, an important question is that, how did you input the data? Did you simply encode original x_i s into token matrices, or you use \\phi(x_i) (where \\phi is a Frourier basis or a second-degree polynomial) to serve as the input? \n\n6. I kind of feel that there are more recent works about the relationship between GD, ICL, OLS and distribution shift [2,3,4]. I am wondering whether their results can somehow explain some of your experiments or to some extent are related to your results?\n\n[1]. What Can Transformers Learn In-Context? A Case Study of Simple Function Classes?\n[2]. Trained Transformers Learn Linear Models In-Context.\n[3]. A Closer Look at In-Context Learning under Distribution Shifts\n[4]. Transformers learn to implement preconditioned gradient descent for in-context learning", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. They have sufficient experiments and the results show some very interesting phenomenon. In some linear and non-linear tasks, the TFs are competitive with the PME or some strong baselines, which showed that TFs are very efficient to learn some algorithms on different linear/non-linear regression tasks. \n\n2. I totally agree with the implication the authors proposed in page 5 that TFs do not need to first recognize the underlying tasks and then solve it  in order to approximate the Bayesian optimal solutions. The authors frankly show that in some cases, the prediction of TFs mimics the PME while in some case the prediction can deviate from that.\n\n3. For  the simplicity bias and the generalization parts, the authors gave many baseline for comparasion and showed that TFs could potentially do something very non-trivial. The experimental evidence is sufficient to convince me that the generalization capacity for OOD tasks increases when the task diversity improves.", "weaknesses": "See 'Questions'.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)