In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization
Authors: Ruiqi Zhang
Jingfeng Wu
Peter Bartlett
Venue: NeurIPS 2024 poster (2024)
Keywords: In-Context Learning
Transformers
Approximation Theory
Optimization
Primary area: learning_theory
TL;DR: We investigate into the in-context learning ability of Linear Transformer Block on linear regression problems and its relationship with one-step gradient descent estimator with learnable initialization.

Abstract: We study the \emph{in-context learning} (ICL) ability of a \emph{Linear Transformer Block} (LTB) that combines a linear attention component and a linear multi-layer perceptron (MLP) component. 
For ICL of linear regression with a Gaussian prior and a \emph{non-zero mean}, we show that LTB can achieve nearly Bayes optimal ICL risk. In contrast, using only linear attention must incur an irreducible additive  approximation error. 
Furthermore, we establish a correspondence between LTB and one-step gradient descent estimators with learnable initialization ($\mathsf{GD}-\beta$), in the sense that every $\mathsf{GD}-\beta$ estimator can be implemented by an LTB estimator and every optimal LTB estimator that minimizes the in-class ICL risk is effectively a $\mathsf{GD}-\beta$ estimator.
Finally, we show that $\mathsf{GD}-\beta$ estimators can be efficiently optimized with gradient flow, despite a non-convex training objective.
Our results reveal that LTB achieves ICL by implementing $\mathsf{GD}-\beta$, and they highlight the role of MLP layers in reducing approximation error.

