Bayes' Power for Explaining In-Context Learning Generalizations

## Abstract

Abstract Traditionally, neural network training has been primarily viewed as an approximation of maximum likelihood estimation (MLE).
This interpretation originated in a time when training for multiple epochs on small datasets was common and performance was data bound; but it falls short in the era of large-scale single-epoch trainings ushered in by large self-supervised setups, like language models.
In this new setup, performance is compute-bound, but data is readily available.
As models became more powerful, in-context learning (ICL), i.e., learning in a single forward-pass based on the context, emerged as one of the dominant paradigms.
In this paper, we argue that a more useful interpretation of neural network behavior in this era is as an approximation of the true posterior, as defined by the data-generating process.
We demonstrate this interpretations' power for ICL and its usefulness to predict generalizations to previously unseen tasks. We show how models become robust in-context learners by effectively composing knowledge from their training data.
We illustrate this with experiments that reveal surprising generalizations, all explicable through the exact posterior.
Finally, we show the inherent constraints of the generalization capabilities of posteriors and the limitations of neural networks in approximating these posteriors.

## 1 Introduction

Neural network training is traditionally seen as parameter fitting for maximum likelihood estimation (MLE) or maximum a posteriori (MAP) estimation.
One performs multiple epochs of training on a single (small) dataset to get closer to the MLE or to the MAP with regularization.

Nowadays, neural networks are commonly trained in a single-epoch setting though, as the community has moved towards larger or infinite data sources (Brown et al., ; Oquab et al., ; Hollmann et al., ).
In these settings, performance is not bounded by available data but compute.
The training happens on unseen data sampled from an underlying distribution in each step, thus it is not very useful to think about it as approximating MLE for a dataset which is largely unseen.
This paper advocates for understanding neural network training and specifically ICL in this setting as an approximation of the true posterior distribution instead (Müller et al., ; Xie et al., ).
We experimentally show that this perspective predicts the generalization behavior of trained neural networks in ICL settings accurately.
Interpreting neural network training as posterior approximation alone cannot fully account for all behaviors of neural networks, given that the architecture, due to its inductive biases, affects how these networks approximate the posterior distribution.
As neural networks improve in approximating the posterior with more scale, these biases become less important though if the posterior is representable by the neural network.

In this paper, we train small neural networks with small budgets on infinite artificial data:
We can see that the true posterior, shaped by the training data, closely matches the network's behavior in most scenarios.
The implications of whether the true data posterior can account for much of a neural network's behavior are significant for the debate over whether large language models (LLMs) are merely pattern matching or reasoning machines.
Under the posterior approximation interpretation, LLMs generalize to new data by constructing a posterior that integrates all observed data, contrasting with the concept of stochastic parrots (Bender et al., ).
The combinatorial power has limitations, though, which we also detail.
These limitations pertain to both the quality of the posterior being subpar for out-of-distribution predictions, as well as the approximation not being tight everywhere, as we detail in this work.

Contributions We show the implications of interpreting ICL as a posterior approximation (Müller et al., ; Xie et al., ) for explaining generalization to unseen data.

- •
We present examples of unexpected out-of-distribution generalizations on simple analytical tasks.
- •
We show how these generalizations can be explained and are even expected through the Bayesian interpretation.
- •
We detail the limits of generalizations in the posterior and the limitations to the posterior approximation quality of neural networks.

## 2 Neural Network Training as Posterior Approximation

Training a neural network $q_{\theta}$ on a dataset $D$ is traditionally viewed as finding an approximation to the parameters $\theta$ that maximize the likelihood of the datasets, $\operatorname*{arg\,max}_{\theta}q_{\theta}(D|\theta)$. This is called MLE.

For the nowadays typical single-epoch training (Brown et al., ; Hollmann et al., ) it is less intuitive to use MLE as an interpretation, however, as large parts of the dataset might not be seen during training.
Thus, we are rather working with the data distribution directly.
The more natural interpretation which we advocate for is as an approximation to fitting the data generation distribution (Goodfellow et al., , Section 5.5).
We approximate

$$ $\displaystyle\operatorname*{arg\,min}_{\theta}\mathbb{E}_{x,y\sim p(X,Y)}[- \log q_{\theta}(y|x)],$ (1) $$

where $p(X,Y)$ is the data distribution from which the examples in $D$ are sampled.

