Bayes' Power for Explaining In-Context Learning Generalizations

• I found the second and third paragraphs in Section 2 really out of place and to contain confusing terminology. For instance $p(y|x)$ is repeatedly referred to as the true posterior distribution. Is this a typo?
• The description of the transformer adapted for ICL feels inadequate. There is a disconnect between presenting Eqn (7) and what the transformer has to do with it. At present, the paper is not intelligible without first reading Muller et al., 2022
• I did not find the related work very adequate. There are in total three paragraphs in Section 7, the last one talks about PFNs. The first one is one sentence describin Xie et al 2022 and Muller et al 2022 as "interpreting neural network training as posterior approximation in ICL". The second paragraph mentions one paper Yadlowsky et al. 2023 as discussing the "generalization behavior of ICL models". What is an ICL model?

The question in practice is not how well the network approximates the PPD, but how well it learns the true generating function. The paper gives examples of successful generalization where the posterior (and the network) can be close to the generating function even when it's outside the training distribution (fig 3), and unsuccessful generalization where the posterior and the network are far from the generating function (fig 2). But then there are cases where the posterior is far from the generating function yet the network is close to it, because the network does not approximate the posterior well (fig 8). In the end one has to question how useful the paper's core PPD thesis is for explaining ICL generalization.

Sec 5.2 shows that, as the number of examples grows from zero, the posterior starts at the prior mean, moves toward true function, but then converges to the KL-closest member of the prior support. So it may pass near the true function but only transiently. I worry Examples like fig 3 are cherry-picked for a nice intermediate point on this trajectory (i.e., a sweet spot in number of examples). The paper would be stronger if it focused on the full trajectory.

Fig 2: wavelength should be wave number

• I am confused by the use of the term “posterior” in this paper. For someone with a background in Bayesian statistics the posterior is a distribution over model parameters, and I would refer to the p(x,y) in this paper as the true distribution and p(y|x) as the true conditional distribution, rather than the posterior. I think the authors are following the terminology of Xie et al, in which marginalising out the latent vector produces a kind of posterior predictive distribution (thinking of the latents as the parameters for a kind of “in-context model”) but this does not mean p(x,y) is the posterior.
• The paper does not situate itself well enough in the literature, including Raventos et al as already mentioned, and also Garg et al “What Can Transformers Learn In-Context? A Case Study of Simple Function Classes”.

1. Previous literature has raised questions regarding the Bayesian nature of ICL. For example, Raventós et al. showed that transformers pre-trained on data with low task diversity struggle to learn new tasks and identified a threshold beyond which ICL emerges. Numerous studies suggest that phase transitions occur with respect to both the diversity of the training data (the cardinality of $L$ here) and the context sequence length (the number of context tokens used both during training and inference). Additionally, some work has observed a simplicity bias, where neural networks tend to “prioritize” learning simpler patterns first. It appears that the authors have not sufficiently addressed these factors in relation to the Bayesian interpretation of ICL.

2. This paper only focuses on the finite discrete prior case. There is extensive literature studying the case of continuous priors (e.g. regression case). For example, in bayesian linear regression one might draw the weights $w\sim N(0,\sigma^2$). It is unclear how this study’s insights extend to such cases -- while it's possible to interpret the $w$'s seen in pre-training as the $L$ set, it does seem quite unnatural and contrived to interpret the unseen $w$'s as some mixture of the seen $w$'s. Additionally, the model architecture requires the data to be exchangeable, which complicates generalization to Markov settings or more complex language tasks.

3. Extending from the common Bayesian linear regression case, it is also a well-known challenge in the field that problem dimensionality correlates closely with the difficulty of training a model to “learn to learn” (ICL). For instance, even Bayesian linear regression becomes challenging when $w$ has dimension on the order of hundreds. In more complex tasks like language modeling, where embedding dimensions often range from hundreds to thousands, it is unclear how useful these insights are relative to the impact of engineering choices and training techniques.

Allan Ravent´os, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression, 2023. URL https://arxiv. org/abs/2306.15063.