In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization
Ruiqi Zhang UC Berkeley rgzhang@berkeley.coo edu
Jingfeng Wu UC Berkeley mujj@bberkelyy.eee
Peter L. Bartlett
UC Berkeley and Google DeepMind
peteroberkeley. edu
Abstract
We study the in-context learning (ICL) ability of a Linear Transformer Block (LTB) that combines linear attention component and a lincar multi-layer perceptron (MLP) component. For ICL of linear regression with a Gaussian prior and a non-zero mean, show that LTB achieve nearly Bayes optimal ICL risk. In contrast, using only linear attention must incur an irreducible additive approxima- tion error. Furthermore, we establish correspondence between LTB and one-step gradient descent estimators with learnable initialization (GD-3), in the sense that every GD-3 estimator can be implemented by LTB estimator and every optimal LTB estimator that minimizes the in-class ICL risk is effectively a GD-3 estimator. Finally, show that GD-3 estimators be efficiently optimized with gradient flow, despite non-convex training objective. Our results reveal that LTB achieves ICL by implementing GD-3, and they highlight the role of MLP layers in reducing approximation error.
Introduction
The recent dramatic progress in natural language processing can be attributed in large part to the development of large language models based Transformers (34,, such BERT [11] LLaMA 332. PaLM [8]. and the GPT series 28 6 25 In some pioneering pre-trained large language models, new learning paradigm known in-context learning (ICL) was observed (see, for example, 6, 24 20)). ICL refers to the capability of pre-trained model to solve a new task based on a few in-context demonstrations without updating the model parameters.
recent line of work quantifies ICL tractable statistical learning setups such as linear regression with Gaussian prior (see [14 3 and references thereafter). Specifically, an ICL model takes sequence (X1, .1.... XM VM Xquery) as input and outputs prediction of Vquery- where (x,, y:) ==1 and (Xquery> 3query) independent samples from unknown task-specific distribution (where the task admits a prior distribution). The ICL model is pre-trained by fitting many empirical observations of such sequence-labe pairs. Experiments show that Transformers achieve ICL risk close to that achieved by Bayes optimal estimators in many statistical learning tasks [14] 3].
For ICL of linear regression with Gaussian prior, the data is generated as
X AO0 H), 31 (33 x, o²).
where is task parameter that satisfies a Gaussian prior, 3 (0, I)) this setup, 14 3 showed in experiments that Transformers can achieve nearly optimal ICL by matching the performance of the Bayes optimal estimator, that is, an optimally tuned ridge regression. Besides, (35] showed by construction that single linear self-artention (LSA) can implement one-step gradient descent with zero initialization (GD-0), offering an insight into the ICL mechanism of the Transformer.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Later, theoretical works showed that optimal LSA models effectively correspond to GD-0 models M. trained LSA models converge to GD-0 models [8 and GD-0 models (hence LSA models) can provably achieve nearly Bayes optimal ICI. in certain regimes
37
Our contributions. In this paper, we consider ICL in more general setup, that is, linear regression with Gaussian prior and a non-zero mean (that is, Biš] = 3 * 0). The non-zero mean in task prior captures common scenario where tasks share a signal. In this setting, we show that the LSA models considered in prior papers 0 [3 37 incur an irreducible additive approximation error. Furthermore, we show this approximation error is mitigated by considering linear Transformer block (LTBB that combines a linear multi-layer perceptron (MLP) component and an LSA component, Our results highlight the important role of MLP layers in Transformers in reducing the approximation error, and they suggest that theories about LSA [ 38 37 do not fully explain the power of Transformers. Motivated by this understanding, we investigate LTB in depth and obtain the following additional results:
We show that LTB can implement one-step gradient descent with learnable initialization, referred to by as GD-B. Additionally, show that every optimal LTB estimator that minimizes the in-class ICL risk is effectively a GD-3 estimator.
Moreover, we show that the optimal GD-B, hence also the optimal LTB, nearly matches the performance of the Bayes optimal estimator for linear regression with a Gaussian prior and a non-zero mean, provided that the signal-to-noise ratio is upper bounded. These two results together suggest that LTB performs nearly optimal ICL by implementing GD-3.
Finally, we show that the GD-3 estimator can be efficiently optimized by gradient descent with an infinitesimal stepsize (that is, gradient flow) under the population ICL. risk, despite the non- convexity of the objective.
Paper organization. The remaining paper organized as follows. We conclude this section by introducing set of notations. We then discuss related papers in Section In Section 3 we set up our ICL problems and define LTB and LSA models mathematically. In Section 4 we show a positive approximation error gap between LSA and LTB, which highlights the importance of the MLP component and motivates our subsequent efforts to study LTB. In Section we connect LSA estimators to GD-3 estimators that are more interpretable for ICL of linear regression. In Section 6 we study the in-context learning and training of GD-B estimators. We conclude our paper and discuss future directions in Section 7
Notations. We use lowercase bold letters to denote vectors and uppercase bold letters to denote matrices and tensors. For a vector x and a positive semi-definite (PSD) matrix A, we write ||x|] A ==
Vx Ax. We write (---) for the inner product, which is defined as (x,y) := xxy for vectors and (), == tr(AB") ) for matrices. We write Ali] as the i-th row of the matrix A, Am, as the (m, n)-th entry, and A-1,- as the right-bottom entry. We write O,,, 0mxn; In for the zero vector, zero matrix and identity matrix, respectively. We denote the Kronecker product as 8. For two matrices A and B, A (9 is linear mapping which operates as (A 3 B) o C = BCA For positive -efinie matrix A, we write A as its principle square root, which is the unique positive A. We also write A A-
semi-definite matrix B such that BB -BB denotes the Moore Penrose pseudo-inverse. For two sets A, B we write A + B for the Minkowski sum, which is defined as {a (a-bb b: E A, b B}. We also write a B {a} + B for element and set B. We write null (-) for the null set of matrix or tensor.
(AA)) where (-)
2 Related Works
Empirical results for ICL in controlled settings. The work by 14 first considered ICL in controlled statistical learning setups. For noiseless linear regression,
showed in experiments that Transformers match the performance of the optimal estimator (that is, Ordinary Least Square). Subsequent works by 3 [& extended their result to noisy linear regression with a Gaussian prior and showed that Transformers can match the performance of the Bayesian optimal estimator (that is, an optimally tuned ridge regression). Besides, (30 showed that the above holds even when Transformers pretrained on a limited number of linear regression tasks. These papers only considered Gaussian

---

prior with a zero mean. In contrast, we consider a Gaussian prior with a non-zero mean. Our setup better captures the common scenarios where the tasks share a signal.
The empirical investigation of ICL in tractable statistical learning setups goes beyond linear regression settings. For more examples, researchers empirically studied the ICL. ability of Transformers for decision trees (see [14, they also considered two-layer networks), algorithm selection [4], linear mixture models (6).. learning Fourier series (2]. discrete boolean functions [5]. representation
learning 13 [6. and reinforcement learning [7 [9]. Among all these settings, Transformers can
either compete with the Bayes optimal estimators or expert-designed strong benchmarks. These
works are not directly comparable with our paper.
Transformer implements gradient descent. A line of work interpreted the ICL of Transformers by their abilities to implement gradient descent (GD) [35 3 10 38
337. In experiments, [35] 10
showed that (multi-layer) Transformer outputs close (multi-step) GD outputs. When specialized to linear regression tasks, [5 constructed single linear self-attention (LSA) that implements one- step gradient descent with zero initialization (GD-0). Subsequently, M showed that optimal LSA models effectively correspond to GD-0 models, (38) proved that trained LSA models converge to GD-0 models, 37 showed that GD-0 models (hence LSA models) can provably achieve nearly Bayes optimal ICL in certain regimes and provided a sharp task complexity analysis of the pre-training. These papers focused on LSA models. Instead, we consider a linear Transformer block that also utilizes the MLP layer.
From approximation theory perspective, 3 showed that Transformers implement multi-step GD under general losses. In comparison, we consider a limited setting of linear regression and show the LSA models can implement one-step GD with learnable initialization (GD-3), moreover, every optimal LTB model is effectively an GD-3 model. Both of our constructions utilize MLP layers in Transformers, highlighting its importance in reducing approximation error. Different from their results, we also show a negative approximation result that reveals the limitation of LSA models.
3 Preliminaries
Model input. We use x E Rd and y E IR to denote a feature vector and its label, respectively. Throughout the paper, we assume a fixed number of context examples, denoted by M >> 0. We denote the context examples by yX E Rx x R&M where each row represents context example, denoted by (x. 3/)). i 1, ....., M. To formalize an ICL. problem, the input of model is a token matrix given by 38
X T
=
E R(d+1)x(M+1),
(3.1)
The output of a model corresponds to a prediction of y.
A Transformer block. Modern large language models are often made by stacking basic Transformer blocks (see, for example, (34)). A basic Transformer block consists of self-attention layer and multi-layer perceptron (MIP) layer 34, E + MLP [ATTN (E) Here, the MLP layer is defined as MLP (E) := w ReLU (W ,E,, where ReLU(-) refers to the entrywise rectified linear unit (ReLU) activation function, and W,, W2 E R x (d+1) are two weight matrices. The self-attention layer (we focus the single-head version in this paper) is defined ATTN(E) = E + W W EM- sfmx( (onx((Wx2 E) E). where sfmx(-) refers to the row-wise softmax operator, W K, Wq [ddex(d are the key and query matrices, respectively, Wr. Wy E (d++1 are the projection and value matrices, respectively, and M is fixed masking matrix given by
IIN 0 E (M)++))x(X+++)
M
(3.2)
This mask matrix is included to reflect the asymmetric structure of a prompt since the label of query input x is not included in the token matrix 21]. In the above formulation, de, de dy are three hyperparameters controlling the key size, value size, and width of the MLP layer, respectively. In the single-head case, common to set d, d, = d + 1 and d = 4(d + 1), where d + corresponds to the embedding size (see, for example, 34 ). Our formulation of basic transformer block ignores all bias parameters and some popular techniques (such as layer normalization and dropout) to focus solely the benefits brought by the model structure.
3

---

A linear Transformer block. To facilitate theoretical analysis. we ignore the non- non-linearities in the Transformer block (specifically, sfmx and ReLU) and work with linear Transformer block (LTB)
defined as
E W
W Q E
fLTB R 1d++1)x((++1) + R, E +
W, W. E WDW, EM
M
(3.3)
where E is the token matrix given by matrix, M is the fixed masking matrix given by (3.2 Here, the trainable parameters are Wp, W E Rd.x(d+1), WK. We E [Rdex(d+1)) and Wi, W2 E xdyx() We use the bottom right entry of the transformed token matrix the model output to form prediction of the label y. The 1/M factor is a normalization factor in linear attention and can be absorbed into trainable parameters. Finally, we denote the hypothesis class formed by LTB models as FLTB == {SLB: WK, Wq. Wv, Wr, WI, W2, d 2 d, de 2 d dj 2 1}, where fLTB is defined in (3.
-1-1-1 refers to the bottom right entry of
linear self-attention. We will also consider linear self-attention (LSA) defined as
E
W W q E
fLSA R(44)x(M+)) R. E+ E W. EM
WWEM
(3.4)
M
-1-11
where W K, Wo, Wp, Wr are trainable parameters. An LSA model can be viewed as an LTB model without the MLP layer (setting dj = + 1 and Wi = W2 = I). We remark that, unlike LTB, the residual connection in LSA plays no role because the bottom right entry of the prompt E is zero (see (.T)). A variant of the LSA model has been studied by [T] (38 37], where W Wq and W.W are respectively merged into one matrix parameter. Similarly, we denote the hypothesis class formed by LSA models as FLSA == {fLAA WK, Wo, Wv, WP,ch 2 d, dp 2 1}, where
fLSA is defined in (3.4).
Linear regression tasks with a shared signal. Assume that data and context examples generated as follows.
Assumption 3.1 (Distributional conditions). Assume that (X,y,x, v) are generated by:
First, a task parameter is independently generated by 3 N (3", v
The feature vectors are independently generated by x, XI, XM i..d NO, H).
Then, the labels are generated by y = (3,x) x) + yi (3,x, + E,, i= I...., M, where e and
i.d. A (0, o2)
ey's are independently generated by e, E1, EM
Here, 0² > 0, H 2 0, T 2 0, and 3 E Rd are fixed but unknown quantities that govern the data distribution. We denote (E1, .c.....5.. EM)
We emphasize the importance of the mean of the task parameter 3 in Assumption 3.1 A non-zero 3 represents a shared signal across tasks, which is arguably common in practice. This assumption is implicitly used in [2] where they assumed task parameters are close to meta parameter. In comparison, the prior works for ICL of linear regression [ 38 377 only considered a special case where 3 = 0. In this special case, they showed that single LSA layer can achieve nearly optimal ICL by approximating one-step gradient descent from zero initialization (GD-0). In more general cases where 3 is non-zero, we will show in Section 4 that LSA is insufficient learn the shared signal and must incur an irreducible approximation error compared to LTB models. This sets our results apart from the prior papers.
ICL risk. We measure the ICL risk a model f by the mean squared error, R())= = EEJ(E) - y)²
(3.5)
where E is defined in 3.. and the expectation over E (equivalent over X, y, and x) and y.
4 Benefits of the MLP Component
Our first main result separates the approximation abilities of the LTB and LSA models for ICI. Recall that an LSA model can be viewed as special case of the LTB model with W₂ W, I. So the

---

best ICL risk achieved by LTB models is no larger than the best ICL risk achieved by LSA models. However, our next theorem shows a strictly positive gap between the best ICL risks achieved by those two model classes. This result highlights the benefits of the MIP layer for reducing approximation error in Transformer.
Theorem 4.1 (Approximation gap). Consider the ICL risk defined by (3.5) and the two hypothesis classes FLTB and FISA. Suppose that Assumption 3.7 holds. Then we have
inf JEFUre R (f) is independent of [3.
sesia R (f) is function of B". Moreover,
(tr (HV) + o²)²
2 inf R 2 max (M + 1)" (M + 1)a tr (HH))) |1 1 FEFLTB
inf (f) SEFLSA
The proof of Theorem 4.1 is deferred to Appendix B Theorem 4. reveals a gap in terms of the approximation abilities between the hypothesis set of LTB models and that of LSA models. Specifically, the best ICL performance achieved by LTB models is independent of 3 while the best ICL performance achieved by LSA models is sensitive to the norm of (3" - In particular, when |1,3* [ = (MM), the best ICI. risk of the former is smaller than that of the latter by at least a constant additive term. So when 9 is large, the hypothesis class formed by L.SA models is restricted in its ability to perform effective ICL. We will show in Section 5 that the hypothesis class formed by LTB models can achieve nearly optimal ICL in this case
We also remark that the (11MM) factor in the lower bound in Theorem 4.1 is not improvable. This is because LSA models implement one-step GD algorithm that is consistent for linear regression tasks (that is, the risk converges to the Bayes risk as the number of context examples goes to infinity), with excess risk bound of (11/M) [38] 337. So the approximation error gap is most (1/MM). Nonetheless, the [2[1 | /M) approximation gap between LSA models and LTB models shown by Theorem 4.1 suggests that 3 is not learnable by LSA models during pre-training. In contrast, will show in Section 6 that L.TB models can learn 3 during pre-training.
We emphasize that the ability of LTB to learn non-zero mean is joint effect of MLP component and a skip connection. Note that LSA also has a skip connection, but its skip connection is inactive. In comparison, the MLP component in LTB activates the skip connection. Therefore, we attribute the ability to learn non-zero mean to the MLP component, which is the only difference between LTB and LSA. Nonetheless, without a skip connection, LTB reduces to LSA with a potential rank constraint on the parameter, which cannot learn non-zero mean as have proved. The above two explanations take different perspectives to interpret the same phenomenon. Finally, we remark that Theorem 4.1 holds even when H and V in Assumption 3.1 not full rank.
attribute the ability to learn non-zero mean the skip connection:
Does scratchpad help? We have demonstrated that employing a single-layer LSA introduces an additional approximation error in the in-context learning problem for the linear regression task defined in Assumption 3.T Nonetheless, are there alternative structures, apart from MLPs, that could potentially reduce this approximation error? One plausible strategy is to include "scratchpad" in the input token. Specifically, we construct the token matrix as follows:
X
x
E
1
RR (M+))
(4.1)
M
0
which we then input into the LSA layer. We discuss the limitation of this scheme in Appendix Notably, this method does not successfully recover the GD-B estimator defined in Section 5 We
leave it as future work to see whether the token matrix with scratchpad could implement other types of estimators that more effectively address the linear regression tasks defined in Assumption well as whether additional structures could help alleviate this approximation error.
as
Experiments on GPT2. Theorem 4.1 shows the importance of the MLP component in LSA models for reducing approximation error. We also empirically validate this result by training a more complex GPT2 model [14] for the ICL. tasks specified by Assumption 3.1 In the experiments, we use GPT2-small model (with or without the MI. component) with 6 Tayers and heads in each layer.
The experiments follow the setting in [4], except that train the model using token matrix defined
5

---

Loss Curve Along Training Process
layer versus LSA layer
(TS Model
LSA Mocel
5.0
Sext 1200
Wn Malalahh Marl Ma
100
150
220
250
200
Eboch
Figure 1: The test loss along the training process for LTB and LSA layer.
in 3.1. We considere two ICL settings, which instantiates Assumption 3.J with 3 = (0, 0, .... 0)T
and Ø (10, 10. .... 10) T, respectively. We set = 20, M = 40, 0, v Н I in both
settings. More experimental details are Appendix [ each setting, we train and test the model
using the same data distribution. The experimental results are presented in Table From Table D we observe that both models (with or without the MIP component) achieve nearly zero loss when task mean is zero. However, when the task mean is set away from zero, the GPT2 model with MLP component still performs relatively well while the GPT2 model without MLP component incurs significantly larger loss. These empirical observations are consistent with our Theorem 4.1] indicating the benefits of MLP layers in reducing approximation error for ICL of linear regression with a shared signal.
Table I Losses of GPT2 with or without MLP component for linear regression with shared signal.
Experiments on LSA and LTB. We trained both the LSA and LTB layers for 3 1,... 1) and found that the trained LTB layer consistently achieved significantly lower ICL risk than the LSA layer (see Figure [T. In these experiments, we adhered strictly to the previously defined LSA and LTB structures, setting the parameters as follows: = 5. M = 5, 0, and A H La- At each training step, we sampled B = 128 new linear regression tasks.
Model
GPT2
GPT2-noMLP
3 0 x
0.003
0.024
B 10
0.013
8.871
In this part, have shown that LSA models, the primary focus in previous ICL theory literature (see, e.g.. [1 38 37) and references therein), are not sufficiently expressive for ICL linear regression with shared signal. In what follows, will show LTB models are sufficient for this ICL problem.
5 LTB Implements One-Step GD with Learnable Initialization
To understand the expressive power the LTB models, we build connection berween FLTB and its subset of models which call one-step GD with learnable initialization (GD-3). We will first introduce GD-3 models and then show that the best LTB models that minimize the ICL risk effectively belong GD-3 models.
The GD-33 models. GD-3 model is defined as
fGD.a R ((++1)((++) R E T-
1 X (x3 x) M
(5.1)
where E is the token matrix given by r E Raxd and ,3 E Rd are trainable parameters. Similarly,
define the function class formed by GD-3 models as FCD-B = {fcD-B B Rd I [dxd]

---

A GD-B model computes a parameter that fits the context examples (X. and uses that parameter make linear prediction of label y on feature x. More specifically, the first step by using one gradient descent step with a matrix stepsize I and an inifialization B3 on the least square objective formed by context examples (X,y).
LTB implements GD-B. A GD-B model (5.1 is a special case of an LTB model 3.3
by setting
W (-L
Od 0 1 1)²
r
W, = W B 1
W
0
where denotes the entries that do not affecct the model output (hence can be set to anything). Note that I and B3 are free provided that d, 2 d + 1, dx 2 d and d 2 (as required in FLTB). In sum, we have proved the following lemma showing that the set of GD-B models belongs to the set of LTB models.
Lemma 5.1. We have FOD-B E FLTB. Therefore, inf JEFoo. R(f) Z inf R(f).
Optimal GD-B models for ICL. We now consider FGD.,s 3 and its optimal ICL risk. We have the following theorem that computes the globally minimal ICL risk over FGD and specifies the sufficient and necessary conditions for global minimizer. The proof deferred to Appendix
Theorem 5.2 (Optimal GD-3 models). Consider the ICL risk defined by
Suppose that
Assumption 3.7 holds. Then we have
The minimal ICL risk of FGD. is
R()) = = tr H HH (I (I-HAUHAD!))
inf fEFos
(5.2)
where Rt = [MM + 1)H (tr (H LaMM
The global optimal parameters for JGD.9 that attain the minimum
take the following form:
B 3= + null (H) r- n (H where r S
(5.3)
and 3 is the mean of the task parameter in Assumption 3.1 Here, null (H) E IR .Hz 0} is the null space of H, and null (H = { z e Rdxe HHH = 0} is the null spuce 3 In particular, when H positive definite, the global optimal parameter is unique, (3,T) (B", r").
Under Assumption 3.1
the global optimal JGD that attains the minimum
5.2 is unique as
function of E (given by 3.1 and takes the following form:
1 -'X (X3" - ,x,,x) M
(*EE
B
(5.4)
Theorem 5.2
characterizes the optimal GD-B models for ICL. In the above theorem, r as the context length M goes infinity (assuming that H is positive definite). In this case, the optimal GD-3 function (5.2) implements one Newton step from initialization , With a finite context length regularized Newton step from initialization 3
M. (52
implements
Optimal LTB models for ICL. Lemma 5. shows that best ICL risk achieved by an GD-3 model no smaller than the best ICL risk achieved by an LTB model. Surprisingly, our next theorem shows that the best ICL risk achieved by a GD-3 is equal to that achieved by an LTB model. Therefore, the hypothesis set FGD-8 is diverse enough match the approximation ability of the larger hypothesis set FLTB for ICL.
Theorem 5.3 (Optimal LTB models). Consider the ICL risk defined by 3.5). Suppose that Assumption 3.1 holds and that rank (H w²) 2 2. Then we have
The minimal ICL risk of FLTB and of FED. are equal,
inf fEFLTB
R(f)
inf fEJco
R(f)- = RHS of 5.2
7

---

Rewrite an LTB model as fLTB in 3.3 with parameters (* denotes parameters that do not affect the output)
W. W = TT : Wk
vI2 (VT
Vii :). W. W.W.W,WW =
V
1-1
Then the sufficient and necessary conditions for fLTB E arg man R(f) are
8-1 U-1#0 00, V- null (H); V21 - null (H); Y B + null (H),
+ null (H
where 3 is defined in Assumption 3./ and I is defined in (5.3). In particular, when H is positive definite, the globally optimal parameter represented by (V V12, V21, U-1,7) is unique up to rescaling of C-1-
Under Assumption 3.1 the globally optimal LTB model (that is, function in arg feJra is unique as a function of E (given by (3.b) and takes the form of (5.4) almost surely.
Lemma 5.1 and Theorem 5.3 together show that FGD-B is representative subset of FLTe that does not incur additional approximation error. In addition, every optimal LTB model effectively an optimal GD-3 model when restricted to all possible token matrices. Note that the optimal model parameters for LTB or GD-3 are not unique because of redundant parameterization. But the optimal LTB and GD-3 models are unique function of all possible token matrices. The above holds even when and V are potentially rank deficient.
Comparison with prior works. A line of papers considers LSA models for ICL. under Assumption 3.1 with 9 = 0 35 38 37). They show that L.SA models can (effectively) implement all possible GD-0 models that specialize GD-B models by fixing 3 = O. In addition, they show that every optimal LSA model is (effectively) a GD-0 model for ICL under Assumption with with 3 = 0 (see, for example, Theorem 1 in []. In comparison, we consider harder ICL problem that allows large shared signal in tasks (that is, large 3 in Assumption 3.1). In this setting, our Theorem 4.] shows that FLSA (hence its subset formed by GD-0 models), subet of FLTB, incurs an additional approximation error propotional || compared with FLTB. In contrast, FGD ,3, subset of FLTB. does not incur additional approximation error according to our Theorem 5.3 Thus the LSA and GD-0 models considered by [35] 38 37 are not capable of learning the shared signal 3 while an LTB model can learn 3 through implementing GD-3 and encoding 3 in the initialization parameter.
Training and In-Context Learning of GD-beta
We have shown that FOD-B is representative subset of FLTB that effectively contains every optimal LTB model. We now examine the ICL and training of GD-B models.
Nearly optimal ICL with GD-3. We will compare the best ICL risk achieved by GD-3 with the best ICL risk achieved by any estimator. The following lemma is an extension of Proposition 5.1 and Corollary 5.2 in (37 (which is based on 33 that characterizes the Bayes optimal ICL risk among all estimators.
X. Lemma 6.1 (Bayes optimal ICL). Given task-specific dataset (X y) sampled according to Assumption 3.7 let g (X,y, be an arbitrary estimator for y and measure the average linear regression risk by f. (g: X) == [X(, y)² X). It is clear that R. (g; X) = R(g)). Then,
The optimal estimator that minimizes the average linear regression risk L(x:)) is g x,) x 3* + V x X + oo" " (33))...
Assume the signal-to-noise ratio is upper bounded, that is, tr (HU) S. 02 then with proba- bility at least 1 exp (-Ω (M)) over the randomness of X, it holds that L (g";; o² ~ Ii min ).), where Mr and (ox)i21 the eigenvalues of & +
The proof is deferred to Appendix Lemma 6. shows that the Bayes optimal estimator a ridge regression estimator centered at B This is consistent with 37 where the Bayes optimal estimator
8

---

is ridge regression estimator as they assumed 3 = 0. The following corollary of Theorem 5.2 computes the rate of the ICL risk achieved by the optimal GD-B model.
Corollary 6.2. Under the setup of Theorem 5.2 additionaly assume the signal-to-noise ratio is upper
bounded, that is, tr (V/H)= & a². Then we have inf JeFeo. R(f) x² = E min (8, where (dz))2o are the eigenvalues of Ap H and Φ = [Lr(V ( + o + 22/M
The optimal (expected) ICL risk achieved by FeD.S in Corollary 6.2 matches the (high probability) Bayes optimal ICL risk Lemma [.T] iignoooin constant factors, provided that the signal-to- toise ratio is upper bounded. Therefore FGD-3 achieves nearly Bayes optimal ICL risk. As consequence, the larger hypothesis set FLTB also achieves nearly Bayes optimal ICL of linear regression under Assumption 3.1
For the simplicity of discussion, we assume fixed context length during pretraining and inference. Our discussions can be extended to allow a different context length during pretraining and inference using techniques in (37] However, this is not the main focus of this work.
Optimization of GD-3 with infinite tasks. We have shown that FGD.g is a representative subset of FLTB that covers the optimal LTB models and achieves nearly optimal ICL risk. We now consider the optimization in the parameter space specified by FGD,p s. For simplicity, we follow (38 and consider gradient descent with an infinitesimal stepsize on the ICL. objective with an infinite number of tasks. That is, we consider the optimization of gradient flow on the population ICL risk under the parameterization of GD-33
d,3(t)
d dt
8 ar
R(JGn.s)
R(JaD.)),
(6.1)
dt 03
where R is defined by 3.5 and JCD.9 is defined by
The following theorem guarantees the global convergence of gradient flow. We introduce some notation to accommodate cases when H is rank deficient. Let Ps be the orthogonal projection operator onto subspace S. Let H Im (H) be the image space of matrix H (viewing H as linear map) and H- == null (H) be its orthogonal complement. Similarly, let 2 == Im (H (
(HZH, Z Rdxd) be the image space of the operator H and 2 be its orthogonal complement. Then we have the following theorem.
Theorem 6.3. Consider the gradient flow defined by 6.1 with initialization 3(0),T(0); We have,
Ph (B(t) + PH (3"), PH (B(1) = PH (3(0)),
Pz (T()) - P2 (T"), Pzi (T(I) P2 (T(0))
as t co. In particular, if H is positive definite, the gradient flow flow converges the unique global minimizer of ICI. risk over GD-3 class, that is, B(1) - 3 and (()) - I as + cc.
The proof, as well as the convergence rate, is deferred to Appendix G We remark that 66. is complex dynamical system on non-convex potential function of 33 and T. We briefly discuss our proof techniques assuming that H is full rank. The rank-deficient cases can be handled in the same way by applying appropriate project operators. To conquer the non-convex optimization issue, we observe that for every fixed I, the potential a function of / is smooth and strongly convex with uniformly bounded condition number. This observation allows us to establish a uniform convergence for B. When (3 is sufficiently close to 3", the potential function of T is approximately convex, allowing us to track the convergence of T.
Theorem 6.3 shows that optimization of GD-3 can be done efficiently by gradient flow without suffering from non-convexity. However, as we have shown previous sections, LTB utilizes a more complex parameterization than GD-3. So Theorem 6.3 does not imply optimization of LTB is easy. We leave it as future work to study the optimization and statistical complexity for directly learning LTB models.
7 Concluding Remarks
In this paper, we study the in-context learning of linear regression with a shared signal represented by Gaussian prior with a non-zero mean. We show that although the linear self-attention layer

---

discussed in prior works is consistent for this more complex task, its risk has an inevitable gap compared to that of the linear Transformer block (LTB), which is linear self-attention layer followed by a linear multi-layer perception (MLP) layer. Next, we show that the effectiveness of the LTB arises because it can implement the one-step gradient descent estimator with learnable initialization (GD-B). Moreover, all global minimizers in the LTB class are equivalent to the unique global minimizer in the GD-B class, which can achieve nearly Bayes optimal in-context learning risk. Finally, we consider training on in-context examples and prove global convergence over the GD-B class of gradient flow on the population loss. Several future directions are worth discussing.
Optimization and statistical complexity. This paper provides an approximation theory of LTB and an optimization theory of GD-3. However, the statistical complexity of learning LTB or GD- is not considered. The work by B7 provided techniques for analyzing the statistical task complexity for pre-training GD-0. An interesting direction to extend their method to study the statistical complexity of learning GD-B. However, their method crucially relies on the convexity of the risk induced GD-O, while we have shown that the risk induced by GD-3 is non-convex. New ideas for dealing with non-convexity are needed here.
From LTB to Transformer block. We focus on LTB in this work, which simplifies vanilla Transformer block by removing the non-linearities from the softmax self-attention and the ReLU activation the MLP layers. This simplification allows us obtain precise theoretical results for LTB (such as its connection to GD-3). On the other hand, non- non-lincarities are arguably necessary for Transformers to work well in practice. An important next step is to further consider the theoretical benefits of non-linearities based on our current results.
Roles of MLP layers. The work by [5] (and references thereafter) empirically found that MLP layers operate as key-value memories that store human-interpretable patterns in some pre-trained Transformers. Their work motivated a method for locating and editing information stored in language models by modifying their MLP layers (see, e.g., (9(23)). ). Our work proves that the MLP component enables LTB to learn the shared signal in linear regression tasks, which cannot be done by a single LSA component. We leave it as future work to theoretically clarify the information stored in the MLP component.
Acknowledgments
We gratefully acknowledge the support of the NSF and the Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMS-2031 10883233883 and #8 #814439, and of the NSF through grant DMS-2023505.
References
[1] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Whirty-seventi Conference on Neural Information Processing Systems, 2023.
[2] Kabir Ahuja, Madhur Panwar, and Navin Goyal. In-context learning through the bayesian prism. In The Twelfth International Conference on Learning Representations, 2024.
[3] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learn- ing algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2022.
[4] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers statisticians: Provable in-context learning with in-context algorithm selection. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[5] Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in-context learning in transformers and LLMs by learning to learn discrete functions. In The Twelfth International Conference on Learning Representations, 2024.
[6] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:18777-901,, 2020.
10

---

[7] Yanda Chen, Ruiqi Zhong, Sheng Zha, George Karypis, and He He. Meta-learning via language model in-context tuning. arXiv preprint arXiv:2110.07814, 2021.
[8] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.023: 2022.
[9] Damai Dai, Li Dong, Yaru Hao, Zhifang Sui, Baobao Chang, and Furu Wei. Knowledge neurons in pretrained transformers. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 8493 8502, 2022.
[10] Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can GPT learn in-context? language models secretly perform gradient descent as meta-optimizers. In Findings of the Association for Computational Linguistics: ACL 2023, 2023.
[11] Jacob Devlin, Ming-Wai Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 04805, 2018.
[12] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In International conference on machine learning, pages 1126-1135. PMLR, 2017.
[13] Jingwen Fu, Tao Yang, Yuwang Wang, Yan Lu, and Nanning Zheng. How does representation impact in-context learning: A exploration synthetic task. arXiv preprint wxXi:22309000044 2023.
[14] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30599 2022.
[15] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, 2021.
[16] 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. In The Twelfth International Conference on Learning Representations, 2024.
[17] Jonathan N Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. arXiv preprint arXiv:2306. 14892, 2023.
[18] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in-context learning. In International Conference Machine Learning, pages 19594. PMLR, 2023.
[19] Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. In The Twelfth International Conference on Learning Representations, 2024.
[20] Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: systematic survey of prompting methods in natural language processing. ACM Computing Surveys, 55(9):1-35, 2023.
[21] Arvind Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations, 2024.
[22] AR Meenakshi and Rajian. On a product of positive semidefinite matrices. Linear algebra and its applications, 295(1-3):33 1999.
[23] Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations gpt. Advances in Neural Information Processing Systems, 35:17359-17372. 2022.

---

[24] Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. arXiv preprint arXiv:2110.55943, 2021. 2021.
[25] OpenAl. Gpt-4 technical report, 2023.
[26] Reese Pathak, Rajat Sen, Weihao Kong, and Abhimanyu Das. Transformers can optimally learn regression mixture models. arXiv preprint arXiv:231 2023.
[27] Kaare Brandt Petersen, Michael Syskind Pedersen, al. The matrix cookbook. Technical University of Denmark, 7(15):510, 2008
[28] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
[29] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[30] Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diver- sity and the emergence of non-bayesian in-context learning for regression. arXiv preprint arXiv:2306. 15063, 2023.
[31] George AF Seber. matrix handbook for statisticians. John Wiley & Sons, 2008
[32] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, MtaieeAAne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 2023.
[33] Alexander Tsigler and Peter Bartlett, Benign overfitting in ridge regression. Mach. Learn. Res., 24:123-1, 2023.
[34] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[35] 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, pages 35151-35174. PMLR, 2023.
[36] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierrie Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma. Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Huggingface's transformers: State-of-the-art natural language processing, 2020.
[37] 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? In The Twelfth International Conference on Learning Representations, 2024.
[38] Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306 09927, 2023.
12

---

A Notation and Variable Transformation
In order simplify the proof, let's first describe variable transformation scheme for our model and data coming from Assumption 3.1
Definition A.1 (Variable Transformation). We fix M as the length of the contexts. Recall that X, x are, respectively, the features in the context and the query input, while 3 and 3 are the true task parameter for the inference prompt and its expectation, respectively. Following the Assumption 3.1 we have B ~ (G", V From the definition for multivariate Gaussian distribution, we know there exists random vector 0 such that
= 3" + 4
(A.1)
where
~ N (0, Ig)
(A.2)
Recall the noise vector is defined as e X,3 and is generated from N (0, IM).
Rank deficient case Note that, even when V is rank-deficient, the variable transformation in A.1 and (A.2) still hold. This can be seen from the definition of the multivariate Gaussian distribution (Def 20.11 and the discussion below in (31)): A vector ] E I with mean 3 and covariance matrix v has a multivariate normal distribution, if it has the same distribution as A + 3 where A is any d x m matrix satisfying AA T y and 0 ~ N (0m, Im). Here, we can recover the variable transformation above if we take = d and A V
Notation Before we delve into the detailed proof, let's repeat the notation part with some additional notations. We use lowercase bold letters to note vectors and uppercase bold letters to denote matrices and tensors. For vector x and a positive semi-definite (PSD) matrix A. we denote | = vx Ax. We denote (:,-) as the inner product, which defined as (x,y) = y for vectors and (A, B) = tr (ABT) for matrices. For a matrix A, we denote ||Allop || as the operator norm and the Frobenius norm, respectively. We denote Ali] as the i-th row of the matrix A, Am,n as the (m, n1)th entry, and A-1,-1 as the right-botom entry, We denote O,,, 0mxns In as the zero vector, zero matrix and identity matrix, respectively.
For a positive semi definief matrix A, we denote A+ for the principle square root A, which is defined as the unique real matrix B such that B is positive semi-definite and BB = BB A. For positive definite A, its principle square root is also positive definite. We denote A as the
Moore-Penrose pseudo-inverse for any matrix A. We also denote A
[A
We denote
& as the Kronecker product. For compatible matrices of proper size A, B, C, B linear mapping which defined by [ A .C = ACB. We denote = & H and
& A is
Φ1 2 62 ... da 0 are its ordered eigenvalues. We also denote \1 X2 ... 2 A 0 as the ordered eigenvalues of H, and X > as its minimal positive eigenvalue. Finally, we denote another important matrix
M + M
(HV) o²
Ω
H H
La.
(A.3)
Note that, under this definition and our assumption on H and L, we have D is invertible.
13

---

B Proof of Theorem 4.1
Proof. The fact that inf feFura R(f) does not depend on the vector 3 is subsumed in the Theorem 5.3 so we do not prove it here. We are going to prove the inequality in the theorem. First, from
Thoorem 5.3 we know that,
R() = = + (r -r ( 2 0-1)
inf
(B.1)
SeFire
where $2 is defined in A.3 FLSA, which denoted
Then, it suffices to lower bound inf yeFisa R(f). So let's take an arbitrary
E
W W
f f(E) = E W. EM
M
where is the token matrix defined in (3.)) Since the prediction is the right- entry of the output matrix, know only the last row of the product W Wy attends the prediction. Similarly, only the last column the on the far right in the above equation atttends the prediction. Since the last column of E is (x 0) calculation (since other parts are multipled by zero). Therefore, we denote
know that only the first d rows of the product W K We enter the
w Wv
W We =
11
U12
u-
u12
where Un E Raxd U12 U21 E R u-1 E R, and denotes entries that do not enter the final prediction. The model prediction can be written as
EM E M
U x
((E) (n u-1)
112
1 Y yTy Ty use x.
=
U M
x
Un
U2i
x y u1 + u-1
M
X X.Un +u
Step 1: simplify the risk function. 3.1 and Definition[A.I we have
We use /3 to denote the task parameter. From the Assumption
= X3+ - (B,x)) +5, B~N(3, v) B= 3" + +
and
X x lid. N H), e[i),e lid. N (0,02), (OLL))
Then the model output be written
((E) = U21
T 1
x
x
Un + uni M
1
XT y uie +u M
1
x U +u- 17 y U1e
M
= (122 + u-13
1 x (Un + BuTo)
M
+ un
M
1 upe + u- M
x
s x Un + u-
M
2 T
xBu
+
M
1
=Te e. u-
To simplify the presentation, we denote
2 U21 + +111
x
1
x (Unr +BuTe)
M
22 un
T T
X
uiz + u-
1
X
Un #u-
2
M
M
M
1
T = E E uuuu
M
14

---

Since x, X, E, are independent, we have
R (f) = E (S(E) (3.x) - -8)
0² + B (S(B) (3.x))
(e is independent from other variables and zero-mean)
E (x) + 22 + 23 -
+ o²
H, EE 21 #2 + 23 /) (21 + z2 23 3 o².
Note that 21 does not contain e, 2g is linear form of E, and ø3 is quadratic form of E. Using Therefore, we have
N(0, o² Ia)),, we have [[(2] - 3) 2 and E|
R (f) 0²
H. E
21 3)
(2 33)
+ HH, E22=)
H. E2zz3
S2
S8
8.
2 H. E
21 1
23
(B.2)
S4
Step 2: compute Sj.
By Lemma
H.4 and that X is independent of all other random variables, we
have
H, E21z1)
EEr
(U11+BB0 + +
1
X x. X.
Uuei + +u_13)
(nnl u-133)
x (U11 + Bu
H
M
M
MM +
UU11 + Bulz) T
U22 + u 13 H.
U21 + u- 13)
H (Un + Bu
.M
1 M
Egtr
tr
U21 + u-1
(nn) + +
H
Un + Bulz)
H (Un + Bubp
We denote
b = U21 + 3-1 E Rd
A Un + Buu E Radd
(B.3)
then applying 3 = 3 + 1 1 get
H. E2)2
M
Ey Bo ir
(A+ H. (b uu v=D (o+m, T
.H A + w v
H
M
1 + M EoLN
tr (o+t (b + u- +
tr
A u +
H A v tūu
H
Note that the first and third moments of are zero. So the above equation only involves the zeroth, second, and fourth moments of Therefore we have
M + tr (A Hbb HAH) + M
1 1 6 H)) tr (ATHAR) +12 + +
M
The leading term
(B.4)
where T2 and T4 are the second order and the fourth order term, respectively. More concretely, the second order term is
T2 = M
M
Etr u- A HbO / ut 0
4 H $ 0b T HAH
15

---

+ +u-uu120 Hbo ı + - A HV ob Hy H
+ u1zn
Hbb HV H 22 A A H0 H
θ H tr A HAH) + Hb tr Bnnس y y H
+ M
.....HHH H tr U11 vzHAH) + 2u 2u_ b tr (A (A }
+ 2u-1b bb
M+1 ++1
M 2u H Hb
Hb + 2u-1 H4HAHu + b HyHH u Hu12
+ u tr (A (A'HHHAH)}
M a tr (wiHw:) tr A HAH
+ tr (W!HVA) b Hb u + 4u-1 b H
2(M + 1)
T (M+ HIH
1
u-1b
[rr (H H HUH AHu
tn(H b 2
M
M
M
M
4
+ , tr
A
HIH
+
tr(H AH
+ u- b HıHAHu (B.5)
M
M
M
The fourth order term is
M +
Etr u U1in 4H wtHy
T1
M
1
4 H 0 tr (nn (.HutouTHH)
+
Etr
u
M
M +
w_1
T Hu
E
v H
M
M +
_², 2tr (HVH))
+ tr (H 1 Hu
(B.6)
M
For the cross term in Si, we have
H. 1Ez13
T 1 x (U11 Buu HB
U8i u-11
(B.7)
10
U21 u- 1
H
(U Bujz)) HB3
(From the distribution of X)
-E
(b+ u- v
H A
0u
H 3r v
(By (B.)
=b -b'HAHS +E U_
H
Ou, H3 + u_1 0 + ot Hu Hy t
=b HAHS + uu tr (HV)
+ u tr (HAHV) + u HUHb T H
(B.8)
where the last row comes from the fact that 0 ~ (0, L,) Moreover, we have
H. E,3
,3 HH3" + tr (HV)
(B.9)
Combining B.4. B.5
B.6
B.8
B.9.. we have
M
S = M
tr (A Hbb
HAH)
M tr 6b T H A (HAH)
+ M +
(2tr
(HDHH) + ( (HV) HigHuI2
M
16

---

2(M + 1)
uuib" [trHH)) HWH AHu
M
M++
4
+ 22
A
HUH
(HH AH +
M n-1 b H
M
M
H
+ -. 10r (HV) H33 + " u-1ir 1rr (HAHV)
M + M
u
HVHHb
+ 3 + ( bT
HVH
(H(H))
u Hu12-
M
(B.10)
Step other terms.
Let us compute S2, S3 and S4. Using definitions, we rewrite 22 as
2 U21
T
X
up +
1
X
Un u-
2
X
M
M
M
1 T x
= -b"
x E. ufu + +
A + u-1
x
(B.11)
M
M
M
Since N IMM N (0, Ia) and they are independent, all terms in S2 vanish except if the term contains even orders of or e. So we have
1
S2 = (H, )
M x
E
Huiz e e bb i XAHA X
+ 24_1 E XAHu e Xb 4w21 E x T H 0 x
0² b Hb u. Huja + tr HAHAT) +u 22 Hb 4u tr (H u. Hu12
M
(B.12)
For the other two terms, we have
S M " 22 B Hu T E)]
1
o0MM + 2)
w² 11 u Hu
(B.13)
M
and
S4 =2 (H,E (21 3 23
= -20B {[(110)
{ [121
+ u-
M MM X (Un+BuT)) - Bu]e) -3
H.
1
T
e
u
M
( T
=202 1F
H.
A v v
- 3
Huj
= 20
HA - 3
Huj + uu
tr H
u Huj2
(B.14)
Step 4: combine all parts.
Combining the four parts BB B.12 BB.13) and B.14 have
R(f) 0² S₁ S2 + S3 S.
We remark that in the 1.
B.13) and (B.14), the parameters are A, b, u and u22, instead of the original parameter U11 U12 U21 u-1- From the definitions of A and b, we know there is a bijective map between (U u1g U21, w_1) and (A, b, u12, u-1), these two parameterizations are equivalent when computing the minimal risk achieved by a model a hypothesis class. From the expression of S1, S2, S3, S4, rewrite the risk as
B.
R (f) o²
17

---

M
1
o²
HYH + (r(HH))H) b Hu +
b HHH Hu12
M
M
+ u² _ tr
TA
M +
HYH+ + (H H
AH
d²
tr (HAHA")
M
M
M
M
o²
+2 1b
(r(H))H+
HVH
H
AHu - 2u_ (HAHV) - HvbHo
M
M
M
b Hb tr A HAH) 44 b HVHAHUR + 442, tr (HVHV) 2 Hu12
M
M + b
22o2b'HAH HAHA H) b 2 20-1tr(HVV) b HAHu + 2 HAH (u_1u12)
M
III
HH 2 (HV) + -3- H (u-iu12) tr
+u²,
M + [(2TT(HHH)++
(H
2+
4 M
o (HV) +
M +
0
M
I+ 11 + IV, where I, II, III, IV are defined as above.
(B.15)
Step 5: lower bound the risk function. Let's first define a new matrix, which by definition is invertible:
M +
H. + tr (H + + I.. M
Ω =
(B.16)
M
So we can write as
= b H Hu - 2r A H AH) 2u_ H AHu
-uu (HAHı)- Hv.HH
H HHuup H H - A A H w 12
- -H ' fn )
(Htupb H -u H H
2 -Ir H
(B.17)
-Ir 2-1
where the last line comes from the fact that Z and H W
commute. Next, we claim II 2 0. To see this, suffices to notice that
4u_1. HęHAH 2 -4 H
H
TH
H
AH
u
F
2
T
AH²
- 4
H W
u12
F
b Hb
tr
A
HAH
- -4 tr (HVH)) 112 Hu (B.18)
where the last line comes from the fact that ||||
= (AA )AAT) for any matrix A.
18

---

Then, let's consider III. Notice that actually, III can be viewed as quadratic function of A Hb, which can be easily minimized. More concretely. have
M
III
8r -tr (H + r u1
H - (HV) + o²) 2 u12
+1
M
3 - (tr (H + 0²) uini1a )
M +
A Hb
H
M +
M
M
A Hb
3 (tr (H + o u-1111/
M ++ 1
20.
(B.19)
Therefore, one has
M
T
M+1 3r (tr (H و²) 1-
H gr -(tr (HV) + o²) (اسسس
Combining the three parts above, one has
R (f) 02
2 IV tr (( 2 )
M
T
8r 33- (tr (HH) + 0²) 1-1112
H 3* (tr (H + o.) 8tnir
M +1
tr (H tr (H:HHz) ) 12-1
2(r (HV) + x²
+
3 H
B H (n_nm1)
M +
M+1
inf yCFLre R(f)-02
2tr (HVHH) +
3M +2
(ir(HV) +o2]² (11 H (u-111))
M(M +1)
Therefore, one has
2 (tr(H))) 02)
R (f)
inf
R(A2
3 H
3r H (1-111))
FEFFLe
M +
M ++
3M 2
(tr (HrP)) + 02)² (w-1 H (w-1U12).
+ 2tr (HOHT)
M(M + 1)
The right hand side the above inequality quadratic function of u-1, so we can take its global minimizer:
(u=HH)++'*)
W- 111
M+1
(Ir (HV) + 02)2/9*
2Lr (HVHH) + M(M+1J
3M42
to lower bound the risk gap as
(In(H))++))?
R (f)
inf R()) 2 JEFLTB
(M+11)²
11
M +1
2tr (HVHV) M[M+1] (tr (HV) + 02)3
Finally, to simplify the results, we notice that on the one hand,
(u(H)))+0')
(rr (H +
(M++1)*
2tr (HVH)) +
M(M+1] (tr (HV) + 22)
3M-22
2(M 1) (HVHV)))
19

---

On the other hand, one has
(u(H))+0*)*
[M+1)]
(M + 1)(3M +2) S 1
M
(tr (H + 02)² S
2tr (HVHV) + M(M+1)
3M-2
3(M + 1)
Therefore, have
2 3(M + 1)" 1 2(M + 1)2r (HUHW)
(tr (H + 00
R()) inf R() 2 max /EFi
||8
(u(H)))+0))
When M+1 2(M (M+
3[++1) have M+T
(M+I))n(HHHW) which implies
(f) inf R()2 2
1
(Ir +
118 2 (M 1) (HVHH) -1
(Ir (HV) + +
FEFuBB
M+1
2(M + 1) (HVHV)
Therefore, finally have
2
(tr (H + +
R (
inf R()) 2 max
110
FEFL
3(M 1)' (M 2 (H
Note that this holds for an arbitrary f E FiSA, so the proof finishes by taking infimum the left hand side.
20

---

C Proof of Theorem 5.2
Proof. Recall that from Assumption 3.1
x~ x~,,0, v= N(,,,( y= y=3B+e
where
B ( v) ( E~N IMt)
Step I: compute the risk function. we have
From the independence of e, with other random variables,
X
X 2 M
RM (3,1) = E
x
B - 3
+ T
B
x
T.
M
ld - r X
2
E
x
(3 - ()
+ M
(x
TX T E) +
M
X X
#2
62
H, E
-
OE
Mr E (x TXT
+
M
(C.1)
where the last line comes from the independence between 33 and X, x, as well the fact that x ~ N (0, H) and the property of tensor product o. Note that, 3 ad I are learnable parameters
here, and we should set them apart from the task vector or the prior mean vector 3 Recall that
B ~ N (3",, MD we have E (B -3 @2
=
+ V. Moreover, using the following
N (), x N~> (0, .IM,
we have that
E Tx e)
Btr XxTTxXXTXX M0² (HTH,
Bridging the two terms above into
we get
82
X
X
RM (8,I)
H. E
r
(B 3
X
X
622
V o²
H, E
La
PHT T
+
X X M
E
La r
H. (3 82002
692
X
H,B
L
PHT
to²
(C.2)
M
M
V2
First, let's compute V1. From the definition above, have
Vi = (B 3²) Hr (3 - )²
where
X
)
82
Hr: [1- -
oH
(C.3)
M
21

---

X
Ls r. X
L
H
(the definition of the tensor product)
M
M
tr HT [HTHI HT
H - -H-HHT+TTT
H+
M +
HTH
M
M
(x ~ N (0, H) and Lemma H.4
tr
HT
HI
[Ia
TH
H [Ia
TH)
H
HT
HIH
0.
(C.4)
M
M
Then, let's compute V22 We have
x
0²
r. X
La - X
σ²
la r-
L
THI
=E
THI
M
M
M
M
(
+
M +
V THE
PHI
+
H
H rT.
M
M
N N,,, H) and the Lemma H.4 we know that it invertible. Moreover, holds that
[x[:] i.d.
From the definition Ω in A.3
X
X
82
v THE THI T
E
La r
M
M
THE
THE
v - WHAN
Therefore, we have
( (H H H st HArH H H
V2 tr
(Hæ) (H wH' - )
Step 2: solve the global minimum. Now we have
R (3,1)) o
tr HH HT
1 Hr
(3 - TT
(Ia Tu-Ii IH HI. IH)+ +
H
HTH
(B 3r)
M
M
[
-Ht - -
(H H
+ tr
+ tr (H tr H wH n H vHb)
(C.5)
tr (HV) [H Ω HH'HH
tr (HV) tr (HH!HH)) 2-1
where the last line comes from the fact that Ω and H w commute. Taking infimum over all B, r in the left hand side, we have
= R -2 (HV) tr ((HtHHt) (2-15 )n-)).
inf R()) fEFeo.
3T
22

---

On the other hand, we take
3- r= = DH S n
(C.6)
where Ha is the principle square root of H and H is its Hoore-Penrone pseudo inverse (see notation part in Section A). Then, we have
R (3" r o² o²
[ H 2 H H. )
+ (H tr H H
Since
H H H - -
Ω VH - - H
and H commute)
Ω H a H Odxd.
H -H²)
Therefore, we have
R of tHV) ( 12-
inf SEFwo
R(f)-o²
inf R (3,
3.
Combining both directions, conclude that
inf R(f) - o² inf (B,,)) - -o tr tt tr ( feFen.
52
(C.7)
3.1
Step 3: identify all global minimizers Above we show that R (B", achieves the global minimum of the ICL risk over GD-B class. Now we will figure out the sufficient and necessary condition to achieve the global minimal risk. From the proof above, we know that for any , E Rd T E Raxd it holds that
V1 = (*) Hr (B - -²) =0
R (B,T) = inf
R() JEFoo
H. (
H H. ss H.rH. H HiwHH
Since Ω is invertible, the second equation is equivalent to
HATH H Hi Odxd:
which is linear system of T. Since I is proved to be one solution, all solutions of this equation is
r z E ddxd H ZH
Oaxd = r + Im (H
where the last equation comes from Lemma H.6 Here, + denotes Minkowski sum. This is defined as A + = {a + b, A, E B} for two sets A, B and B. Under this condition, we know that
B {a} B for entry and set
HI M
HT
M
Hr H H +TM H+
HI"H.
M
M
Consider the following inequality:
tr HH M
HT
I H (I' + r
H
I +
M
H.T:T
M
M
23

---

La H
T T
H +
M +
H HI
M
M
Ia
M +
H
M (VM
M
la
M +
H
+
M
Ia
M M+1
M
M
M +
Odxd-
Therefore, we have
(3 - 3*) Hr (B - (8) --==0 =0
tr
HT M
HI
- [B0 -3") H
Ia I- H² + H
Ia
M
M +
+
H.T M
H
H (B -3)
M
H (3 3r)
0₂
B = B" + null (
( 3 3" null (H).
where null (.) denotes the null space of matrix and denotes the Minkowski sum. The last equivalence comes from Lemma H.6 Therefore, we conclude that the farr achieves the minimal ICL risk in GD-B class if and only
B = 3a + null (H);
-= Im H 22)
Specially, if H is positive definite, there is unique global minimizer in GD-3 class with parameters
β 3 r=r.
Step 4: equivalence the hypothesis class Finally, we show when H is rank-deficient, any global minimizer actually corresponds to single function in FGD-3 almost surely. For arbitrary global
minimizer, assume B 3 + h r + Z, where h e null (H) Z Im (H02). Suppose
we have prompt X,y, x, which follows the Assumption 3.1 and the token matrix is formed by we have
3.11
r
fa. (E)
=
B
x
(X3 - y),
M
F
X
(X,3* -)),x
M
+h x M
x zx X,8
M x
ZX T Xh
x x Xh 1h+0 +x ZXTx y.
M
Notice that X(1),x i.d. N (0,, H) have
E (h ThTx)² h Hh 0
since h E null (H) Therefore, we know h T x almost surely, and similarly, Xh 0,2 almost surely. Finally, we have
(
(+x22XT)
E
Etr HZH
0,
M
M
which imples x1 Z O, almost surely. Therefore, conclude
fa. (E) =s (E) almost surely.
24

---

D Proof of Theorem 5.3
Proof. Recall the definition of Linear Transformer Block class:
fLB ((++))xM++) R
E
W WQE
E
W W1 B Wv EM
M
where E is the token matrix defined by 3.1 Let's take an arbitrary function in the LTB class with trainable matrices
WK, Wi E Rxxx((1))
Wp. Wv (4+(( W1. We E R(x((++))
Similar to the proof of the Theorem 4.1 we know that only the last row of W Wi and the first d-columns of W Wi attend the prediction. Therefore, we denote
W.W, W.
W W
( 21
*
W
V11
12
(D.1)
;- R, Vn E Rdxd, and denotes entries that do not enter the prediction. Then, the prediction of LTB function can be written
where 7, V12, V21, E Rd w-1
EM ET M
Vn 11
f(E) = Txx+ x (v2r
(221 v-1)
X
V12
-TTT
V21
x X
Vn
+ V21
X
y
V12
T
v-1
1 M
y X
Vn
1
yyy
v12
T
M
M
M
Step 1: simplify the risk function. 3.1 and Definition A.I we have
We use B to denote the task parameter. From the Assumption
y XB++. y (3.x)
(=(xx)++ tE, B~N(8", 1) 3= B +
and
X Lid. N H), efil,e Lid. N (0,00), 0 (~.().
Then the model output can be written
1
1 M v12
f(E)
V21 M
x X
Vn
+ 221 M
X y v12 +
M
Tx. Vn tu-
yT + (v21 + v-13)
1 x
x (V11
(3T()) xx
M
V21 M
T x vI tui
M
x Vn +
T X + "Me" v x.
M
1
To simplify the presentation, we denote
1 X (V + +Bzx)))
2 V21 + v-13)
M
1
2 M
2 Vil x g vi + w-
x. Vn w-1
M
M
T Tee UVIV
23 M
Since x, X, e, 3 are independent, we have
R (ʃ) = E (S(E) (3.x) - -8)
2
25

---

0² +E ({(E) - (3,x)) )²
(e is independent from other variables and zero-mean)
E
x2 + B. +
H E
23 + + - to? 21 22 2 + B) (21
Note that 2 does not contain e, 22 is linear form of E, and 23 is a quadratic form of e, Using
(0, o Ia), we have [1 + () 2 0 and [[ = 0. Therefore, we have
R (f) - 0² = H. 21 + 7 - 3) 21 + 7 3)
+ H. Ezg22
H, Ezyz3
54
8
o
2 H.E E
21 + - 3) 23
T
(D.2)
Step compute the risk function. Let's compute the risk function. Note that, the risk function is in the same form as the risk function of LSA layer, which is in the proof of Theorem 4. (see Section and equation (B.2). There are two differences: one we replace U11 U12 121 u- here with 11, V12: V21, V-1- The second difference is that we replace 3 in (B.2) with 3 -. This equivalent to replacing 3 with 8 7, since / - N (B* y, 1) Therefore, similar to (B.15). the risk function in (D.2) can be written as
RJ) 0²
M +1
1 HIH + (r(H)) b Hv12 +
0² M
=
b Hb Vig
.
M
22 tr
AT
M +
HWH (r(HV)H
AH
o²
tr (HAHA")
M
M
M HvH
o²
+2v_1b' b
H
H
AHv12 20-1tr (HAH Iv
M
M
b Hb tr ( HAH) + 41-11 -T H + 40 tr (HIHD) vig
+
M
VI
M
TT
HAHA H b- 26 HHH (3 7) (rr1)((HDD b HAHv12 20 HAH M
2THH (F-1V12)
VII
(0* 7) H 7) - (tr (HT) o²) -7) H (F-IV12) (H))
4u2, [( (2t (2t
M + (HW)2
4
M + 2 od vI
-VI_ + 2
oer (HV)
M
VIII
V + VT VII VIII,
(D.3)
where V, VI, VII, VIII are defined as above and Ω M+1 H WH
to
LI and b
M
M
221 + 3- E Rd = Vn + BVv E Rdxd
26

---

Step 3: solve the global minimum of the risk function. This is very similar to step 5 in Appendix The V term in (D.3) is actually equal to the term (B.15). so we have
- A TH -H -Htwtt ]
H H T H - H A H H s2
H i TH
H
= tr
H
2 -t (H
= -r (HHWH))) 1
where the last line comes from the fact that Ω and H commute. For the same reason, we know that the VI term above is equal to the II term in (B.15). which implies VI 2 0, since
4v-1 b H 4 b H
H
THE
H
AH
H
2 V12
F
F
6 H
H
AH
H 4
V22
b Hb. tr
HAH
4v
(HPHV)
V Hv12, (D.4)
where the last line comes from the fact that | tr (AAT) for any matrix A. The term VII above is equal to III term in (B.15). except that replace 3 with 3 7. Therefore, we have
VII
M
Br
(tr (HV) + + U-1V12
H H 1- (tr (HQ) o²) 0- V12
M +1 1
M
(3 - (ir (HV) + + V12 )]
M (M++)) H
=
A Hb
M 1
M
M
8 -7- 1 (HW) 2-1112)
A Hb
M + 1
2 0,
(D.5)
which implics
M
H H((8 - 7 - (tr (H + +o7)e. 1- 1
VII M M+1 --r (Hr o²) -- 1V12
Combining the three parts above, one has
R (f) o²
> VIII tr (HHHHH) 2 1 )
M
T
- - (tr (H + 0²) e-iV12
- (tr (H + o²) v_ V12
M +1
tr (H(() - tr (H.HH)) 2 1-1
1
(0r (0" -7) T H ( - 7)
+
M +
inf R(
(tr (HV) + +02)
(3" -7) H (V-1VV12
M ++
3M +
+ 2tr (HVH))+ +
M(M (tr (HQ) + (H²)) (V-IV12) T H (V
27

---

Here, the global minimum of GD-3 class is taken from Theorem 5.2 whose proof does not reply on the proof here. Therefore, one has
inf R()) 2 M+1 (8" 7)) H -7) 2 2 (HV) + + feFoo.
(f)
-3)
H (U-1VV12)
M ++1
2ir (HVHH)) + M(M +1) (Ir (H(V + =2) (V-IV12) H (V-1V12).
3M + 2
The right hand side in the above inequality is quadratic function of U-1 so we can take its global minimizer:
(o=H))+t*))
V_
V12
M+1
(tr (H 02)3 (3" - -7)
(HVHV) + 3M+2
M(M+1]
to lower bound the risk gap
(I=(H)))=022
R (
inf R()2 2
(M+1))
18--710 2 0.
M+1 2tr (HУHИ) + [(M+1] (tr (H (HV)+ 02)
3M+2
Note that, taking infimum on the left hand side, we get
inf R (f)
R(f) 2 0. fEFan.
inf
feTum
On the other hand, in the main text we have showed that FGD-B C FLTB, which implies
inf
(f)
inf fEFo..
R(f) I 0.
SEFLTB
Therefore, we conclude
2 Ω xr)) R(f) = (H tr ((H WH
inf
R(f)
inf feFco.
fEFLBB
where the last equation is from Theorem 5.2
Step 4: sufficient and necessary conditions for global minimizers. Let's verify the conditions for the global minimizers. For a function in LTB class, the sufficient and necessary condition for it to be a global minimizer that inequalities in the above step all hold, which are
H
H
+
H
A
H
H H
(D.6)
H
AH
=
H
V12
(D.7)
F
F²
0-1
H
6 H
|0 H W
H
AH
|
(D.8)
M
H
A Hb-
8 3-- (tr (HV) + o²) 2-IV12
0axds
(D.9)
M+1
(Ir(H))+6*)
1-1
2hr (HNPH)) + M(MFIJ (Ir (HV) + 02)² (0" 7)
M+1
(D.10)
3M42
118- - 0.
(D.11)
Let's first verify the necessary conditions of the system defined above. From [.111 and Lemma
H.5. know 7 E 8 null Hb
3 + null (H). Then, we have 8-1V12 null (H). Then, in
D.7)
we know
b
H
=4 141-H
H H 1ivia 1|
H
DH
1- H H
V12 12 O.
28

---

which implies either H b 02 or H AHE Oaxd. If we assume H АH = Oaxd, then D.6
implies H v12b H H H From assumption, know rank(H v 2, which
implies rank(H H WH ) 2 2, since for any matrix Z, holds that rank ( rrank Z Z Since H WH Ω 2 2,
multiplication by an invertible matrix does not change the rank, we know rank
while H H is a matrix of rank at most one, which contradicts with D.6. Therefore, we
have
H 0,
(D.12)
H H H
(D.13)
Recall in Theorem 5.2 we have defined
r= T
Simple calculation shows
H H H H Ω
DH H Ω H H I
where the second and the last equalities come from the fact that Ω and H WH commute, and the third equality comes from the property Moor Penrose pseudo-inverse. Therefore, one solution of
DD133
u 1 r
Since (D.13) a linear equation to 1-1A, we know its full solution is
0_ AEr
H ZH
0axd
P Im H
The solution (D.12 is
null H U-1 + null (H).
V21 ---1
Therefore, we know the necessary conditions for (.6 to (D.II) are
1- #
V-1V12 E null (H)
V21 = - + null (H).
(D.14)
VI1 E T 10-13"v + Im HH
== 3" + null (H)
It is easy to verify these equations above are also sufficient conditions of (D6) to (.11I) by directly replacing each variable with its value and validating equaions from (D.6) (D.11.
Specially, H is positive definite and hence, invertible, the global minimizer is unique up to a scaling
:
1- + 0, V12
0a.
-U_18" V
T
3=
(D.15)
Step 5: equivalence in the hypothesis class. Finally, we will prove that any global minimizer of FLTB is actually equivalent to one single function in FLTB almost surely. More concretely, let's take function FLLB with parameters W K, Wa Q: W p. Wv, W1, W, and recall the parameter
transformation in (D.I. We assume equations in D.14 and Assumption 3.1 hold. Then, for any
vector a E null (H), one has x x a = since x, X, N (0,, H) We denote
i..dd
V21 -0-1 + a1,
U- Vn T 0 Z, 7 3 +
where a1, a2 E H. H Z = Odxd- Then, we have
W WoE
((E) = W w. E WT W.EM W.EM
M
29

---

x x Xyy)
V 11 v12
= (3" + +a2) x+ (-0-13T +aj v-1)
M
y X y-y)
x
3 x + (-3rT 1) M
X X
V
(Xa: = 0, 0, x = 0)
y T X yJy)
U-IVIex
3r x (-3xT 1)
1
X x
xx
(v X,XV 0)
M
y yTy
0
= -(3r m x (XB y),x = S (E),
M
where /3 r (.) is the GD-B function defined in 5.1
30

---

E Proof of Corollary 6.2
Proof.
We denote Φ1 02 2 2s 2 0 are ordered eigenvalues of A Hv?. From Theorem
we have
(HHWH)) tr ( s2-))
inf / fefcn.
R(f) - 0²
where
M +
Ir (HYV + o²
Ω
H WH
M Id.
M
Therefore, have
inf R(ʃ) 0² H (n (2- H 12-) feFoo.B
n H (H + trr (HAHHH) + +o") 1a))
M tr
Since
tr
H
+ +
Is H WH + (HA + +02) Ia 32 tor (HiHHt) + + It:
we have
tr
(HHHHH))
to²
inf JEFco,
R(f) 0²
(2 (n-'Hi'' wH!) tr
M
d
Φ
M+1
Ø1 +
1=1
min {0i x)
1
Therefore, we conclude.
31

---

Proof of Theorem 6.1
F.1
Proof of first equation
Proof. From the classical bias-variance decomposition, we know
C (gg == EE (g(X. x) y)² X
E (g(X, E Tv x.y.x))
X
+E
(E X,y.x] y)²
x
since the cross term vanishes. The second term does not depend on g and hence, the Bayesian optimal estimator is given by the posterior mean, i.c.,
Bayes = E Eyy X. y,
B x, y, x x).
Since 3 ~ N (3" ) we have there exists a random vector Ø N (0s, Ia) such that 3
3 + T almost surely. Therefore, one has
Daayes 1G" xx + (B 01x0, x,,, vxxx)
To compute E Ø x,y, suffices to solve the posterior distribution ,3 given X, From
N (0a, La). we have
y y-x x (3- ( + 1
P. (@) X.y,x)
x P
P
y
x exp
1~T 0
0
202
X + )L 0 2 - x
x exp
20
Note that, the function above matches the probability density function of a multivariate Gaussian distrbution. Therefore, the posterior mean is given by
0 x,y,x = (v+ X + +ooIa) -1 x (y -x9').
Therefore, we conclude
TByyes = x TV! x X + +oola)= ViXT ( (x-xx))
+ +T
(F.1)
F22 Proof of second equation
Proof. Let's first do a variable transformation. We denote
X = X
X = vixx
y y A= H.H.
(F2)
iid.
Then, know X x
N (O,, A).
We can write the Bayesian optimal estimator in
F.I
Bayes =
x x x
xx
The true label is Therefore, have
y x=x =x T
-1 x5 3 -x xT(B* (3++ vtt
2
C (DBayesi X) - 0² =E X X x (
32

---

o 0 E Xy R x (xTX (xTx+ +o
where
0: xx= Xy.
Note that, this is equivalent to the risk of estimating the ground true linear weight 0 using 0 under the Gaussian prior 0 / (0, Ia) The estimator 0 is function of transformed input-output pairs in the context X y and the transformed query input x, and takes the form of standard ridge estimator with regularization coefficient being o². We then revoke the standard results for the risk of ridge estimator. Applying the Theorem and in [3], we know for a fixed weight vector 0, with probability at least 1 - exp (-I((M) we have
0² + Ei>k. di
k- M
M E of Lisk Lish Φi
0
0
0
+
H
M
H
Ha
0²
where Φ1 2 22 2 2 δa 2 0 are ordered eigenvalues of A.
d² + E
k* min k Φk 2
and c is an absolute constant. Here, Ho. SVD approximation with respect to the largest k- singular values and Ha is the SVD approximation in the remaining singular values. Namely, if
we have the decooposooooiiiionnn of = Q diag (01,02,... Qe) Q where Q an orthogonal matrix, then Ho:k- and :a are given by
Ho = Q - diag (01, Φ2, .... ok. 0, .... 0) QT,TT H = H. Hork.
Taking expectation over 0 ~ N (0, La), we have
L (Bayys X) -2 0² = 1oLL
Bayes
0
H
- Σisk. Øs
2
2k" ht-
M MEss of
Σ
Σ
0i +o²
M
M
0²
Ei>k
iSk
i>k*
Now we simplify this expression. First, we define
δ = 2 Eisk 05
M
From the definition, we see / >
2a M. On the other hand, from the assumption that tr (H)
V ( = -i di S 0z, we know that & x co?
Mr Combining two parts, we get
o²
(F.)
M
Therefore, we have
σ²
Eiskr of
L (iBaaes: X) -2
0²
Σ
+
Σ
④ +
K²
Φi
M
x²
isk
1>k
2
min
0
+
Φ
Σ
min
(from FB
})
min
Φi
+
min
Φ
min ( 5)
This finishes the proof.
33

---

G Proof of Theorem 6.3
G.1 Training dynamies
Now we consider doing gradient flow on the risk function, or equivalently, on the excess risk function which differs by only constant:
1 Ø
[R(3,T) min R(;
(G.1)
df 203
dr
1 0
[R(3, r) min R(:
(G.2)
dt
ar
First, we have the following corollary of T'eorrm 5.2 min R(.-). Corollary G.1 (Excess Risk). We fix M as the context length. Consider the ICL risk in 3.5
which computes the excess risk R(B,T)
assume
the data is generated following Assumption 3.7 Then, we have
R(B,T) r) min R(-,.)
tr HT HT
(3 B-8 (Ia TH) H IH)+
HIH
(B (3)
M
M
(HiTHt ( (HHrHH - -
(G.3)
Proof. This is obtained directly from equation (C.5) and (C.7
Now, we write out the differential equations explicitly in the following lemma.
Lemma G.2 (Dynamical system). The dynamical system of gradient flow described in
G.1 and
G.2 is
d3 (4-TH)
HT
HT
( TH)
+
+
HIH
(B (3- 33) (G.4)
dt
M
M
dr
HTHH HDH)
M+
HIH
(B) B3) H
dt
M
M (a- 3*) H 3*) HTH H (B- H.
(G.5)
Proof. This can be obtained by directly calculating the derivatives over the excess risk G.3 To
write out the dynamics of B3, it suffices to notice that ,3 only attends the first term of the RHS of G.3 which is standard quadratic function. For the derivatives of T, we first have
0 20r
(H'TH: - H s2
HHTHH HAwHt T
tr
=H
H H Htwnte s2 D
(the sixth equation in Lemma H.3
HHH HVH.
Now it suffices to compute
Hr HT
1 HT HIH
0r (B3 T3)T
(I. TH) H (H) +
+ M
(B - -3))
M
34

---

Let's compute the partial derivativess separately. From Lemma H.3
we have
19- (9-
0
--H H (B (B- 8*)
à (B 8r)
HHH (3-
0r
är 2
- - (²) (3 -3) H:
2
0 1
M +1
(B 8*)
Hr HTH (B - -8))
M - HT 133 3*) H:
0r
2
M
tr
HT
HI
33)
T H (B (
3*) -5)
= M (3 -)) H (33 - 3*) НIH.
ar
M
Summing over the three equations above and applying the definition of gradient flow, we conclude the dynamics of r in G.5).
G.2 Proof the global convergence
Let's now prove Theorem 6.3 Since the first part of Theorem 6.3 is directly implied by the second part, we only deal with the case with general H. In this section, we denote A_ the minimal non-zero eigenvalue of H. As in the main text, we define
H = Im (H)
and 7-- as its orthogonal complement. First, let's prove the convergence of B.
Lemma G.3 (Convergence of (). Under the dynamical system (G.4) and (..5) one has
-21-11
|Px (B(t)) - PH (B*) s exp
|P (B(0)) Px (8))
(G.6)
M 11+1
which implies
H² (B(t) -800) λ1 exp
-2A-1t [|8(0) 33
(G.7)
M +1
and
Px (B(1)) + Pn (3")
when + from arbitrary initialization ((0) and [(0). Moreover, one has for any 0, it holds that
PH (3())) = PH (3(0))
(G.8)
Proof. We first consider the orthogonal projection operator P. From Lemma H.5 and equation (..4 we know
( --HHHH Hr (3 - ))),
dt
where
HT HT
Hr = (I, TH) H PH) From the property of pseudo-inverse, we know
H +
Hr HTH.
M
M
dP* (B(t)) 33 3r) Hr (B 3r) HH H H
HH (3 Br)
Therefore, we have
d 1 Px(((())
201iee PM (B- HH PH (3)
dt
2
Notice that
T
M
M
M
M+
Hr 2
TH
H
TH
+
M +1 H.
M+11
M
M
11+1%
M
M
+1
35

---

This implies
1 dt 2 (PH()( - -3>113]
1
PH (B -3))²
. H.PH (3 - 33)
M +
A-1
Prr. (B 3')1
M ++1
where the last line comes from Lemma H.5 and A-1 the minimal non-zero eigenvalue of H. Via standard integration method in ODE, we know this suggests
-21-1t M +1
pP (3(1) - -3 < exp
|P (3(0) B²2 +
when + 00.
This indicates that Pi (B(t)) - PH (3*) when t. + co. Finally, we have
n4 (3(1)- -3") (Ia HH+ Hr (B - -3") 0d
which implies PH4 (3(t) - (3)) = PH= (3(0) - -3") and hence, PH1 (B(t)) = PH1 (3(0)) for any
t > 0.
Next, let's consider the convergence of r. As defined in the main text, we have
2 Im
H & H
and Z its orthogonal complement. We have the following lemma.
Lemma G.4. Under the dynamical system G.4) and G.5 one has
Pzi (T r=) Odxd:
(G.9)
dt
d
Pz (T -') - P2 (T) H. - M+1 1
H. P (T- -3 TT) H
dt
M (B 3*)
H (*) H Pz (T r*) H
M (B- B*)
H (3 (*) HI*H
HI HI H (B - 3*) (8 3*) 3²) H
M
E -1
H + H + + I4 H (B 3r) ) (B 3²) H.
+ M
(G.10)
Proof. The first ODE is trivial since the RHS of G.5) always lies Z Im H H) so its
projection on -+ always vanishes. To prove the second equation, rewrite (G.5)
dr H TT) H HWH + HH
d
M+1
M HT -T--T")) 33)) (B 3r)
M ((-))) H (8 8)) H I*) H H (B 3²) H
M
M (B B*) H (B B*) HITH
HI ( B¹) (3- T H.
M
For term A, recalling I" wH H 4, we have
H H H H
- H
(9 and H. VH commute.)
36

---

H H
-
HHHHA-TOH
(Q and H TH commute.)
НУи.
This suggests A = 0. For B + C, we have
1 HI
B+C=
H (3 3²) (3 3*)
H (Ia (3- B*) H.
M
We compute (Ia
HT
H as
LL.
HI
H
H H-
Н Q H H
(Q and H WH commute.)
H²
H Ω H
H
n nH - H
( (tr H (H* Lo)) H².
Therefore,
B+C HI (3 (B-
M
M H WH + (ur (H + +o²) Ia) H (B- (B- H.
Bridging A = 0 and the result of B + C into the ODE, we have
(T - r")
--H (TT M +
H H (0- 3a) "T H
M
M (3 - 3) H
H 8) H HT -") H
1
M (B 3r) H 3¹) HITH HI (3- (B B3) H
H H T tr ( 2² L H (B 3r) (B B*) H. (G.11)
M
Note that, all terms in the RHS above lie in 2 = Im H2 so this summation also equals di Pa [T - I'). Moreover, since
H (T-r") H = H² P mm ( (T-T') + null (Ht (H+ehH) (T - r*) H
H. Im (HH (T- (T-
(G.12)
Bridging (.12) into G.T1 we conclude.
Then, let's upper bound the dynamics of [| (T (T" in the following lemma. Lemma G.S (Dynamics of Frobenius norm). Under the dynamical system (G.4 and G.5 one has
d [ Pz (T 1 < 8-A (T - 1 Ag
-22
Pz IT
l')l
dt
M M+1
(G.13)
37

---

where
A X Amin (St)
2
(G.14)
A 2
> [(0
-
M
are two positive constant. Here, X-1 > is the minimal positive eigenvalue of H. X is the maximal eigenvalue of H, Xwin (St) is the minimal eigenvalue of 22 (defined in (A.3) and is strictly positive.
Proof. From the dynamics of PzT in G.10). we know
d
Pz T r]
d = tr
Pz (T- P2 (T-))))
dt
dt
= G + G2 + G3.
(G.15)
where
H Pz (T --) H Pz (T - -
G1 tr
M+1
H H.Pz (T H (3) 33) H Pz P2 -r
G2 -r
M
1 [
Pz (T (T-)
tr
(B - Br) H - 3r) H P2 (T -") H-P2
M
1
G3 -r
M (B 3*) H (B -3) HITH Pz (T (T-
+ tr
HI (3 3r) (B 3*) H P2 ( (--'
M
H
H . + (Ir (H + +o²)
Ia] H (3 (3) B*) H H.Pz (T rm)
+ir
M
(G.16)
From Lemma H.2 we know that G2 Then, we consider G1 and G3- For G1, we have
G = tr (HiPz r") Hb) (HiPZ -r- 2
5 t2 Pz r - r P2 (T- s2
((3) Lemma H.6 and (2) in Lemma H.2
- =>X (2) P2 (T (--r")
where (-in(-) denote the minimal eigenvalue. Since Ω is positive definite, we know Amin (Ω) > 0.
Finally, let's upper bound G3- First, have
1 (g B¹) (B- 3t) HI'H P2 -r')"
M
Vd
(B 3T3 H -3)) H H Pz (T (T")"*
H
(Lemma +.
M
Vad1
(22-_1t)
exp
||3(0) - 33, Amax
(H
|||||| I| (T-F")ll
M
M +1
(from equation (..7
vdx²
-22-1t
||3(0)) - 13, Amax H || (T-''))
M
M +1
For Auax H
we recall r" wHtΩ 'H + and obtain
H
H w - Ω H ' (9 and H 4H commute)
= = -
H = Hb)
38

---

Since Q 1 and H H commute, they are simultaneously diagonalizable. The eigenvalues of
Ω- H THE are
2 = 1, 1,2,.... d.
M+ 0 +
M
Note that, every eigenvalue upper bounded by so we simply get
Amax
H 1
(G.17)
Therefore, have
(g (3-8") H33- 3¹) HH P2 (T (T-)"
tr
vdd2
(-2A-11)
exp
[13(0) 3 P2z (T (--r") M +1
(G.18)
M
Similarly, we have
HI (33 (3) ( 33) H P2 (T r
P2 (T ---)
H2 (8 -3") H
|||r H
(Lemma H.1
M
Vax}
(22--16)
exp
[13(0) 33 |Pz (T - r->
(G.19)
M
M+11
and
tr
H Ω
H + ttr (H + +0") I1 H (B -B-3") 3²) H. ..Pz (T -T')
M
trr H wH) + + La) 1a)l111 +
H (3- (B- 8r) H Pz TT - r*) H
M
(Lemma H.I
Vd.
.
H - (3-3*) (B B²) H:| |||| |Pz (T - r->Ill
M
-2A-1t [|@(0)) 11 PPz (T
< Vdyi exp
l')
(G.20)
M +1
Bridging
G.18),
G.19
and G.20 into G.16), we get
2
-2A-1t
G3 S
2 +
x
exp
||3(0)
331 1P (T r')l
(G.21)
M
M +1
Finally, we finish this section by proving the convergence of
Lemma G.6 (Convergence D). Under the dynamical system G.4 and G.5
one has
Pz (T()) - P2 (T"), Pz- (T(t)) = Pz (T(O))
when t + from arbitrary initialization (0)) and ((0).
Proof. We observe that
1 1
(T
P2
P2 (T-
P2 (T - -
di
2
di
Combining it with Lemma G.5 we have
d di | (T- S --Ai
-2X_1t
(T- Ax
M ++
39

---

where A > 02 > 0 are defined in (..4 Simple calculation shows that
exp (A1t) PP (T ll - A exp [
2 M +
(G.22)
di
When A1 *
2A M41 we integrate both sides from to t = T, then divide them by exp (A1T).
This gives
P2 (T(O) - -I
A2
exp
(A1
2h-1 M+
T
|Pz (T(T) -') S
exp (AIT)
A
2A
exp
(A1T")
M+1
|P (T()) - -I')l
A2
22-1T
+
exp
M exp (-A17)
+0
exp (A1T)
A
2A-1
M+1
when T + 30. Otherwise, if A 2A-1 the right hand side of G.22 reduces to a con-
stant, which implies exp (A,4) P[ (T(() M+1
1 grows at most at linear rate, which implies
[Pz (T(t) - I + when t - 00. Together, in all cases, we have |Pz (T)))-r')l +
From G.9
soon get Pz (I()) Pz (T())) foo all t 0. Therefore, we conclude.
O
The Theorem 6.3 is proved by simplly combining Lemma G.3 and Lemma G.6
40

---

H
Technical lemmas
Lemma H.l (Von-Neumann's Trace Inequality). Let U, V E gdxx with S We have
V) S i=1
oi(U)o,(V) |U'l Eo.(v) S Vd. ||V|lop||V|I=
i=1
where 01(X) 2 (x 2 2 (X( are the ordered singular values of X E Bd
Lemma H.2 (22)
For any positive semi-definite matrices A, B E Bdxd have
tr[AB] 2 O.
AB ž 0 and only if A and B commute.
Lemma H.3 (Derivatives, (7)). We denote A, C,D.X X matrices and a, b, vectors. Then, we have
① 8x tx
X
BXC
BXC + BB X
8x Bx Bx
[B B x.
ix
da ab T.
0x
da LaTx" Taxxb ba
IX
8b x DXc = DJ Xbc + DXcb
OX
tr (AXX+C))
2A (AXB C)B
Lemma H.4 (Lemma D.2 in [38] Lemma 4.2 in 377. If Gaussian random vector of d dimension, mean zero and covariance matrix H, and A E Rxdd is a fixed matrix. Then
xx Axx
(A+ A H (rHH)H.
If is symmetric and the rows in E Md are generated Independemllfrom
X|i] ~ H), i= 1....,, M.
Then, it holds that
X x = -r (HA) H + M(M + 1) HAH.
Lemma H.5 (Linear Algebra). Suppose H a (non-zero) positive semi-definite matrix Raxxd and H is its principle square root. We denote A_1 as the minimal non-zero eigenvector of H. Then, we have
I. null (H) null (H&) Im (H) = Im (H:)
2. HH H H, where (-)+ denotes the Moore-Penron pseudo-inverse.
3. For any vector E Rd, the orthogonal projection operator Im (H) and null (H) are respectively
PIm(H)) (a)= = HH a H*H Pouli(H) (ar) = -(Ia -I.--HH -Ia--HH) (I. H H+H) a.
* 4. For any vector a E Rd, we have
PIm(H) (a) HP (a) 2 A-1 (( (a)
41

---

Proof. We consider the sigom-decomposittonn of matrix H, which is
H = QDQ",
(H.1)
where D diag (i1, A2,.., Aς) is a diagonal matrix with diagonal entries being the eigenvalues of H, and Q is an orthogonal matrix. Then, from the definition of principle square root, we know
H QDIQT
(H.2)
where D diag (VλI, VA,, (ad) We denote columns of Q as Q1.... We know they are the eigenvectors of H. From the eigen-decomposition of H and H we know they share the same set of eigenvectors. Without loss of generality, we assume rank(H)
and Å 2 À2 ... X > 0, A, i 1,... d. Then, we denote H. as the linear vector space spanned by q1: .... 9 and H as its orthogonal complement (which is the subspace spanned by 9r+1, qa). For any vector
ar E Rdd we can write its orthogonal decomposition as a = Σi= Li=1 a,qi so we have
Ha
Σ 1=1
Hq.
Σ a; E H. 1=1
which implies Im (H) C H. On the other hand, we know for any a H, we can write it as Ei= a;q, for some a1 Or, then we have
bay H
Σ
ai qi
Im (H),
A
which implies H C Im (H). This shows
Im (H) = 7 = span (q1,..99).
and
null (H) = H = span (9r+1,-..aa)
since null (H) the orthogonal complement of Im (H). Then, (1) proved by noticing that and Hà share the same set of eigenvectors. To prove (2), it suffices to notice that
HH+ Q diag 1.... 1,0,0, 0) Q H*H.
To prove (3), it suffices notice that actually HH+ E Im (H) and
(HH a, (Ia H-H+ a) (H Ho, (L. HH* a) a (HH+ HH HH*HH**) a 0.
Finally, to prove (4), we can write a E a;q; for some aı,.., as Then, by orthogonality we have
d
d
Pim(H)) (a) HP HPImHH) (a)
Σ a:q;
H
Σ aqq
Σ af Le qi 21 Σ aंq q:
= (PImHH)
(a)
Therefore, we conclude.
Lemma H.6 (Tensor product). Suppose H is (non-zero) positive semi-defiite matrix in Rdxd and Hà is its principle square root. We denote & Kronecker product, which is defined
(A A A
We define H
Im
{(H H) Z: Z E Rdxd] null HH H (Z E R HH H) Z = =0)
and define Im
3 @ H and null
( @ H
similarly. Then, we huve
42

---

I. Im H @ H:)
Im (H @ H).
null H Hb)
null (H @ H).
. 2. We denote 2 = Im H @ H Im (H s H) and z null H & H
null (H H). For any matrix Z [ddd
we have
P2 (Z) [HH*) 802
= HH ZH H H -
(H.3)
Pzi (Z) - - (HH+) 82
Z = Z HH* ZH H = 2 H 2H H
(H.4)
3. For any matrix E Rdxd, , holds that
( ( Pz (2)) (2))
Pz T 2 (2)
(H.S)
where A-1 is the minimal positive eigenvector H.
Proof. Let's first prove the second part. From the definition of tensor product and the fact that HH H H (see Lemma H..)), we know the second equation in H.3 holds. To prove the first equation, it suffices to notice that for any matrix Z E Rdxd the matrix HI ZH H is Im HH together with the fact that
HH ZH H. Z HH H
HH
ZH HZ
HH ZH HHH Z H H
HH
ZH
HZ -
HH HH ZH HH'
(HH+ = H
0.
The proof of (H.4) is similar. Then, from Lemma H.5
know Im (H)
Im H which
implies the projection operator onto those two subspace are identical, indicating HH+ Therefore, have
H H 1
PI (Z)
092 Z
HH ZH*H H
ZH H
P
H @H (Z).
Im
Since this holds for any matrix Z E Rdxd this suggests Im H
H
Im HH H)
Sim-
ilarly, we also have null (H's @ H²) null (H & H). Finally, to prove (3), we use the eigen- decomposition of H and H
H = QDQ H =D=QT
where Q (q q qa) is orthogonal matrix and D diag (A1, >d): where X 2 d2 2 ....As 2 0 are ordered eigenvalues. We assume rank(H) = and X > 0 = Ay+1- Then, from the property of Kronecker product (eg, see (27). the eigenvalues of H @ H are {v : i.j < d}. The eigenvectors are {q: @ q, S i.j < d} which forms an orthogo- nal unit basis of dd We define
= span {q & qj 15 1, <r)
the eigenspace spanned by all eigenvectors corresponding to positive eigenvalues. For any matrix Z, we can decompose it as
di.j qi qi Z
So
H x
H)
Z
Σ
al
(H & H) (q (g)== Σ
(Hq.) X (Hq;)
1
i
43

---

p Σ
a
anddddaa eqes,
w-
which implics Z c S. On the other hand, for any Z E S, we write it as Z Ei-=l bug q qig so we have
bis bis (19))) = b
bis X1dy
Z
Hqq 2 Hq,) = (H @ H)o
(q: @q))
E EZ,
Waml Mrdy
Xxdj
which implies c Z. Combining two directions, we have S 2. Therefore, for any matrix Z E Rdxd we can write it as Z = Eij=t C,, qi @ qj for some C,, Then, we have
H & H Pz (2) = H & H
Ps (Z) = H H
Σ
Ci.
9: 8 qj
1,==1
Σ W=
cu
(H²q:) @ (HiQ))=
Σ
au Vody 9..
$,1=1
Therefore, we have
( H (H & Hz) Pz (Z)) ( H²) o PZ) T
Σ ao S 00a
Σ
CKIV Andi qu & qi
k1==1
L (4))44))))))))) &q)
(By orthogonality)
2 x3 Σ
E. (q. 2 (q 8 s1)
1,1=1
= 2 Pz (Z) (2)
Therefore, conclude.
I Experiment Details
Our experiments on GPT2 mostly follow the setting in [14, except that we use the token matrix
defined in (3.1). In experiments, use d 20, M = 40, V = H = L and 0. We train the GPT2 with and without MLP layers on two settings: 3 = (0, 0, ..., 0) T and 3 = (10, 10, .... ..,) 10) For GPT2 with and without MLP layers, we initialize the Wv and Wa by normal distribution with standard deviation 0.02/ number of residual connections, where the number of residual connections equals the number of layers for GPT2 without MLP layers. For GPT2 with MLP layers, this is 2 times the number of layers. Initialization for other matrices follow the default setting in (36). We use Adam with learning rate 0.0001. We train the model for 200000 steps and we sample a batch of 256 new tasks for each step.
Does scratchpad help?
In this section, we show the limitations of adding a scratchpad to the token matrix. We will show that by using a single LSA layer and the token matrix with scratchpad, one cannot recover the GD-B estimator defined in Section 5 We leave it as future work to see whether the token matrix with scratchpad could implement other types of estimators that more effectively address the linear regression tasks defined in Assumption 3.1] well as whether additional structures could help alleviate this approximation error. We follow the notations previous sections. The token matrix with scratchpad is defined as
44

---

X T
x
E
E
R
[2(++))x(M++))
(I.1)
where 1 R. refers to the vector filled with ones.
LSA model f. is defined by
((E) = E W.WWENNNN W W
E W W
WSW, EM
M
M
where the second equality because the bottom right entry of is zero (see (3. prediction is the bottom right entry of the output matrix. So only the last row of W last column of E attend the prediction. Denote
Note that the
Wy and the
Q
*
W Wv
W Wq
=
T
b
#
10
01
a12
92
b2
where
w E Rdd 0, E R, E b, 91. 12 E R.d, b b2 E IR,
and denotes entries that not enter the final prediction. Then we have
Q
b
EM E
x
S(E)
(w) a a2)
q
1
1
M
0
Q2
b2
Q
b
EM E
x
=(wT
a1 a2)
91
b1
1
92
b2
X X
X
1 M
X y
Q
b
=wwT a
az)
1
I
X
M
1 I
q
b
x
M
1
T
X
y
1M
Ty
T
b2
Following the Assumption 3.1, we use @ to refer to the task parameter, then
3 3:= 3 + v where A~ (0a: Ia)
We then decompose (E) as
X X
X T 1M
yy
Q
f(E) (
a1
02
1 1
X
M
1 yy
q
x + II,
M
T
X
y
lM
y
y
q
where terms and II are independent of x. Therefore, we have
R (f) = E (E()) y)²
E (S(E) - (B,x) x)
+ +
since y/8,x ~ NN(,0²)
(I. [[.+ (,,x) 2
+ o
since /(E) I x
E I - 2 +
since x x~ (0,,H)
45

---

3 +
-
Note that the above equation holds II 0, which can easily achieved by setting b 0g. bj b2 = O. Therefore, without loss of generality, we can consider the LSA function which takes the following form:
X
X
X
1M
X
Q Ii
1 M
(EE = 1- Tw a
M
X
M
MY
x.
(J.2)
y
X
1M
y
q2
Let's then try determine whether such function effectively represent GD-3 function, which takes the form of
r'X X -y)
fGD (EE) = (3" x)
(J.3)
M
where 3 is the prior mean in Assumption 3.1 in our submission, and F W and S H WHA + (H#)
one simple way is to let M Ig is defined in (5.3) in our submission. In order achieve this,
a= Q -T"), qi = q2 Oa-
(J.4)
However, this will incur some additive terms and will potentially enlarge the ICL risk. Inserting the parameters above, have
I M
EE) = (E)
+
(X(T*)T= - (B*) +x())) x
This shows that the extended token with a single LSA cannot easily implement the GD-B function class and may ineur additive ICL risk depending on 3 (as shown in Theorem 4.1
46

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We believe our introduction and abstract are factually accurate in describing the contributions of the paper and its result.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitation in Section 4 and the future work section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
47

---

Justification:
We provide the assumption in Assumption 3.1
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide all detailed experiment setup in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
48

---

Answer: [Yes]
Justification:
The data we use are simulated. We plan to release our code upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://niis.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https Cooeeeoodeesssseeeeee //nips.com / ssiniiiiccypyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
We include all details and hyperparameters our experiments Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: The cost for producing the experiments precludes from reporting error bars. Moreover, the experiment section is not the main contribution of this paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
49

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [NA]
Justification: There no strict requirement for the computer resources for our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:p/pw...eeepoocco/ooooooo/oa
Answer: [Yes]
Justification:
None to report.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work focuses on the theory of in-context learning so it is not expected to have direct societal impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
50

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: The data used in this paper are fully simulated, so there is no data or model that have a high risk.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification:
We do not relase any asset.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt paperswi thcode. cow/ datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
51

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not relase any asset.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We do not involve humans in our research.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing.
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
52


Reviews:
[{"review_id": "4kdjeTaRnU", "reviewer": "Reviewer_mBt4", "summary": "This paper studies the in-context learning of linear regression with a Gaussian prior that has a finite, non-zero mean across tasks. Specifically, the authors show that a linear transformer block (linear attention with an MLP layer) can achieve near-optimal Bayes risk for this task, whereas a linear attention-only block suffers from an approximation error. They further demonstrate that an LTB hypothesis class can be well-represented by a subset class of one-step gradient descent estimators, GD-$\\beta$, for the ICL task.", "questions": "Have the authors seen experimental evidence that trained LTB models converge to GD-$\\beta$ models, similar to how LSA converges to GD-$0$? I believe the correspondence established between the two is excellent, and theoretically studying the convergence is probably another work in itself, but seeing some experimental evidence would have been really cool.\n\nI believe the restrictions imposed on $d_k, d_v$ for the hypothesis classes are to ensure sufficient overparameterization, which is pretty common for theoretical works. However, I think in practice these dimensions are typically set to be less than $d$. Are these necessary conditions, or do the authors believe there are ways to override them for their analysis?\n\nThe constructions for Lemma $5.1$ and Theorem $5.3$ both operate on pairs of matrices composed together (if not more), which is also the common practice in the analysis that I’ve seen. Any reasons why the authors didn’t opt for analyzing with merged matrices, especially when the inner dimension of the product $W_p^TW_v​$, that is $d_v$​, is assumed to be larger than $d$? It doesn’t even enforce any additional rank constraints on the product matrix.\n\nIt was interesting to see that in your GPT-2 experiments, the full transformer model (with all the non-linearities) also relied heavily on the MLPs. How much of this do you think can be attributed to the presence of softmax only? Did you try experiments where you used LSA in GPT-2 training with MLPs?", "limitations": "Yes, the authors have addressed the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work delves into theoretically understanding the very relevant and hot topic of ICL. The problem motivation is clear and concise, and I really like how the authors essentially redo the previous results of LSA and GD-$0$ estimators with LTB and GD-$\\beta$, demonstrating the effective correspondence between them and showing that they achieve near Bayes optimal ICL risk. Overall, I think the authors do a great job of portraying the benefits of the MLP component in transformers, and it’s a good theoretical contribution, as most papers involve theory only from the LSA perspective.", "weaknesses": "I don’t see any major weaknesses in the paper as such, but I do have a few questions (see questions).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for supporting our paper! We answer your questions below. \n\nQ1: Have the authors seen experimental evidence that trained LTB models converge to GD-beta models, similar to how LSA converges to GD-0?\n\nA1: Yes, empirically we also see that trained LTB models converge to GD-beta models. Similar results are also obtained in [2] for LSA converging to the GD-0 (see their figure 2). We will try to add the figure for LTB converging to GD-$\\beta$ in the next version.\n\nQ2: I believe the restrictions imposed on $d_k$, $d_v$ for the hypothesis classes are to ensure sufficient overparameterization, which is pretty common for theoretical works. However, I think in practice these dimensions are typically set to be less than $d$. Are these necessary conditions, or do the authors believe there are ways to override them for their analysis?\n\nA2: In multi-head attention, $d_k$ and $d_v$ are set to be $(d+1)/m$, where $m$ is the number of heads and $d+1$ is the embedding size. Since we only consider single head attention, this reduces to $m=1$ and $d_k = d_v = d+1$, which satisfies our assumptions.  \n\nThe assumption that $d_k, d_v \\ge d$ simplifies the discussion, but we believe our results can be generalized under relaxed assumptions using standard techniques. For instance, if $d_k < d$, then the global minimum of LSA is achieved when $W_K^\\top W_Q$ is the best rank-$d_k$ approximation of $\\Gamma^*$. Similar results are obtained in [1]. We expect a similar result can be proved for LTB, where the global minimum is achieved by certain low-rank approximators. \n\nQ3: The constructions for Lemma 5.1 and Theorem 5.3 both operate on pairs of matrices composed together (if not more)... Any reasons why the authors didn’t opt for analyzing with merged matrices, especially when the inner dimension of the product $W_p^\\top W_v$​, that is $d_v$​, is assumed to be larger than d?...\n\nA3: We choose to present our results in terms of separate matrices just to be aligned with the practical Transformer setup. As you mentioned, our construction results also hold under merged matrices. We will clarify this in the revision.  \n\nQ4: It was interesting to see that in your GPT-2 experiments, the full transformer model (with all the non-linearities) also relied heavily on the MLPs. How much of this do you think can be attributed to the presence of softmax only? Did you try experiments where you used LSA in GPT-2 training with MLPs?\n\nA4: We think the gap observed in our GPT-2 experiments (w/ MLP or w/o MLP) is caused by MLPs instead of softmax since softmax is applied in both cases. We thank the reviewer for the proposal to train deep LTB and LSA models. However, training deep linear models is extremely hard without normalization techniques like softmax, and we observe unstable training performance when we train deep LTB and LSA. In practice, the normalization is important for training a deep model.\n\n[1]. Y Li et al. Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond. 2024\n[2]. JV Oswald et al. Transformers Learn In-Context by Gradient Descent. ICML2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MqAPQZvPrK", "reviewer": "Reviewer_D6mY", "summary": "In this paper, the authors investigate the in-context learning (ICL) of a linear transformer block (LTB), i.e., a single block comprised of a linear self-attention layer and a MLP layer. The task considered is ICL of linear regression with a Gaussian prior. Unlike earlier works that studied this problem, a Gaussian prior with a non-zero mean is used for generating the ICL tasks. For this setting, it is shown that the MLP layer is essential, as using only a self-attention layer results in a weaker approximation. Furthermore, the LTB at minimum is demonstrated to have a specific structure akin to one-step gradient descent with a learnable initialization, referred as GD-$\\beta$. The optimisation and statistical properties of the GD-$\\beta$ estimator is also widely studied.", "questions": "1)  What happens if we include positional encoding (P) are also included?  It is possible to get an additional $P^{\\top} W_K^{\\top} W_QP$ term can come from the attention and can this overcome the need of positional encoding?\n\n2) Does the optimisation dynamics of GD-$\\beta$ offer any insights on the optimisation of the LTB ?   \n\n3) Does the results qualitatively change if the input sequence is generated from uncentered Gaussian $\\mathcal{N}( \\mu, H )$ instead of zero mean Gaussian prior ?\n\n\nMinor corrections: \n\n1) In l.251, it should be $ \\Tau^{\\top} $ in the expression of $W_K^{\\top} W_Q$.\n\n2) In l.263, an additional bracket in the expression of $\\Omega$.", "limitations": "NA", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "a) In-context learning of statistical tasks particularly the linear regression has received particularly interest. The authors have identified an interesting and relevant modification to this problem setting when the ICL tasks are generated from a more general setup Gaussian prior with non-zero mean. \n\nb) In this setting, an MLP layer and a skip connection are needed for approximation which is a novel finding. Even with these additional parameters it is shown that minimum of the task can be still be characterised and it corresponds to GD with learnable initialisation. \n\nc) The statistical and non-convex optimization aspects of the GD-$\\beta$ are interesting from a technical standpoint.", "weaknesses": "The paper attributes that property of learning the MLP layer, however the skip connection is also equally important. Hence it is a bit of an misleading claim to attribute this to the MLP layer. Overall, the model restricted to only linear attention and also a single layer.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your comments. We will make sure to fix the typos in the revision. We will address your concerns below. \n\nQ1: The paper attributes that property of learning the MLP layer, however the skip connection is also equally important. Hence it is a bit of an misleading claim to attribute this to the MLP layer. Overall, the model restricted to only linear attention and also a single layer.\n\nA1: We emphasize that the ability of LTB to learn non-zero mean is a joint effect of an MLP component and a skip connection. Note that LSA also has a skip connection (see also [1,2]) but its skip connection is inactive. In comparison, the MLP component in LTB activates the skip connection. Therefore, we attribute the ability to learn non-zero mean to the MLP component, which is the only difference between LTB and LSA. \n\nNonetheless, we agree one can attribute the ability to learn non-zero mean to the skip connection — as you have mentioned, without a skip connection, LTB reduces to LSA with a potential rank constraint on the parameter, which cannot learn non-zero mean as we have proved. The above two explanations take different perspectives to interpret the same phenomenon. We believe the way of interpretation does not diminish the significance of our contribution. We have discussed this in Lines 198-206 in our submission, but we will clarify this further in the revision.\n\nQ2: What happens if we include positional encoding (P) are also included? It is possible to get an additional $P^\\top W_K^\\top W_Q P$ term can come from the attention and can this overcome the need of positional encoding?\n\nA2: This is an open question. Although we show evidence that a scathpad hardly helps, it may be possible that a positional encoding can achieve the effect of a MLP component. There are several ways of implementing positional encoding. For example, GPT 2 uses an additive trainable matrix $P$ as positional encoding. Note that this significantly increases the amount of trainable parameters in both LSA or LTB, and also makes optimization harder. A more popular method for positional encoding is to use sinusoidal functions, which are only partially trainable. A fully rigorous analysis of the effect of all kinds of positional encoding is an interesting question, which we will comment on as a future direction. With that being said, we believe our current results, which clarify the benefits of an MLP component, have already made significant contributions. \n\nQ3: Does the optimisation dynamics of GD-beta offer any insights on the optimisation of the LTB ?\n\nA3: Yes, we think the optimization of GD-beta offers insights for the optimization of LTB. First, as we have shown that the global minimizer of LTB is equivalent to that of GD-beta, so efficient optimization of GD-beta offers a way to find the global minimizer of LTB. Second, GD-beta can be viewed as LTB with restricted parameters. So optimization of GD-beta offers insights of optimization of LTB in certain subspaces. Finally, we emphasize that, although GD-beta is easier to optimize compared to LTB, its landscape is still non-convex, which captures part of the non-convexity challenges of the latter problem class. We will include these discussions in the revision. \n\nQ4: Does the results qualitatively change if the input sequence is generated from uncentered Gaussian $N(\\mu, H)$ instead of zero mean Gaussian prior ?\n\nA4: Part of our results still hold when covariates are from uncentered Gaussian instead of centered Gaussian (see a detailed discussion below). Examining all of our results under uncentered Gaussian is left as a future work. In fact, we can derive an analog to Theorem 5.2 in our submission when the features $x$ are sampled from an uncentered Gaussian distribution $N(\\mu,H)$. Without loss of generality, we assume $\\Psi$ and $H$ are invertible. Under assumption 3.1 in our submission, we have the global minimizer of the GD-$\\beta$ class is $\\beta = \\beta^*$ and \n\n$$\\Gamma = \\Psi (H + \\mu \\mu^\\top) \\cdot A^{-1},$$\n\nWhere $A = \\frac{1}{M^2} \\mathbb{E}[X^\\top X \\Psi X^\\top X] + \\frac{\\sigma^2}{M}(H+\\mu \\mu^\\top)$ is a positive definite matrix and $X \\in \\mathbb{R}^{d \\times M}$ is the feature matrix. The optimal ICL risk is $\\sigma^2 + tr(\\Psi (H+\\mu \\mu^\\top) - (H+\\mu \\mu^\\top) \\Psi (H+\\mu \\mu^\\top) A^{-1} (H+\\mu \\mu^\\top) \\Psi),$ where $A$ is defined above. Moreover, since the LTB class can always implement a GD-$\\beta,$ we know there exists an LTB function that can achieve a constant level of ICL risk regardless of how large $\\beta^*$ is.\n\n[1]. R Zhang et al.  Trained Transformers Learn Linear Models In-Context. JMLR 2024.\n[2]. K Ahn et al. Transformers learn to implement preconditioned gradient descent for in-context learning. NIPS 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Enu0n20eSl", "reviewer": "Reviewer_xFNi", "summary": "his submission extends the earlier theoretical analyses of in-context learning with linear transformers by taking into account the data covariance and a non-zero mean on the task parameters. It is shown that a linear transformer with a linear MLP block (LTB) can implement the algorithm of one-step gradient descent with non-zero initialization. This additional flexibility of learned initialization allows LTB to match the Bayes optimal risk for linear regression, whereas the linear transformer studied in prior works cannot achieve such optimality at finite $M$. The authors also showed that the optimal one GD step solution can be obtained by training the parameters with gradient flow.", "questions": "See Weaknesses above.", "limitations": "See Weaknesses above.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "In the past years, there have been many works that examined the ICL ability of linear transformers in the linear regression setting, where it is typically shown that transformers implement one preconditioned GD step on the squared error in-context. This submission goes beyond the prior results by showing that linear transformers can adapt to the non-zero mean of the task distribution with the aid of a linear attention block. This yields a separation in the ICL efficiency between transformers with or without an attention block, which is an interesting message in my opinion.", "weaknesses": "I have the following concerns and questions:\n\n1. Firstly, the obvious limitations of the analysis: the problem setting is still linear, and the authors only showed adaptivity of ICL to the task mean but not the full prior (see point below). Given the additional linear block, the approximation gap is not really surprising. The optimization analysis directly assumes the GD-$\\beta$ parameterization, and it is not clear if gradient descent on the standard LTB reaches the same solution. Also, convergence is only shown at a population level with no finite-sample guarantees.\n\n2. The Bayes optimal estimator in Lemma 6.1 is a generalized ridge regression estimator where the $\\ell_2$ regularization is anisotropic and adapts to the prior. On the other hand, when $\\beta^*=0$, the GD-$\\beta$ solution in Theorem 5.2 reduces to the one-step GD estimator studied in prior works, which does not incorporate the prior covariance.  \nIt is not clear to me why these two estimators always have the same statistical efficiency, as it is well known that anisotropic shrinkage provides a substantial advantage (Wu and Xu 2020) (Li et al. 2023). I guess the conclusion in Corollary 6.2 relies on $tr(H\\Psi)$ being bounded. Can the authors comment on the optimality of GD-$\\beta$ if this assumption is not satisfied?  \n(Wu and Xu 2020) *On the optimal weighted $\\ell_2$ regularization in overparameterized linear regression*.  \n(Li et al. 2023) *Transformers as algorithms: generalization and stability in in-context learning*.\n\n3. Some minor questions: \n* On page 5, why does the \"scratchpad\" take the specific form of (3.1)?\n* For the experiment reported in Table 1, what is the size of the trained model? Is the curriculum learning strategy in (Garg et al. 2022) also used?\n\n---\n**Post-rebuttal update:** the authors’ response addressed most of my concerns. I found that one of my criticisms is based on a misunderstanding of the optimal solution that the authors analyzed.  \nI have increased my score accordingly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments. We will address your concerns below. \n\nQ1: Firstly, the obvious limitations of the analysis: the problem setting is still linear, and the authors only showed adaptivity of ICL to the task mean but not the full prior (see point below). Given the additional linear block, the approximation gap is not really surprising. The optimization analysis directly assumes the GD-beta parameterization, and it is not clear if gradient descent on the standard LTB reaches the same solution. Also, convergence is only shown at a population level with no finite-sample guarantees.\n\nA1: Although an additional MLP component improves approximation power is to be expected, we think it is quite surprising to rigorously prove an approximation gap between LTB and LSA. Because such a result formally justifies the role of an MLP component in LTB for learning the mean signal. \n\nWe agree that our results are in a simplified setup, where the task is linear and the optimization is under the GD-beta parameterization and at a population level. Yet even in the simplified setup, there is no such kind of result before our work. Without a full understanding of LTB in the simplest possible setup, it seems unlikely that one can theoretically understand LTB in more practical setups. As our work is the first work on studying LTB for ICL, we believe we have taken an important step in this direction and the contribution of our work is significant. \n\nQ2: The Bayes optimal estimator in Lemma 6.1 is a generalized ridge regression estimator where the $\\ell_2$-regularization is anisotropic and adapts to the prior. On the other hand, when $\\beta^*=0$, the GD-beta solution in Theorem 5.2 reduces to the one-step GD estimator studied in prior works, which does not incorporate the prior covariance. It is not clear to me why these two estimators always have the same statistical efficiency…. Can the authors comment on the optimality of GD-beta if this assumption is not satisfied?\n\n\nA2: We would like to point out a potential misunderstanding of our result. The GD-beta solution in Theorem 5.2 does incorporate the prior covariance. This is because $\\Gamma^*$ is a function of $\\Omega$ defined in Theorem 5.2, which is a function of the prior covariance. We hope this also clarifies the follow-up comments made by the reviewer. \n\nQ3: On page 5, why does the \"scratchpad\" take the specific form of (3.1)?\n\nA3: Initializing the “scratchpad” with a fixed constant is a natural choice. \nHere the constant is set to be $1$ for concreteness but this number is not special. \n\nQ4: For the experiment reported in Table 1, what is the size of the trained model? Is the curriculum learning strategy in (Garg et al. 2022) also used?\n\nA4: In our experiments, we use a small GPT2 model with 6 layers and 4 heads for each layer. The key dimension is 32 and the inner dimension in the FeedForward Network is 128. (Since we use a multi-head attention in the experiment, this does not directly correspond to the d_k and d_v in our theory part). We use the curriculum learning strategy in (Garg et al,2022). We will clarify this in the revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TvPhn32Tft", "reviewer": "Reviewer_XdPR", "summary": "This paper studies the in-context learning (ICL) in the linear regression setting with a Gaussian prior with a non-zero mean. They first prove that linear Transformer block (LTB) will enjoy a smaller approximation error than linear self-attention (LSA) in this setting, where the non-zero mean of Gaussian emerges. Second, they show that the core benefit brought by linear MLP is that LTB can implement the one-step gradient descent with learnable initialization (GD-$\\beta$). Besides, all global minimizers in the LTB class are equivalent to the unique global minimizer in the GD-$\\beta$ class. Finally, they investigate the non-convex dynamics of the GD-$\\beta$ class and prove that it will converge to the unique global minimizer. Simulations on the GPT-2 support their theory.", "questions": "1. In the LTB class, the paper does not merge $W_K^TW_Q$ and $W_P^TW_V$ into one matrix parameter, respectively. Are there any benefits? Does it will improve the approximation ability of the LTB?\n2. In line 272, I think the $\\Gamma$ is a precondition but not equivalent to the Newton step. Is it a typo?\n3. The pretraining goal of this paper is different from the autoregression training in practice. Are there any potential approaches to study the autoregression pretraining and few-shot regression inference?", "limitations": "The authors have adequately addressed the limitations.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Clarity: The paper is presented very clearly and is easy to follow.\n2. Significance: The paper characterizes the benefit brought by the linear MLP in the ICL for linear regression setting, which is undoubtedly an important topic in the ICL theory community.\n3. Quality: The paper investigates the additional effect of linear MLP from approximation and optimization perspectives in detail, and the derivation is solid.", "weaknesses": "1. In this paper, all non-linear parts in the transformer block (softmax and ReLU) are dropped. However, similar works on the non-linear transformer block have emerged in the ICL theory. [a] considers the ReLU MLP followed by a linear attention layer in the linear regression setting. [b] considers the standard transformer block in the in-context classification regime. Authors are suggested to discuss the potential impacts of these studies about non-linear transformers on the setting in this paper.\n2. Though the LTB class enjoys the same global minimizer as the GD-$\\beta$ class and the training dynamics of the GD-$\\beta$ class can converge to the global minima, whether the training dynamics of the LTB class can converge to the global minimizer is not discussed by the authors. The authors should clarify this point.\n\n[a] Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape, ICML, 2024\n\n[b] Training Nonlinear Transformers for Efficient In-Context Learning: A Theoretical Learning and Generalization Analysis, ICML, 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for supporting our work! We will address your concerns as follows. \n\nQ1: In this paper, all non-linear parts in the transformer block (softmax and ReLU) are dropped. However, similar works on the non-linear transformer block have emerged in the ICL theory. [a] considers the ReLU MLP followed by a linear attention layer in the linear regression setting. [b] considers the standard transformer block in the in-context classification regime. Authors are suggested to discuss the potential impacts of these studies about non-linear transformers on the setting in this paper.\n\nA1: There are several differences between our work and [a,b]. Specifically, [a] considered an LSA layer after a network for ICL over the Barron class. In contrast, we consider a linear MLP after an LSA layer. In the former setting, the network plays the role of feature learning, while in the latter setting, the MLP plays the role of improving approximation power. [b] considered ICL of a stylized binary classification problem, while our focus is ICL of linear regression. The loss landscape of these two settings is significantly different. In the former setting, the authors use hinge loss for classification while we use square loss for regression. We will cite these two papers and discuss their relationship with our work in depth in the revision. \n\n\nQ2: Though the LTB class enjoys the same global minimizer as the GD-beta class and the training dynamics of the GD-beta class can converge to the global minima, whether the training dynamics of the LTB class can converge to the global minimizer is not discussed by the authors. The authors should clarify this point.\n\nA2: We agree that an optimization guarantee for GD-beta does not imply an optimization guarantee for LTB. This has been noted in Lines 362-366, but we will emphasize this more in the revision. However, the optimization of GD-beta offers insights for the optimization of LTB. First, as we have shown that the global minimizer of LTB is equivalent to that of GD-beta, so efficient optimization of GD-beta offers a way to find the global minimizer of LTB. Second, GD-beta can be viewed as LTB with restricted parameters. So optimization of GD-beta offers insights of optimization of LTB in certain subspaces. Finally, we emphasize that, although GD-beta is easier to optimize compared to LTB, its landscape is still non-convex, which captures part of the non-convexity challenges of the latter problem class. \n\n\nQ3: In the LTB class, the paper does not merge W_K^\\top and W_P^\\top W_V  into one matrix parameter, respectively. Are there any benefits? Does it will improve the approximation ability of the LTB?\n\nA3: We choose to present our results in terms of separate matrices just to be aligned with the practical Transformer setup. Our construction results also hold under merged matrices. Merging the matrices does not affect the approximation ability of LTB in our setting. We will clarify this in the revision.  \n\nQ4: In line 272, I think the \\Gamma is a precondition but not equivalent to the Newton step. Is it a typo?\n\nA4: In general $\\Gamma$ is a precondition. When $M$ goes to infinity, the global optimal $\\Gamma^*$ converges to $H^{-1}$, that is the inverse Hessian of the population MSE. Plugging this into equation (5.4), we get one step of Newton step under population MSE from $\\beta^*$. We will clarify this in the revision. \n\nQ5: The pretraining goal of this paper is different from the autoregression training in practice. Are there any potential approaches to study the autoregression pretraining and few-shot regression inference?\n\nA5: Good question. We believe the few-shot regression inference results in Theorem 5.3 in [1] can be extended to GD-beta. However, we think autoregression training would not improve the few-shot regression inference for GD-beta, since the model capacity is limited. We believe a more comprehensive model is required to show the benefits of autoregression training. \n\n[1] Wu et al. ​​HOW MANY PRETRAINING TASKS ARE NEEDED FOR IN-CONTEXT LEARNING OF LINEAR REGRESSION? ICLR2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper provides the benefit of an MLP layer on top of a linear transformer on an in-context learning task. This paper generalizes the problem setting to a non-zero mean task distribution while previous work mainly focused on mean 0 task distributions. The MLP layer is beneficial to deal with such a non-zero mean of the task distribution. Finally, it shows the convergence of gradient flow on the population risk .   

- The paper is well written. The message of the paper is very clear and thus the readers can easily grasp the main contribution.  
- Although there is not so much mathematical difficulty, this paper gives an interesting contribution to the line of research on this topic. Especially, it is nice that the paper provides not only statistical advantage but also optimization guarantee on the population risk.     

Indeed, the reviewers are overall positive on this paper. I also agree with their opinion and would like to recommend acceptance of this paper.
Decision: Accept (poster)