The idea for the posterior approximation interpretation, can now simply be found by rewriting Equation 1 as a Kullback-Leibler (KL) divergence minimization with respect to the posterior $p(y|x)$

$$ $\displaystyle\operatorname*{arg\,min}_{\theta}\mathbb{E}_{x,y\sim p(X,Y)}[- \log q_{\theta}(y|x)]$ (2) $\displaystyle=\operatorname*{arg\,min}_{\theta}\mathbb{E}_{x\sim p(X)}[\mathbb {E}_{y\sim p(Y|X)}[-\log q_{\theta}(y|x)]]$ (3) $\displaystyle=\operatorname*{arg\,min}_{\theta}\mathbb{E}_{x\sim p(X)}[\mathbb {E}_{y\sim p(Y|X)}[\log p(y|x)-\log q_{\theta}(y|x)]]$ ( $\log p(y|x)$ is indep. of $\theta$ ) (4) $\displaystyle=\operatorname*{arg\,min}_{\theta}\mathbb{E}_{x\sim p(X)}[D_{ \text{KL}}(p(y|x)||q_{\theta}(y|x))].$ (KL divergence definition) (5) $$

Thus, by minimizing the cross-entropy loss in Equation 1, we implicitly seek a model distribution $q_{\theta}(y|x)$ that closely approximates the true posterior distribution $p(y|x)$ on data supported by $p(x)$.
The exact posterior is recovered with infinite training data, a neural network that can model the distribution $p(y|x)$ and a training that finds the optimum, as previously derived by Müller et al. ()(^1^11They focused on in-context learning settings, but the proof can be applied to other input modalities, too.).

In an ICL setting, our input $x$ is a concatenation of a training set $D_{train}=\{({\bm{x}}_{i},y_{i})\}_{i\in\{1,\dots,n\}}$ and a query input ${\bm{x}}_{query}$ for which we are predicting the output $y=y_{query}$.
In this paper, we consider ICL problems where the dataset is sampled by first sampling a latent $l\sim p(l)$ and sampling each example in the dataset based on it $(x,y)\sim p(x,y|l)$.
This construction is commonplace in Bayesian modelling and might even be seen in language, where a particular person starts to write a document with a particular intent.
Thus, the posterior distribution, also called posterior predictive distribution (PPD) in this setting can now be written as

$$ $\displaystyle p(y_{query}|{\bm{x}}_{query},D_{train})=\int p(y|{\bm{x}}_{query },l)p(l|D_{train})dl.$ (6) $$

In the above representation one can already see the mixture over different latents $l$, that makes the predictions of our models interesting.

The optimization from Equation 1 becomes

$$ $\operatorname*{arg\,min}_{\theta}\mathbb{E}_{({\bm{x}}_{query},y_{query})\cup D _{train}\sim p(\mathcal{D})}[-\text{log}\ q_{\theta}(y_{query}|{\bm{x}}_{query },D_{train})]$ (7) $$

and matches $p(y_{query}|{\bm{x}}_{query},D_{train})$ in the limit (Müller et al., ).

## 3 In-context Learning with Priors Over Finite Sets of Latents

In this work, we look at multiple priors and their posteriors.
To allow easily computing the posterior for our priors, we focus on priors with a finite set of latents.
That means that the considered priors have latents $l$ that are drawn from a large but finite set $\mathbf{L}$ with equal probability, if not specified otherwise.
We built an easily extendable framework to efficiently compute the posterior $p(l|D_{train})$ and the PPD $p(y|{\bm{x}}_{query},D_{train})$ on CPU and GPU to allow further research in understanding in-context learning through the Bayesian lens.
In the discrete setting, the PPD, as defined by Equation 6, simplifies to a weighted sum

$$ $\displaystyle p(y|{\bm{x}}_{query},D_{train})=\sum_{l\in\mathbf{L}}p(y|{\bm{x} }_{query},l)p(l|D_{train}).$ (8) $$

For each setup, we define a distribution $p({\bm{x}},y|l)=p(y|{\bm{x}},l)p({\bm{x}}|l)$ over examples given a latent.
If not specified otherwise, $p({\bm{x}}|l)=\mathrm{U}(0,1)$ is the uniform distribution and we define $p(y|{\bm{x}},l)$ via a deterministic mapping $f:X\to Y$ with output noise $y=\mathcal{N}(f({\bm{x}}),0.1^{2})$.