The paper provides interesting toy examples that suggest natural research questions to explore, but does not address any of them in depth. I truly appreciate the visualizations and the discussion in the submission. However, the paper currently reads as if it is a well-written report on toy simulations at the beginning of a research project, or a lecture note surveying well-known facts about sequence models and open research questions. It is unclear to me what the authors' main contributions are, if it is to be peer-reviewed as a conference proceeding.

Having said this, I see ample opportunities for the submission to grow into a fully fleshed paper.

First, I suggest authors study the relationship between the capacity of the model and its extrapolation behavior in-depth. A sufficiently flexible model can exhibit any arbitrary smoothing behavior, and the inductive bias built into the model will govern its extrapolation behavior. It is thus unclear to me why the generalization/smoothing behavior is explained by the Bayesian posterior approximation view alone. I see a couple of different ways to proceed. On the empirical side, one can build scaling law on how model capacity impacts ICL extrapolation behavior could be a valuable contribution to the literature. On the other hand, the authors could also analyze extrapolation behavior theoretically, and provide insights as to how transformers introduce particular behavior on unseen latents.

A Bayesian model provides a way to combine the observed data points to infer the latent structure (l in the paper's notation). To my understanding, the current paper does not provide substantive insights on how sequence models (e.g., PFNs) perform this synthesis process in an implicit manner. The phenomenological discussion on the posterior predictive distributions are true tautologically but in my opinion, does not add much to our understanding of the robustness of ICL (or lack thereof). Furthermore, as the authors also observe, Bayesian models with misspecified priors can lead to incorrect predictions, and increasing data can worsen performance under such conditions. It is obvious that sequence models will suffer the same fate, but a detailed understanding of how their behavior under misspecification will be valuable.

Another research question of interest is a careful study of different sequence modeling architectures. The authors solely focus on PFNs but as the authors also note, this particular modification to the attention masking mechanism suffers limitations. My understanding is that there are alternative masking approaches such as Nguyen et al. or the more recent work by Ye et al., as well as recent works such as Sun et al. that propose new state space modeling architectures (see below for proper citations). The impacts of different modifications to attention or SSMs appears poorly understood in the literature, warranting further analysis.

Finally, the external validity of the observations in the paper are not clear and seem to apply specifically to the particular experimental setting studied. Even with resources available in typical academic labs, for ICLR I would expect a careful study of scaling behavior up to 100M-1B of parameters across different environments.

As a minor feedback, I see some gaps in the authors' discussion of the literature. As the paper notes, the interpretation of autoregressive models as learning posterior predictives is well-established. More broadly, the meta-learning literature has long taken this view of sequence modeling.

In addition to Xie et al. and Muller et al. cited in the submission, several authors have contributed to this growing literature, e.g., see Jeon et al. or Nguyen and Grover below. Going beyond this observation, a number of research groups have explicitly drawn connections between autoregressive generation and posterior inference, e.g., see Zhang et al. Ye et al., and Falck et al..

• Jeon, H. J., Lee, J. D., Lei, Q., & Van Roy, B. (2024). An information-theoretic analysis of in-context learning. arXiv preprint arXiv:2401.15530.
• Nguyen, T., & Grover, A. (2022). Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. arXiv preprint arXiv:2207.04179.
• Zhang, L., McCoy, R. T., Sumers, T. R., Zhu, J. Q., & Griffiths, T. L. (2023). Deep de Finetti: Recovering topic distributions from large language models. arXiv preprint arXiv:2312.14226.
• Ye, N., Yang, H., Siah, A., & Namkoong, H. (2024). Pre-training and in-context learning is bayesian inference a la de finetti. arXiv preprint arXiv:2408.03307.
• Falck, F., Wang, Z., & Holmes, C. (2024). Is In-Context Learning in Large Language Models Bayesian? A Martingale Perspective. arXiv preprint arXiv:2406.00793.
• Sun, Y., Li, X., Dalal, K., Hsu, C., Koyejo, S., Guestrin, C., ... & Chen, X. (2023). Learning to (learn at test time). arXiv preprint arXiv:2310.13807.