Paper: In-Context Learning through the Bayesian Prism (ICLR-2024)

Reviewer comment (Novelty):
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 the 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?

Author response:
**Q1. HMICL and MICL**: While the reviewer is correct in a formal sense, we believe it’s clarifying to use our naming. Roughly speaking, the distinction between MICL and HMICL is reminiscent of the distinction between the usual supervised learning and meta-learning. As an example, in the MICL setup, consider linear regression. Here all tasks are instances of linear regression and each weight vector $w$ defines a task. Now if we consider a mixture of two “meta-tasks” (what we call function classes in the paper), say linear regression and decision trees, then 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). The formal definition in [2] which allows general hypothesis classes can similarly model HMICL. 

Thus, we agree with the reviewer that on this interpretation, MICL is equivalent to HMICL. However, this view is too coarse-grained for our purposes. What is of interest, both from the application and theory perspectives, is a more fine-grained view about whether and how models learn to perform these different “meta-tasks”. And the hierarchical structure is central for even talking about this. Previous work, e.g.  [1, 2], considers MICL with only a single “meta-task” and is thus not suitable for the type of analyses we perform. Compared to MICL, HMICL is arguably closer to the ICL for LLMs which can perform a vast variety of “meta-tasks.” 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 [5] and FLAN-T5 [6] (with a caveat for the latter being true for fine-tuning and not pre-training) or training of multilingual models [7] which involve pre-training corpora in different languages are some examples of the hierarchical nature of the training distribution in real-world LMs. 

We welcome further feedback and any suggestions for better naming.

Reviewer comment (Experiments):
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 an 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.

Author response:
**Q3. GMM plots with OLS for over-determined region**: We thank the reviewer for suggesting these experiments. We provide the answers to the two questions below:

(1) Figure 21 in Appendix plots the errors for the transformer model from Figure 1, 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 $\approx 0$.

(2) As shown in Figure 21, transformer’s errors are smaller than OLS errors and agree with the PME(GMM) errors for all context lengths. This shows that Transformer is indeed simulating the mixture PME and not OLS.

Reviewer comment (Experiments):
For the experiments on multiple linear models (DR, SR, SVR, Skew-DR, etc.), even though the PME is not analytically tractable, 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 strong evidence to say whether the TFs are really mimicking the Bayesian optimal estimators. Although this could be hard in the 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.

Author response:
**Q4.1. Sampling based approximations of PME**: Thank you for asking this question. We should first clarify that for Dense Regression and Skewed-Covariance Regression, we can compute Bayesian predictors exactly and we are able to show transformers for these two problems perform very close to the Bayesian predictor (Figures 6 and 9 in Appendix). However, for other linear inverse problems where PME seems intractable, based on the reviewer’s suggestion, we have tried Markov Chain Monte Carlo sampling method (No U-Turn Sampler (NUTS)) to approximate the PME. Particularly for Low-Rank and Sign-Vector Regression problems, for small problem dimensions i.e. 16 and 8 respectively, we were able to approximate the PME. We provide the plots in Figure 11 in Appendix and TF can be seen to follow NUTS-based PME in both cases. In fact, transformers seem to perform slightly better than the NUTS solution, which we believe could be due to the errors in the approximations due to sampling. We had less success in approximating the PME for sparse regression, which we believe might be due to the complex nature of the prior which involves a mixture of discrete (to select which features to keep) and continuous distributions (for the values of selected features). However, the alignment between transformers and exact PMEs for Dense and Skewed-covariance regression as well as numerically approximated PMEs for Low-Rank and Sign-Vector regression do provide strong evidence of transformers behaving like Bayesian predictor during ICL for linear inverse problems.

Reviewer comment (Experiments):
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 did you use \phi(x_i) (where \phi is a Frourier basis or a second-degree polynomial) to serve as the input?

Author response:
**Q5. Feature Encoding for non-linear regression tasks**: Thanks for this question. We just provide $x_i$ as the inputs, so the model needs to figure out $\Phi$ on its own.

Reviewer comment (Experiments):
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.

Author response:
**W1. About synthetic setting**: We agree with the reviewer that our work is mostly concerned with synthetic datasets to study how in-context learning works in transformers. The reason for choosing the stylised setup for our investigation stems from the fact that the real-world datasets used to train LLMs are difficult to model theoretically. Our work builds upon a long line of prior works, e.g., [1, 2, 3], that make simplifying assumptions on the pre-training data distribution to theoretically understand in-context learning in transformers. Our hierarchical meta ICL setup is also arguably closer to in-context learning in practical LLMs which can realize different classes of tasks (sentiment analysis, QA, summarization etc.) depending upon the inputs provided. Our work provides evidence on a diverse set of problems that the behavior of transformers during ICL agree with the Bayesian predictor. How well this holds for LLMs in practice is an open question, but we believe our work takes a step towards answering this question.

Reviewer comment (Experiments):
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.

Author response:
**W2. Effect of Model Size on Deviations**: As per reviewer’s suggestion, we conducted experiments with transformers of varying sizes. For the forgetting phenomenon, bigger models for larger task diversities have larger OOD losses as the training progresses ⇒ more forgetting. For smaller task diversities, there is no clear trend, however the forgetting phenomenon continues to hold. Hence, the deviations are not alleviated by increasing model size. (Figure 40, Appendix).

Reviewer comment (Reproducibility):
As an empirical investigation paper, I would suggest that the authors consider providing the code to facilitate further investigation and replication of their findings.

Author response:
**W3. Code Availability**: We will be releasing our code publicly upon the publication of our work. For now, we provide the anonymous link to the code repository for reviewer's reference: https://anonymous.4open.science/r/icl-bayesian-prism/README.md.

Reviewer comment (Experiments):
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.