For all experiments, we use a transformer (Vaswani et al., ) adapted for ICL, introduced by Müller et al. () and called a Prior-data Fitted Network (PFN).
PFNs are particularly well suited for ICL thanks to their permutation invariance with respect to the set of context/training as well as query examples.
For each experiment, we report the mean of the predictions and provide a notebook to reproduce it easily(^2^22 For each experiment, we performed a grid search for the best final training loss. We searched across $4$ and $8$ layers, batch sizes $32$ and $64$, Adam learning rates $0.0001$, $0.0003$ and $0.001$, embedding sizes $128$, $256$ and $512$, as well as $100\,000$, $200\,000$ and $400\,000$ steps. The sizes for our training sets were uniformly sampled between $1$ and $100$. Code can be found at [https://github.com/SamuelGabriel/BayesGeneralizations](https://github.com/SamuelGabriel/BayesGeneralizations) .)

Figure: Figure 1: The model is only trained on step functions (left), still it learns to make smooth predictions (right) just like the true posterior for the step function prior.
Refer to caption: x1.png

## 4 Generalizations Explainable By Posterior Approximation

In the following, we will show three examples of interesting generalizations of neural networks when performing ICL, each of which is explainable by the ground truth posterior.

### 4.1 Training on Step Functions Yields Smooth Predictions But Not Everything Representable

In our first experiment, we show that even when a model was not trained on a smooth function, only on step functions, it becomes a smooth predictor as it approximates the posterior.

During training, we sample step functions, see Figure 1 (left), that start on different heights $\Delta y$, have different step positions $\Delta x$ and step sizes $h$. We can define the set of functions in our latent set as

$$ $\displaystyle\mathbf{L}^{(1)}=\left\{f^{(1)}_{\Delta x,\Delta y,h}:\begin{ aligned} \Delta x\in\{-1.,-0.98,\ldots,1\},\\ \Delta y\in\{-1.,-0.98,\ldots,1\},\\ h\in\{0,0.02,\ldots,2\}\\ \end{aligned}\right\},$ (9) $\displaystyle\text{where }f^{(1)}_{\Delta x,\Delta y,h}(x)=\begin{cases}\Delta y &\text{if }x<\Delta x\\ \Delta y+h&\text{else}.\end{cases}$ (10) $$

The predictions in Figure 1 (right) are smooth, as there are multiple step functions that might have produced the line, and the PPD now averages all of these step functions as in Equation 6.
Thus, even though our model was trained solely on non-smooth step functions, its predictions are (approximately) smooth as they are an average of many step functions.

Figure: Figure 2: Training a model on sine curves of a single amplitude, frequency and different offsets (left), the model does not only learn to model these curves (center), but also models the posterior for a sine that has a wavelength of $2$, instead of $3$. The posterior is flat, as the model is very uncertain about the offset $\Delta x$ of this curve.
Refer to caption: x2.png

### 4.2 Training on Sine Curves Can Yield Flat Line Predictions

Similar to the above experiment, we observe that training on sine curves with different offsets can generalize to an unseen function, a flat line, if the data comes from a sine curve that has a different frequency from the sine curves observed during training.

Our prior is defined as follows

$$ $\displaystyle\mathbf{L}^{(2)}=\left\{f^{(2)}_{\Delta x}:\Delta x\in\{0,2\pi/10 0\ldots,2\pi\}\right\},$ (11) $\displaystyle\text{where }f^{(2)}_{\Delta x}(x)=0.2sin(3\pi x+\Delta x)).$ (12) $$

In Figure 2 (right) we show that the model closely matches the ground truth posterior for a sine curve $sin(2\pi x)$ that is outside the training distribution.
This again is explained by a plethora of latents being averaged in the predictions of the model.

Figure: Figure 3: We train a model on two distinct classes of functions, sines and sloped lines, only (left). It not only learns fit both function classes well (center), but also learns to model slightly sloped sines, when prompted with a data from a sloped sine.
Refer to caption: x3.png

### 4.3 Training on Sloped Lines and Flat Sines Teaches Predicting Sloped Sines

We now show that models can even mix two different classes of functions.
We show that models trained on sine curves at different heights and on lines at different heights and slopes, generalize to sine curves that have a slope.

This was previously thought not to be possible, as elaborated by Yadlowsky et al. () for example.
It turns out that this is not only possible, but to be expected according to the posterior approximation interpretation.

Our latents are a uniform mixture of the previous prior $\mathbf{L}^{(2)}$, and a prior over sloped lines

$$ $\displaystyle\mathbf{L}^{(3)}=\left\{f^{(3)}_{\Delta x,m}:\begin{aligned} \Delta y\in\{-1.,-0.98,\ldots,1\},\\ m\in\{-1.,-0.98,\ldots,1\},\\ \end{aligned}\right\}$ (13) $\displaystyle\text{where }f^{(3)}_{\Delta x,m}(x)=mx+\Delta y$ (14) $$

In Figure 3, we can see that the generalization to a slightly sloped sine curve is possible, and that it even is in agreement with the true PPD.
Here, the PPD is a linear mixture of sine functions and lines and can thus represent a wave function that has a slope.

For this prior, we performed an ablation that considers different training times to see whether models tend to converge to the Bayes optimal prediction as training goes on.
In Figure 4 we see that the approximation of the posterior becomes better as we perform more training steps, but seems to suffer diminishing returns as we increase the number of training steps.

Figure: Figure 4: We see that the approximations of the true posterior become better with more training steps and a lower cross-entropy loss, like we expect for a powerful model as outlined in Section 2. The losses are negative, as we are in a regression setting, where the density can be above $1$.
Refer to caption: x4.png

## 5 Limitations of the Posterior

While we show interesting generalizations that are explainable via posterior approximation above, in this section we will focus on the limitations of generalizations supported by the posterior itself.
That is, in which scenarios can't we expect the posterior to yield useful and intuitive predictions.

Figure: Figure 5: While the model could make the optimal prediction (left) using a posterior mixing just the two latents in the center, it predicts differently as the latents in the center have a low likelihood to have generated the data.
Refer to caption: x5.png

### 5.1 Being representable is necessary but not sufficient

A combination of latents that perfectly fits the data does not guarantee the posterior uses this or an equivalent combination. We demonstrate this by extending the step prior $\mathbf{L}^{(1)}$ from Section 4.1 to include negative steps, i.e., $h\in\{-1,-.98,\dots,1\}$.

The prior consists of simple step functions, but combining two yields a mean modeling a step up followed by a step down (Figure 5, left).
This can be achieved by many posterior distributions, such as one assigning $50\%$ probability to each of the latents in Figure 5 (middle).
Thus, it could theoretically fit the data perfectly. Figure 5 (right) illustrates the actual posterior and its approximation, though.
The prediction is much flatter than it should be.
This is because the latents that would yield the right distribution do not get enough probability mass in the posterior, as they predict parts of the data very poorly (in our example the left- and right-side).
Thus, while there is a combination of latents that optimally approximates the function, the posterior does not choose that combination; this does not even change when providing more data (as this would also lead to more data points being modelled poorly by each of the latents that could be mixed to fit the function optimally).

Figure: Figure 6: We evaluate a Bayesian model's performance in estimating coin bias, using latent options $p=0.3$ and $p=0.6$. The posterior deteriorates with more data as the likelihood for $p=0.6$ dominates. Results are averaged over all potential coin toss outcomes.

### 5.2 Bayesian Models With Misspecified Priors Become Exponentially Worse

Bayesian ensembles tend to concentrate on a smaller subset of the prior as one conditions on more data.
While this is beneficial if the to-be-fitted function has support in the prior, it is detrimental if it does not.
We say a function has support in the prior, if there is a latent that shares the mean prediction with the function.
In general, certain functions lie outside the support of any prior distribution. Even Gaussian processes lack support for specific functions, such as step functions.

In cases where the function does not have support in the prior, the posterior will still concentrate.
It will concentrate on a close-by latent in terms of KL-divergence (Burt et al., ), but a necessarily wrong latent.
And as the likelihood for a latent has a multiplicative form ($\prod_{i}p(x_{i},y_{i})$), it becomes exponentially more confident in the wrong latent as more data is provided.

In Figure 6, we demonstrate this effect analytically on a simple coin flip prior.
Here, we use two discrete latents: either the coin has a head probability of $p=0.3$ or $p=0.6$.
The coin flips we condition on are actually fair ($p=0.5$), though.
On the left, we show that while the posterior on average is close to the actual true probability after observing only a few coin tosses, it diverges towards the closer latent ($p=0.6$) as more coin tosses are observed.
On the right, we see this is due to the likelihood of $p=0.6$ growing exponentially compared to $p=0.3$.

We can further see this effect for more interesting priors, too. In Figure 11 of the appendix, we show that when we feed more data to our prior that combines sine curves and lines (Section 4.3), we see the same effect of converging to a wrong solution while starting from a good prediction.
Finally, in Figure 12 of the appendix we show this effect even happens for a Gaussian process, when modelling a step function.
The Gaussian process becomes increasingly worse at modelling the step as more ground truth data is provided.

## 6 Limitations to the Posterior Approximation Interpretation

While the posterior approximation interpretation is useful to predict generalizations in many cases, it has limitations which we showcase experimentally in the following section.
We first look at the unreliable behavior of neural networks when we are outside of the support set of the prior, then we look at cases where we are inside the support but with low probability.
Finally, we discuss model limitations that might hamper the ability of models to approximate the posterior.

### 6.1 Approximating an unknown distribution, yields unknown outcomes

The intuition of posterior approximation breaks down when the input data $x$ has no support ($p(x)=0$) in the training distribution (extrapolation).
The posterior $p(y|x)=p(x,y)/p(x)$ is simply not defined in this setting.
In this section we show that neural networks tend to behave unpredictably in this setting.

In Figure 7 we provide evidence that neural networks are less reliable in extrapolation settings.
Here we train a standard MLP on a balanced binary classification problem: classify 0 0s and $1$s, i.e., the examples all have the form $(x=0,y=0)$ or $(x=1,y=1)$.
While the models are all able to predict well on the support set {0 0, $1$}, in between these values, where we are outside the support, each model behaves differently.
As soon as we add a bit of Gaussian noise (standard deviation of $0.1$) to the inputs during training, though, the behavior of our models becomes much more predictable on the $(0,1)$ interval(^3^33All models are 3-layer MLPs with ReLU activation and hidden size $64$. They were trained with Adam Kingma & Ba (), a learning rate of $0.001$, and a batch size of $1024$. The experiment can be reproduced using the notebook [https://github.com/SamuelGabriel/BayesGeneralizations/blob/main/Tiny_MLP_Generalization.ipynb](https://github.com/SamuelGabriel/BayesGeneralizations/blob/main/Tiny_MLP_Generalization.ipynb). .)

Figure: Figure 7: Predictions on the $[0,1]$ interval of 100 MLPs (x-axis) trained with different seeds, sorted by their prediction at $0.5$. The noiseless training data (right) results in less reliable outcomes, with models exhibiting greater variations in their predictions based on their seed.
Refer to caption: extracted/5896040/figures/one_dim_generalize_pred_figure/one_dim_generalize_pred_dist-000.png

We posit that extrapolation ($p(x)=0$) is typically unreliable in neural networks and should thus be avoided, e.g., by augmenting training data to ensure a well-defined posterior.
The definition of support is only precise up to invariances and equivariances of the neural network though, e.g., a translation-invariant model will behave predictably when translating inputs, even outside the support set.

### 6.2 The threshold of support

While for the exact posterior the support set is well defined, for the neural network-based approximation there is no apparent difference between a likelihood very close to zero and exactly zero, both are typically not in the training data.
It might have included a very similar one, though.

We can see, however, that prediction quality decreases with the data likelihood.
We showcase this in Figure 8 for the prior from Section 5.1, where this phenomenon is easy to see.
For the high likelihood data, the approximation still is close to the true posterior, but as we move the lines further apart, thus making our dataset less likely, the neural network falls back to a nearest-neighbor-based prediction.

While in the limit, we expect to have correct predictions wherever we have support, we see that in practice the predictions deteriorate before that.

Figure: Figure 8: This illustrative example of PFN behavior with decreasing data likelihood shows that the model aligns with the posterior for well-supported datasets but reverts to shortcut predictions with less likely data. Above each plot we provide the data likelihood according to the prior described in Section 5.1.
Refer to caption: x6.png

### 6.3 Architectural Limitations

Neural networks are constrained by the limited complexity of functions they can model due to finite computational resources, thereby unable to accurately represent certain posterior distributions.
Consequently, predictions do not converge to the true posterior despite minimizing the loss function.

We illustrate a recognized limitation in architectures akin to PFNs, which we employ: encoder-only transformers lacking positional embeddings fail to count repeated inputs (Barbero et al., ; Yehudai et al., ).
This is due to the nature of self-attention involving weighted averages in a permutation-invariant fashion.
Thus, a posterior that involves counting the number of identical inputs will not be well approximated by our models.

To demonstrate, we construct a simplistic prior focused solely on estimating the probability of a coin landing heads, devoid of any additional features. Coins are sampled each with a distinct head probability $p\in\{0.01,0.02,\dots,0.99\}$ for different datasets, and data samples are generated by flipping the coin with the outcome assigned as the target $y$.
The setup is illustrated at the top of Figure 9 (left).
The posterior corresponding to this prior must incorporate the decrease in uncertainty as additional samples are observed. Solely heads samples were inputted, as depicted in Figure 9 (left, bottom), to estimate the probability of obtaining heads on a subsequent coin flip. Despite the influx of identical samples, the neural network’s prediction remains unchanged, as shown in Figure 9 (right), irrespective of the actual variation in the posterior.

The inductive biases inherent in neural network architectures manifest less critically in specific phenomena, such as the consistent errors observed across different models in our experiments on modeling sloped sines.
These models consistently overestimate the sine magnitude relative to the true posterior as the slope increases. The predictions of the eight best models are presented in Figure 10 in the appendix.

Investigating the precise modeling constraints and inductive biases of contemporary neural networks constitutes an active research domain (Bhattamishra et al., ; Weiss et al., ), potentially augmenting the probabilistic insights obtained from the posterior.

Figure: Figure 9: On the left (top), we outline our prior, sampling a probability $p$ for heads (green) and generating samples by coin flips. At test time (left, bottom), we condition on varying counts of coins displaying heads. On the right, we can see that the transformer-based PFN model is not able to approximate the posterior, as it would need to count the number of examples, which all are identical, in the context set.

## 7 Related Work

While interpreting neural network training as posterior approximation in ICL was established in previous work (Xie et al., ; Müller et al., ), this study investigates its explanatory power in single-epoch ICL setups and identifies its limitations.

Yadlowsky et al. () previously discussed the generalization behavior of ICL models. They found that transformers do not have particular inductive biases making them especially strong on out-of-distribution ICL. What we have added to their analyses is that data alone, defining a posterior, enables generalization to new input-output mappings.

Finally, there is a line of work using the interpretation of neural network training as posterior approximation, namely PFNs (Müller et al., ; Hollmann et al., ; Müller et al., ; Rakotoarison et al., ; Adriaensen et al., ), which are trained on data sampled from a prior as an approximation to the posterior.
PFNs are a form of amortized inference, where the PPD (mapping from training data and test example to test output) is learned directly.
Thus, they heavily rely on the posterior approximation interpretation.

## 8 Conclusion & Future Work

In this paper, we demonstrate that transformers trained for in-context learning (ICL) can achieve notable out-of-distribution generalizations, that are unlike a nearest-neighbor matching, and involve the composition of training examples.
We found that these generalizations can be explained by the posterior that the transformer learns to approximate implicitly in many cases.
We did also find, though, that these generalizations exhibit clear limitations and may not function intuitively in all scenarios.
If the data-generating function does not have support in the prior distribution, model predictions tend to converge to incorrect solutions.
Moreover, when conditioning on data that lies outside the support of the prior, neural network behavior becomes less reliable and predictable.

Open future work are the following points.
i) Develop a practical definition of support in neural network training, considering that most data within the support is not sampled during training.
ii) Understand what types of priors models prefer to approximate. This is interesting because the data sampled from the prior typically does not identify the prior.
iii) Finally, we would like to learn if this interpretation can be useful for understanding more about the behavior of language models in general sequence settings.

