Towards Better Understanding of In-Context Learning Ability from In-Context Uncertainty Quantification
Authors: Shang Liu
Zhongze Cai
Guanting Chen
Xiaocheng Li
Venue: Accepted by TMLR ()

Abstract: Predicting simple function classes has been widely used as a testbed for developing theory and understanding of the trained Transformer's in-context learning (ICL) ability. In this paper, we revisit the training of Transformers on linear regression tasks, and different from all the existing literature, we consider a bi-objective prediction task of predicting both the conditional expectation $\mathbb{E}[Y|X]$ and the conditional variance Var$(Y|X)$. This additional uncertainty quantification objective provides a handle to (i) better design out-of-distribution experiments to distinguish ICL from in-weight learning (IWL) and (ii) make a better separation between the algorithms with and without using the prior information of the training distribution. Theoretically, we show that the trained Transformer reaches near Bayes optimum, suggesting the usage of the information of the training distribution. Our method can be extended to other cases. Specifically, with the Transformer's context window $S$, we prove a generalization bound of $\tilde{\mathcal{O}}(\sqrt{\min\{S, T\}/(n T)})$ on $n$ tasks with sequences of length $T$, providing sharper analysis compared to previous results of $\tilde{\mathcal{O}}(\sqrt{1/n})$. Empirically, we illustrate that while the trained Transformer behaves as the Bayes-optimal solution as a natural consequence of supervised training in distribution, it does not necessarily perform a Bayesian inference when facing task shifts, in contrast to the \textit{equivalence} between these two proposed in many existing literature. We also demonstrate the trained Transformer's ICL ability over covariate shift and prompt-length shift and interpret them as a generalization over a meta distribution.

Published in Transactions on Machine Learning Research (05/2025)
Towards Better Understanding of In-Context Learning Ability from In-Context Uncertainty Quantification
Shang Liu Imperial College Business School Imperial College London
s.liu21 @imperiall.com
Zhongze Cai
2. 2.0a222@mmmriial..ccco
Imperial College Business School Imperial College London
Guanting Chen
guanting@unc.ccm edu
Department of Statistics and Operations Research University of North Carolina
Xiaocheng Li
riaocheng. Ii@imperial..c.. uk
Imperial College Business School Imperial College London
Reviewed on OpenReview: https: // opt///veenrrr...e/ foru?? id= Jutpbhheey
Abstract
Predicting simple function classes has been widely used as a testbed for developing theory and understanding of the trained Transformer's in-context learning (ICL) ability. In this paper, revisit the training of Transformers lincar regression tasks, and different from all the existing literature, we consider bi-objective prediction task of predicting both the conditional expectation [X] and the conditional variance Var(Y|X). This additional uncertainty quantification objective provides a handle to (i) better design ont-of-distribution experiments to distinguish ICL from in-weight learning (IWL) and (ii) make a better separation between the algorithms with and without using the prior information of the training distribution. Theoretically, we show that the trained Transformer reaches near Bayes optimum, suggesting the usage of the information of the training distribution. Our method can be extended to other cases. Specifically, with the Transformer's 's context window S, prove generalization bound of (S, analysis compared previous results of Ö(V1/n). Empirically, the trained Transformer behaves the Bayes-optimal solution natural consequence of supervised training in distribution, it does not necessarily perform a Bayesian inference when facing task shifts, in contrast to the equivalence between these two proposed in many existing literature. We also demonstrate the trained Transformer's ICL, ability over covariate shift and prompt-length shift and interpret them as a generalization over a meta distribution.
on n tasks with sequences of length T, providing sharper illustrate that while
1 Introduction
A particularly remarkable characteristic of Large Language Models (LLMs) is their ability to perform in-context learning (ICL) (Brown el al., 2020). Once pretrained on vast corpus of data, LLMs can solve newly encountered tasks when provided with just few training examples, without any updates to LLMs parameters. ICL has significantly advanced the technique known as prompt engineering (Ekin, 2023), which has achieved widespread success in various aspects daily life (Oppenlaender et al., 2023; Heston Khun, 2023; Li et al., 2023a). Behind the empirical success of ICT. this method has captured the attention of the

---