Author response:
**Q1. Impact of order of examples**: In our initial experiments, we did evaluate the influence of order of examples on the ICL performance. In particular, for dense linear regression, we found that the ICL performance is fairly robust to the order of examples. We have added Figure 10 in Appendix which shows a box plot for loss@k values at different prompt lengths, and for each prompt length we consider 20 permutations. As can be seen the standard deviation of the errors is close to zero across prompt lengths. Hence, the order of the in-context examples do not impact the ICL performance here. We did expect such behavior to hold in our setup, since we deal with continuous functions and things like majority label bias and recency bias [4], which are a source of variance in performance due to the order of examples for NLP tasks, do not apply here.

Reviewer comment (Experiments):
I am curious if increasing the model size of transformers can alleviate the forgetting phenomenon.

Author response:
**Q2. Effect of Model Size on Forgetting**: As mentioned in response to W2, increasing the model size does not alleviate forgetting. Instead, it leads to more forgetting for larger task diversities, as the model fits the pretraining distribution better over the course of training and hence has higher OOD losses. (Figure 40, Appendix).

Reviewer comment (Presentation):
The figures are very hard to understand, making it challenging to summarize the conclusion from the figures.

Author response:
**Readability**: We will take the reviewer’s comments into account and will add more descriptive captions so that the figures are easier to understand. About references to the appendix, we believe that given the contents and the scope of our work, it would have been difficult to fit all the technical details in the main paper given the page limit. However, the details in the appendix are not primary for understanding the core theme and contributions of our work. However, we will take this feedback into account and move some of the details to the main paper. About references to “some specific” paper, we are not sure about which paper the reviewer is referring to here. If it is Akyurek et al. 2022 [1], we cite their work while describing the weight recovery method in page 3, since we use their method for our experiments, but we do provide all the details about how the method works after citing their work. In case they are referring to Raventos et al. (2023) [2], while we do provide all the important details from their work relevant to our experimental setup and findings, we will add some more details from their work in Section 6 to improve readability.

Reviewer comment (Novelty):
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.

Author response:
**Readability**: We will take the reviewer’s comments into account and will add more descriptive captions so that the figures are easier to understand. About references to the appendix, we believe that given the contents and the scope of our work, it would have been difficult to fit all the technical details in the main paper given the page limit. However, the details in the appendix are not primary for understanding the core theme and contributions of our work. However, we will take this feedback into account and move some of the details to the main paper. About references to “some specific” paper, we are not sure about which paper the reviewer is referring to here. If it is Akyurek et al. 2022 [1], we cite their work while describing the weight recovery method in page 3, since we use their method for our experiments, but we do provide all the details about how the method works after citing their work. In case they are referring to Raventos et al. (2023) [2], while we do provide all the important details from their work relevant to our experimental setup and findings, we will add some more details from their work in Section 6 to improve readability.

Reviewer comment (Writing):
The "Bayesian predictor" is not properly explained in the context of "in-context learning".

Author response:
**Technical definitions**:
- **Bayesian Predictor in context of ICL** : Informally, we have a probability distribution on tuples $(x, f(x))$ induced by the priors on $x$ and $f$. This in turn, induces a distribution on the prompts, namely $(x_1, f(x_1)), (x_2, f(x_2)), …, (x_i, f(x_i)), x_{i+1}$. Conditioned on the prompt we would like to sample $f(x_{i+1})$---this is the posterior predictive distribution (PPD). We use the mean of the PPD as the Bayesian predictor since the nature of our setup requires a point estimate. The formula for PME in the paper formalizes this. We provide the formal definition of Bayesian predictor in context of ICL in Section 2 (under the PME paragraph) and discuss the relationship between ICL and Bayesian inference in the Bayesian predictor paragraph in the introduction Section.

Reviewer comment (Writing):
The "deviations from Bayesian prediction" are not clearly explained; the author should explain that clearly.

Author response:
**Technical definitions**:
- **Deviations from Bayesian prediction**: “Deviations” here refer to predictions of transformers during ICL that do not match with the Bayesian predictor given the pre-training distribution and the observed data.

Reviewer comment (Writing):
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.

Author response:
**Technical definitions**:
- **Simplicity Bias**: Simplicity bias is by now a well-established term ML literature; for example, [3, 4, 5, 6].

Reviewer comment (Writing):
If you mean "generalization is out of distribution (OOD) detection", then just using OOD detection would reduce a lot of confusion.

Author response:
**Technical definitions**:
- **Generalization**: We agree that a more descriptive term should be used here. We will change it to "OOD Generalization" in the final version for clarity.

Reviewer comment (Theory):
The correctness of Equation 2 is questionable. What is the exact definition of "$df$", the differentiation with respect to a function $f$?

Author response:
**Differential of a function $df$**:  This is standard notation in Bayesian inference. We can view a function as a point in a higher-dimensional Euclidean space defined by the function parameters. E.g, for linear regression, this reduces to the weight vector $w$ and hence $df$ is simply $dw$.

Reviewer comment (Novelty):
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?

Author response:
**Why HMICL?**: Prior work on theoretically understanding ICL, considers the simpler case where functions are always sampled from a single function class (which we term as MICL). We present a more advanced setup where the functions can come from different classes. Further, our setup can serve as a testbed for investigating the origins of multi-task learning in large language models (LLMs) because HMICL is arguably closer to ICL in real-world LLMs, which can realize different classes of tasks (sentiment analysis, QA, summarization etc) depending upon the inputs provided. We show that the ability to solve a mixture of tasks arises naturally as a consequence of Bayesian prediction, which can be viewed as computing the posterior when the prior comes from a mixture of function classes. This challenges the existing intuition of multi-task learning where the model first identifies the task and then solves it.