## Appendix A Appendix

Figure: Figure 10: We show the predictions of the 8 strongest models from our grid search. We see that all models tend to make similar mistakes on this task, which is likely due to a bias in the architecture.
Refer to caption: x7.png

Figure: Figure 11: The prior combining sines and lines introduced in Section 4.3 deteriorates in accuracy as more data points are conditioned upon, as detailed in Section 5.2. Notably, the neural network approximation makes an erroneous approximation here and still predicts a sloped sine. This model was trained on dataset with up to 100 examples, thus this is not due to length generalization.
Refer to caption: x8.png

Figure: Figure 12: We show a simple Gaussian Process (GP) conditioned on 10 to 400 linearly spaced context points of a step function. The GP, with an RBF kernel (lengthscale $0.4$, outputscale $1.0$), a constant mean function, and Gaussian noise ($\sigma=0.1$), mostly models the step within its $95\%$ confidence interval at first but increasingly becomes over-confident in an over smooth function as more points are added. Reproduce this experiment with our notebook [https://github.com/SamuelGabriel/BayesGeneralizations/blob/main/GP_fitting_a_step.ipynb](https://github.com/SamuelGabriel/BayesGeneralizations/blob/main/GP_fitting_a_step.ipynb).
Refer to caption: x9.png

## Peer Reviews

confidence: 5
contribution: 1
initial_rating: 3
presentation: 4
questions: N/A
soundness: 2
strengths: The paper offers examples that illustrate how sequence models generate posterior predictions on latent factors unseen at training time. I thought the experiments were thoughtfully designed to highlight the basic Bayesian interpretation. While widely known in other contexts, the paper's discussion on the limitations of the posterior approximation interpretation is useful and adds color to the authors' narrative. I found the discussions on how PFN-style attention cannot capture repetitive patterns is illuminating. 

Overall, the authors' illustrations are informative and interesting and suggest immediate research questions.
summary: Going beyond the usual MLE interpretation, this work takes a Bayesian interpretation of in-context learning. The authors highlight that autoregressive probabilities in a sequence model are posterior predictives. They demonstrate how sequence models extrapolate to unseen latent factors (e.g., predicting smooth outputs from step functions). The authors note  sequence models may generalize poorly when the data is not supported in the model’s prior.
weaknesses: 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.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: While the central claim of the paper is interesting, I already believe in this conclusion based on Raventos et al and do not currently see what additional evidence this paper provides.

Things that might change my views:

- A more thorough comparison to related work including Garg et al and Raventos et al, and clarifying the additional contributions made by this paper. What do step functions tell me that linear regression does not?
- In this paper it is somewhat implicit that the population loss for the training objective does have (6) as its global minima, something that is e.g. close to proven in Raventos et al. Is it possible to argue for this theoretically? Arguably then this setting of step functions is more “minimal” than linear regression.

Some more technical notes:

- Could you more carefully define the “posterior ground truth” and “NN posterior approx” in Fig 1-5? I think “NN posterior approx” just means the output of the transformer and “posterior ground truth” just means the true function, in which case it is unnecessarily confusing terminology.
- Also the authors use “prior samples” to mean samples from the training distribution. While it might be appropriate to think of the transformer as doing in-context Bayesian inference with the training samples as the “prior” and outputs as the PPD, this is an empirical hypothesis to be proven, not something that should be labelling plots a priori (in my opinion, and if I am understanding the setting correctly).
soundness: 2
strengths: - The study of step functions as a simple class of functions for in-context learning is novel, as far as I know.
summary: The authors point out that a common informal way of thinking about neural network training, namely that training is approximately searching for the maximum likelihood estimate of the parameter given the data, is not appropriate in a setting like large language model training where models are trained in a single epoch. They advocate instead focusing on approximation of the “posterior” by which they mean the true distribution of data from which the “infinite” stream of examples is drawn.

In order to explore this hypothesis they train a series of transformers on in-context learning (ICL) tasks, where (x,y) pairs are presented in context y = f_l(x) + epsilon for some latent l fixed in each context but varying between contexts and normally distributed noise epsilon. This is a common testbed for understanding the nature of ICL in recent literature.

They show that transformers trained in such a setting, where the f_l are either step functions or sin functions, appear to learn to predict according to a kind of Bayesian posterior predictive distribution, which marginalises out the latent l. This is a hypothesis for ICL that has been explored extensively in recent years, in works including Xie et al “An explanation of in-context learning as implicit Bayesian inference” and Raventos et al “Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression” NeurIPS 2023 and the overall situation is (as far as I know) still unclear.
weaknesses: - 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”.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Just making sure, the posterior ground truth you are plotting is the posterior predictive mean? i.e. $E_{\text{true posterior predictive}} [y_{query}\mid x_{query}, D]$. And these were analytically solved?
soundness: 3
strengths: 1. The experiments demonstrating the interpretative power of ICL are clear and compelling, with ideas presented in a straightforward and illustrative manner.

2. The authors also examine inherent constraints that lead to common failures, such as analyzing the X shift (out of support).
summary: This paper studies ICL as doing bayesian inference by approximating the true posterior defined by the data-generating process. The authors propose function classes that make this Bayesian interpretation of ICL evident, with experiments showcasing surprising generalizations that can be fully explained through the exact posterior. Additionally, the authors highlight the inherent limitations of posterior-based generalization, as well as the architectural constraints of neural networks.
weaknesses: 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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: It would help to include a few sentences describing the PFN.

One of the paper's main points is the posterior is a weighted mixture of functions in the prior, where the weights are the posterior on the latent l. It would be useful to see the posterior on l, especially in fig 3 and 5. 

The asymptotic posteriors as in figs 11-12 are helpful in connection with my comment above about the full learning trajectory. It would be useful to see the same thing for fig 5.

The idea that generalization performance might be nonmonotonic (i.e., best for some intermediate number of examples) is interesting. Can this happen for language tasks? That seems like it would be a significant result, assuming it hasn't been demonstrated before.

Fig 10: how is strongest defined?
soundness: 3
strengths: The paper advances an interesting ideal-observer perspective on ICL, that we can understand generalization just by looking at the posterior implied by the data (without reference to the NN architecture).

More specifically, it points out that out of distribution testing leads to posteriors that are mixtures of functions in the prior.
summary: Recent work has argued that ICL can be interpreted as Bayesian inference, where the set of examples are taken from a randomly sampled function, and the goal is to produce the posterior predictive (PPD) for the test example. This paper shows how the PPD interpretation can help to understand generalization behavior in ICL. A series of experiments train a modified transformer on some prior distribution of functions and test its out-of-distribution generalization, comparing its output to both the generating function and the exact posterior.
weaknesses: 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

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: * Section 2: Would anything be lost if you removed paragraphs 2 and 3, "For the nowadays..." and "The idea for the posterior approximation..."?
* Section 3: "We built an easily extendable framework to efficiently compute the posterior $p(l|D_{train})$" Where do you do this...? I thought your framework is for computing the PPD via Eqn (7) which actually bypasses the posterior entirely. 
* Section 3: "we define $p(y|x,l)$ via a deterministic mapping $f: X \to Y$ with output noise $y=N(f(x),0.1^2)$". Where did the dependence on $l$ go?
soundness: 3
strengths: * This work leans on the strengths of Muller et al., 2022 which showed that Transformers can be trained using an ICL-objective to give us a good approximation of the Bayes PPD. Because there is a rich literature of the generalization behavior of the Bayes PPD, we can explain the ICL behavior of transformers trained for ICL by appealing to this classic theory. 
* There's a good discussion of the limitations of their conjecture
summary: This paper studies transformers that are adapted for ICL in the sense of transformers trained a la Muller et al., 2022. This involves a meta-training type of objective that leads to Eqn (7). The majority of the hard results center on trying different classes of latents $L$ and deterministic mappings $f: X \to Y$ which go onto define the meta-training datasets in Eqn (7). Transformers trained in this way are seen to generalize to unseen tasks and this can be explained because the transformers trained in this way are trying to approximate the Bayes posterior predictive density (PPD).
weaknesses: * 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?

Meta-review: This paper proposes a Bayesian interpretation of in-context learning (ICL) in large language models, arguing that neural network training can be viewed as an approximation of the true posterior defined by the data-generating process, rather than maximum likelihood estimation. The reviewers praise the paper's interesting and well-designed experiments, which demonstrate the interpretative power of ICL and highlight the limitations of posterior-based generalization, but criticize the paper for lacking depth and clarity in its contributions, with some reviewers feeling that the paper reads like a report on toy simulations or a lecture note surveying well-known facts, and that the authors do not sufficiently address related work or provide substantive insights into how sequence models perform implicit Bayesian inference. Overall, all reviewers lean towards rejection and the authors have not provided any rebuttal. We would still like to encourage the authors to resubmit an improved version of the paper in the future.

Decision: Reject