Published in Transactions on Machine Learning Research (05/2025)
theoretical machine learning community, leading to considerable efforts into understanding ICL from different theoretical perspectives (Xie et al., 2021; Akyürek et al., 2022; Von Oswald et al., 2023; Zhang et al., 2023a).
This work aims to enhance the theoretical understanding of ICL by examining the Transformer's context tradeoff. Although we obtain the results
window and showing its effects on the axproximation--stiaatie for the case of uncertainty quantification where the model is asked to predict both the mean value and the uncertainty of its prediction, our analysis is applicable across various ICL tasks and provides sharper bounds compared to previous works. In addition to developing theories, empirically demonstrate the effectivencss of Transformers to in-context predicting the mean and quantifying the variance of regression tasks. We design scries of futt-f-ddisribution (OOD) experiments, which have generated significant interest within the community (Garg et al. (2022); Raventós et al. (2024); Singh et al. (2024)). These experiments provide insights in designing the pretraining process and understanding the ICL capabilities of transformers.
Our contributions are as follows:
We theoretically analyze the problem of in-context. uncertainly quantification. We consider the case when Transformers only process the contexts within context window capacity S and derive generalization bound of O( (tt/min(SS,,Tii result can be easily extended to other cases under the assumption of almost surely bounded and Lipschitz loss functions. As far as we know, our generalization bound is the first of its kind and provides a tighter bound compared to the existing analyses (Li et al., 2023b; Zhang et al., 2023b) when S < T. In particular, we use the context- wintxxt-window structure to establish a Markov chain over the prompt sequence and construct an upper bound for its mixing time. We also examine the extra approximation error term due to a finite context. window S (Section B.2). Combining those discussions together, we quantify the convergence of the trained Transformer's s risk to the Bayes-oplinal risk. Moreover, note that all the theoretical results only show that the trained Transformer achieves a near-optimal in-distribution risk compared to that of the Bayes-optimal predictor. It is incorrect to draw (from the theory or the in-distribution numerical results) either of the conclusions that (i) the Transformer that achieves the near-optimal risk exhibits a similar structure as the Bayes-optimal predictor by performing Bayesian inference (Zhang et al., 2023b; Panwar al., 2023) or (ii) the Transformer performs as the Bayes-optimal predictor for dut-o--distribution tasks.
for pretraining over n tasks with sequences of length T (Theorem 3.2). Our
Numerically, for the uncertainty quantification problem, provide comprehensive study of the in-context learning ability of the trained Transformer under three scenarios of distribution shifts: task shift. (Section 4.1), covariate shift (Section 4.2), and prompt length shift (Section 4.3). We find that transformers are capable of in-context learning of both mcan and uncertainty predictions, even under a moderate amount of task distribution shift, provided that the task diversity the training data is relatively large. Additionally, we find that increasing the task diversity with a meta-learning approach helps the transformer learn in-context robustly under covariate shift. Lastly, we observe that removing positional encoding from the embedding vector massively helps the generalization ability, enabling it to better learn tasks in-context with unseen prompt length.
We defer more discussions on the related literature to Section A.
2 Problem Setup
Consider training & Transformer for some regression lask f: X' x from a function class F. The covariates X Rd are generated from distribution Px, and the output variable y ff(x) + for some function f F, noise level a, and some random noise with E[c] = 0 and Var(c) = 1. The Transformer performs a sequential prediction task over the following sequence
(o), Dh..8...., TT. UT)
where T is the total number (in-context) samples. For a Transformer model with parameters 0 E Q we denote it as TFg. At each time = 1, ..., T, the model TF, observes H, = (#1, Vi. Tt- 1t--111/-1130 11-1. 2;) (which is called history or prompt) and makes a bi-objective prediction of y to both predict the mean with go and quantify the uncertainty of the prediction with (H.H)). With slight abuse of notations, denote the
2

---

Published in Transactions on Machine Learning Research (05/2025)
output of the model by (H, == (fo(H.) (H)).. The pretraining dataset consists of n sample sequences
{
(0) (1)
(8)
10
1 TT. 9T
(1) FT
}
a i=1
D = 21 Vi 22 12
To generate each sample sequence in D, a function fi is sampled from a distribution Py supported on F and a noise level oj is sampled from a distribution Pa supported on [0,5] c R. Then each i; and yt is generated pairwise by
lid. 5t
Px, = So (zt +o, oo 4 Ltd. P.
where 67 ' s are ii.dd noise of mean zero and unit variance.
The Transformer is trained by minimizing the following empirical loss
@ERM = arg min
e TFo (H{ (H{²)..0 vv.))
(1)
dEG nT
1= t=1
where H²(0) (x) (
;1
"., ...., 2-11, 1t-1: and ((())),) (R x R*) R - R denotes the loss function. We
use (1) (3) (7, to demote the samples in the training dataset and It, yt, He to denole an arbitrary feature, label, and history. Throughout this paper, we assume that each probability distribution is contimuous and has a probability density function (p.d.f.), and we also assume the conditional distribution ye on observing He exists almost surely.
The loss function is accordingly defined by
(y 0)
((U, == =o
28
Definition 2.1 (Bayes-optimal predictor). The Bayes-optimal predictor under the distributions Pr, Px, Pa and P, is defined by
())
mig (G)
E (((((())))OH)))) u.)]
(2)
where Ge is the class of all measurable functions of H, E 74. The expectation is taken with respect to the following dynamics:
Px, & Pe, f ~ P7, Pa, y1- f(xx)) + to and H, = (f1, "..,,,,.......... D.).
The loss the right-hand of equation is the expectation of the empirical loss equation 1. With rich enough function class and an infinite amount of training samples, the trained Transformer TF Fggas converges to (yi of) as will be shown in Theorem 3.2.
2.1 Motivation for the uncertainty quantification objective
We first give semi-formal definition for in-context learning and in-weight learning of the bi-objective linear regression task considered in this paper.
In-context learning refers to that the Transformer gains the ability to learn from the in-context samples (samples in H₂) and behave as an algorithm. For example, the Transformer exhibits in-context learning when it behaves like a ridge regression model when performing the linear regression task. The in-context learning ability should persist under the out-of-distribution setting such as shifting Px, Pr, and P..
In- weigtttht learning, for the NLP tasks, generally refers to that the Transformer relies the information stored in its weights to make predictions, rather than the in-context samples. In this light, it memorizes the training samples and uses the memorization to make predictions. This memorization mechanism is more aligned with supervised learning setting; in particular, the mechanism is sensitive to distribution shifts and thus is not considered a desirable outcome of training Transformers.
Then the question is when we train the Transformer according to the objective equation 1, does it exhibit in-context learning in-weight learning? For the classic single-objective linear regression task, we note two
3

---

Published in Transactions on Machine Learning Research (05/2025)
facts: (a) the trained Transformer is near the Bayes-optimal predictor (Xie et al., 2021; Zhang et al., 2023b; Panwar et al., 2023); (b) under a Gaussian prior, the Bayes-optimal predictor is exactly a ridge regression model (with proper choice of the regularization parameter, (Wu et al., 2023)). With these two facts, it is tempting to draw the conclusion that the Transformer gains in-context learning ability and behaves as a ridge regression model, but this has to be done with caution or may even lead a wrong conclusion, given that the trained Transformer does not exhibit out-of-distribution ability for full generality as noted in the numerical experiments (Garg al., 2022). This motivates us to consider the bi-objective tasks. First, for the uncertainty quantification objective, when the number of in-context samples is fewer than the dimension of Z4 is, a near Bayes-optimal predictor must utilize the information in the training procedure, whereas there is no algorithm that can optimally do this via only in-context samples. Thus it well distinguishes the Bayes-optimal predictor from any possible algorithm and therefore gives a clearer picture of whether the trained Transformer indeed behaves as an algorithm. Second, the uncertainty quantification objective provides us an casy handle to designing numerical experiments, such as the flipped experiments (Wei et al., 2023; Singh al., 2024), distinguish in-conlext learning from in-weight learning. Third, the uncertainty quantification objective is of independent interest as it indicates whether the trained Transformer knows its uncertainty not. Furthermore, through this bi-objective task, we illustrate that while the trained Transformer is near the Bayes-optimal predictor, it does not necessarily perform Bayesian inference when making predictions, which contradicts the arguments in (Zhang et al., 2023b; Panwar et al., 2023; Jeon et al., 2024).
3 In-Context Learning when In-Distribution
In this section, we focus on the in-distribution property of the trained Transformer. We provide a finite-sample analysis of how trained Transformers reach near Bayes-oplimum. While our analysis is made on the case of uncertainty quantification, can be easily adapted to other loss functions such as mean squared error. To procoed, first provide the exact form of the Bayos-optimal prodictor defined in equation 2 for the mcan and uncertainty prediction.
Proposition 3.1 (Bayes-optimal predictor for mean and uncertainty prediction). The Bayes-optimal predictor of the step-wise population risk defined in equation 2 is given by
ye H.) = H.y, (H,, = B[(v - yi (H.)) = =((())))))))(((()))))))))))))))) oi (H.)) + B(o2)H,)
The optimal mean predictor shares the same form as the Bayes-optimal predictor for single-objective mean prediction task. The additional uncertainty prediction task does not change the nature of the mean prediction part. The two terms in the optimal uncertainty predictor can be interpreted follows. The first term is epistemic uncertainty, which indicates the uncertainty (of identifying the that governs the history He) due to lack of information. The term decreases as the samples accumulate, i.c., as the number of in-context samples increases. The second term is aleatoric uncertainty also known as intrinsic uncertainty.
Recall that the empirical risk estimator defined by equation 1. Now we define the population risk as
R(TFo) = ] T
1 R.. Σ (TFo vn)
where He is another sampled sequence that is independently and identically distributed as H( in the training data. We denote the population risk minimizer as 0":
0" E arg min (TFa))
(3)
8E0
Now we present our main thcoretical result.
Theorem 3.2. Let gERM denote the ERM estimator as defined in equation over the function class of the L-layer, M -hedss Transformer models. Suppose that at each time t, the Transformer has context window of making predictions based on xe and previous S pairs of (I ys) for s max(1, -L 1. Then under
4

---

Published in Transactions on Machine Learning Research (05/2025)
16
Bayes Optimal
Bayes Optimal
14
didge Regression
3.0
Ridge Regression
Linear Regression
Ave. of predicted uncertainty
Linear Regression
Mean prediction error
12
TF-4096
2.5
TF-4096
TF-63536
TF-65536
10
# in-context samples
# in-context sa mampls
(a) Mean prediction
(b) Uncertainty prediction
Figure 1: Transformer behaves close to the Bayes-optimal predictor for in-distribution tasks. Details of the distributions in data generation are given in Section G.1. The numbers 4096 and 65536 refer to the number tasks (configurations of 0.)) used in the training, which is formally defined in Section G.2. The Bayes-optimal predictor is stated in Proposition 3.1 and calculated analytically error in predicting yı, For the right panel, the y-axis gives the average of the predicted uncertainty test samples (average of ((H,) or o² (H.) on test samples). In particular, we note that ridge regression and linear the of
Section G.3. For the left panel, the y-axis gives the mean squared
all the
regression (ordinary least squares) do not naturally produce measurement of uncertainty, residuals on the in-context samples as their estimates of uncertainty. More visualizations are deferred to Section C.1.
some boundedness assumptions of the Transformer's porameters (Assumption B.5 and B.6), we have with probability at least 6,
R(TFgSRH)) R((Fa-) S <O
where Ô omits poly- logarithmie terms that depend n, T, 1/8 and boundedness parameters.
Proof sketch. First, prove that (a slightly redefined version of) the truncated history forms up Markov chain conditioned observing the full hidden information f(i) and σ()), and upper bound the mixing time by min(S, T} to enable the concentration arguments. Second, we prove that the loss function is almost surely bounded (Lemma E.3) in preparation for McDiarmid-type concentration inequalities (Lemma F.2, (Paulin, 2015)). Third, we show that the loss is almost surcly Lipschitz to control the difference between loss functions with respect to the change of the parameter (Lemina E.7). Fourth, we prove that there exist. distributions PýERR and over parameter space O satisfying number of properties as constructed in Lemma E.11. Lastly, we use standard PAC-Bayes arguments Paeas and i and conclude the proof. The detailed proofs are deferred to Section D.2.
Comparison with previous results. There are also other theoretical results that characterize the outcomes of the (pre-)training on Transformer models (Zhang et al., 2023a; Wu et al., 2023; Xie et al., 2021; Li et al., 2023b; Bai et al., 2024; Zhang et al., 2023b; Lin et al., 2023a). Our analysis differs from theirs in terms of both the conclusion and the techniques. One stream of results examines the property of the gradient flow (or gradient descent) over the loss function for linear regression problems. The exact quantification of the gradient flow entails a simplification of the Transformer's architecture to the case of single-layer attention mechanism under linear activation even simpler settings (Zhang et al., 2023a; Wu et al., 2023). While their analyses provide insights into the learning dynamics of Transformer models, the lcarning of the single-layer attention Transformer can be very different from multiple-layer Transformers (Olsson et al., 2022; Reddy, 2023). Another major line research uses statistical learning arguments (Xie et al., 2021; et al., 2023b; Bai et al., 2024; Zhang al., 2023b; Lin al., 2023a) such as algorithm stability, chaining, PAC-Bayes arguments. Bai al. (2024) focus on making predictions after observing a fixed length of variables under the i...d. setting (which is more aligned with the standard supervised learning setting), which differs from the

---

Published in Transactions on Machine Learning Research (05/2025)
more practical setting of making predictions at every position as in Theorem 3.2. Xie et al. (2021) prove the convergence between the Bayesian inference and the true underlying distribution rather than the trained model and the Bayesian inference. Lin et al. (2023a) consider a sequential decision-making problem and use covering arguments to derive generalization bounds, while their analysis does not adopt the concentration arguments inside each sequence, resulting in an T(TT/n) upper bound for the average regret. The most. related works ours are Li et al. (2023b); Zhang et al. (2023b). The major difference is that they all consider the only case of S T. Li et al. (2023b) use the algorithm stability arguments to give generalization bound over |R r| of order O(vI/(nT)). They prove the loss difference caused by perturbing one imput pair over history of length & is controlled by (11/t). Averaging those differences leads to a (log((T)TT) = 0 (T) inside each sequence (see their equation (15) in their Appendix C), which appears in the argument to prove that the loss per sequence is ((7--/2)-sub-Gauus -sub-Gausian.. However, in the case of S < T (which is more often the case in practice), the algorithm stability term is of O(1/S). Averaging these terms inside cach sequence lcads to a difference of order (1)))) If we stick to the original Azuma-Hocffiie arguments, the sum of squares of these terms is of O(T/S²), leading to a far worse sub-Gaussian norm of O(T'/25-1), ) resulting in a final generalization bound of order 0(T ''n 11/2) that is clearly suboptimal compared to our D(SS/nI]) Besides, such a bound also grows with T, which is undesirable. Similar to ours, Zhang et al. (2023b) also use a concentration argument for Markov chains. However, their Theorem 5.3 has two limitations: The first is that their result is of the order P(VTTinn((xTT) but they do not specify Tmin- Since they do not consider the truncated history but the full history, the Markov chain (which is not verified by them) will never mix inside each task sequence (see our discussions in Section D.2). Thus, the term Twin in their result is actually T, leading to an order of Ó(VI/n), which is suboptimal compared to our O(VS/(nT) when the context window S < T. The second limitation is that their error dccomposition is not light: their excessive risk bound (measured by the total variation distance between the distribution induced by å and that by 8*) has a term Da(PPruee Po-) - (V(eeree which means their result has an extra term of the approximation error since the Kullback-Leibe variation distance (Polyanskiy & Wu, 2024). Our work is the first theoretical analysis showing the effects of the context window S on the performance of the Transformer up to our knowledge. The construction of the truncated history serves two-fold: not only does the truncation fit the practical model of finite context window but it also gives an upper bound on the mixing time. Concentration inside each sequence makes it possible to analyze the training dynamics broader than fixed-length sequences and prove the convergence to near Bayes-optimum. The context window S also captures a novel dimension of the approximation--sting tradeoff in the Transformer model.
divergence is stronger than the total
Extension of Theorem 3.2 to other problems. We remark that the result and its derivation do not pertain to the uncertainty quantification setting, but hold for more general loss functions and are of independent interests. In particular, our analysis still holds as long the loss function is almost surely bounded and Lipschitz with respect to the change of parameter 0, as we can see from the proof sketch. We note here that to enable the Markov chain's concentration arguments, the almost surely bounded loss requirement cannot be relaxed to other tail properties such as sub-Gaussian (see the counter example in Theorem 4 of Fan et al. (2021)).
We defer discussions on the approximation error to Section B.2.
4 In-Contex: Learning under Distribution Shifts
In Section 2, we describe in-contex, learning ability as algorithm-like that predicts based on the learning from in-context samples, and such an ability should be generalizable to an out-of-distribution (OOD) environment. In this section, we differentiate the OOD scenarios into task shift, covariate shift and length shift, and examine the Transformer's in-context learning ability in each scenario. As far as we know, we provide the first comprchensive group of numerical experiments (for the lincar regression task) that demonstrates the Transformer's ability to handle these three types of distribution shifts. We provide preliminary theoretical discussions for such abilities and hope this points directions for future theoretical research.
6

---

Published in Transactions on Machine Learning Research (05/2025)
4.1 Task shift
When the trained Transformer performs well on the OOD data, it means that the Transformer gains algorithmic ability that learns to make predictions based on the in-context samples, because such an ability is not restricted to the distribution of the inputs. Comparatively, the mere observation that the Transformer works well on the in-distribution data does not demonstrate its in-context learning ability as a traditional supervised learning model also has such ability and generalization performance over in-distribution data.
In the previous section, when we show the in-distributton performance of the Transformer, the variance parameter o² is generated by the prior of the unverse-Gamin distribution o² Inv- Gow--amma(T,, 7) with parameters and F. The details of the other generation distributions are deferred to Section G.1. For the in-distribution setting, we set F = 20 which leads to a prior mean around 1. Now we consider three out-of-distribution (OOD) settings for the
S-OOD (small OOD): I = 80, f = 20. The prior mean of a is around 0.5.
M-OOD (medium OOD): I = 100, 400. The prior mean of is around 2.
L-OOD (large OOD): 100, F 1600. The prior mean of is around 4.
Bayes optimal (L-0OD)
Bayes-optime (L-0OD)
Bayes-optimal (M-OOD)
Bayes-optima (M-0OD)
Ave. of predicted uncertality nky
Rayes optimal (S 00D)
Ave. of predicted uncertainty nky
Bayes optima (S COD)
1F-4096 (L-OOD)
1F-65536 (L-COD)
TF-4096 (M-0OD)
TF-65536 (M-00D)
TF-4096 (S-OOD)
TF-65536 (S-00D)
20
60
80
100
0
20
60
80
100
# in-context samples
# In-context sa mples
(a) Transformer trained w/ small pool size
(b) Transformer trained w/ large pool size
Figure 2: OOD performances of Transformers and the Bayes-optimal predictor. The g-axis gives the average of the predicted uncertainty over all the test samples (average of à (HH) o (H,) test samples), and ideally, they should converge to the expected uncertainty level 0.5 (S-OOD), (M-OOD), and 4 (L-OOD) as in-contexi samples increase. There three OOD environments: small (S-OOD), medium (M-OOD), and large (I-OOD) that reflect the intensity of the OOD. Two versions of the Transformer model are trained with pool size of 4096 and 65536. The Transformers and the Bayes-optimal predictor are the same the ones Figure The only difference is that they are evaluated on OOD data here.
We make following observations based on Figure 2: First, the Bayes-optimal predictor predicts well. We note that the Bayes-optimal computed based on the in-distribution prior distribution (with respect to o²). Thus when the Bayes-optimal predictor tested under the OOD environment as in Figure 2, the prior used by the Bayes-optimal predictor is wrong. But we note from Figure 2 that the Bayes-optimal predictor has the OOD ability to correct the prior the in-context samples accumulate (noting that the three Bayes-optimal curves converging to the correct mean of 0.5, 2, and 4). This is also known as the washing out of priors in Bayesian statistics. Second, Transformers deviate from the Bayes-optimal on these OOD lasks. For both plots in Figure 2, we note that the predicted values from the Transformers deviate from those of the Bayes-optimal predictor when the OOD intensity is large. This tells that the trained Transformer does not conduct Bayesian inference under task shift. other words, is incorrect to conclude that the trained Transformer behaves the Bayes-optimal predictor just. from the matching in-distribution loss (as Figure 1). Moreover, the

---

Published in Transactions on Machine Learning Research (05/2025)
Transformer achieves near-optimal loss for in-distribution tasks (as Figure 1) but it does so via different avenue than the Bayes-optimal predictor (as Figure 2). This is in contrast with the findingscclamss in the previous papers (Zhang et al., 2023b; Panwar et al., 2023). Third, the deviation of the trained Transformer from the Baycs-optimal is smaller when the task diversity large or the OOD intensity is small. This is aligned with the findings (Raventós et al., 2024) for in-distributton performance, while the OOD setting is not studied therein.
The theoretical evidence only states that the trained Transformer has near-optimal in-distribution loss the Bayes-optimal predictor. But it does not give any evidence that these two have structural similarity that persists for OOD tasks. In particular, we note that the trained Transformer may take statistical shortcuts: When evaluated under in-distribution tasks or some simple task shifts (e.g. scaling the weights vectors or changing the signal-noise ratio), Zhang et al. (2023a); Wu et al. (2023) show that Transformer will construct shortcuts using the statistical property of the training distribution. More specifically, Transformers (can, and will) encode the information of the covariance matrix into their model parameters reach ncar-optimal in-distribution performance. Such statistical shortcuts are beneficial to the in-distribution performance but. can hurt its OOD ability. Increasing the training task diversity, such as larger training pool size, may remove some of these statistical shortcuts to obtain ncar-optimal empirical loss, and thus better enable its in-context learning ability.
We defer more discussions and visualizations on this OOD experiment to Section C.2.
4.2 Covariate shift
For all the numerical experiments so far, the covariates are generated from (00, Ia)- This follows the standard setup of the existing literature (Akyürek et al., 2022; Von Oswald et al., 2023; et al., 2023b; Raventós et al., 2024). Tt also noted from the literature (Garg et al., 2022; Zhang et al., 2023a) that the trained Transformer in this way lacks in-context learning ability under covariate shift. However, if the Transformers trained on different tasks (different fi's) can generalize to unseen tasks during the test phase, one may wonder if the Transformers trained on different covariales (different distributions of at 8) can also generalize. In this subsection, we give a positive answer to that question by proposing meta-training procedure that improves the trained Transformer's ability to handle covariate shifts. Specifically, we consider generating the covariates in the training data as follows:
For each training sequence (say, the i-th), we first sample a vector (Ai...., Aa) where each X3 is i.idd. Uniform(0, 2]. Then all the x( ( for = I..., T are sampled from (0, sense, the covariance matrix of x(())s 's is also random variable, and the x!!)s can be viewed being sampled in a hicrarchical manner from meta-distribation.
Aa))). In this
We examine the performance of such training procodure under four OOD test settings. In other words, the x(00s the test data is generated from the following four distributions where
Large covariance (I-cov): x.(0s s are sampled from V(0,4Ta).
Decreasing diagonal (Dec.): x (1) 's are sampled from (00 diag( diaag(dd/i]] [))
Shrinking diagonal (Shr.): x (i))ss are sampled from diag(d/ (
Rotation (Rot.): (6) 's are sampled from (00,
where U, is an orthogonal matrix
independently generated for each sequence.
Figure 3 gives the evaluation result under the 4 OOD settings. We nole that the mela-distribution used is still significantly different from the four OOD test environments. Thus the results show the effectiveness of the meta-training approach. We evaluate the prediction error of models ordinarily trained, and models trained by the vecaattrainng process. For both mean and uncertainty, the models trained by the meta- mota-training procedure have a smaller prediction error.
8

---

Published in Transactions on Machine Learning Research (05/2025)
Deviation from Bayes Optimal μ
Deviation from Bayes Optimal o
0.8
static_K X_ x_model_on Lo__eoo_oooo L-cov meta x model _n___-ovv
1.75
0.7
Er. Mean Prediction
1.50
Er. Uncet Prediction
0.6
1.25
0.5
1.00
0.4
0.75
0.3
0.50
0.2
0.1
0.25
20
40
60
80
100
20
40
60 sampl
8D
100
samppee es
0.9
static model
0.8
M
meta, x w_weel on Dec.
Err. Moan Prodiction
0.7
Ert Uncer. Prediction
0.25
0.6
MI WWw WM Mal Nhv
0.20
0.5
0.15
0.10
0.3
0.2
in-context samples
in-context samples
MM
static x. m_mdeel on_ SSr. model _o__oo_S Shr..
1.2
Etr. MMaa Prediction
1.0
FT. Uncer: Prediction
0.30
W M
Ahh oith mn
MM Yow M
0.8
yow
0.20
0.6
mil
0.4
0.10
0.2
20
40
60
80
100
20
60 in-context sample
40
80
100
text. sampl
static x mod_lllloooo on Rot.
1.2
0.5
mela_x x model on Rot.
Br. Mean Prodiction
1.0
Err. Uncor. Predicton
0.4
M
0.8
0.3
0.6
0.2
0.1
100
20
100
in-context samples
in-context samples
Figure 3: The errors of the and uncertainty prediction where the is measured by the absolute difference against the Bayes-optimal predictor. The static x model corresponds to models trained with the standard way in generating X,'s, while the meta__ model corresponds to the new approach of drawing X''s from the meta-training procedure. In all OOD settings, mela-trained models have better performance.

---

Published in Transactions on Machine Learning Research (05/2025)
4.3 Lenghh shift and positional embedding
Existing work (Dai et al., 2019; Anil et al., 2022; Zhang et al., 2023a) have pointed out the failure of Transformers to generalize to longer contexts than the ones they have seen during training. It is worth mentioning that the code implementations of some previous works (Zhang ct al., 2023a; Garg et al., 2022) are based on the "transformers" package of Hugging Face. Although these works have not included positional embedding explicitly, the GPT2 module imported from this package adds a built-in positional encoding implicitly. We suspect that some unxxpected behaviors (like the "unexpected spikes of prediction error" mentioned in Zhang et al. (2023a)) are due to that the built-in positional encoding is not disabled. In this subsection, we investigate the length generalization ability of the trained Transformer on the uncertainty quantification task. Specifically, we control the prompt lengths that the model is trained on. Previous experiments train the model on prompts with lengths (number of in-context samples) ranging from 1 to 100. In this experiment, we control the training prompts such that the lengths are either shorter than 44 or longer than 45 (the choice of 45 as the cutoff point is not essential). We specify these two configurations below.
Trained on I 44: the model is trained on prompts with length ranging from 1 to 44, and is evaluated with prompt length from to 100
Trained on 2 45: the model is trained on prompts with length ranging from 45 to 100, and is evaluated with prompt length from 1 to 100
We regard this difference in prompt length between training and testing as length shift. We evaluate the effect of removing positional encoding under this prompt length generalization task. If positional encoding is added to the embedding, samples at unseen positions will be associated with an unseen positional encoding vector in the embedding space. This requires the model to handle not just an unseen number of in-context samples, but also a possibly unseen embodding distribution, and generalization ability will likely deteriorate. As mentioned previously; the built-in positional encoding of GPT2 model use a positional encoding which is set to be (t, O,....... _0) for the t-th token, and the encoding will then be concatenated to the embedding vector. We validate the above intuitions with the following 4 training configurations.
No positional encoding (w/o Pos.): the model is trained without positional encoding.
Add positional encoding (w/ Pos.): the model is trained with GPT2's built-in positional encodings.
Add segment encoding (w/ ): the positional encoding is added with a random amount offset. For the i-th training sequence, a random offset & is first uniformly sampled from 10, 22}. Next, for each token in this prompt at position t, the positional encoding is set to tt.
Add full range encoding (w/ F-Pos.): similar to the S-Pos. configuration, the positional encoding is added with a random amount offset. But here the offset is uniformly sampled from 1.... 100).
For the model trained with the "w/o Pos." configuration, it is also tested without positional encodings. For the models trained with the rest configurations, they are all tested with the "w/ Pos." way of encoding.
The results are shown in Figure 4. The models in the left figure are trained on prompts shorter than 44, and the models in the right figure are trained on prompts longer than 45, We make the following observations. The pre-trained transformer in general can generalize to prompts with unseen length, under the condition of using/ reing/remooing the positional embedding properly. The "w/o Pos." curve in the left figure shows that even at positions larger than 44, the model can still produce predictions close to Bayes-optimal. Adding positional encoding hurts the generalization ability. From the "w/ Pos." curve in the left figure, find that the model's performance drops significantly at positions larger than 44. The main cause of the failure of length generalization is due to the distribution shift in the positional embedding space. As given in the "w/ S-Pos." and "w/ F-Pos." curves in the left figure, if the model has seen the positional encodings for a ocrtain position during training, then its performance at this position is significantly improved, even the corresponding prompt length is never seen. The length generalization ability is not unrestrictively strong, and such generalization ability for smaller lengths is generally weaker compared to that for larger lengths. The
10

---

Published in Transactions on Machine Learning Research (05/2025)
1.0
0.40
on
Trained
Trained on w/ Pos.
0.35
Trained Pos.
Trained on S-Pas,
Br.. of Uncer. Prediction
Trained on w/ F-Pos.
Ert. of Uncer: Prediction
0.30
0.25
0.6
0.20
0.4
0.15
0.10
0.2
0.05
100
40
# in-context samples
# in-context sa mples
(a) Generalization from long to short
(b) Generalization from short to long
Figure 4: The effect of removing positional encoding prompt length generalization. The y-axis records the average error of uncertainty prediction, which is the difference between the uncertainty predicted by the transformer and the Bayes-optimal estimator. (a) For models trained with prompt lengths < 44, the figure the left shows that positional encoding has the worst generalization capacity with a larger length, and removing positional encoding could effectively euhance the length generalization power. (b) For models trained with prompt lengths 45, removing positional encoding can help generalize to smaller lengths, although the generalization ability for smaller lengths is generally weaker compared to that for larger lengths.
right figure shows that even for the "w/o Pos." configuration, performance still degrades when the prompt length is shorter than 20.
Wu et al. (2023)'s Theorem 5.3 points out that under the case of the single-layer iinearatteention-nnly Transformer model on a linear regression task with Gaussian priors (without position encoding) if we train the model to only predict one single label after observing 7' context exemplars, the optimally trained model under T = T, also performs well at the case T = T2 (compared to the Bayes-optimal predictor for T = T₂) if Ti and T2 are close. In particular, such a property holds because the single-layer linear-attentin structure constructs "statistical shortcut" to achieve a near-optimal solution to the lincar regression problem (of which the Bayes optimal predictor is an optimally tuned ridge regression). The ridge regression's regularizing parameter is related to the sequence length and the signal-to-noie to-noise ratio (SNR). If SNR is fixed and T1 and To are close, the Bayes-optimal predictors for Ti and T2 are close, implying that the "statistical shortcut" learned at Tj still works in the case of T2. This result implics the possibility of context length gencralization by a simplified Transformer model due to shared structures in the attention matrices. The model of Wu el al. (2023) is also a simplified Transformer that skipped the position encoding part. Another point is that our empirical results also show that as the training and test context lengths get more different, the generalization becomes worse, which is of the same spirit as Wu et al. (2023).
5 Conclusion and Limitations
In this paper, we study the in-context learning ability of the trained Transformer through the lens of uncertainty quantification. In particular, train the Transformer for bi-objective task of mean prediction and uncertainty prediction. We develop new results both theoretically and numerically, The takeaway messages are: First, the Transformer can perform in-context uncertainty quantification. Second, the trained Transformer is only guaranteed to achieve a near-optimal in-distribution risk against the Bayes-optimal predictor. This does not imply that the Transformer behaves as the Bayes-optimal predictor either in n--distribution or out-of- distribution. Third, the major concern of restricting the window size has been computational due to the quadratic growth of the computational consumption with respect to the context length, while our theory shows that limiting context length/ hhigth/wwndww window sizes also benefit the generalization when the data size is limited.
11

---

Published in Transactions on Machine Learning Research (05/2025)
Finally, the Transformer has the in-context ability for out-of-distribution tasks, but this in-context ability is contingent on a proper training method, such as sufficient task diversity, meta-training for covariate shift, and effective removal of the positional encoding. Two important future directions are as follows. First, we believe our method for deriving the gencralization bound has implications for a scope much larger than uncertainty quantification and can be used to improve the existing bounds for various tasks using Transformers. Second, all the numerical experiments in the paper are conducted for the linear functions fi's. We believe the same results still hold for nonlinear functions well; and such results can further consolidate the in-context ability for uncertainty quantification of the Transformer.
Our work also has limitations both theoretically and numerically, First, although Theorem 3.2 gives a tighter generalization bound when the context window S is limited, it might not adapt to those LLMs that are of very large window sizes nowadays. However, might still be valuable for those situations where smaller Transformers suffice. For example, for online decision-making problems where the action space and the sequence length are much smaller than the natural language processing problems, smaller models such as GPT-1 structure have been applied to construct. Decision Transformer (Chen et. al., 2021). Second, our work criticizing the Bayesian explanation of ICL ability itself does not give positive explanation of the ICL dynamics; how ICL. works remains unclear for future study. Finally, our discussions are limited to the regression tasks, while whether and how Transformers perform ICL in other tasks is beyond this work's scope.
References
Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, Rajendra Acharya, et al. review of uncertainty quantification in deep learning: Techniques, applications and challenges. Information fusion, 76:243-200 2021.
Gustaf Ahdritz, Tian Qin, Nikhil Vyas, Boaz Barak, and Benjamin L Edelman. Distinguishing the knowable from the unknowable with language models. arXio preprinl arXiv:2402.0335 2024.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. orXiv preprint ar orx::2211.15661, 2022.
Com Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowyez, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. Advances in Neural Information Processing Systems, 04138546-35556 2022.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural informalion processing systems, 36, 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, al. Language models are few-shot learners. Advances in neurol information processing systems, 33:1877- 1901, 2020.
Michele Caprio, Souradeep Dutta, Kuk Jin Jang, Vivian Lin, Radoslav Ivanov, Oleg Sokolsky, and Insup Lec. Credal bayesian deep learning. arXiv preprint wxXi:2302.096556 2023.
Michele Caprio, Maryam Sultana, Eleni Elia, and Fabio Cuzzolin. Credal learning theory. arXiv preprint 0020072402.00057, 2024.
Chao Chen, Kai Liu, Ze Chen, Yi Gu, Yue Wu, Mingyuan Tao, Zhihang Fu, and Jieping Ye, Inside: Lims' internal states retain the power of hallucination detection. arXiv preprint NFf:2402.007444 2024.
12

---

Published in Transactions on Machine Learning Research (05/2025)
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 34:15084- 15097, 2021.
Herman Chernoff. measure of asymptotic efficiency for tests of a hypothesis based on the sum of obscrvations. The Annals of Mathematical Statistics, pp. 493-507, 1952.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoe V Le, and Ruslan Salakhutdinov. Transformer- xl: Attentive language models beyond fixed-length context. arXiv preprint arXio:1901. 000600 2019.
Monroe Donsker and SR. Srinivasa Varadhan. Asymptotic evaluation certain markov process expectations for large time. iv. Communications on pure and applied mathematics, 36(2):183-213 1983.
Sabit Ekin. Prompt engineering for chatgpt: quick guide to techniques, tips, and best practices. Authorea Preprints, 2023.
Fabian Falck, Ziyu Wang, and Christopher C Holmes. Are large language models baycsian? a martingale perspective on in-context learning. In ICLR 2024 Workshop on Secure and Trustworthy Large Language Models, 2024.
Jianqing Fan, Bai Jiang, and Qiang Sun. Hoeflding's inequality for general markov chains and its applications to statistical learning. Journal of Machine Learning Research, 22( 2021.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes, Advances in Neural Information Processing Systems, 11330583-35999
2022.
Jakob Gawlikowski, Cedrique Rovile Njieutcheu Tassi, Mohsin Ali, Jongseok Lee, Matthias Humt, Jianxiang Feng, Anna Kruspe, Rudolph Triebel, Peter Jung, Ribana Roscher, et al. A survey of uncertainty in deep neural networks. Artificial Intelligence Review, 56(Suppl 1):1513-1589, 2023.
Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? case study on learning with representations. arXiv preprint ar arX:::2310.10616 2023.
Thomas F Heston and Charya Khun. Prompt engineering in medical education. International Medical Education, 2(3):198-20) 1011198--0,.200 2023.
Hong Jun Jeon, Jason D Lee, Qi Lei, and Benjamin Van Roy. An information-theoreti analysis of in-context. learning. ar preprint 0rXiv:2401.15530 2024.
Lorenz Kuhn, Yarin Gal, and Sebastian Farquhar. Semantic uncertainty: Linguistic invariances for uncertainty estimation in natural language generation. arXiv preprint orXiv:2302.0966, 2023.
Beibin Li, Konstantina Mellou, Bo Zhang, Jeevan Pathuri, and Ishai Menache. Large language models for supply chain optimization. arXiv preprint ar arX1::2307.0387, 2023a.
Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms:: Generalization and stability in in-context learning. In International Conference on Machine Learning, pp. 19565-19599 PMLR, 2023b.
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. arXiv preprint arXiv:2310.0855 2023a.
Zhen Lin, Shubhendu Trivedi, and Jimeng Sun. Generating with confidence: Uncertainty quantification for black-box large language models. arXiv preprint arXiv:2305111881 2023b.
Potsawee Manakul, Adian Liusie, and Mark JF Gales. Selfeheckgpt: Zero-resource black-bor hallucination detection for generative large language models. arXic preprint arXi::2303.08896 2023.
13

---

Published in Transactions on Machine Learning Research (05/2025)
Colin McDiarmid et al. On the method of bounded differences. Surveys in combinatorics, 141(1):148-188, 1989.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXie preprint w:Xiv:220911180) 2022.
Jonas Oppenlaender, Rhema Linder, and Johanna Silvennoinen. Prompting ai art: An investigation into the creative skill of prompt engineering. arXiv preprint arXiv:2303. 13534, 2023.
Madhur Panwar, Kabir Ahuja, and Navin Goyal. In-context learning through the bayesian prism. In The Twelfth International Conference on Learning Representations, 2023.
Daniel Paulin. Concentration inequalities for Markov chains by Marton couplings and spectral methods. Electronic Journal of Probability, 20(nonc): - 32, 2015. doi: 10.1214/EJPP.20-4083. URL https:
//doi. wrg/ 1214/EJP...... +20-4039
Yury Polyanskiy and Yihong Wu. Information theory: From coding to learning. 2024.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAT T blog, 188):9, 2019.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in Neural Information Processing Systems, 36, 2024.
Gautam Reddy, The mechanistic basis of data dependence and abrupt learning in an in-context classification task. Xix preprint or arXiu:2312.03002 2023.
Yusuf Sale, Michele Caprio, and Eyke Höllermeier. Is the volume of a credal set a good measure for epistemic uncortainty? In Uncertainty Artificial Intelligence, pp. 1795-1804. PMLR, 2023.
Aaditya Singh, Stephanic Chan, Ted Moskovitz, Erin Grant, Andrew Saxe, and Felix Hill. The transient. nature of emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 36. 2024.
Aviv Slobodkin, Omer Goldman, Avi Caciularu, Ido Dagan, and Shauli Ravfogel. The curious case of hallucinatory (un) answerability: Finding truths in the hidden states of over-confident large language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp- 3607-3625, 2023.
Ralph C Smith. Uncertainty quantification: theory, implementation, and applications. SIAM, 2013.
Timothy John Sullivan. Indroduction to uncertainly quanlificntion, volume 63. Springer, 2015.
Alexander Tsigler and Peter L Bartlett. Benign overfitting in ridge regression. Journal of Machine Learning Research, 24(123):1-766 2023.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174. PMLR, 2023.
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently. arXio preprint ar WwX.::33330038/6 2023.
Lisa Wimmer, Yusuf Sale, Paul Hofman, Bernd Bischl, and Eyke Hüllermeier. Quantifying alcatoric and epistemic uncertainty in machine learning: Are conditional entropy and mutual information appropriate measures? Tn Uncertainty in artificial intelligence, pp. 2282 2292. PMLR, 2023.
14

---

Published in Transactions on Machine Learning Research (05/2025)
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? ar preprint or orXix:23100089 2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint (xxi::211.000000 2021.
Fengzhuo Zhang, Boyi Liu, Kaixin Wang, Vincent Tan, Zhuoran Yang, and Zhaoran Wang. Relational reasoning via set transformers: Provable efficiency and applications to marl. Advances in Neural Information Processing Systems, 35:35825-35838 2022.
Ruiqi Zhang, Speneer Frei, and Petor L Bartlett. Trained transformers learn linear models in-context. arXiv preprinl ar Xiv: 09927, 2023a.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. are preprint orXiv:2305.1942 2023b.
A Related Works
Thcorctical Understanding of In- In-Context Learning. There are two streams of rescarch the theoretical understanding of ICL: the first tries to give sharp approximation error bounds on different. tasks, while the second focuses on how the trained Transformer approaches the potential optimum. For the approximation error, following the pioneering empirical investigations on simple function classes (Garg et al., 2022), Von Oswald et al. (2023); Akyürek et al. (2022) conjecture that the Transformer is doing ICL via gradient descent, and verify it both empirically and theoretically. Based on the mechanism of layer- wise gradient descent construction, Bai ct al. (2024) show that Transformers are able to behave (approximately) as well as some well-known algorithms on some statistical problems. Some following works generalize the layer-wise gradient descent construction to other settings such as rocision-makingg (Lin et al., 2023a) and lincar regression under representations (Guo al., 2023). Apart from the layer-wise gradient descent, some other works consider the one-step gradient descent reached by a single-layer linear-activated Transformer (Zhang et al., 2023a; Wu et al., 2023) and curve the excessive population risk of the optimal model compared to oracle or the Bayes-optimal predictor. Ahn et al. (2024) give a set of global optima for some specific one-layer or two-layer attention-only models with lincar or ReLU activation. Aside from characterizing where the Transformer can reach, another group of works is making elforts towards understanding where the Transformer will reach. One typical way is to study the simplified attention-only Transformers. Zhang et al. (2023a) start the analysis of the training dynamies of the gradient flow over the population risk on the linear regression task and show that single-layer ineea--attnntion-nnl model converges to some specific sets with suitable initialization. Wu et al. (2023) keep the same spirit and give sample complexity bound based on a certain gradient descent scheme. For general Transformer models, technical tools from the statistical learning theory are applied. As for the task of predicting the next token in natural language tasks, Xie et al. (2021) provide a viewpoint from the Hidden Markov Model (HMM) and prove the asymptotic consistency under the regularity condition. Bai et al. (2024); Lin et al. (2023a) use chaining arguments with covering numbers for generalization, where Bai et. al. (2024) consider the training under fixed length and Lin el al. (2023a) consider the problem of sequential decision-making Li et al. (2023b) adopt algorithm stability arguments obtaining a bound of Õ(1/VnT). As is discussed in the main text (sce discussions after Theorem 3.2), their analysis will result in suboptimal Ö(S/VnT) for the case S < T. Zhang et al. (2023b) adopt similar concentration inequality for Markov chains to get a bound of ( ////mi//nni (nT)) Since they do not consider the limit of context window S, their derivation ends up with Tmix 2 T, which suboptimal compared to our case. In short, our paper is the first theoretical analysis on the limit of context window S and gets a tighter generalization bound than previous works on the generalization bound when S <.
Bayesian Behavior of In-context Learning. Due to the complex structure of transformers, showing the theoretical properties of ICL without proper assumptions are challenging. There has been growing interest. in
15

---

Published in Transactions on Machine Learning Research (05/2025)
developing experiments to test various properties of ICL, leading to new observations and insights. Some of the earliest works that show transformers behave like Bayesian estimator can be found in Akyürek et al. (2022); Garg et al. (2022), and this argument is supported in follow-up works including Li et al. (2023b); Wu et al. (2023); Bai ct al. (2024). However, there is also increasing empirical evidence demonstrating transformers' non- Ban-Bayeii behavior. Singh et al. (2024) design flipped experiment and show transformers' Bayesian behavior could be transient. Raventós al. (2024); Panwar et al. (2023) demonstrate that the Bayesian behavior of transformers dependent on the task diversity in the pretraining dataset, and transformers could deviate from the Bayesian predictor if number of different training tasks is large. Falck et al. (2024) design experiments based on the martingale property, necessary condition of Bayesian behavior, and provide evidence that transformers exhibit non-Bayesian behavior from a statistical perspective.
Transformers for Uncertainty Quantification. Uncertainty quantification has seen significant develop- ment within the general machine learning and deep learning domains (Abdar et al. (2021); Gawlikowski et al. (2023)), generating considerable interests within communities working on transformer- based large language models (I.LMs). Soe Kuhn al. (2023); Manakul ct al. (2023); Lin et al. (2023b) for uncertanty quantification using black-box LLMs, and Slobodkin et al. (2023); Chen et al. (2024); Ahdritz et al. (2024) for that of white-bor LLMs. Most of these works focus on natural language processing tasks that have less statistical properties. Indeod, uncertainty quantification has traditionally been developed from a more statistical and probabilistic perspective (Smith (2013); Sullivan (2015)). By adopting transformer models to study more statistics-relatdd problems, our work aims to bridge and contribute to both fields.
We thank the reviewers for pointing out some literature for discussion. Although the problem setup seems to be similar to domain generalization/ /adaptation credal set learning theory (Caprio et al., 2024), the main purpose of our work is different. The learner may face different types of distribution shifts between the training data and the test data, and those works in the previously mentioned field are trying to curve the learned results in those settings. However, our work aims to show that there is a gap between Bayes optimal and Bayesion. When infinite training data is available and facing ID test data, the Bayesian learner can behave as Bayes optimal. But the Bayes optimal may just be a result of ERM training as shown in our work. We test the Transformers on OOD data to see if they coincide with Bayesian predictions. The answer is suspicious especially when the OOD issue becomes severe. Another line of research focus on quantifying the uncertainty, especially factorizing the alcatoric uncortainty (AU) and the epistemic uncortainty (EU) (Sale et al., 2023; Caprio et al., 2023; Wimmer et al., 2023). In this paper, our work differs from the aim of the inove-menttooed uncertainty quantification literature. We do not manage each term accordingly (the AU and the EU). We are designing simple yct illustrative uncortainty quantification examples to test the ICL ability of Transformers rather than proving Transformers should be (one of) the most appropriate approach(es) to do uncertainty quantification.
B Transformer Model
Following Radford et al. (2019), we consider a decoder-only L-layer Transformer model that processes the input sequence H, by applying multi-head attention (MHA) of M heads and multi-layer perceptron (MLP) laser-wise. Without loss of generality, we assume I; E Rd for some d 2 2. We concatenate each ye with d 1-1 1 zeros so that it matches the format of cach It, while we still denote the concatenated vector by y/ with a slight abuse of notations. We denote H, by a matrix in (2 for I,...,T, T, where H, = [21, v.. x.]. We may also refer to F, by h2--1 and ye by hz. In practice, the attention mechanism has maximum dependence length, and therefore the Transformer model can only produce an output based on the most recent tokens up to a context window size S. Hence we assume that at cach time step t, the Transformer model has a maximum capacity of making predictions based on It and previous S pairs of (Ie, V&) observations for s = t -- S,..., - 1. In other words, the Transformer has a maximum capacity of processing 25 + tokens, and it is making predictions (H,, = (Hş) based on the truncated history H&S where H == (Imax (1.t |max(1,t-S)..... ..)). In the following, we formally describe the architecture of the Transformer used in this paper.
22).
16

---

Published in Transactions on Machine Learning Research (05/2025)
Definition B.1 (Multi-Heal Head Attention). A multi-head attention layer with M heads and activation function act(.) can be defined as function (+) for any sequence Z4 E (2dxx2(-1) and = 1,...,
M
MHAW (Z,) Ze +
Σ (WV Zit ((WP)) (WQ Zi)).
where W ((W@) WT. ) denotes all the parameters, W W E Rd xd, W E Rdxd for cach m ,,.., M, and act: RR (2+ + 1(2-1)x((-11 is the activation function.
Here we merge the residual connection into the multi-head layer and skip the layer normalization case the notations and simplify the analysis. The activation function is usually set to be columns-wise softmax in practice: for each vector R
exp(221-1) T
exp(21) Di=1 exp(2;)
softmax(2) =
2r exp(21)
Some theoretical results also consider alternative choices for act. For example, Akyürek et al. (2022); Ahn et al. (2024); Zhang al. (2023a) consider the linear activation (that is, to entry-wis divide by the sequence length 2/ 1). Bai et al. (2024); Guo et al. (2023) also examine the ReLU activation (that is, to entry-wise apply ReLU function ReLU(2) ma 2} and later divide by the sequence length 2t -). 1).
Definition B.2 (Multi-Layer Perceptron). multi-layer perceptron layer with hidden dimension da can be defined as a (token-wise) function MLP M (-)) for any sequence Ze E pdx(2t-1) and t ,,...,5+ 11,,
MLP (Z,) = Ze + AeReLU(A,
where A (2) denotes all the parameters, A1 E Rdd. xd A2 E Rdx,, and ReLU is the entry-wise ReLU function.
We merge the residual connection into the multi-layer perceptron layer and omit the layer normalization to simplify the theoretical development.
Definition B.3 (Transformer). Transformer model with L layers can be defined as function TFo(-) for any sequence Zz Rdx((1-1) and t = S + 1. For the 1-th layer, the model receives z (!l-1) the input. and processes it by an MHA block and an MLP block, such that
2 MLP AOD (M (Z)-")), VI =I=1...,,
where 2 = Zt. After the L-th layer, the model lincarly maps the z111) E Rdx(2t-1) onto (2r-1) via a matrix E Ixd and we process the second dimension by a softplus function to get the final prediction as
30 (0())) (PZ(4)) 11.2111
and
Go softplus (P
Here 0 = (((WO,,0)]LL P) encapsulates all the parameters and the function softplus(z) log log(1+xxp(22) is introduced to avoid negative output. The output is sunmarized as (Z. (2o(Z.), 6n Z.)).
Remark B.4. To enable parallel training, the decoder-only Transformer receives a full sequence in the training phase. The model has a masking component that prevents the model from sceing into the "future". However, such masking unnecessary in our setting as the Transformer model receives exactly what it should "see" each time and the full dynamics identical to those in the masked setting.
17

---

Published in Transactions on Machine Learning Research (05/2025)
Miscellaneous notations. Denote the set {I....,K) K by [K]. Denote the consecutive sequence {i, + 1,....5) by i: j. Denote the matrix A's entry at the i-th row and the j-th column by Aij- Denote the vector r's i-th element by (z);. Define the d-dimensional vector z's p-norm as (E (x))) 1/2 for p [o], where |/ maxisiga(z)i- Define the m n-sized matrix A's (p, f)-morm as (2j=1 |A.allg)// Denote the d-dimensional diagonal matrix by di... Ad). Denote the d-dimensional identity matrix by La- Denote the total variation distance between two probability distributions P and Q by TV(P, Q). Denote the Kullback-Leibler divergence between two probability distributions such that P & Q by Dai(P))). Denote the product measure of P and Q by P Q or P 8 Q. Denote the Cartesian product by x V. Denote the tensor-product o-algebra of two r-algebras Σn and Σg by Σn @ Σg. Denote the limiting behavior of being upper (lower, both upper and lower, respectively) bounded by up some constant(s) by O (12, O respectively). Denote O be the but omitting some poly-logarithmic terms.
two spaces X' and y
B.1 Assumptions
Based on this setup of the Transformer model, we introduce the following bounded assumptions used for Theorem 3.2. Such assumptions are common the analyses of the Transformer model (Bai et al., 2024; Zhang et al., 2023b) by either assuming an extra clipping operator or explicit upper bounds.
Assumption B.5. Assume = B(0, Brr), where the norm is defined as
|||| = max
[P] 1 1.....A L}.
The corresponding norms are defined
|| 22. m I,...,My |||| max || 2, W 12,
||||| == max || [ ||P|| = || 12,,0
where we omit some suporssriptss/uubsrriptss of the layer number (1) for simplicity.
Assumption B.6. Assume |HHlll, bounded by Bи almost surely. Such a regularization is equivalent to assuming | S Bи and ly, S Bи almost surely.
B.2 Approximation Error
In Section 3, we provide the generation bound in Theorem 3.2. Now we give analysis for the approximation error. We define the Bayes-optimal risk obtained by the Bayes-optimal predictor in Proposition 3.1: for each T,
R =
(4)
However, Transformers only have access to the truncated history H.S which prevents them from reaching R; By using Proposition 3.1 for the H&, we denote the truncated Bayes optimum for each t:
vi (HH = H[
and
(o St (HH))² El(f(x.) (H:))) H Elo² H
We denote the truncated Bayes-optimal risk
R - 1 (H,),, (HH))
(H{)),,{(HHH)) v1)
(5)
It is straightforward to check that
R R₁ for any < S.
(6)
However, the equality is generally not true for S. We give example illustrate the gap.
18

---

Published in Transactions on Machine Learning Research (05/2025)
Example B.7. Consider the case where one has oracle access to the noise level a. Note that the oracle knowledge only reduces the risk R2", since we use information that is not a measurable function of H.. The problem is reduced to regression problem.
Suppose the function f is linear and its weight vector has a prior distribution of N da),, and the noise 1). Suppose Zt NN0 Ta)- Then the optimal estimator is an oplimally luned Ridge regression.
Tsigler 8 Bartlett (2023) show that with high probability, the optimal Ridge regression estimator has on average risk of 1/2 ((1/S), where the term 2 is due to E[(y ( ((2)))) But the length t approaches infinily, the average risk of the optimal Ridge regression over the full sequence He will converge to 2 with high probabilily, meaning that the estimated I will converye to true I for every sequence. Hence one can always construct an uncertainty estimation by averaging all the residuals, and such an estimation & will converge to the true a. Thus, we have Ryy approaching as t grows to infinity, leading to the conclusion that
R - R 2 201/5))
for sufficiently large t.
Example B.7, together with Theorem 3.2, shows the approximation-cstiaattin tradcoff in selecting the context window S of Transformer models. Previous works (Wu et al., 2023; Bai el al., 2024; Zhang al., 2023b; Guo et al., 2023) consider the case where S S, and establish the upper bounds for the approximation error. In other words, these existing results are all made with respect to the gap between R25* and I(TFa-). To our knowledge, we are the first work to point out the extra term of approximation error due to truncation.
C More Numerical Results and Discussions
C.1 In-distribution performance
In Figure 1, we provide a comparison of the in-distribution performance of the trained Transformer v.s. the Bayes-optimal predictor. subtle point that for the uncertainty prediction, we only plot the average predicted uncertainty, this does not fully imply that the Transformer gives a similar prediction as the Bayes-optimal predictor. To this end, Figure plots the difference between each of the models and the Bayes-optimal predictor in terms of uncertainty estimation.
C.2 Out-of-distribution perfomance
In Figure 2, we plot the Baycs-optimal predictor under three OOD settings, and we note that though the Bayes-optimal predictor uses a wrong prior, it has the ability to work as algorithm to correct the prediction with the in- in-context samples. Now in Figure 6 (a), we compare the Bayes-optimal predictor that uses the wrong prior with the Bayes-optimal predictor that uses the correct prior (which replaces the in-distribution prior with the correct OOD prior o²). The figure is based on the large OOD setting. We observe that the Bayes-optimal predictors with the ID prior or the OOD prior both converge to the true uncertainty level. For Figure 6 (b), we plot the performances under the same large OOD setting. As reference line, we copy and paste the Bayes-optimal predictor's curve Figure (b) here. We note this reference line computed based the in-distribution (ID) data and is not comparable at all to the predicted uncertainty level on the OOD data. Yet, we note that when the number of tasks is small when training the Transformer (say N = 4096), it tends to make predictions on the OOD data by treating the OOD data just as ID data, and this means the trained Transformer is doing in-weight learning and has no in-context learning ability. As the number of tasks increases, the Transformer gradually gains the in-context ability and moves towards the Bayes-optimal predictor on the OOD data.
C.3 Training dynamics and task shift OOD performance
Now we zoom into the training dynamics further investigate the OOD performance under task shift. In the following example, we derive a theoretical result based on Theorem 4.1 in Zhang et al. (2023a). Specifically, R and R' (following the notations therein) denote the in-distribution and out-of-distribution expected risk.
19

---

Published in Transactions on Machine Learning Research (05/2025)
Ridge Regression
Linear Regression
TF-4096
TF-65536
5
El El-log(lō o
2
-1
5
10
20
40
80
# in-context samples
Figure 5: In-distribution performance of the uncertainty prediction against the Bayes-optimg predictor. The y-axis gives an estimate of B log (H,)) (H(H)] where the expectation is taken with respect to H.. Here o(H,) is the uncertainty estimate produced by an algorithm (ridge regression, linear regression, or transformer). and o² (He) is the Baycs-optimal predictor given in Proposition 3.1 and calculated by Section G.3. The figure shows that the Transformer and the Bayes-optimal predictor produce similar uncertainty predictions. Tn addition, the Transformer trained on a larger pool of tasks (larger N) produces better approximation of the Bayes-optimal predictor.
w/ OOD prior
Payes-optimmn w/ ID prior
Dayes-optimal prior
4.0
ID Ref. Line
Ave. of predicted uncartainly
Ridge Regression
TF-4096
TF-4095 TF-65536
TF-16384
TF-65536
of predicted
2.5
1.5
0.5
0
20
40
60
80
100
0
20
40
60
80
100
context samples
in-context samples
(a) Bayes-optimal w/ or OOD prior
(b) Moving from ID to OOD
Figure Performance under L-OOD setting, For both (a) and (b), the y-axis gives the average of the predicted uncertainty over all the test samples (average of ((H)) o (HH) test samples), and ideally the curves should converge to the true uncertainty level of as the number in-context samples increases. In (a), we compare the Bayes-optimal predictor that the wrong prior with the Bayes-optinal predictor that uses the correct prior (which replaces the in-distribution prior with the correct OOD prior of o²). Both work well in that the curves converge to the mean uncertainty level of around 4. The Transformers deviate from both Bayes-optimal predictors due the large OOD intensity. In (b), we observe that as the training task diversity increases. The transformer gradually moves from the ID reference line the Bayes-optimal predictor.
The result says that while the in-distribution risk continues docreasing over time, the out-of-distribution risk may keep increasing or may first decrease and then increase. Importantly, the out-of-distribution risk may depend the initial point of the training procedure.
20

---

Published in Transactions on Machine Learning Research (05/2025)
Reaching the Bayes optimum requires prior knowledge of the underlying distribution. We provide simple example of where the Transformer stores its prior knowledge and how it hurts the OOD performance even under a mild distribution shift.
Example C.1 (A corollary that can be derived based on Theorem 4.1 in Zhang et al. (2023a)). Consider one-layer allemtion-only Transformer model with linear activation and one allention head on the linear regression task. We now concatenate each of the inputs to be [x ye] E Rd+1, Suppose we focus on the linear regression task on the (T' 1)-th sample after observing context ezemplars, where each w ~ (0,, La), each "!!" NO, Ta), each (8) ~ (0, 1), and vi" wli) Zhang et al. (2023a) (with details referved to therein), then for any - -ol 2 A for some A > 0, if we train on the distribution w.r.l. on but lest on the distribution w.r.l. o (denoted by R2), then for C = d/(16(2 + -o)) and any sequence 81 < <d2- <.... ...CAA, there erists a non non-decreasing sequence 0 S T S T S ....
(8) Vi + ao 1o.e, _E2") If we adopt the same training selup
such that
(T())) Rig.. 8,, for each i 1,2,...,
while the parameter (W K WQ Wo) (Wy)&11d-1 converges to L/(1+ (2+1+22+00))) (2+ (o) (which is the corresponding part of some ). Here 6* and 6" minimise the population rish: R and R', accordingly.
We design experiment to show that training proceeds, the model's OOD performance is improved abruptly in the starting phase, but then degrades steadily after too many steps of training. We introduce the experiment settings below. A visualization of the setup is given in Figure
Testing:
In Distribution:
Training:
ID:
OOD:
>>>00 0
0
better
better
a E 92
0 E 11
0
0
< 0 a E 9
Out of Distribution:
Training Too Many Steps
OOD:
G G2
slightly better
worse
G2 [0.3, 0.5)
00.1,0.3]
0.0.7,0.9]
0.7]
Figure 7: The settings of the OOD experiment. (Left) The in-distribution (ID) tasks are sampled from regions denoted by the green blocks, and the OOD tasks are sampled from the red blocks. (Right) In the starting phase, training improves both the ID and OOD performance. But if training for too many steps, the ID performance is only marginally improved, while the OOD performance steadily degrades.
Each linear task in our uncertainty quantification setting is characterized by parameters (w, o). We define two regions for w, denoted by Wi and W2- And two regions for denoted by G1 and G2. When w is sampled from W1, w follows the following distribution
w = 13,, B~ N(0, 1s).
When w sampled from W3, w follows the following distribution
w = -18, 8 ~~~ No(,
For a, define 9i = [0.1, (0.5, 0.7] and G2 = [0.3, (.3] 0.9.. Define G1 and G2 to be the "complementary" group of each other. We sample a independently from w, and we always sample a uniformly from either group Gn Gz. As marked in Figure 7, the "ID" tasks sample its parameters from (w, σ) E W1 &1 Uw2 & G2 and the "OOD" tasks sample its parameters from (w,a) E Wi & G2 UW2 G1. The training is on "ID" tasks, and
21

---

Published in Transactions on Machine Learning Research (05/2025)
the trained model is tested on both "ID" tasks and "OOD" tasks. The metric evaluate in this experiment is the "prediction accuracy" of uncertainty. The accuracy denotes the probability that the model predicts the a into its "right" group. (for a prompt generated from (w, with a E g we say that the model makes a "right" prediction the predicted a falls into g).
Figure presents the experiment result. The prediction accuracy on the ID dataset peaks after 20k steps of training. At the same time, the prediction accuracy on the OOD dataset also increases 80%. After that, the ID performance remains unchanged, but the OOD accuracy keeps dropping.
1.0
Uncertainty prediction accuracy
0.8
Mo fo M
0.6
0.4
0.2
In distribution Out of distribution
0.0
0
40k
80k Training steps
120k
160k
200k
Figure 8: The accuracy denotes the probability that the model predicts the a into the "right" group. For example, the sampled tasks take a from group 91, then accuracy denotes the probability that the model predicts a into G1. The data is collected for the 100-th token in order to climinate the epistemic uncertainty due insufficient in-context samples. The x-axis denoles the training steps. This figure shows that when training too many steps (> 40k in this case), the generalization ability of the model steadily declines.
In order to verify that the degradation of OOD performance is due to the increasing confidence in the prior information of the training data, check for the OOD distribution whether the model has predicted o into the complementary group. The result is presented in Figure 9, which verifies that after training too many steps, the model tends to predict a following the training prior. A more concrete way to explain it: consider an OOD sampled task (w, a) where w > 0. According to the sampling rule of OOD tasks, it must have o E G2- If the model has the OOD ability, it should prodict a E G2- But if it has too much confidence in its training prior, will predict o into G1. the complementary group of Gz. Figure 9 shows that for the misclassified OOD tasks, the model has predicted them into complementary groups
22

---

Published in Transactions on Machine Learning Research (05/2025)
OOD Prediction
1.0
Accuracy
Prior Percentage
0.8
percentage
0.6
0.4
0.2
0.0
20k
40k
80k
160k training steps
200k
240k
320k
400k
Figure 9: This figure validates that the decline of OOD ability is due to increasing confidence in the training prior. The blue bars correspond to the OOD accuracy, and the red bars give the probability that the model predicts uncertainty into the complementary group (L.e. the training distribution of σ). As the training proceeds, most of the misclassified
predicted following the training prior.
D Proofs of the Results in the Main Paper
D.1 Proof of Proposition 3.1
Proof. Recall that the population risk is
log (H,) + (ye - (( (H.)) 28 28 (H,)
L(g, 0) = B
We first prove that for any (H,), the choice of = Ely, H.) minimizes the population risk. With any fixed > 0, when o, dp. then minimizing the population risk reduces to minimizing E[(y. (H.))) Using Fubini's Theorem and the fact that the conditional distribution exists, we have
E B [(ve ((H))) = EH,, [ - 9 ((H)))
EH. [[(((z)) ((H+))) H.]] + EH. [I[[2|
(7)
where the last equality follows from the fact that 6, is independent of H, and a and of zero mean and unit variance. Since the second term on the right-hand-side of equation 7 does not depend on i we only need to
locus on the first term. For each realization of H,, the prediction ((H.) is a single point; combining with the fact that the squared loss minimized with respect to one single point prediction if and only if that point is the expectation (in this case, the conditional expectation HH),), prove that for any do, the population risk's minimizer
0. (ao) (B (x)())) + (.H Ely, [H..
23

---

Published in Transactions on Machine Learning Research (05/2025)
where the second equality follows again from the fact that e, is independent of H2 and 0, and e, is of zero mean. Since this equality holds for arbitrary do. we conclude that
y/ = El He].
Now we have confirmed the optimal choice of regardless of whatever â is. We can thus find the optinal choice of & by fixing ğ = yf and minimizing the population risk. Similarly, we can change the integration order so that we only need to minimize Ellog (H,) 26 (1%)
+ (M- (0-KHH)
[H)] for any realization of Ht. Calculations
show that
aI
log Ō (H+) + (ue ((H)
H4
0 (logo + B[ 9
(8"(H.)
(He(H))
(H,)
0G(H.)
(H,))
[E( 0 H,))
(H,)
where the first equality follows from the fact that on observing He, (HH,) is fixed value, and the second equality from the calculus. Thus, the risk is minimized if and only if 6 (H.) [ (HH))) Substituting for yi have
(H,) = Bl( vf (H.)) H.] B((((=) - Vt (H.)) H. H|
where the last equality follows again from the fact that e is independent of He and is of zero mean and unit. variance.
D.2 Proof of Theorem 3.2
Proof. (ERM Before start the detailed proof, we define another flattened sequence (ãx, Us)) for k 1, nl' where for k = iT + t we have
(TeT = (xt -//2
(8)
Here, we merge all the sequences [(x{1) Mi ) for i= I,..., n into one sequence (3 d=) Similarly, we can define a flattened truncated history H as
(1)
1 = (z (6)
max Mmax(r-s..
(2)
99)
Note that HS k = (H S(0) (/1),, since we have added the target label yi (()) into the flattened truncated history Eg for notation simplicity. With a slight abuse of notations, we have
to ( = ((TF((H)) v1) =
(10)
where the equality holds since we are making predictions based on most pairs of (I+, yc). We can similarly replace the l function in the definition of empirical risk and population risk RG obtaining
r(TF(O)) = nT 1=1 1=1
8(TFa (TFo(H))) (HS). v))
(11)
nT
k=1
and
1
R(TF(O)) T E.,
8 (TF(
+ E
Ea S (1
nT
Enng
Lo (
(12)
24

---

Published in Transactions on Machine Learning Research (05/2025)
where H is another flattened truncated history that is ii.d. to Hy. For notation simplicity, we define
H$ = (HS HT)
(13)
Then we simplify the notations as
ro (H( r((FF()))
(14)
and
Rg = Ens (ro(i))) R(TF(O)).
(15)
To control the difference between Rg and ro (+( for any 0 (which could potentially depend on training data D), we use PAC-Bayes argurentts for simplicity.
All the following arguments are made with the conditional distribution knowing each fO) and a(), for cach i = 1,...,, We omit the conditional dependencics in notations only for simplicity.
By our definition of data generation, the flattened truncated history H naturally forms up Markov chain on the space 2 (verified in Lemma E.13), since the newly generated (æt, ye) are conditionally independent of all previous observations. Here D4.k=I++ = (X' x y) (i,().
Fix 0 that does not depend on the training data D. We now bound the difference between Ro and (P())) via concentration inequality for Markov chains. From Lemma F.2, we know that if the Markov chain's mixing time is small enough (which means it quickly converges to the stationary distribution), the concentration properties over the Markov chain would be good enough to enable the standard PAC-Bayes arguments. We also know from Lemma E.15 that the flattened truncated history has a mixing time no greater than min Th since all the histories S pairs before the current time would be truncated from the input, and the history H. restarts every time a sequence reaches length T. With these observations, we start our detailed derivation.
Since the function is almost surely bounded by C2 is shown in Lemma E.3, have almost surely for any H. and H
nf' Σ nf
2C2
ro ro(H)) -ro ro(
1 # M
(16)
k=1
We can use McDiarmid type's inequality for Markov chains (Lemma F.2, with the mixing time upper bound no greater than min(S, T) (spocifiod in Lemma E.15), such that for any R,
Ep exp (A(ro(f()) Ro (T( < exp 2x* min mii((,TT)
(17)
nf
Set x to be the distribution over O defined in Lemma 11 Since x is chosen independently from D, we can integrate equation 17 with respect to ~ л such that
ED exp (A(ra(F(S) Ro((is)]]] exp 22 min(S,
Eg
(18)
nf
Using Fubini's Theorem, we can exchange the order of integration, such that
Ep Bown exp (A(ra(F()) - exp 22 min(S,
(19)
nf
By applying Donsker-
Donsker-Varadhan'ss Varadhan's formula (Lemmn F.3), we derive from equation 19 that
C2*C3 min (S,
(Famp[(rro(f Ro Dia(pll1)))] S sexp OEP(e)
ED
exp
sup
nT
Rearranging terms, we have
{Eerg (Xro((()) R -a((plT)}} - 2XºC3 min (S,T) (S,T)) 1
Ep
exp
sup (EP(()
(20)
nf
25

---

Published in Transactions on Machine Learning Research (05/2025)
Using Chernoff's bound (Lemma F.4) with probability 8/4, we have with probability at least 1 - w.r.t. D,
(Te (M(o((H) - Ro(H) D 2x²C2 min (EP(O)
sup
S log(4/6).
(21)
nF
Since this bound equation 21 holds for any distribution p over Q we can set p to be PâERS as defined in Lemma E.11, resulting in a high-probabblity bound
Eg Ed-PPERse (ra(H()) [
(ppeeM (T) + 2AC3 (iin) + (4/6)
(rearranging terms)
A
(nd'))
C2 /in (DalPeERM |I=) +2 log(4/6) (by setting À VnT/ (in -3,/2)))()
(1/C2))
< 50 [ (TTT)))))))))
(by Lemma E.11)
(22)
By Lemma E.12, the loss function is Lipschitz. Since for any 0 supp(PgERm), 0 is up to O(I(nTT) away from @ERM, we can control the difference between the risks of any E supp(PSERM) and jERM as
|na (FL - -'gano (H+) ((L((TT))
(23)
Rs - ROENSE| S {((1/(nTT).
(24)
Thus, we have
Tianna (Ht)) RgEns S ( (iImmmS,,T//(11)))
(25)
Applying the above arguments again for the negative of r, we have with probability least 1 8/2,
IREER (HL RIERRS Ö Oiv (S,T)/(TT))
(26)
For 8" we can repeat the above steps and get
|re- (71 Ro- I O min(S,T)\/(n)))
(27)
The probability that all these bounds hold simultancously is at least 1 - w.r.t. D.
Hence with probability least 5,
R(TF R((FF.)
Rijuns - Ro.
(by definition in equation 15)
IggERs (FLS) - rg- (H+)) O(/miinS,,T)/(mn
(by equation 26 and equation 27
=r(TF (ERM) - TTFa.)) (VV11It vS/T)
(by definition in equation 14)
< Õ (S,T(//(((nTT)
(by definition of ERM equation 1)
(28)
We now take the expectation over each (() and (1) to conclude the proof.
Remark D.1 (Why truncation). Previous analysis (Zhang et al., 2023b) to derive a similar Bayes-optimal argument does not truncate the history and treats the whole history as an inhomogeneous Markov chain. Then they apply the concentration inequalities on Markov chains (for example, Lemma F.2) to control the difference between Rt and r. However, their arguments have two limitations: the first one is that their model is assumed to make decisions based the full history, which clearly exceeds the Transformer's model's capacity. The second limitation is that such concentration argument for Markov chains often relies upper bounding the mixing time or lower bounding the spectral gap (for example, Fan et al. (2021)). But Zhang et al. (2023b) do not specify this the mixing time. Furthernore, in each sampled task sequence (assume we know the task f()),, the mixing time of the (untruncated) history H, is infinity: two sequences start with different initial pairs of (x1,8) Vi), then they will never become identical no longer what comes consccutively. Thus, their mixing time will be T, leading to Ö(1//n) generalization, which is suboptimal if S < compared to result..
26

---

Published in Transactions on Machine Learning Research (05/2025)
E Proofs of Lemmas
In this section, we prove these lemmas based on the choice of the activation function act = softmax. Similar results for other options act ReLU can also be found in many existing literatures (for example, see Bai et al. (2024)).
E.1 Boundedness of Transformers
Lemma E.1 (Layer-wise boundedness). Suppose at the I-th layer of the Transformer, we have (WV.0)) || S Ba- Then for any inpul H we have
By for any m 1, ................................ M, || 12,2 [|
ll2.2
| S (1 (1) MBv) H 112000
Proof of Lemma E. 1. For notation simplicity, we denote sof toftmax((s (W (1) H(-1)) T H((-1)) as s Note that every column of s is of unit 1-norm. Denote each column of s by st" For any input H, we have
K
|HA w 12,00
M
$ 21 2,
Σ
HS| 2,00
(by triangle incquality)
m=1
M
=|| 2.
max
W (1) Hspl1a
(by definition of || l12,,)
m=1
||
2,00
max W m.m.(1) (1)
H 12.00 |2,.001187711 |1 (by Lemma F.5)
M
2,,00 (1)
(since sf is of unit 1-norm)
M
$ H 2, +
Σ
W
2,2 22
(by Lemma F.6)
m=1
S (1 + M MB) 2,.00
(by assumption of bounded norm)
(29)
For any input H, we have
|MPP A (H) (H)
< H 2,00 + || | ReLU( (A() H)
(by triangle incquality)
|H| 2,00 + ||4!") 2.2 |ReLU( A H)
(by Lemma F.6)
2, + ||4 (1)
2. max |ReLU(A]] H)alll
(by definition of || 1|2.00)
< 2, + |14 (1
2.2 max (1(1 H),ellz
(since [ReLU())| S |2| for any R)
1 + || ||2 (14 H 12,00
(by definition of || 112,00)
< H 2, + 1144 ||2 (1 ||2 12,00
(by Lemma F.6)
< (1+ B ||
(by assumption of bounded norm)
(30)
Combining equation 29 and equation 30 yields the conclusion.
Lemma E.2 (Transformer's boundedness). Suppose [|WTh.(0 By for any 1, 1,..., M, ( 12,2 (| 12,22 < BA for any I E [L]. We further assume the projection matria P is of bounded norm PT 2,300 & Bp. Then the Transformer's oulputs satisfy that
(0(H)) < C| and exp(-Ci||H) (H(H) $ 1 G. |2
27

---

Published in Transactions on Machine Learning Research (05/2025)
where C1 = 1: (1) MBv) is specified constant.
Proof of Lemma E.2. By Lemma E.1 and "pecling" argument, easily prove that
|1 (1+ B3) MBB)
Thus,
|1 S (H S (1+ B1) + 1 1,,,00
Denote P by P [PI.Pa] where pı and p2 are vectors of dimension d. Then the first output
i Pi H!
where we have (by Cauchy-Schwarz inequality),
M S 11illall S Bp(1 + (1) MBy)* 12,.00
The other output & can be proved similarly as long as one notices
log(1 + exp(-x)) 2 exp(-x), and log(1 + exp(x)) 1 11+0. x,
for any x 2 0.
Lemma E.3 (Boundednoss of loss). Under Assumption B.5 with ||0|| S BTF and Assumption B.6 with |1 I Bh almost surely, we have
(H.F) S C2
almost surely, where C2 = (C1+ + 1) exp(2C,Br) + max BH, log(C,Bh)) is a specified constant, and Cj is constant defined in Lemma E.2.
Proof of Lemma E.3. By Lemma E.2, we have
(y:- =(H.))² (H,)
exp(2C) BT)
S (y:: dr (H.))
2
s (n? + (.(H)) exp((C,BB)
s (C1 + 1)² exp(2C;BH))
where the second incquality follows from Cauchy's incquality. Combining with a triangle incquality, we have the desired result.
E.2 Lipschitzness of Transformers
Lemma E.4 (Lipschitzness of multi-head attention). Suppose we define the oulput's 's norm as || 12,00, the norm of W as
||]| == max([] 12 || 112,, [ lla2 m = L...,,,,MM
and the imput H 's norm 1,,,00 Suppose at the 1-th layer of the Transformer, have |I Bw for any m = 1,...,M, M, and (1 & Bll-1) almost surely. Then MHAww: (H()-1)) is (1) -Lipschitz with respect ((I) and C2- Ci-Lipcchttz with respect to (l-1)) almost surely. Here c(O 2B (B) p + (B(1"1) and Ca 1= M Bw specified constants.
Proof of Lemma E.4. We first prove the Lipschitzness result for W. To ease the notations, omit the dependence on 1 and sometimes abbreviate W Wa WKQ- For any W and W using triangle inequality twice, we have
|MHHW (H) MHAW (H)
28

---

Published in Transactions on Machine Learning Research (05/2025)
M WP Hsoftmax( H WKQ H) WV Hsoftmax( H W KQ H)
2,00
M
WV ((sottmax(H W KQ H) softmax( H WKQ H))
2,00
m=1 1
Σ
(wv "WW Hsoftmaxx H WKQ H)
(31)
2,00
We now deal with two terms in equation 31 separately. Since conclusion will be made for arbitrary m E [M], omit the dependence for notation simplicity from now on.
For the first term, we have
Wy (softsaoamaaa H WxQlf) softmax(H T H))
||
2,00
= mmx W| H (softmax H W WxQh) - soft H Wkoh
(by definition of || l02,
< [Wy M| 2, max
(softmax H WkQh)) sof toft H II
(by Lemma F.5)
< | H 2,
max (softmax H WkQh)) - softmax H Wk) ll
(by Lemma F.6)
< [/]Wvll2.22
12,00
max
11 W kkh H W kk
(by Lemma F.7)
< |2.22/22
12,00
max
|| W Ixw W hellle
(by Lemma F.5)
s 2/
max
Wxq - 1 11
(by Lemma F.5)
< 2/
1,,00 ||Wxe= W 22 12.12,00
((by definition of || 1|2,00)
(IIW>-W2 Wx W/ 1ol + WK W - W W 222))
x 12
((by triangular inequality)
< 2/
2, 12.2ll22/WWW W2ll2.2 20 + IIWx - Wk 12.21 W/ (122)..
((by sub-multilicaaiiiiyy of matrix norm)
< (B
(IIWa Woll2.2 + I|xx - Wkl12.2)
((by bounded norm assumption)
(32)
For notation simplicity, we denote soft Way H) by Note that every columi of unit 1-norm. KQ
Denote each column of S by 8t. For the second term, we have
(Wv Wv) Hsoffmax(( H W H)
2,00
I1(Wy
(by notation substitution)
max llWwv H'+-)=ssle
(by definition of || l12,00)
S max 11Wwv Wi- 12 s
(by Lemma F.5)
Il(Wy Wl.)
(since 8t is of unit 1-norm)
< IWy 12 H 2,,00
(by Lemma F.6)
< (44 -I W 1la.2.
(by assumption of bounded norm)
(33)
Substituting equation 32 and equation 33 into equation 31, we can conclude that MHAw (11(1-1) is c.0 Lipschitz with respect to ((()) for c/S)) = 2B (B - + +(B (Bll1)))
As for the second Lipschitzness conclusion (the one w.r.t. H),, it is straightforward one replaces H with H' in the proof of equation 29.
Lemma E.5 (Lipschitzness of multi-layer pereeptron). Suppose we define the output 's norm as || 112,00 the norm of W as
|||| max (1141 || ||22|222)
29

---

Published in Transactions on Machine Learning Research (05/2025)
and the input H's norm as 112,00 Suppose at the 1-th layer of the Transformer, have ||4 - Ba and 2, I B(l-1) almost surely. Then MLP (H) is C -Lipschitz with respect to A and Ci-Lipschite with respect almost surely. Here C BABH and Co = + B3 are specified constonts.
Proof of Lemma E.5. We first prove the Lipschitzness result for A. To ease the notations, we omit the dependence on For any A and A', we have
|MLP MLP x 22)
$ |[(A2 42) ReLU A 12, + (R) ReLU(41)))2. (by triangle inequality)
I ||42 4/2 2.2 2,2
ReLU HT) 12,00
+ ||4 22.2 ReLU A1 H - ReLU(A) H) 2,20
(by Lemma F.6)
||Az - A 2.2 max [
+ ||4 2.2 max [ H)LU(A
(by definition of || 112,,00)
< 1| A2112,2 max || (A1 H) ||1 ||A2l12.2 max H), H:,0l2
(since ReLUU()) ReLU |2) - |د| 22l for any 21, 22 E IR)
|| 1/ ( H 00 || A A H 2,00
(by definition of || 1111,,0))
< $1122 Abll2.2 26112.2||A1112.22 H 2, || A1 A' 2 H 1,,00
(by Lemma F.6)
< BABH 11) (IIA1 1/112,2 +1A12 1-A2-122122
(by assumption of bounded norm)
(34)
As for the second Lipschitzness conclusion (the one w.r.t. H), it is straightforward if one replaces H with
H H' in the proof of equation 30.
Lemma E.6 (Lipshitzness of Transformer). Suppose we define each output's 's norm as | for 9 and & the norm of as
|0|| == max [ [|||| IPII,
where ||W|| || is as defined in Lemma E.4, ||A|| is as defoned in Lemma E.5, and ||||| = input norm as ll- 112,00 Suppose we have ||0|| S BTF, and || S Bи almost surely. Then go (H) is Cr-Tipschilz with respect to 0, and (H(H) is Cs- Tsppsssc with respect 0.
12,,00, and the
Proof of Lemma E.6. First we quantify the constants B(t1)) Lemma E.4 and the constants B (!-1) in Lemma E.5 via Lemma E.1. As is shown in the proof of Lemma E.1, we can define
B(-1) = (1+ MBBM)BBD + B 1= 1=1....,4
and
B B(-11) = (1 1L+MBBF) B- 1= 1=
such that all requirements in Lemma E.4 and Lemma E.5 are met almost surely. Thus, we bound the gap between H() (the output of TFy after I layers) and H'()) (the output of TF after I layers) by induction. We claim that if (1)) H()0), then there exists a constant c!.l) for any I= l- L that do not depend on
H. such that
12,00 < C 0'11
We prove it by induction. For = the case be verified by calculation: by Lemma E.4,
|MHA (H))))
MHA (H =c3"110- 0'1]-
30

---

Published in Transactions on Machine Learning Research (05/2025)
Similarly, by Lemma E.5,
| 112,0
|MLP AND (M (1( MLP APCI) (T (11 (0) |]
S MILP AND (M (H())) MLP ACt) (MHAW (H 1)
|MLP AND (MHA (H MLP (r( W (H 1)
Co|IMHAWorr (11 MHA (H1 (0) 1122,00 C 0'll
- (C6C3) cg 6'1l.
(35)
where define Cj as C.2 = CoC3 + C Suppose our conclusion holds for any / to 1. Then for ( lo. we have
W (H10--1)) - w (11 - Ro) (H110-1)) - MHA w (11' 1|2000 + @M (HH)0-1)) MHAw (H))
S Call (1-1) 11 11, + c 110 8'll
- (CaC (11-1) - 111,,
by applying Lemma E.4. We can again compute the difference between ( and (* similar to what we
do in equation 35 as
(()))) 11,000 I (ColCaCy (1- 1 + C (10)) ) 110 0'1
Hence the induction holds if we define dgla) = Co(C,Chlo1) + G ( + Coaa), Now we have proved
[| H'L)) 1|,,0 S 1 1lll.
We shall see from Cauchy-Schwar Schwarz inequality that
10 8'1 < [p - - P - ( - H²(L) 12,00
$ ||0 (1 + . MBrr) (1 (1 + B+ ) + BTr 1 0' (1 (1+ B + BT
((1 MBF
8'11.
where the second inequality follows from the proof of Lemma E.6. We can now define
Cy = (1 MBF) Bir) + BTFC!
and conclude the proof for ÿ. As for 0, we can see from the fact log(1 exp(-)) is 1-Lipschitz that the Lipschitzness also holds for Cs = Cr.
Lemma E.7 (Lipschitzness of loss). Suppose we have ||@|| < BTF and || & Bи almost surely, where the norm of is the same as defined in Lemma E.6. Then (HT)((H)))) y) is C1o-Lipschitz with respect to B almost surely.
Proof of Lemma E.7. Based on the Lipschitzness of the Transformer w.r.t. (Lemma E.6), we only need to prove that both partial derivatives
Øc and de are bounded. For the first partial derivative, we have Bý
ae Ö5
1
(y 0)
<1 + BH (x) (by Lemma E.2)
(36)
For the second partial derivative, we have
- (y - 1)2 + 0²
06
Φ3
((1 B3 + (1+ C1BH)) exp(3C)BR). (by Lemma E.2)
(37)
31

---

de By Lemma E.8, we Published in Transactions on Machine Learning Research (05/2025)
Combining inequalities equation 36 and equation 37 with the Lipschitzness of and & w.r.t. 8, we conclude the result with
Co = (1 + B1) exp(2C,B:)C+ + ((1+ (1) (1+C))H)²) exp(3C1 (xp
where C7 and Cg are constants that appear in Lemma E.6.
E.3
Constructing Distributions over Parameter Space
In this section, we formally define two distributions over the parameter space Q The first distribution Pa may depend on the empirical distribution, while the second distribution Ny should be independent of the training dataset. We control the Kullback-Leibi divergence between Pg and my in Lemma E.11. For notation simplicity, we may use some notations of different meanings from the main text.
For any dimension d, we denote the Lebesgue measure over Bd by Ad(-). Then we have the following lemma. Lemma E.8 (Upper bound for p.d.f.). Suppose p is the uniform distribution over B( 3r) B(0, R) for some zo E [0,, R) Rd, where the Lebesgue measure is defined as >a(-), and R > 3r. Then the p.d.f. Pol.) ezists and
Po(x) S AB(,(O,r)
Proof of Lemma E.8. Denote the set to be S = B( 3r) n B(0, R). Since p is the uniform distribution, we just need to prove that
Aa(S) 2 Aa (B),,))) r)).
This is true because there exists some x' E Rd s.t. B(x', r) C S. In fact, we can construct the small ball
B
T0
Tu
1.5r, r) C.
Lemma E.9 (Upper bound for KL divergence). Suppose the probability space is defined on B(O, R). Suppose p the wniform distribution over 3(zn,3r)nB() R) for some xo E B(0, R) C Rd where the Lebesgue mensure is defined >a(:), and 3r, Suppose π is the uniform distribution over B(0, R). Then
Dal(pl))) O(Ca log (Rg())))))
where Ca == log(Aa(B(O,, 1))) is some constant related to d.
Proof of Lemma E.9. Since & we can define the Radon-Nikodym derivative de can upper bound the RN derivative by
dp (z) 1/>a(B(zo, BB(,, R)) < OCCa log(R/r)).
da 1/A,B(0, R)
Hence,
Da((p(T) [f
dp (x)) dp(x) dr
log
JrEB(0.m)
OCCa log (
2EB(0.R)
O(Ca log (Rrr))
Remark E.10. Note that C2 = circumference to its diameter, and r is the Gamma-function
2 +[ is uniformly upper bounded. Here т denotes the ratio of a circle's
32

---

Published in Transactions on Machine Learning Research (05/2025)
Lemma E.11 (Upper bound for Dal(palTTT), Suppose space specified by Assumplion B.6 (thal is, θ = B(0, BTF)). ). For each layer and each m, suppose we define the norm over each Wo, WK E Rd xd Wv E Rdddd A1 E Rd igduxd, A2 E Rdxdd, to be the Frobemius norm (that is, || 112222. Suppose P [Pi.Pal", and we define the norm over Pr, P2 E Rd to be the Euclidean norm. For each loyer and each m, suppose we have the probability mensures Pwr' PWxx Pwy? PA PA, as the uniform distribution over B(0, (mTT) n [O,, BrF)), and the probability measures TWa TWx; RWv, RA TA as the uniform distribulion over B(0, BTF)). Suppose we have the probability measures Ph1 Pfe as the uniform distribution over 1(nTT/(nT) n Bi0, BrF)), and the probability measures Tpi Tp as the uniform distribution over B(0, BTF)). Suppose we define
considering probability measures over the
Pg
PWrm.(t))
X
Pirm.t) 13) K
X Pimm.tt) m.
x
A
Pil
X Pp. @ P#z1
(38)
and
To =
X
2
X m.d
m
m
π
A
2 Tp.
Tp2
(39)
where @ represents the product of measures. Then we have
DalPgl(Ta) OCin (og(nBBTE))
where Cn is some specified constant that depends polynomially L, M, d, dms dh
Proof of Lemma E. 11. By setting and I = 1,.... L,
1 for Lemma F.9 and R = BTF, we have for each 1,..., M 3xT
DalP wm.. < (Cau,,, log(nTBrrp
Dulp I w= m - 0 (Caam, log(nTTBrF)
Dalpa mmm)))) s (Co Iog() BTF)):
Da(P (w) s (Caa, log(nT
DulP || [ < (Car, log(n' BTF));
Dai(p, [mT - O log(n'T Daa(Ps - (Ca (oggnBTFF))
By Lemma F.8, we can sum up the above inequalities and get the final result.
Lemma E.12 (Bounded difference). For any Ô O suppose we construct the distribution Pg as in equation 38. Then for any 0 supp(pa), under Assumption B.6 and Assumption B.5, we have
((TFo(H),)) ((TFg(H). w) S (CCio/(TT)).
almost surely. Here C1o is the same as defined in Lemma E.7.
Proof of Lemma E. By construction shown in equation 38, we can see that for any 0 E supp(pg),
| 0 < 1/(nT).
Then from the Tipschitzness of the loss function w.r.t. (Lemma F.7), we conclude the proof.
33

---

Published in Transactions on Machine Learning Research (05/2025)
E.4 Markov Chain's Property
Lemma E.13 (H is Markov chain (conditioned on knowing I and σ)). Suppose we have HS defined equation 13. Then H is Markov chain conditioned on knowing each f() and σ( for each i i= ,,...,
Proof of Lemma 1.1 By definition, the state of H will restart and does not depend all previous histories once H index k reaches the point of k = iI' 1. Therefore, only need to verify that inside each task's sequence, the state H is also Markovian.
Suppose == i7 + & for some i, and we considering k == + +T for each = 1,...,, We write ,r,
T.
H and (I Mmaxx (1.txx(111)).......... ...,,,,.. ft,, interchangeably for notation simplicity.
Each pair of (#r, Br) is now independent conditioned on knowing the underlying f() and g(0) We omit the conditional dependencics on f() and g(t) for notation simplicity. The p.d.f. of Hg conditioncd on observing ++1)+=++11 and knowing ((0 and ((1) is
(M$) (
p( 5). (1. ................................ Zt. Vi (H Ir= (H, +1 f() . o())
= =xx () = L "max(111-5p)
w6-1 = b } p(z,, f(0) = o!0)
(by conditional independence of each pair of (xr.)r))
= p(2max(1,,-S5) Wmax(1, - Tt, Ve - = Hrr = f(0) o= oo!)
Thus the Markovian property holds.
We present the definition of miring time as used in Paulin (2015).
Definition E.14 (Mixing time for inhomogeneous Markov chains), Let XI...., Xx be a Markov chain with Polish state space 2 x D2 (that is, X, 24)). Let L(X++e|X; x) be the conditional distribution of Xit+ given X; x. Let us denote the minimal & such that L(X x) and L(X y) are less than away in total variational distance for every S i < N and v, Q. by r(e), that is, for 0 < e < 1, let
d(t) = max sup TV = r), C(X++o/x =)), ISIEN- a.yen
r(E) = min(t E d(() et.
We now upper bound the mixing time of (H, (
Lemma E.15 (Mixing time for truncated history). Suppose we are considering the conditional distribution knowing each (() and o(i), Then for the Markov chain H.S we have
r(e) min(
for any e E [0, 1).
Proof of Lemma E5, We first consider the case when S < T. The mixing property inside each sequence H for h = iT +1....... 1,..., iT + T. Since each (x, is i.i.d. distributed conditioned on knowing f(i) and (()) the conditional distribution of the consecutive sequence (I++1 (e+1)...... (xT, (иT,YT) is never affected by previous T.
pairs (11,))).... (It, yt) for any 1 1 I T. We consider the conditional distribution on knowing (() and ((i) from now on and omit the dependencies for notation simplicity.
For any 1 < T S, for any two points H + HSS the distribution of H is independent of previous & pairs of observed samples. In other words,
C(U ( 1 (1 C( 10 i1
for any l' 2 S. Hence,
d(t) 0, for any 12 S.
34

---

Published in Transactions on Machine Learning Research (05/2025)
We have
(() S S, for any E [0, 1).
When S > T, note that the flattened (truncated) history H restarts every time it meets the end of sequence generated by some f() and (()) Since the length of those sequences is T, we have
T(E) s T, for any E [0, 1).
F
Technical Lemmas
In this section, we present some technical lemmas. Note that all the notations in this section are chosen for simplicity and may have different moanings than those in other sections.
Lemma F.1 (MeDiarmid's inequality (McDiarmid et al., 1989)). Let X = (XI...., XN) be vector of independent. random variables laking values in a Polish space = A x ... x AN- Suppose that f A + satisfies
N
f(x) - (()))(()) - Σ calf wh
1=1
for any x. y E A. Then for any E R,
X 1e1e - -B((X))((((()))) S
exp
2
Lemma F.2 (Corollary 2.11 in Paulin (2015)). Let = = ,N be Markov chain taking values in Polish space A = A1 x ... ... x AN, with miring time T(C) for 0 S < 1. Define
2 -2
Tmin == inf r(t) cE(0,1)
1
Suppose that f: A + IR satisfies
N
f(x) f(0) -
Σ c c1ff + vi}.
i=
for any y E A. Then for any X E R,
E exp (A(S(X) - E[((X))) Tmin |c11
Lemma F.3 DDosker-Varadhan variational formula (Donsker & Varadhan, 1983)). Let P and be two probability distributions over (O.F). I < P, then for any real-valued function h integrable w.r.t. P.
EpgEplexp
sup [EQth] Du((QIP)).
QEP
Lemma F.4 (Chernoff's bound (Chernoff, 1952)). For any random variable X, if E[exp(X)) s 1, then for any E (0, 1),
P(X < log(1/8)) 2 1 5.
Lemma F.5 (Lemma M.7 in Zhang et al. (2022)). Given any two conjugate numbers p.g E [1, oo] s.t. 1/p+1/c = 1, for any r E [1, col, we have
|xl| < ||A|l=/|=d and ||x]|| S [|I 1pr- ||2ll
for any matric A E RRxx and vector x E R
35

---

Published in Transactions on Machine Learning Research (05/2025)
Lemma F.6 (Lemma M.8 in Zhang al. (2022)). Given any two conjugate numbers p.g E [o] s.t. 1/p + 1/q 1, we have
|AB p,00 1 ||
Blpp.o
for any matric A E m and matria B E Raxx,
Lemma F.7 (Lemma M.9 in Zhang al. (2022)). Given any two vectors x, Rd, we have
||oofmaa(()) sof tmaxdy)]l S 2|| y|100
Lemma F.8 (Property of Kullback-Leibler divergence, Proposition 7.2 in Polyanskiy & Wu (2024)). Given any two probability distributions μr and μ2 over F) and any two distributions and is over (D'FF)) if Pi < #2 and in < V2, then we have
Dai @ un 222 & v2) D2((11)22)) + Da((i
G Experiment Details
G.1 Training data generation
We first describe basic setup of all our experiments. For some experiments, we change some part(s) in below to design the corresponding "flipped" experiment to examine the OOD ability of the trained transformer. In particular, the i-th thread the training data
(1) T VT
is generated by the following distributions:
Px: the feature vector It Lid. NO, Ia) where La is d-dimensional identity matrix.
Pe: the noise & lid. 1)..
Pa: the noise intensity a, sampled i.i.d. from
T4 ~ Gamma(I, F), or =
where the parameters I == 20 for the basic setup of the experiment. We change these two parameters for some OOD experiments.
Pf: The function f(() = w.' where w, is generated from
w o N(w, of Ia)
where La the d-dimension identity matrix and ã set to be an all-one vector of dimension d. The covariance matrix of w, related with the noise intensity og to control the signal-to-noise ratio.
Finally, the target variable is calculated by
: xt (8)
(1) asc
y1
Throughout the paper, we consider the dimension d = 8.
36

---

Published in Transactions on Machine Learning Research (05/2025)
G.2 Number of Tasks N and Training Procedure
In the previous Section G.I, we define how generate the training data. As in the previous work, we introduce the notion of task where each realization of (w,, 0,) is referred to as one task. The rationale is that each configuration of (w; o;) corresponds to one pattern of the sequence (It, ys)ss. While the distribution of (w,, σ.) corresponds to infinitely many possible task configurations, we use a finite pool of tasks for training the Transformer. Specifically, we generate
T= [= o
from the distributions discussed above. Throughout the paper, we use N refer to the total number of tasks the pool size.
Training the Transformer for our setting is slightly different from the classic ML model's training, We do not use fixed set of training data. Rather, we generate new batch of training data freshly for each batch.
The batch size b = 64. For each batch, we first sample with replacement tasks from the task pool T.
And based on cach sampled (w;, dt), we gencrate a training sequence (x) Vi I22 :22 TT. .T following the setting in the previous Section G.1.
(6) (i) (3)
(8) (8)
All the numerical experiments in our paper run for 200,000 batches.
The validation and testing sets are also randomly generated instead of fixed beforehand. But unlike the training phase which draws the task configuration from the task pool T, the validation and test phase samples (wy, o.) directly from the original distribution described in the previous Section G.1. This is aimed to validate or test whether the trained model has learned the ability to solve a family of problems, or it only just memorizes fixed pool of lasks
G.3 Derivation of
Predictor
In Proposition 3.1, we state the Bayes-optimal predictor in the form of a posterior expectation. Now we calculate the Bayes-optimal predictor explicitly under the generation mechanism specified Section G.1. Conditional on history H1, (æ1, ,l...., the posterior distribution of (w, o) that governs the generation of H, can be calculated based the Bayesian posterior
P(r IT,, VF
P(wla, H:) = NWw,, Ex),
where
Σ = 1a + Σ E
Wy = Eg
w +
T.
#=1
==1
t-1
r
F = 7
+
w
10
w T
Σ, w.)
~'
2
s=1
Accordingly, the Bayos-optimal predictor bocomes
vi (HH,) = Elyo H.] = w. xt.
and
of (H, El(ve V (H.)) H..
B(o
(tr ( Σx) + 1).
Ie
These formulas used to generate the Bayes-optimal curves in the figures.
37


Reviews:
[{"review_id": "I3N1gJiWIo", "reviewer": "Reviewer_5xWM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths \n\n*  This work presents a new shaper generalization bound of the in-context learning for in-distribution regression task.\n\n*  This work investigates the capability of the in-context learning on out-of-distribution task revealing that proper training tricks is additionally necessary for the transformer to perform in-context learning capability on out-of-distribution tasks.\n\n\n### Weaknesses\n\n*  The claim of the in-context learning capability on out-of-distribution task seems less solid because the claims on out-of-distribution task are only empirically demonstrated and the necessary training tricks for Transformer is different depending on the types of out-of-distribution task. On the other hand, the claims on in-distribution tasks are supported by empirical results and theoretical analysis.\n\n* Some parts of Chapter 4 are not clearly explained. For example, in Section 4.2 on covariate shift, the meta-distribution approach is introduced abruptly, making it unclear why it is considered. Additionally, the results of Section 4.2 are reported in the appendix, even though they could potentially be included in the main section. Furthermore, in the last paragraph of Section 4.3, the justification of context length generalization using Theorem 5.3 from Wu et al. (2023) is unclear, as the theorem is introduced without sufficient explanation of its relevance in this context.\n\n* The experiments in this work are conducted only on regression tasks, raising doubts about the generalizability of the claims to classification tasks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work analyzes the in-context learning capability of Transformers on both in-distribution and out-of-distribution regression tasks. Specifically, for in-distribution regression tasks, the authors present a generalization bound that exhibits a sharper decay rate using the window size $S$ and in-context size $T$. Empirically, they demonstrate that the in-context learning capability of a trained Transformer achieves a near-optimal risk comparable to that of the Bayes-optimal predictor.  \n\nFor out-of-distribution regression tasks, the authors note that existing analyses of in-context learning have been established only for in-distribution settings, and thus investigate its effectiveness on out-of-distribution tasks. Empirically, they confirm that when appropriate training strategies are applied based on the specific out-of-distribution task such as task-shift, covariate-shift, and length-shift, the in-context learning capability of Transformers remains valid.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* I recommend explaining the motivation behind the meta-distribution approach in Section 4.2 and reporting the corresponding results in the main section. This result appears to be one of the key findings, demonstrating that in-context learning with the appropriate training trick can be applied to covariate shift datasets.\n\n*  I recommend providing a clearer explanation of Theorem 5.3 from Wu et al. (2023), described in last paragraph in section 4.3, to better support the justification of context length generalization. In the current draft, it is unclear how the absence of position encoding for Transformer, which improves context length generalization, relates to Theorem 5.3. If this connection cannot be adequately explained, it may be better to omit the use of Theorem 5.3 from Wu et al. (2023).", "broader_impact_concerns": "This work investigates the in-context learning of the transformer and does not seem to raise the concerns of ethical issue.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "FZTpWnxz1U", "reviewer": "Reviewer_zqB3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. The theory considers a very general multi-layer Transformer architecture with multi-head nonlinear activations, tackling a challenging theoretical setting. The limited context window design aligns well with practical scenarios. \n2. The paper provides comprehensive and detailed comparison with previous theoretical literature, making it clear to understand the theoretical contributions of this work.\n3. The findings in Sec. 4.3 that the failure in length generalization of ICL is mainly due to the distributional shift in the positional encodings is insightful.\n\n**Weaknesses**\n\n1. The main theoretical finding of Theorem 3.2 that ICL converges to a Bayesian-optimal predictor lacks novel insights. Many previous works [1-2] have empirically and theoretically revealed that ICL achieve nearly optimal performances on the i.i.d. linear regression task, and the performance improves with longer context length and more training examples.\n2. The practical value of the theory remains unclear. Although the theory considers a more general Transformer architecture and provides a tighter generalization bound, I'm still concerned about the practical value of the theory. Specifically, what does a tighter generalization bound imply for real-world applications? How can these theoretical insights inform the design of more effective ICL algorithms or paradigms?\n3. The results in Sec. 4.1 is somewhat less surprising, since previous works [1-3] have highlighted the limited OOD performance of ICL. Similarly, the result in Sec. 4.2 also seems not surprising to me either, since in my opinion, the ICL performance will get better as the test distribution becomes closer to the training one. Given that the proposed meta-distribution training enlarges the training distribution, it's plain to infer that it will help generalize to new distributions.\n\n[1] What Can Transformers Learn In-Context? A Case Study of Simple Function Classes, NeurIPS 2022\n\n[2] Trained transformers learn linear models in-context, JMLR 2023\n\n[3] A Closer Look at In-Context Learning under Distribution Shifts, Arxiv 2023.05\n\n[4] Pretraining Data Mixtures Enable Narrow Model Selection Capabilities in Transformer Models, Arxiv 2023.11\n\n[5] Can in-context learning really generalize to out-of-distribution tasks, ICLR 2025", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "1. This work theoretically proves that, for i.i.d. scenarios, when a nonlinear Transformer is trained on ICL data, the downstream ICL performance will converge to a population risk minimizer given long enough context.\n2. This work investigates three OOD generalization scenarios: task shift, covariate shift (input distribution shift), and length shift. It empirically reveals that, for task shift, the ICL prediction deviates from Bayesian-optimal predictions; for covariate shift, training on a larger meta-distribution can improve generalization; for length shift, it is mainly caused by the distributional shift in the positional encodings.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Add more discussions about how the practical value of the theory.\n2. In Section4.1, in the sentence \"Comparatively, the mere observation that the Transformer works well on the in-distribution data does not demonstrate its in-context learning ability as a traditional supervised learning model also has such ability and generalization performance over in-distribution data\", should \"... generalization performance over in-distribution data\" be \"... generalization performance over OOD data\"? Is this a typo?", "broader_impact_concerns": "No concerns on the ethical implications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "6PjdzegOZ7", "reviewer": "Reviewer_corT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is generally well-written, and the theoretical claims are largely correct. I would recommend for an acceptance after the requested changes (mostly dealing with links to existing literature) are implemented.\n\nTheorem 3.2 is a really nice one.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors revisit the training of Transformers on linear regression tasks: they consider a bi-objective prediction task of predicting both the conditional expectation and the conditional variance.  This additional uncertainty quantification objective allows the authors to (i) better design out-of-distribution experiments to distinguish in-context learning from in-weight learning (IWL) and (ii) make a better separation between the algorithms with and without using the prior information of the training distribution. Their theoretical findings are corroborated by compelling empirical experiments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Q1** Is it always guaranteed that the Bayes-optimal predictor in Def 2.1 is unique? Is it assumed? If not, can the authors prove it? If not, it'd be enough to substitute $:=$ with $\\in$\n\n**Q2** The problem setup considered by the authors is similar to the one put forth in https://openreview.net/forum?id=AH5KwUSsln. It would be nice to see a discussion on how they relate (even in Appendix A). More in general, credal sets or second-order distributions may be of interest for the authors when modelling distribution misspecification or drift.\n\n**Q3** The authors discuss epistemic uncertainty, and measure it via conditional expectation. The latter has been abundantly shown to be a suboptimal way of quantifying EU (see e.g. https://openreview.net/forum?id=4NHF9AC5ui, Section 2.1). While we agree with the authors that it could be computationally preferable to work with such a measure in the context of LLMs, we at least encourage them to add a discussion on EU measures involving credal sets and second-order distributions (arguably, the most principled way of quantifying EU) by referencing works such as https://proceedings.mlr.press/v216/sale23a/sale23a.pdf, https://proceedings.mlr.press/v216/wimmer23a.html, https://openreview.net/forum?id=MhLnSoWp3p, https://openreview.net/forum?id=1Wg0J2WtEU, https://proceedings.mlr.press/v235/sale24a.html, https://arxiv.org/abs/2401.00276, https://arxiv.org/abs/2308.14815\n\n**Q4** Am I correct if I understood that the experiments in Sec 4.1 are somehow similar to those in https://openreview.net/forum?id=4NHF9AC5ui, Section 4.2?\n\n**Q5** Does the in-context learning/in-weight learning problems discussed by the authors relate to continual and/or active learning (see e.g. https://arxiv.org/abs/2308.14815, https://arxiv.org/abs/2305.14782)?", "broader_impact_concerns": "There are no concerns on the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The AE is not an expert in this area, so relies on the judgment of the carefully selected expert reviewers.  The AE's assessment is that the reviewers performed careful reviews and provide balanced feedback on the strengths and weaknesses of the paper, as well as following up with author responses. 

In final recommendations, the reviewers agree that the paper provides useful analysis and insights about downstream performance of in context learning and would be of particular interest to the UQ community.  Two of the reviewers note limitations in the thoroughness/rigor of theoretical analysis and potential for practical impact.  These opinions are elaborated in the reviews and comments that are viewable by the author.

Based on the comments and the recommendations of reviewers, the AE feels comfortable to recommend the paper for acceptance to TMLR.   The "as is" is meant to indicate that changes indicated in author responses should be incorporated, but additional significant changes are not required.  The AE does not recommend ICLR journal-to-conference track, in accordance with the majority recommendation of reviewers.
Decision: Accept (poster)