Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context
Authors: Kazusato Oko
Yujin Song
Taiji Suzuki
Denny Wu
Venue: NeurIPS 2024 poster (2024)
Keywords: Transformer
in-context learning
feature learning
single-index models
Primary area: learning_theory

Abstract: Transformers can efficiently learn in-context from example demonstrations. Most existing theoretical analyses studied the in-context learning (ICL) ability of transformers for linear function classes, where it is typically shown that the minimizer of the pretraining loss implements one gradient descent step on the least squares objective. However, this simplified linear setting arguably does not demonstrate the statistical efficiency of ICL, since the pretrained transformer does not outperform directly solving linear regression on the test prompt. 
In this paper, we study ICL of a nonlinear function class via transformer with nonlinear MLP layer: given a class of \textit{single-index} target functions $f_*(\boldsymbol{x}) = \sigma_*(\langle\boldsymbol{x},\boldsymbol{\beta}\rangle)$, where the index features $\boldsymbol{\beta}\in\mathbb{R}^d$ are drawn from a $r$-dimensional subspace, we show that a nonlinear transformer optimized by gradient descent (with a pretraining sample complexity that depends on the \textit{information exponent} of the link functions $\sigma_*$) learns $f_*$ in-context with a prompt length that only depends on the dimension of the distribution of target functions $r$; in contrast, any algorithm that directly learns $f_*$ on test prompt yields a statistical complexity that scales with the ambient dimension $d$.  Our result highlights the adaptivity of the pretrained transformer to low-dimensional structures of the function class, which enables sample-efficient ICL that outperforms estimators that only have access to the in-context data.

Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context
Kazusato Oko!.3, Yujin Song2., Taiji Suzuki
Denny Wu1.5
University of California, Berkeley, University of Tokyo, 'RIKEN AIP "New York University, Fllttion Institute
oko@berkeley.edu, song-vujin1399g eerssg.ecc..ee.. u-tokoo ac. mel.jimmitt...uttooe. jc.jp, dennynudnyu.edi
Abstract
Transformers can efficiently learn in-context from example demonstrations. Most existing theoretical analyses studied the in-context learning (ICL) ability of trans- formers for linear function classes, where is typically shown that the minimizer of the pretraining loss implements one gradient descent step on the least squares objective. However, this simplified linear setting arguably does not demonstrate the statistical efficiency of ICL, since the pretrained transformer does not outperform directly solving linear regression the test prompt. this paper, we study ICL of nonlinear function class via transformer with nonlinear MLP layer: given a class of single-index target functions f- (x) a. ((x, B)), where the index features 3 R2 are drawn from r-dimensional subspace, former optimized by gradient descent (with a pretraining sample complexity that depends on the information exponens of the link functions σx) learns f- in-context with prompt length that only depends on the dimension of the distribution of target functions r; contrast, any algorithm that directly learns f. on test prompt yields a statistical complexity that scales with the ambient dimension d. Our result highlights the adaptivity of the pretrained transformer to oow-dimesiional struc- tures of the function class, which enables sample-effiche ICL that outperforms estimators that only have access to the in-context data.
show that nonlinear trans-
Introduction
Pretrained transformers [VSP* 17] possess the remarkable ability of in-context learning (ICL) [BMR+20], wherehy the model constructs a predictor from a prompt sequence consisting of pairs of labeled examples without updating any parameters. A common explanation is that the trained trans- former can implement learning algorithm, such as gradient descent on the in-context examples, in its forward pass [HvMM* 19, DSD+22 22, VONR + 23). Such explanation has been empirically studied in synthetic tasks such as linear regression [GTLV22, ASA #22), which has motivated theoretical analyses of the statistical and computational complexity of ICL in learning simple function classes.
Many recent theoretical works ICL focus on learning the function class of linear models using linear transformers trained via empirical risk minimization. this setting, it can be shown that minima of the pretraining loss implements one (preconditioned) gradient descent step the least squares objective computed on the test prompt [ZFB23, ACDS23, MHM23J. This implies that the learn linear targets with n Z d in-context examples,
forward pass of the pretrained transformer hence matching the sample complexity of linear regression on the test prompt. Subsequent works also studied how the distribution and "diversity" of pretrained tasks affect the ICL solution in similar problem settings [WZC+23 23, RPCG23, ZWB24, LLZV+
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

The motivation of our work is the observation that the simple setting of learning linear models with linear transformers does not fully capture the statistical efficiency and adaptivity of ICL. Specifically,
A linear transformer has limited expressivity and hence only implements simple operations in the forward pass, such as one gradient step for least squares regression. This limits the class of algorithms that can be executed in- in-context, and consequently, the pretrained transformer cannot outperform directly solving linear regression on the test prompt. We therefore ask the question:
With the aid of MLP layer, can a pretrained transformer learn a nonlinear function class in-context, and outperform baseline algorithms that only have access to the test prompt?
A key feature of ICL is the adaptivity to structure of the learning problem. For example, [GTLV22] empirically showed that transformers can match the performance of either ridge regression or LASSO, depending on parameter sparsity of the target class; [RPCG23] observed that transformers transitions from weighted-sum estimator to ridge regression the number of pretraining tasks increases. Such adaptivity enables the pretrained model to outperform algorithms that only have access to the test prompt, which cannot take into account the "prior" distribution of target functions. Hence a natural question to ask is
Can pretrained transformer adapt to certain structures of the target function class, and how does such adaptivity contribute to the statistical efficiency of ICL?
1.1 Our Contributions
1.1.1
Function Class: Gaussian Single-index Models
To address the above questions, we study the in-context learning of arguably the simplest nonlinear extension of linear regression where a nonlinearity is applied to the response. Specifically, we consider the following class of tinll-index target functions, where the t-th pretraining task is constructed as
(1.1) where of : :R + I is the unknown link function, and & E Rdd is the index feature vector which is drawn from some fixed r-dimensional subspace for some r < d. The task is to learn this input-output relation by reading the context (Æ1, "/,......................... NN, VN) and predict the output corresponding to the query x (See Section 2 for details). This problem setting is based on the following considerations.
sjj x .......... x'y x i.d. (0, La), vf = =0'((a),, 3²) +
Nonlinearity of function class. Due to the nonlinear link function, single-index targets cannot be learned by linear transformers. For this class of functions, the statistical efficiency of simple algorithms has been extensively studied: given link function with degree /' and information exponent Q (defined as the index of the smallest non-zero coefficient in the Hermite expansion), we know that kernel methods require least n Z d' samples [GMMM21, DWY21J, whereas two- layer neural network trained by gradient descent can achieve better sample complexity 2 dd((Q) [BAGJ21, BBSS22]. Hence these algorithms, if directly applied to the test prompt, require a context length polynomial in the ambient dimensionality d, which arguably deviates from practical settings of ICL where the pretrained transformer learns from a few in-context examples which may come from high dimensional data space. These algorithms serve as a baseline for comparing the statistical efficiency of ICL.
Two types of low-dimensional structures. Prior works on gradient-based feature learning high- lights the adaptivity of neural network to low-dimensional functions, i.e., single-index model a. ((x,,, 3)) that depends on one direction (index features) / in Rdd [AAM22, BES+ BMZ23]. In such setting, the complexity of gradient descent is dominated by the search for direction B3 E Rdd, the difficulty of which can be characterized by various computational lower bounds [DLS22, AAM23, DPVLB24].
Importantly, our pretraining problem setup introduces another notion of low dimensionality: the "distribution" of target functions is diw-dimensional, since the index features for each task 3r are drawn from a rank-r subspace. This wwwdeeeniionalttyy of function class cannot be exploited by any algorithms that directly estimate the target function from test prompt, but as will see, transformers can adapt to this additional structure via gradient- -aased pretraining, which reduces the search problem (for the index features ( to r-dimensional in the in-context phase. Therefore, when r < d, we expect pretrained transformers to outperform baseline algorithms on the in-context examples (kernel methods, neural network, etc.).

---

Empirical observations. We pretrain GPT-2 model [RWC* 19] (with the same configurations as the in- context linear regression setting in [GTLV222 to learn the Gaussian single-index task (1.1) with degree-3 link
10
prediction risk
function, and compare its in-context sample complexity against baseline algorithms (see Section 4 for details).
10
In Figure I we observe that the pretrained transformer achieves low prediction risk using fewer in-context ex-
-O-
d= 16
d= 32
kemel
amples than two baseline algorithms: kernel ridge re- gression, and neural network trained by gradient descent. Moreover, observe that unlike the baseline methods,
twe-layer
pretrained TF
10
10a
10"
in-context sample size N'
Transformer achieves an in-context sample complexity (almost) independent of the ambient dimensionality d.
Figure 1: In-context generalization error (with standard deviation) of kernel ridge regression, neural network + gradient descent, and pre- trained transformer. The target function is a polynomial single-indes model. We fix r = and vary d = 16, 32.
The goal of this work is to rigorously establish a d- independent in-context sample complexity in an ideal- ized theoretical setting for a shallow transformer opti- mized via gradient-basee pretraining.
1.1.2 Main Result: Learning Single-index Models In-Context
We characterize the sample complexity of learning (I.1) in-context, using nonlinear transformer optimized by gradient descent. Each single-index task is specified by an unknown index feature vector [ IR drawn from some r-dimensional subspace, and a link function σ, with degree P and information exponent Q S P; we allow the degree and information exponent to vary across tasks, to model the scenario where the difficult of pretraining tasks may differ. We show that pretraining of the nonlinear MLP layer can extract the oow-dimensional structure of the function class, and the attention layer efficiently approximates the nonlinear link function. Our main theorem upper bounds the in-context generalization error of the pretrained transformer.
Theorem (Informal). Let| f (æ1, (01,3)...... .I....,, NN YN, X) + y be transformer with nonlinear MLP layer pretrained with gradient descent (Algorithm 1) on the single-index regression task (1.1). With probability at least 0.99, the model f achieves in-conter prediction risk f() f-(x)]-T oa(1). where T the noise level, if the number of pretraining tasks T, the number of training examples N. the test prompt length N and the network width m satisfy (we omit polylogarithmic factors)
T E (QQ)
NT >> T Z (Q)
Z r ((P)
T7L Z
pretraining
inference
opproximation
where Q, P are the information exponent and the highest degree of link functions, respectively.
To the best of our knowledge, this is the first end-to-end optimization and statistical guarantee for in-context regression of this nonlinear function class. We make the following remarks.
The required sample size for pretraining T, N scale with the ambient dimensionality d. In particular, the number of pretraining tasks T is parallel to the complexity of learning single-index model with information exponent Q using two-layer neural network. On the other hand, the sample complexity for the in-contexi phase only depends on the dimensionality of the function class r < d.
Note that any estimator that only has access to the in-context examples requires Z d samples to learn the single-index model, as suggested by the information theoretic lower bound (e.g., see [MM18, BKM 19]). Therefore. when r & d, we and algorithms that directly learn from the test prompt, such as linear/kernel regression and neural network + gradient descent. This highlights the adaptivity (via pretraining) of transformers to low-dimensional structures of the target function class.
separation between pretrained transformer
Our analysis of pretraining reveals the following mechanism analogous [GHM+23]: the nonlinear MLP layer extract useful features and adapt to the low- dimensionality of the function class, whereas the attention layer performs in-context function approximation on top of the learned features.
1.2 Related Works
Theory of in-context learning. Many recent works studied the in-context learning ability of transformers trained by gradient descent. [ZFB23, ACDS23, MHM23, WZC+23, 23, ZWB24] studied
3

---

the training of linear transformer models to learn linear target functions in-context by implementing one gradient descent step in the forward pass. Similar optimization results have been established for looped linear transformers (GS) transformers with SoftMax attention [HCI.23, NDL.24, CSWY24a, CSWY24b] or nonlinear MI.P layer [KS24, L.II++44]. Our problem setting resembles [KS24], where a nonlinear MLP block is used to extract features, followed by linear attention layer; the main difference is that we establish end-to-end guarantees on the optimization and sample complexity for learning a concrete nonlinear function class, whereas [KS24] focused on convergence of optimization. [CCS23] showed that transformers can learn nonlinear functions in-context via a functional gradient mechanism, but statistical and optimization guarantees were given. If we do not take gradient- based optimization into account, the function class that can be implemented in- context by transformers has been characterized in many prior works [BCW 23, GHM 23, ZZYW23, JLLVR24, SGS 24, KNS24). These results typically aim to encode specific algorithms (LASSO, gradient descent, etc.) in the forward pass directly analyze the Bayes-optimal estimator.
Gradient-based learning of low-dimensional functions. The complexity of learning low- dimensional functions with neural network has been extensively studied in the feature learning theory literature. Typical target functions include single-index models [BAGJ21, BES 22, BBSS22, MHPG 23, DNGL23, BES BE3] and multi-index models [DLS22, AAM22, AAM23, BBPV23]. While shallow neural network can efficiently approximate such oondiinenioonal functions, the effi- ciency of gradient-
training is governed by properties of the nonlinearity a². In the single-index setting, prior works established a sufficient sample size Z d(((K), where K E N is the infor- mation exponent for algorithms utilizing correlational information [BAGJ21, BBSS22, DNGL23, MHWSE23), or the generative exponent for algorithms that employ suitable label transformations [DPVLB24, LOSW24, ADK JMS24J. Moreover, n d samples are information theoretically necessary without additional structural assumptions; this entails that estimators that only access the test prompt inevitably pay sample size that scales with the ambient dimensionality. As we will see, pretrained transformers can avoid this "curse of dimensionality" by exploiting the lww-imeesiionality of the task distribution. oiw-dimeniional structure of the function class similar to our setting has been assumed to study the efficiency of transfer learning [DLSS22] and multi-task learning [CHS+23) 23].
2 Problem Setting
denotes the l norm for vectors and the &g + f2 operator norm for matrices. For a vector w, we use Wabb for a S b denote the vector [wa, Wa+1...... wy] 1n denotes the all-one vector of size N. The indicator function of A is denoted by II. Let N be a nonnegative integer; then [N] denotes the set (n E Z
Notations. ||
1 n S N). For a nonnegative integer i, the i-th Hermite polynomial defined as (x) (-1)feT d'
For a set S, Uni denotes the uniform distribution ||ж||| = 1} by g Ō(-), represent O(-) and I2(-) notations where polylogarithmic terms are hidden. We write S when there exists constant c such that < cb. If both S b and b S a holds, we write a x b.
dz
over S. We denote the unit sphere {æ E IRd
2.1 Data Generating Process
2.1.1
In-context learning
We first introduce the basic setting of ICL [BMR+ of simple function classes as investigated in [GTLV22, ASA+22]. 22]. In each task, the learner is given sequence of inputs and outputs (æ1, V1.......................... NN, UN, æ) referred to as prompt, where æ; x E R2 and yi E IR. The labeled ex- amples X (x)
yN) E R are called context, and is the query. Given input distribution æ ..,... NN, Lid. Drr, the output y/ is expressed
ΣN) E IdxxN y (V1
y = fm, + Si. i E [N].
where f. is the true function describing the input-output relation and & i..d. D. is i.i.d. label noise. Note that f. also varies across tasks we assume f. is drawn i.i.d. from some true distribution D,.- In the pretraining phase, we optimize the model parameters given training data from T distinct tasks {(=) .f....,,a0 Fhy: Vis. x', ,v''') ) which is composed of prompts {(a), ".....,,hrr vhr. a. ) and responses ( Ye=ı for queries {x))T where y' fl(") + s and 5² Lid. De.
Je=1
4

---

We say model learns these functional relations in-context, when the model can predict the output fx( corresponding to query by solely examining the context (X, X.y) y), without updating model parameters for each task. Given the pretrained model f(x, y. ax 0) with parameter 0 which predicts the label of query æ from context (X; y), we define the expected ICI. risk as
R (f) = (XI I:N. DINN- x; 0) -yl,,
(2.1)
where y = (2) + & and the expectation is taken over the in-context data: Ii...., N , x ~ Da, fo ~ SN S D.. Note that we take the expectation with respect to contexts (X (Xi:N= V1:N-) IRxNN" x xRN of length N in order examine the behavior of ICL specific context length.
2.1.2 Gaussian single-index models
We consider the situation where the true onput-output relation is expressed by single-index models, i.e, functions that only depend the direction of the index vector /3 in the input space. An efficient learning algorithm should adapt to this feature and identify the relevant subspace from high- dimensional observations; hence this problem setting has been extensively studied in the deep learning theory literature [BL19, BES 22, BBSS22, MHPG 23, MZD 23] to demonstrate the adaptivity of gradient-based feature learning.
Assumption 1. Let be the noise level. The prompt (æ V1. ................................ NN. UN 2)) is generated
L.d. De = -NO, Ia), y: fx + St. fa(x,) = (((2.. .N,E x
31, 22,
Lid. Unif({ (TT,))) and the distribution D,. of the true function fa is specified as:
where 51,
Index Features. Let S be an r S d-dimensional linear subspace of R.. We draw B uniformly from the unit sphere S(S) in S, i.e., from 8(5) (a B E S, =
Link Function. o2(2) = Σ $ (z), where S Q=Q S P. We draw the Hermite coefficients q from any distribution satisfying
P
= Oa,r # 0,
a s Re (a.s.) and (cq..... cp) + (0...., (a.s.).
(2.2)
Throughout the paper, we assume that P < d, r and < d: specifically, we take p= (1, and r I d2/2, Note that the condition r < entails that the class of target functions is low-dimensional, and as we will see, such structure can be adapted by the transformer via pretraining.
Remark 1. We make the following remarks on the assumption of single-index function class.
For each task, the target is a single-inder model with different index features drawn from some rank-r subspace, and different link function with degree at most P and information exponent (defined as the index of the lowest degree non-zero coefficient in the Hermite expansion of the link function, iee, minçi 6; A 0} in this case; see [BAGJ21, DH18J) at least Q. This heterogeneity reflects the situation where the difficulty of learning the imput-output relation varies across tasks. Note that we allow for different distributions of the Hermite coefficients {c:}: for example, we may set (cq....,ccp)
(. o =1] (manifold of coefficients satisfying E [x))(()))= = = 1), Uni 10--1 (O,...,))), or Unif(((1....,)),, (0......)
Unif (cQ....cp)) Ei-q
2
The condition Q 2 (i.e., the Hermite expansion of a does not contain constant and linear terms) ensures that the gradient update detects the entire r r-dimensional subspace instead of the trivial rank-1 component. For generic polynomial o.. this assumption preprocessing step that subtracts the low-degree components, as done in [DLS22].
be satisfied by simple
2.2 Student Model: Transformer with Nonlinear MLP Layer
We consider a transformer composed of a single-layer self-attention module preceded by an embed- ding module using nonlinear multi-layer perceptron (MLP). Let E B xdxxxxx be embedding
5

---

VN. Φ). A single-layer SoftMax self attenttnn matrix constructed from prompt (x1, (21,31.... $1........................... module [VSP* 17] is given as
(W E) W E
fAuta( W WV W WO) = E W softmax
(2.3)
where p is the temperature, and WK W E Rdxxd, W Rd.xd, and W E Rd. are the key, query, value, and projection matrix, respectively. Following prior theoretical works [ZFB23, ACDS23, MHM23, KNS24], with p = N; it has been argued that such simplification can reproduce phenomena in practical transformer training [ACS* 23). We further simplify the original self-attention module (2.3) by merging WPWV as WPV E Rd and (Wk) W as WKQ E Rdd.dd,, and consider the following parameterization also introduced in [ZFB23, ACDS23, WZC+23) 23].
remove the SoftMax and instead analyze the linear attention
K
WPV
WKQ
(2.4)
1x
0 01x
where E R and K E ()x f E: W WKQ)
Then, the simplified attention module is written as r PV w KQ E) and we take the right-bottom entry as the
prediction of y corresponding to query
Prior analyses of linear transformers [ZFB23, ACDS23, MHM23] defined the embedding matrix E simply as the input-output pairs; however, the combination of linear embedding and liner attention not sufficient to learn nonlinear single-index models. Instead, we set de = m + 1, dy = N + 1 for m N and construct E using an MLP layer:
(ww x1 + b1)
aww EN b1) alw x +b)
E
: o(w + bm)
E & (R(mtD)(N+1))
(2.5)
o ( aN + bm)
aw +bm) 0
y1
UN
where Wi,..., Wm E _dd and bl. ...... bm E R are trainable parameters and o: IR + R is nonlinear activation function; we use (z)) = (z) = max(z, (2,0, 0}. This to say, we define embedding as the "hidden representation" o(w x + b) of width-m two-layer neural network. For concise notation write W (wi,-.., wm)
(bn.-..,b,,).
Remark 2. We make the following remarks the considered architecture.
The MLP embedding layer before attention has been adopted in recent works [GHM* 23, KS24, KNS24]. This setling can be interpreted as an idealized version of the mechanism that lower layers the transformer construct useful representation, top of which upper atlention layers implement the in-context learning algorithm.
Our architecture is also inspired by recent theoretical analyses of gradient-be basedd feature learning, where it is shown that gradient descent on the MLP layer yields adaptivity to features of the target function and hence improved statistical efficiency [AAM22, DLS22, BES 22).
Combining the attention fAu and MLP embedding (2.5), we can express the model prediction y as
(w a + bz)
f(X, y1 W,T,D) = To X b1
(2.6)
N
о ( T x + bm))
where I - T.T and ( W X + +b1)) denotes the m x N matrix whose (i.j)-th entry is (w, æz- b:); see Appendix for full derivation. We refer to I as the attention matrix.
2.3 Pretraining: Empirical Risk Minimization via Gradient Descent
We pretrain the transformer (2.6) by the gradient-based learning algorithm specified in Algorithm 1, which is inspired by the layer-wise training procedure studied in the neural network theory literature [AAM22, DLS22, BES+ 22]. In particular, update the trainable parameters in sequential manner.

---

Algorithm 1: Gradient-based training of transformer with MLP layer
Input :Learning rate ?/1, weight decay λ1, A2, prompt length N1, N2, number of tasks Ti, T2, attention matrix initialization scale 7.
Initialize w Unif(sd-1) (j E [m]];: 8/ Unif([-1,, (i e [m]);
p (nni(( (i (ml) and L (i + [m]). Stage I: Gradient descent for MIP layer
Draw data {(x11 vi aNx uN,, x', v² )) with prompt length Ni.
W (1) - w (0)
E s(x', yy x'; W()), p b(o)))2 X \mw)y (0)
7
5 Initialize by ~ Unif log d, log d)).
Stage II: Empirical risk minimization for attention layer
Draw data {<f vi.... ................................ ..... N' VNg x', yd +T+ M=Ti+4 with prompt length N2. argminr 1 El 2 (v/ - f(X', " y _ W(), r. b))² + 2 |
I t
T2
Output Prediction function x + f(X; y. a. W r b)).
Stage I we optimize the parameters of the MLP (embedding) layer, which non-convex problem due to the nonlinear activation function. To circumvent the nonlinear training dynamics, we follow the recipe in recent theoretical analyses of gradient-based feature learning [DLS22, BES+ BEG 22]; specifically, we zoom into the early phase of optimization by taking one gradient step on on the regularized empirical risk. As we will see, the first gradient step already provides reliable estimate of the target subspace S, which enables the subsequent attention layer to implement a sample-efficient in-context learning algorithm.
Stage II we train the attention layer, which is a convex problem and the global minimizer can be efficiently found. We show that the optimized attention matrix I performs regression on the polynomial basis (defined by the MLP embedding), which can seen as an in-context counterpart to the second-layer training (to learn the polynomial link) in [DLS22, AAM23, OSSW24).
3. Transformer Learns Single-index Models In-Context
3.1 Main Theorem
Our main theorem characterizes the pretraining and in-context sample complexity of the transformer with MLP layer (2.6) optimized by layer-wise gradient- based pretraining outlined in Algorithm 1.
Theorem 1. Given Assumption 1 and I d',, We pretrain the transformer (2.6) using Algorithm I with m = ),T,
(1(82+199). N,T, (((d22+1r), 7 mārbag x m2
(log df -C, for constant C,,. Then, for appropriately chosen regularization parameters \1 A2 > with probability at least 0.99 over the data distribution and random initialization, the ICL prediction risk (2.1) with test prompt length N for the output f of Algorithm can be upper bounded as
3P
1
1
Rx-(f)
+ 44P
T2
N2
N
Theorem 1 suggests that to achieve low in-context prediction risk, it is sufficient to Ti, Ni ((((91))) and m, T2, N2, N ((P)). Observe that the pretraining complexity Tı, Nj scales with the ambient dimensionality d, but the in-context sample complexity N only scales with the dimensionality of the function class r d. This illustrates the in-context efficiency of pretrained transformers and aligns with our observations in the GPT-2 experiment reported in Figure I.
Remark make the following remarks.
The sample complexity highlights different roles of the two sources of low dimensionality our setting. The low dimensionality of single-index f. entails that the pretraining cost scales as N Z (8))), which consistent with prior analyses on gradient- masient-based feature learning [BAGJ21, DNGL23), On the other hand, the low dimensionality of function class (i.e., B lies in r-dimensional
7

---

subspace) leads to an in-context sample complexity that scales as N 2 r)(P), which, roughly speaking, is the rate achieved by polynomial regression or kernel models on r-dimensional data.
The multiplicative scaling between N1 and T1 in the sample complexity suggests that one can tradeoff between the two quantities, that is, pretraining on more diverse tasks (larger Tı) can reduce the required pretraining context length N1.
Similar low-dimenn dimensional function class has been considered in the setting of transfer or multi- task learning with two-layer neural networks (DLS22, CHS" 23), where the first-layer weights identify the spun of all target functions, and the second-layer approximates the nonlineurity. However, the crucial difference is that we do not update the network parameters based on the in-context examples; instead, the single-index learning is implemented by the forward pass of the transformer.
Comparison against baseline methods. Below we summarize the statistical complexity of commonly-usee estimators that only have access to N in-context examples, but not the pretraining data. Note that the in-context sample complexity all depends on the ambient dimensionality d > r.
Kernel models. Recall that target function is degree- polynomial. and hence kernel ridge regression requires N Z d in-context examples [GMMM21, DWY21].
CSQ learners. The correlational statistical query (CSQ) lower bound suggests that an algorithm making use of correlational information requires N* Z (9()) samples to learn single-index model with information exponent Q [DIS22,, AAM22]. This sample complexity can be achieved by online SGD training of shallow neural network [BAGJ21, DNGI.23).
Information theoretic limit. Since the single-index model (1.1) contains unknown parameters. we can infer an information theoretic lower bound of N* 2 d samples to estimate this function. This complexity can be achieved (up to polylog factors) by tailored SQ algorithms [CM20, DPVLB24] modified gradient-base training of neural network [LOSW24, ADK 24, JMS24).
3.2 Proof Sketch of Main Theorem
We provide a sketch of derivation for Theorem 1. The essential mechanism is outlined as follows: after training the MIP embedding via one gradient descent step, the MLP parameters (woft) align with the common subspace the target functions S. Subsequently, the (linear) attention module estimates the input-output relation f. (which varies across tasks) this r-dimensional subspace. We explain these two ingredients in the ensuing sections.
3.2.1 Training the MLP Layer
We first show that the first gradient descent step W results in significant alignment with the target subspace S, using a proof strategy pioneered in [DLS22]. Note that under sufficiently small initialization and appropriately chosen weight decay, we have
2 Etn v'V m f( (X', y. x:; ( p()), b(0)).
The key observation is that this correlation between y and the gradient of model output contains information of S. We establish the concentration of this empirical gradient around the expected (population) one, which determines the required rates of Tj and Ni. For the population gradient, we make use of the Hermite expansion of σb, (2) = o(2 + b;) and show that its leading term is
(0)
proportional to
w wir which is contained in the target subspace S; see Appendix B for details. 0a-r
3.2.2 Attention Matrix with Good Approximation Property
Next we construct the attention matrix P which satisfies the following approximation property: Proposition 2 (Informal). There exists such that with high probability,
[(X' y¹, x (L), I.b) v² -r=== h 722 /N2
for all t (Ii + 1,... T2). Moreover, we have ||F||5 Ö(f2P/mm))
8

---

We provide an intuitive explanation of this construction. Recall that the target function can be written in its Hermite expansion (x) -i= a He: ((x, (3x,)) (3)). We build an orthonormal basis for f+ as follows. Let (,1,...,,BB Br) be an orthonormal basis of S. Then, any f. can be expressed
linear
combination of functions in H = [[ Hep, ((3). 1) Q S p1 + ... + Pr - P.Pi 2 0,..., Pr 2 0}, where Eg Ea [h(2)) (a)) I== holds for h. h' E H. We write H = {h ..... ................................ h with Bp |H| = (r and observe that two-layer neural network can be constructed approximate each hug Specifically, there exist vectors a'...., a Br 1mm such that ش3= az a ( ((w))) f) x) + bj) = hn (x)) for each n E [By].
Consequently, we can build the desired attention matrix using coefficients a², a Br, Let (a a BB) E mRmxBr and T AA the attention module can recover the true function
N₂ To (W + b1 y (W + +b)
دز= ai ..) + h1) U/) (E,- aj (1) ,2) + bb))
n=1
No
ii=
2 وي E(h,(()) (z)) f=(x). No hn (zx))) hm.()) E N2
Bp
(3.1)
Roughly speaking, the self-attentio architecture computes the correlation between the target function and basis element h; to estimate the corresponding coefficient in this basis decomposition. We evaluate (a) the approximation error of two-layer neural network in Appendix C.1, and (b) the discrepancy between the empirical and true correlations in Appendix C.2.
Note that the approximation errors and the norm of r at this stage scale only with r up to poly- logarithmic terms; this is because ( already identifies the wiw-dimesiionll target subspace S. Consequently, the in-context sample size N only scales with the target dimensionality r < d.
3.2.3 Generalization Error Analysis
Finally, we transfer the learning guarantee from the constructed I to the (regularized) empirical risk minimization solution r By the equivalence between optimization with L2 regularization and norm-constrained optimization, there exists A2 such that || < ||]]| and the empirical ICL loss by I* is no larger than that of T. Hence, we can bound the generalization error by I using standard Rademacher complexity bound for norm-constrained transformers provided in Appendix D.I. One caveat here is that the context length N* at test time may differ from training time; hence we establish a context length-free generalization bound, which is discussed Appendix D.2.
4 Synthetic Experiments
4.1 Experimental Setting
We pretrain a GPT-2 model [RWC+19] to learn the Gaussian single-index function class (1.1).
Specifically, we consider the 12-layer architecture (with 22.3M parameters) used in [GTLV22] for in-context linear regression. The pretraining data is generated from random single-index
of models: for each task t, the context ((*), W iid. Unif{ (CQ.....).. Ex = and
is generated as of iid. NOO Ia) and
vi Eie He:((ar; B²)), where (cb.... ..)ppp
ge ltd. Unif(,8 1B....,
..... Br. .........1.......... 0J ||8|| = 1}. See Appendix F for further details.
4.2 Empirical Findings
Ambient dimension-free sample complexity. In Figure 2 we examine how the in-context sample complexity of the GPT-2 model depends on the ambient dimensionality d and the function class dimensionality For each problem setting the model is pretrained for 100, 000 steps using the data of degree P and information exponent Q = 2 (see Appendix for details). In Figure 2(a) we observe that for fixed r = 8, varying the ambient dimensionality d = 16, 32, 64 leads to negligible change the model performance for the in-context phase. In contrast, Figure 2(b) illustrates that for fixed d, the required sample size N scales with the dimensionality of the function class 2, 4, 8. This confirms our theoretical finding that transformers can adapt to dow-dinensional structure of the distribution of target functions via gradient-based pretraining.

---

1.4.
1=
1.2
6=
1.2
d= 8
d=
1.0
d= r=8
d=
prediction risk
1.0
prediction risk
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
50
150 in-context sample size
250
100
150
200
250
in-context sample size N
(a) risk varying d. fixed r.
(b) ICL risk - fixed d varying
Figure 2: In-context sample complexity of GPT-2 model pretrained Gaussian single-index function (see
Section 4.1 for details) of degree-4 polynomial. Observe that (a) the ICI. risk curve overlaps for different ambient dimensions d but the same target (subspace) dimensionality r, and (b) the required sample size N' becomes larger as increases.
Superiority over baseline algorithms. In Figure we compare the in-context sample complexity of the GPT-2 model pretrained by data of Q 3 and /' 2 against two baseline algorithms that directly learn f. on the test prompt: (t) kernel ridge regression with the Gaussian RBF ker- nel k(x, x²) = exp(-11x x'1l x'133/o²) and (ii) two-layer neural network with ReLU activation fNN(F) = Em ( w.)) trained by the Adam optimizer [KB15]. We observe that for 16,32, 32, the pretrained transformer outperforms both KRR and two-layer NN; more- over, the performance of these two baseline algorithms deteriorates significantly as the ambient dimensionality d becomes larger.
Conclusion and Future Direction
We study the complexity of in-context learning for the Gaussian single-index models using a pre- trained transformer with nonlinear MLP layer. We provide an end-to-end analysis of gradient-based pretraining and establish a generalization error bound that takes into account the number of pre- training tasks, the number of pretraining and in-context examples, and the network width. Our analysis suggests that when the distribution of target functions exhibits low-dimensional structure, transformers can identify and adapt to such structure during pretraining, whereas any algorithm that only has access to the test prompt necessarily requires larger sample complexity.
We outline a few limitations and possible future directions.
The in-context sample complexity we derived (P)) corresponds to that of polynomial regression kernel methods in r-dimensional space. This rate is natural as discussed in Section 3.2.2, where the linear self-attention module extracts the coefficients with respect to fixed basis functions (of size ,(()), An interesting question is whether transformers can implement more efficient in-context algorithm that matches the complexity of gradient-basee feature learning in dimensions. This can be achieved if the pretrained model learns features in in-context.
Our pretraining complexity is based makes use of the correlational information; hence the information exponent of the link functions plays an important role. We conjecture that the sample complexity can be improved if modify the pretraining procedure or training objective, as done in IDTA+24, LOSW24, ADK+24, JMS24].
GD step analysis similar [DLS22, BES 22] which
Acknowledgement
KO was partially supported by JST, ACT-X Grant by ISPS KAKĚNHI (24K02905) and JST CREST (JPMJCR2015), This research unrelated to DW's work at XAI.
TS was partially supported
10

---

References
[AAM22] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewiez. The merged- staircase property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory. pages 4782-4887. PMLR, 2022.
[AAM23] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. SGD learning on neural networks: leap complexity and Madllett--sddle dynamies. In Conference on Learning Theory, pages 2552-2623. PMLR, 2023.
[ACDS23] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2023.
[3C Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention (maybe) all you need (to understand transformer optimization). arXiv preprini arXiv:2310.01062. 2023.
[4DK*24] Luca Arnaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iuvant: Data repetition allows sgd to learn high-dimensional multi-index functions. arXiv preprint arXiv:2405. 15459, 2024.
[ASA+22 22] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint 2rXiv:2211.15600 2022.
[BAGJ21] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensionn inference. The Journal of Machine Learning Research, 22(1):4788-445 2021.
[BBPV23] Alberto Bietti, Joan Bruna, and Loucas Pillaud- Pillaud-Vivien. On learning Gaussian multi- index models with gradient flow. arXiv preprint arXiv:2310.19793, 2023.
[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems, 35, 2022.
[BCW+23] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statis- ticians: Provable in-context learning with in-context algorithm selection. Advances in Neural Information Processing Systems, 36, 2023.
[BEG+22] Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit, Advances in Neural Information Processing Systems, 35, 2022.
[BES+ Jimmy Ba, Murat Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35, 2022.
[BBS+23] 23] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu. Learning in the presence of tow-dimensional structure: A spiked random matrix perspective. Advances in Neural Information Processing Systems, 36, 2023.
[BKM+19] Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions in high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, (6(12):5451-546) 2019.
[BL19] Yu Bai and Jason D Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. arXiv preprint 1rXiv:1910.01619, 2019.
[BMR 20] 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 Neural Information Processing Systems, 33, 2020.

---

[BMZ23] Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.0000 2023.
[CCS23] Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.06528, 2023.
[CHS+23] Liam Collins, Hamed Hassani, Mahdi Soltanolkotabi, Aryan Mokhtari, and Sanjay Shakkottai. Provable multi-task representation learning by two-layer relu neural networks. arXiv preprint arXiv:2307.06887, 2023.
[CM20] Sitan Chen and Raghu Meka. Learning polynomials in few relevant dimensions. In Conference on Learning Theory, pages 1161-1227. PMLR, 2020.
[CSWY24a] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamies of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.19442 19442, 2024.
[CSWY24b] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. arXiv preprint arXiv:2409. 10559, 2024.
[DH18] Rishabh Dudeja and Daniel Hsu. Learning single-index models in gaussian space. In Conference On Learning Theory, pages 1887-1930. PMLR, 2018.
[DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference Learning Theory, pages 5413-5452. PMLR, 2022.
[DNGL23] Alex Damian, Eshaan Nichani. Rong Ge, and Jason D. Lee. Smoothing the landscape boosts the signal for SGD: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems, 36, 2023.
[DPVLB24] Alex Damian, Loucas HilluddViviens Jason Lee, and Joan Bruna. The com- putational complexity of learning gaussian single-index models. arXiv preprint arXiv:2403.05529 2024.
[DSD+222 Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers. arXiv preprint arXiv:2212. 10559, 2022.
[DTA+24] Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborová, and Florent Krzakala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and lapp exponents. arXiv preprint arXiv:2402.03220, 2024.
[DWY21] Konstantin Donhauser, Mingqi Wu, and Fanny Yang. How rotational invariance of common kernels prevents generalization in high dimensions. In International Conference on Machine Learning, pages 2804-2814, PMLR, 2021.
[GHM+23] Tianyu Guo, Wei Hu, Song Mei, Huan Wang. Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? case study on learning with representations. arXiv preprint arXiv:2310.10616 2023.
[GMMM21] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Lin- earized two-layers neural networks in high dimension. The Annals of Statistics, 1925),1029-1054, 2021.
24] Khashayar Gatmiry, Nikunj Saunshi, Sashank J Reddi, Stefanic Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi-step gradient descent for in-context learning? In International Conference on Machine Learning, 2024.
[GSS21] Friedrich Götze, Holger Sambale, and Arthur Sinulis. Concentration inequalities for polynomials in e-sub-exponential random variables. Electronic Journal of Probability, 26:1 22, 1.21.
12

---

[GTLV22] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. Advances in Neural Information Processing Systems, 35, 2022.
[HCL23] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprini arXiv:2310.05244 2023.
[HMM* 19] Evan Hubinger, Chris van Merwijk, Vladimir Mikulik, Joar Skalse, and Scott Garrabrant. Risks from learned optimization in advanced machine learning systems. arXiv preprint arXiv: 906.01800, 2019.
[JLLVR24) Hong Jun Jeon, Jason Lee, Qi Lei, and Benjamin Van Roy. An information- analysis of in-context learning. arXiv preprint arXiv:2401.15530 /55300 2024.
[JMS24] Nirmit Joshi, Theodor Misiakiewicz, and Nathan Srebro. On the complexity of learning sparse functions with statistical and gradient queries. arXiv preprint arXiv:2407.05622, 2024.
[KB15] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
[KNS24] Juno Kim, Tai Nakamaki, and Taiji Suzuki. Transformers are minimax optimal nonparametric in-context learners. arXiv preprint urXiv:2408.12/41 2024.
[KS24] Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Noncon- mean-field dynamics on the attention landscape. arXiv preprint arXiv:2402.01228, 2024.
[LLZV+24] Yue M Lu, Mary I Letey, Jacob A Zavatone- Veth, Anindita Maiti, and Cengiz Pehle- van. Asymptotic theory of in-context learning by linear attention. arXiv preprint arXiv:2405. /1751, 2024.
[LOSW24] Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with sgd near the information-theorrtie limit. arXiv preprint 4rXiv:2406.01581,200 2024.
[LWL+24] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. How do nonlinear transformers learn and generalize in in-context learning? In International Conference on Machine Learning, 2024.
[Mau16] Andreas Maurer. A centar-conttaation inequality for rademacher complexities. In Algorithmic Learning Theory: 27th International Conference, ALT 2016, Bari, Italy, October 19-21, 2016, Proceedings 27, pages 3-17. Springer, 2016.
[MHM23] Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:230..03576, 2023.
[MHPG+23] Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with SGD. In International Conference Learning Representations, 2023.
[MHWSE23] Alireza Mousavi-
Denny Wu, Taiji Suzuki, and Murat A Erdogdu. Gradient- based feature learning under structured data. Advances in Neural Information Process- ing Systems, 36, 2023.
[MM18] Marco Mondelli and Andrea Montanari. Fundamental limits of weak recovery with applications to phase retrieval. In Conference On Learning Theory, pages 1445-1450. PMLR, 2018
[MZD+23] Arvind Mahankali, Haochen Zhang, Kefan Dong, Margalit Glasgow, and Tengyu Ma. Beyond ntk with vanilla gradient descent: A mean-field analysis of neural networks with polynomial width, samples, and time. Advances Neural Information Processing Systems, 36, 2023.
13

---

[NDL24] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint WxXiv:4422.14755, 2024.
[OSSW24] Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Learning sum of diverse features: computational bardness and efficient gradient-based training for ridge combi- nations. arXiv preprint arXiv:2406. 0 11828, 2024.
[RPCG23] 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. Advances in Neural Information Processing Systems, 36, 2023.
[RWC+19] 19] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
[SGS+24] Michael E Sander, Raja Giryes, Taiji Suzuki, Mathieu Blondel, and Gabriel Peyré. How do transformers perform in-context autoregressive learning? arXiv preprint urXiv:2402.05787, 2024.
[Ver18] Roman Vershynin. Bigh--iimnsional Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.
[V3NR+23] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexan- der Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174. PMLR, 2023.
[VSP+17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, F. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[Wai19] Martin J Wainwright.
Statistics: A Non-Asymptotic Viewpoint.
Cambridge University Press, 2019.
[WZC+23] 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? arXiv preprint arXiv:2310.0839), 2023.
[ZFB23] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint 2rXiv:2306600927,,, 2023.
[ZWB24) Ruiqi Zhang, Jingfeng Wu, and Peter I. Bartlett. In-context learning of a linear transformer block: Benefits of the mlp component and one-step gd initialization. arXiv preprint 2402:22402.195.1 2024.
[ZZYW23] Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305. 200.::205./9420, 2023.
14

---

Table of Contents
1
Introduction
1.1
Our Contributions
1.2
Related Works
2
Problem Setting
2.1
Data Generating Process
2.2
Student Model: Transformer with Nonlinear MLP Layer
2.3
Pretraining: Empirical Risk Minimization via Gradient Descent
6
3
Transformer Learns Single-index Models In-Context
3.1
Main Theorem
3.2 Proof Sketch Main Theorem
Synthetic Experiments
4.1
Experimental Setting
4.2
Empirical Findings
Conclusion and Future Direction
10
Preliminaries
16
A.1 Definition of High Probability Event
16
A.2 Tensor Notations
16
A.3 Hermite Polynomials
16
A.4 Other Auxiliary Lemmas
18
Proofs for the MLP Layer
19
B.1 Calculation of Population Gradient
20
B.2
Concentration of Empirical Gradient
24
Construction of Attention Matrix
29
C.1 Approximation Error of Two-layer Neural Network
................................
29
C.2 Concentration of Correlation
34
C,3 Proof of Proposition 2
36
D Generalization Error Analysis and Proof of Theorem
37
D.I Rademacher Complexity Bound
D.2 Prompt Length-free Generalization Bound
................................
40
D.3 Proof of Theorem
41
E Derivation of Simplified Self-attention Module
42
Details of Experiments
43
F.1 Detailed Experimental Settings
43
F.2 Experiment on Our Architecture and Algorithm
44
15

---

A Preliminaries
We consider the high-dimg diteemiieniin setting, i.e., our result holds for all d > D where D is a constant which does not depend on d and r. Throughout the proofs, we take
(11,..... Tr, 0. . 0)
Fi,.... 2r ]R} and Unif(S(S)) Unif([(B),..... Br, 0, ...., 0) 8 + + 3 1}). Note that this assumption is without loss of generality by the rotational invariance of the Gaussian distribution.
A.1 Definition of High Probability Event
We say that an event A occurs with high probability when there exists sufficiently large constant C which does not depend on d, r and
Definition 3.
1 - Pr[A] S )(d-c')
holds.
We implicitly assume that we can redefine the constant C* be sufficiently large as needed. The lemma below is basic example and will be used in the proofs:
Lemma 4. Let x ((0, 1). Then, |a|| S Vlog d holds with high probability
Proof. From the tail bound of Gaussian, P([æ] 2 t)
eex (-t2/2) holds. Hence we get
P[[x 2 V2C log d) < ((
Note that C* can be redefined by changing the hidden constant in |x]] X Vlogd. d.
If A....., ....... AM occurs with high probability where M Opoly then A ... n AM also occurs with high probability (by redefining C*). In particular, throughout this paper we assume that m, N1, N2, Tı, T2 O(poly(d)), which allows us to take such unions.
A.2 Tensor Notations
In this paper, a k-tensor is multidimensional array which has h indices: example, matrices are 2-tensors. Let A be a k-tensor. Aq....f denotes (ig....,44)-t entry of A. Let A be a k-tensor and B be a I-tensor where 2 A(B) denotes a tensor whose (11,..., ii-t)-th fi......k..)-t entry is
A
Σ A B B
and is defined only when sizes are compatible. If = I, we sometimes write A(B) as A B or (A, B). Let u E Rd be a vector and be a positive integer. Then, S E Rxdx denotes a k-tensor whose .i4)tth entry v,, .... U.g
Let fxx) Rdd + R be variibll differentiable function. k-tensor V f(x)) defined as
θ
(V* (x=))
f(r)
Oxi
Dos,
The following properties can be verified easily.
Lemma 5. For tensors A and B, [A) < A B holds.
Lemma 6. For vector v, $
k || (1)+ holds.
A.3 Hermite Polynomials
We frequently use (probablists") Hermite polynomials, which is defined as
Hc,(2) = (-1)'c di dz
where i is a nonnegative integer. We often make use of the orthogonality property: Ez~N(0,1) (2) il6.,-- The Hermite expansion for o R - IR is defined as
a(2) = Σ.2 $ (z) where a, = =E Similarly, the multivariate Hermite
16

---

expansion for f R' - R is defined as f(z) = En d000 (1))1-14))
Aiunrd He,, (21) ... Her (za).
where Oii..nia E (io.. (z1) (21)....ue Hei (za)l. The coefficient ai can also be 0n 827
obtained by Oi R P
are a=
f(z)
Consider f. in our problem setting (Assumption 1). We can bound (z)) and [ (x) as follows.
Lemma 7. Under Assumption 1, Em....f.(ar)) (1, and Er (f.())) = Oa,- (1) holds.
Proof. From (2.2) and P = (1),
Σ i!
(f, = BE-N(O.1)(a)
Hes(z)
r
(1)
We can bound Ex [f (x)4 = Bs- (Ei + @ =))e naively as follows: let C$
maxosisp E. Noo (e.(()))). Sp is an O(1) quantity depending only on P and Q. Then,
P
Σ
Ho)(2)
BewN((.1)
Bllacc
E (,
(2))
a!j!!!!!
QsigkigP
Σ
",
QSIAKIE Sp-R
The lemma below is useful to find a basis of the set of single-index functions.
Lemma 8. Suppose B S(S). Then,
Pi 4,
p!
Hcp((2,13)
Σ
Pi!-.. pr! 8 SP......
Hepp (x1) ... Hep, (Ir)
21.0.... 20
holds.
Proof.
Note that Ex,...,, (0011)
841
del He, (3)) is nonzero only when
Oz
+ + in p. because g
atr
He, (x,
p(p 1)... (p (is +... + i+) +
d 1) ... ... Hc (1r (61tm.ii)))) (.) (x,,)) B)) and E
Dx.
I (((01) Hc,((x, (xx,))) = EsmN(O.)) [Hc,(2]] (s.NNo.))|Hc(a())) =
if > O. When i +... i, = p. then Ex, . (((11)
81
gir Da.--
Hep (xc, Bi) = =p!8 ... 3
ax"""
holds. Thus, from the multivariate Hermite expansion we obtain the claim.
Corollary 9. Let (z) Ee $ He,((2, B)). Then, Eg f=())) 3 f00.kff$$ S and
otherwise it is the zero tensor.
17

---

A.4 Other Auxiliary Lemmas
Lemmas random vectors. Let x(d) be the distribution of v .... 2 where 2 1), Note that the Gaussian vector v ~ (N La) can be decomposed as v = 2w where x(d) and (n because the norm and the direction of the Gaussian vector are independent of each other.
Lemma 10 (Example 2.11 in [Wai19]), For z x(d) and <t<>
P[|2² d 2 dt - 2 expl dt² /8)
holds.
Lemma 11. Lef C > O. For u Uni | 2 v
4Cr d
log d holds with probubility at least
1 - -2xxp(-d/33) 2rd-C (then with high probability).
Proof. Let ~ x2 (d) be independent from w. We observe that the distribution of 2² w 1 the d((r)-distibbttinn First, from Lemma 10, with probability least 2 exp((d//3
=² 2 d/2
holds. Moreover, we can say that (z')² ~ (r) satisfies (z')² I 2Cr log d with high probability for sufficiently large constant:
us decompose (z')² uf + ... + r? where e; ~ (0, 1); from
Lemma 4, u² I 2C log d holds with probability at least - 2d 2d-C
C and thus (z')² - 2Cr log d holds
with probability at least - 2rd C, Taking the union bound yields the result.
Lemma 12. For k 2 and r S d,
Ew--gd W1: 2 = Od.r (())
A
holds.
Proof. Since || polynomial of w. Lemma 24 in [DLS22] yields E 1 |1 || x
Ew-gs II W1r [1 ||// Here, Euw--Uiit(S 1) syll|w1:rl |1 3x. 1 [[o] +1..) +1,311 3r+r(r-1)
x
d(
(r²² /d2) yields the result.
Lemma 13 (Sub Gaussian vector). Let x be a d-dimensional random vector and suppose that (u, x)
zero-mean -u-sub Gaussian for any u E Rdd. such that || = 1. Then, ||>|| = Ö(avd)) holds with high probabillity.
Lemma 14. Let X be a oaa02-uu sub Gaussian random variable and Y be a bounded random variable satisfying YY| S R almost surely. Then, XY - E[XY] is O >-su Gaussian.
Proof. XY is ( Rub Gaussian from Lemma 2.6.8 in [Ver18).
by examining the tail probability. Hence, XY E[XY] is
Polynomial concentration. We frequently use the following bounds on output data y (which is a polynomial x:).
Lemma 15 (Theorem 1.5 in (GSS21]). Lef N( In) and f be degree- D n-variable polyno- mial function. Then, for all it holds that
2/h
S 2exp
min 1Sk5D
CDM²
where M is an absolute constant and Cp is a constant depending only on D (if [V (x)) = =0 then we ignore the term with k).
Lemma 16. a N Ta) and y fax) + Ei @ He:((a 33) + where e; and B3 are drawn according to Assumption 1, and ~ Uni Then, for any fixed {c,) and B on the support of their distributions, it holds that
=
2/h
1
min 1
P(lvl > t 2
CpM2 2
Re
where M is an absolute constant and Cp is constant depending only on P.
18

---

Proof. Note that from Corollary 9, = [V = 0 and Eg f+ = c, (3) "
= 0,
by (2.2), Lemma 6 and ||3|| 1, [ (z)]llr R. holds. Lemma 15 yields the result.
Corollary 17. Fix any {ci and 33 on the support of their distributions. Then, If.(2=)), ly] X (log d) P/2 holds with high probability over the distribution of x and s.
B Proofs for the MLP Layer
This section analyzes the behavior of MLP weight w after one gradient descent step (line 4 in Algorithm 1), as sketched in Section 3.2.1.
We set À 1/1 -1 From line 4 of Algorithm 1, for each j E [mn].
w (1)
1 2m
f(x yn x. W ((0)),(0))
Ti
(0)
Ti
2m Ti
[(X y
x ,x'; w ) b(o) / u s(x', y a w r (())
,',
T
N Ni
= -2m 10
10)
Oby (w)
(0)
x²)
Σ
(0) at)ac
y o ( (0)
x²)x'
(w,) (0)
ai)
(B.1)
Ti
N1
1=1
1 Ti
2m
Σ
f( fx y x: ( I b mn f(x', yy s:; ( r (())..
Ti
holds where σb, (z) = a( + bj). Here, of and yf is the input and output i-th context at I-th task, respectively. First we show that the final term is sufficiently small, if we set 7 13 sufficiently small.
(0)
Lemma 18. With high probability over the distribution of ((x y xt y Di and {w) (0) pm
15=1
T
2m
1
Σ
f(X',
V. (0 ((0) (( 10) f(xt ,y x ( p(O) ((0))
TH
-
-ō
Proof. Recall that
N;
f(x' 3. ,i; (0) (() (00) = Σ
(())
yoo (w (0) xt))
oae (oo) (0) 2²),
yi
jij
N
1=1
wo f(X' 00,x" w()), r()) -
vio6, (w))
(0)
or (w)) (0) xt)
N₂
N
(0)
(w) (0) 2²)) oh, (wo
T x')x'.
+
N
Fix i, f and consider the inner product between w (0)
Unif(Sa-1) and x ~ NO Ia). From
the rotational invariance, without loss of generality, we can assume that w = [1, 0, .,.0....... _0 and
of ~ (o, Ia). Therefore, (10)0)
(0) ac! ~ NO 1). From Lemma 4, (1)
(0) x!!
2 S log d holds with
high probability. Moreover, from Lemma 10 and Corollary 17, we know that 1v/1, 1y'1 x (log P/2
19

---

and || æ E Vd holds. We can take the union bound with respect to all the i, t. retaining the high probability bound (see Appendix A). Therefore, noting that |b;| - 1, we obtain
[f(' y. x.; W p b( S my(log d) P/2+1
and
V go: f(x' v. x'; W p b b(
S >(og d)
Thus obtain the assertion.
From now on we focus on analyzing (B.1) which expresses the correlation between the output label y and the gradient of model output W Let
9T NT (w, b, {(x', y", x', x",")" y')}?=n>
N w x²)
v xf)
N
T y o6www a')x²
Σ
viooww xt),
(B.2)
N1
and (wv, b) = ElgT,, (w, b. {(x' vi x² )'))?))) where the expectation is taken with respect
the data {(X y" xt ) Note that
w 2m
(0) j.j 9T NI (10g (0) (100 [(x' D a. ())))))) (
from Lemma 18.
In Appendix B.1. we analyze the explicit form g b) via the Hermite expansion: we can check that its main term aligns with the low-dimensional subspace S. In Appendix B.2, we show the concentration of Jr,,N, around g.
B.1
Calculation of Population Gradient
B.1.1
Hermite Expansion
First, note that
g(w, b) = E[yor(w' x)yoo(( ( + E[yo x z)xyoo((]T
((.. [L)
+ si)o((w a;)=;:]
Ey. (f z) (z.) s)ool xx)]
= =E,, (())))))
[f (z;) + w))o(w Txt)x:]
+ Ry, [R,,
x) E ((.)) (m:) + 5:)o [x]]]]]
x)) [f (z) x)]]
Here E,, denotes the expectation with respect to the distribution the true function f- (i.e., the a:(b)
distribution of {c and (3) specified in Assumption 1. Now let cT((=) Eizo i! Hey(z) be the Hermite expansion of student activation (2(2) ReLU( + b). Lemma 19. It holds that
Ex (2)) x)æ]
(+) (x) +
('))
(B.3)
i!
Q=Q11
(b(b) V (x
Ex Ef x)]] Σ
i!
20

---

Proof. Consider a function g(w x)
Eiso ês He x). Here (x) x)] can be
calculated as
En [f. (z))( x)
Σ i20
ai
Ex (x)
x)]
ti
dd
ai i
1
Exls- (2) (x1)
He (ea)( (: Lemma 8)
id!
120 42
11
wd d
Ex
0i
8e
So (2))
id!
Oxy
dxp
120 20....... $2
a. Sa (x) @1)
i!
20
Where the final line obtained noting that the tensor (x) has iiliia!
entries whose value is
gi1
8th
fx)).
Dr.
Note that Ex (x)) E V Σド= / He, (ar, B3) is nonzero tensor only when Qsis P.
from Lemma 8. We can also show (B.3) as
(z)) x))=
Ex (2) =)] + (f
x)] w (: Stein's 's lemma)
a;+1 Ez Vi+1 fo (x) 00))
0=t2 (b) V f. (x) [ 84
w
i!
i!
120
220
P-1
a;+1 Em /i+1 fo fx(x)]
a,42 (b) V f. (x) (w 48)
Σ
w
i!
i= Q-1
i=Q
Using Lemma 19, g(w, b) can be expanded as
(b)1
V
(
g(w,b) 6) = EE, 2
i!
a 21)2(6) ( fo V
P
(b) V f.(x))
+ w
i!
i=
(b)(b)
(0-1) aq (b) (b) V (a:) (w)
Ey, (Q - 1)!
Ex V fo
(x)
+ ssw b)
Q!
(b)
E,, E2 (0--1))
E. (V) 8Q)1 + s(w, b), (B.4)
Q!(Q - 1)1
where
s(w, b)
P-1
+ (''
-2B,.
i!
i!
(=Q-1
a, (b)E V/L.
u2 (b)² (Q - 1)!
Ex
49
(2
F.n
(Q0-1)))))))
Q!
expresses terms in the asymptotic expansion which are not the leading term. Note that, in this explicit formulation, a;(b) (e appears. It will be helpful to show that this coefficient is well-controlled.
21

---

There are constants Cи and Ct. depending only on Q and P satisfying the following condition. Let (2)2) = o2 b) and Unif([-1,1)). Then,
Lemma 20.
(Ez-wroonloo(((Heeecc S CH (2 - i - 2) and Bznwvoon/ooo(())eee 2 C2 (B.5) holds with probability 1/2. Especially, [B=~Nio.1) (z) S Ch (2 22 holds with probability I.
Proof. From Lemma 15 in [BES 23], we know that for i 2 2, 1
= V2x [(6) holds. First, as Unif([-1,1]. by setting sufficiently large Ch. [He)(b) - Cи holds for all 0 I S P. Therefore, (Bz=)V(..)lo((()Hee(e))
5//2
(E=-)(roonloo(((]Hee(]]]] s (b))) On holds for all 2 S i P +22 2 with probability one.
From continuity we know that there exists C2 such that [E.. probability 1/2.
Cz with
In the following, we first calculate the main term and then upper bound the residual terms.
B.1.2 Calculation of the Main Term
Let us calculate the main term in (B.4) explicitly. Lemma 21.
2a (b)
Q!(Q= - 1)! E. EE [ f+ ((Q-1)) Ex [<)
200 [6 (2Q - 1)!!
||0
2Q-2 2 Wirr
Q!(Q - 1)! B [220]
0a-r
holds.
Proof. From Corollary 9, it holds that E. [c)] = (3) Then, we obtain
E, EL / fa (x) 0 ((0-1))) Exl (
-Eag 1 E& (8
80 w (Q-1)) (300 o w 09]
-Eng 16 E B 020 8(2Q-1)).
Next, let us show
E [ (10 ( (122)))1))
(2Q 1)!!
11:rr 11
22Q-2 2 Wi.r
Ex 20
0d-r
First, note that by letting B' ~N((0, where
I
0,
Σg =
0_
Oa-r
and = Xr independent from B, B ~ Az holds (recall that we assumed that B ~
Uni O,, TTT 1 + +... 82 = 1}) see Appendix A). Then,
Eg[
Esaxe [2 Byg -N
620
It suffices show that (o (=020(00020000000000000000000)00 (=0120-11) (2Q - 1)!1=|1 2Q-2 for Rr. As an ex
ample, we show that the first entry of E (= (= equals to (2Q - 11] the other components can be calculated similarly. First, we expand (00.0((0.)))0020000 (2 1))) as
(Ev-N10.L.) (orta.r.///0000 1202 2(20-1))))
Σ
E E(ur%; ... Uiag 12i, 2iz
ig....fog--Eh
Note that Eluie, ... 1 I + 0 if and only if the degree of fi1n ... ... Uie-r with respect to fj is even, for all j E [r]. Then, we obtain
(Ben. 020 (2120)1)))
22

---

q1 U1+ +g.
Con 1.
E
201
1/2
241
22
2.22
2/2
242
24
21 2..20.....0
01 to +.
-1
Cq
Elu 2
242
...
U 12
22r
2211
20z
24r
912 1.02 10:0000000000000
where Cr.....rr
(11(())2))) (2..)) is the number of ways to assign (291 - 1) "1"s, (2g2) "2"s, ...
(2Q
(241 (9:1-11((2g)))))
among 8 12Q-1-
Note that Ev 1200 = (2q 1)!!. Therefore, we obtain
BowN(a,.))le02g 120 (=0022-))))1
Q1 0.41-g.-0 A.
(2Q - 1)!
Σ
(2q1- 1)!(222) .. (29,)) (2g) - 1)!!... .2.....(....) - 1)1112 24: 2,222
20r
=11
2
Q1 92
@1 +
(Q-
Σ
(11 - ()!(2)) ... (q,)1 (2Q - 1)) (2,))r
=21
q1 2 1.2....... 20
=(2Q 1)!l23 (2) + 2,))
multinomial theorem)
which concludes the assertion.
By Lemma 21, we observe that the main term of the expected one-step gradient aligns to the dimensional subspace S. In other words, the MLP layer captures the information of S via pretraining, which will be useful the later stage.
B.1.3 Bounding Residual Terms
First, similarly to the Lemma 21, obtain the explicit formulation of the entire g(w, b).
Lemma 22.
g(w, b)
Σ Q-1Si<P-1, QSISP,
24 (6)(((1c)
(i + j)!! Rznx. [2(+++1]
Wi.r
a!j!
02-r
i+] 15
20,+2 (b)E (1 + j - 1)!!
+ w Σ
||
(B.6)
Elj!
E Exxx [2(+5]
QSISP QZ/SP. ++3 even
holds.
Proof. Note that
(n-1)!
n-
Wirr
(n is even)
Eg (w0
F F [2n]
0a-r
(n is odd)
and
Ex-ep [sn] | 1" (n is even)
(n-1))) 1)!!
Eg 1 (w m0)
0 (n is odd)
which can be obtained similarly the proof of Lemma 21.
Let us upper bound the non-leading terms; we need show that the moment of first r components in residual term s(w, b) are sufficiently small.
23

---

Lemma 23. Ler = 4k where & E [P]. then,
E Ew-ga- (2) 1r = =ōar
(Varē+)
d20+1
r
(B.7)
holds uniformly over all b E 1-1
Proof. Note that ||s(w, 1) [ is polynomial in w; from Lemma 24 in [DLS22], it suffices to show the case & = 1.
s(w, b)r
20,- 1 (U)), G (6)
+ j)M!
Ezax- [2++++]] w1 | 1 wir
ilj!
Q-1SiSP-1, QSISP.
(4.J odd. ()7(Q-1.))
201=2(0)a,()) (i j j 1)!!
Σ QSiSP
ilj! E [2tt]] ||11r i+5 (+/ Wir-
QSJSP. (4)
By Minkowski's inequality E|I + 1 < E|]|æ| 4 /////, suffices to show that (B.7) holds for each term. Note that B[c+11,,] S R and Elc,cyl Rs from Assumption 1, [b,( < Cи from Lemma 20, and Ee [rrl known that E. E_~x, [2 = (rr)). Putting these things together yields the assertion.
((r/d)2k) from Lemma 12. Moreover, it is
is also needed in later stages to give high probability upper bound the ( b)ll
Lemma 24.
(V
[g b) =
sup 6EI-1,1]
holds with high probability over w 1
This can be shown similarly to Lemma 23 we use Lemma 11 instead of Lemma 12 to obtain a high probability bound (note that this bound includes not only s(w, b) but also the main term, and it is for ( b), instead of g b)
Proof.
B.2 Concentration of Empirical Gradient
Now we control the deviation of Ir, (w,, b, {(X', ,",",, "a,"""))))))) ( from g(w, b).
Lemma 25. Under Assumption
[f. (x)x=l1 S (" log Q-1
sup 1_11
(B.8)
(f x)<|1a|1 S / ([[oog) log
Q-1
sup f. 1e1111)
(B.9)
E (x))oo T x)]l S / loge)
sup
(B.10)
f. [_,-
holds with high probability over w Unif() respectively. Here sup,. denotes the supremum over the support of f. whose distribution is specified in Assumption 1.
Proof. We bound the leading term aq aq(b) 9 f- [w o0 ")
in (B.3) of Lemma 19 to show (B.8).
(Q-1)I
From Lemmas 5 and 6, we have
aq (b Ex 4 fz())] ( 1(Q-()))
aq (b)cq (9) ( (w)
(Q - 1)!
[Q 1)!
24

---

ag (g) 8Q (w 1:
0(Q-1))
2
1Q 1)!
ag(b)co
||3| 2Q
W 11:l ((Q-1))
(Q - 1)!
(b)c
W1::> (2Q-1))
(Q - 1)!
Moreover, Lemma 20 tells us that [ao(b)
I C# always holds when b E [-1, and I R.. Therefore, from Lemma 11, we can show
Assumption ensures that lool
that ao(b)E=|V9 1, (a]]( 010 1) 2
CH (Q-1))
)(( [crr log with probability at least
S
(Q-1) 13
1 1 exx((_d/32) 2rd (B.9) and (B.10) can be obtained similarly.
The following lemma is parallel Lemma 19 [DLS222].
Lemma 26. Fix any t [Ti), E [-1,1] and w S satisfying the conditions (B.8) (B.10). Then, with high probability over the distribution of (X y .,"''"",
v²), it holds that
viol aa)aa (2) x)
S
N
N1
v xa)aaaa Eolf x)œ]
r
N1
N1
Σ vlooww Ta) Em (x x)]
1
N
N₂
Here the right-hand sides do not depend on t.
Proof. We only present the proof of the first inequality, as other bounds can be attained similarly. From Corollary 17, [ff())) & where R S (log d)p/2 holds with high probability for all E [Ti]. Conditioned on this,
N1 H( at)xa F. R. ( x)aa)]
N
N
(a Il (o ( -
x)
N
Ex If (x) (x) (2()) Ro (w x)=|||
Note that ft
(w (wT a))a u)
x) u\) is Rub Gaus-
sian for any || = 1 from Lemma 14. Hence the first term can be upper bounded by Ö N. with high probability from Lemma 13. For the second term, we observe that
[B <>)xl]
< E. (= T<x xxx[|]
< Eellf("))""""" Tx) (I 9//44
< [|," [[]{)))))))
||
W o(Vd.d 0/4)
for sufficiently large C; hence this term can be ignored. Finally, so (t) u) is --b Gaus-
sian for any || = and then || 1 ΣN, so T=) = Ö(v d/N,) with high probability. II
By similar procedure, we can obtain bound of [g (w, b, ((X' y¹ x ... ()) g(w, b)]l-
25

---

Lemma 27. Fix any b E [-1,1] and w E S satisfying the conditions (B.8)-(B.10. Then, with high probability over the distribution of {(X), .y", xx, vv') )
d
1,, ((X', , yy, x., y )))))))) - g (w, b)ll =0
dd-1
T
Ti
(W, ((x y', x², ","""''')) )1))))))>)))) g(m. b()=1)))==l=
11--1
r
r²
+
Tj
Ni
holds.
Proof.
In this proof, the term "with high probability" indicates the probability with respect to the distribution of {(X', y' x', yyy) b²)) Let
Tj
z1 (w b, ((x y )))))
Σ
1 x²)
y ( T
1
y ( æ))a
Ti
Ni
and
Ti
N1 y vor T x!)
22(w, b, {(X', yy at, y ッ
y (
T
N1
Note that JT,N, = =1 + 22 holds.
Let us consider z1 (w, b, {(x', y', x², v²) ) Since we assume (B.8)-(B.10) (0) hold, by taking union
bound for Lemma 26 with respect to E [Ti), with high probability |l |, ΣN utol (w xx)aa[] s
R=
(V Ni
+
+ प for each t. Moreover, from Corollary 17, |v'| S R' = Õ(1) for each t.
Here, we notice that the procedure in the proof of Lemma 26 can be applied, yielding
(wo, b, {(x', y², x², 3,, ",,,"))), M (w, b, ((x), y,, a &
=ō RR' VJ/Ti
-ō 19-1 do-r
d
N1
Tj
with high probability. Similarly we can obtain
|| -| = 00((1
+
T.
(1(zı R[21)))=7 0
(((v
19- JQ-T
and
12 ((V
[1(2 - E[zz])a>> = Ō
T5
which concludes the proof.
We need obtain a bound uniformly over b E [-1, We first introduce the following definition.
Definition 28. Fix any w E gd-1 and X = U
U (zi..... I'y x²} c Rd. Then, define a finite
disjoint partition B(w, X) = (B....., BN(w, of 1-11111, iee U,B; = [-1, 1] and B; n B3 @ (i + j) as follows: b and 1/ belong to the same B, if and only if siiig +b) = sign( (w, +b) for all x X.
It is clear that |B( (w, X)I S | < (N₂ +1TT, holds because it suffices to divide 1- 1] at the point satisfying (w, x) + b = 0 for each x E X.
Lemma 29. Fix any
Then, with high probability over the distribution of {('',)) y² x², "e".y")) > the following holds:
gd-1
26

---

g(w,b) b) is L1 Lipschht contimuous with respect to b in the interval [-1,1].
is Lo-Lipschitz con- Let X = 1 ,.y, w²) Then, .T..N; (w, b, {(X', y², æ², "/.,"'))) (1))) 17 tinuous with respect to b in each interval B, E B(w,, X).
Here L1 011 and L2 Ö(Vd) do not depend
Proof. By the explicit form of g(w, b) (B.6), we notice that each term is made by taking the product
between the term as ay (b)az (b) and vector independent of b whose norm is (1). Thus, it suffices to
12 show that ag (b) (b) ((1)-Lipscchiz This follows from the fact dg V2r He, H (see the proof of Lemma 20) and (b)e 62/2 is O(1)- continuous on the bounded set [-1, 11.
Next, let see 9TI,N. (w, b, {(X" y", xt, , 1/))) Recall that (w,, b, {XX, ,y',""" y', x', y )
Ti Σ v'osww x²) :=1
(lo)(wTT))a + Σ Tj y v x')a' 1 Ni
Σ i=1
N1 vloo(w xt).
Σ
N1
From Lemma 10 and Corollary 17, we know that |o! |. 1/ I (log d) P/2 and || ||=*|| S vd holds for all i,t with high probability. Assuming this and noting that is 1-Lipschitz, in each interval in ((w, X), MN,N, (w. b, {(X', , ,y",", ",,"""""" ,')) is L2 ((v))
Corollary 30. With high probability over the distribution of w and {(X' ,',,"' yyy
(0,,"))
sup I b, {XX, ,,,x" x." ( g(w, b)ll = hE[-1,1]
/-0-1 d
d²
do-1
Ti
N1T1
r2
sup 11 (w.,b (XX y 2., ,/))))) g(w, g(w,,b)==11== 6E[-1,1]
(V
+
T1
N1T
(B.11)
holds.
Proof. Let X = U (a)..... a x²) and consider B(w. X) (B..... _ B Then, let /Vd -coverings of B, E B(w, X). Since [B] is
B the union Ô polynomial in d, Lemma 27 holds with high probability uniformly over b B.
d2
M.T.
Moreover, from the construction and Lemma 29, for b [-1, B, there exists (bb) B such that
119T..,. (w, b, {(X', y x" )) -
|19T, .N. (w, (b),, {(X', y', (o',")))/i)) g(w, ग (b))
<ō
V 79-1 Ti
+
N.
holds.
Lemma 31. With high probability over the distribution of (Xx, 1 y', x², wy ?
It holds that
(1g(w, b) -TINNI (XX y', x., "'.'") v'))
sup bE(-1.1]
d
+
(B.12)
do-1
Ti
NITi
with high probability over w Undf
27

---

It holds that
sup b
E
Moo-Unit(3t-)))lll(e (w, ((X', y', x', "',,' +.)))=))
r2-1
r
x²
(B.13)
do-1
V Ti
NiTi
for j E [4P].
Proof. If an event A occurs with probability at least 1 d-C over the simultaneous distribution of (w, (IX',y'," x', >>>))) then we can say that
with probability at least 1 d over {(x', y. x², (²)) Pw [A ((X','' u'','', ,'))=",))) 2
1- d-2 holds.
(B.12) follows from this remark and Corollary 30.
For (B.13), with high probability over the distribution of {(X', y", x.."""')))) g")) Pw Fwf(B.1))
{(X', y" æt ('y'") 2 1- c is satisfied for some Ö. Also, from Lemma 10 and Corollary 17, with high probability over the distribution of {(X', yy x., y )1= we have ||a]| ||ad"|| O(vi) and ly"| = O()). From the definition of 9rs,N, (w, b, {(X), ,',',, yy ) (B.2) and expan- sion of g(w, b) (B.6), sup.. an llg(w, b) 9T, .N. (w, b, {(Xt, y", x² ".yy)) ME)il I sup,, ( b)ll sup llg,,NN (w, b, {(x y" x", " bility
1 Illl is upper bounded by Ö(d). Therefore, with high proba-
E Unw (yllg( b) gT,N, (w0 {(X' y x. " 6E|-1.1]
sup
(
27 11g(w, b) IT.NN (wo, b, ((x' yy ,¹, a +))))))
E uni 1)
sup bE|-1,11 [
((1-0
=ō
rQ-1
22
+ +d00)
dQ-1
Ti
NiT
By setting sufficiently large Ö, we arrive at the claim.
Summary. We combine the obtained results and show that llgT.,,, (w, b, ((X*, y xt, "et."")) ( and the moment of the residual terms are bounded with high probability.
Corollary 32. With high probability over the distribution of {(xt, y² at, y²) ?
(I) Mr,N (w,br b, {(X. y², x' ,3))))))) m( b) + r(w, b, {(x yn x ) holds where
m(w,b) b) = (B) &eg (2)) (2Q - 1)!!
W
2Q-2 2
Wirr
Q!(Q 1)!
E 22
0a-r
and
(
( {(x', ,,))))))))))
1/5
Enw 1) 1
sup
1:r
=Ōar
V
/1Q-1
r +
p2
r
Vd4-IV
Ti
V
NrTi
d2Q+I
28

---

(II) Ifmry < CrdQ-1//R /C1 (log for sufficiently large C₁ and C2. Ti 2 r and NiT1 2
rd2Q+1, then
1 (w, (XX y', x' M 1
sup
6E1-1.11
holds with high probability over w 1 Moreover {Xx, y x', ")')), S log
sup bE(-1.1]
22
holds with high probability over w Unif(Sd-1) and the training data for the second stage
{(X' ,', ,/) where U l x²}.
Proof. Combining Lemmas 23 and 31 yields (I), For (II), by the condition for Ti and N1, we obtain
d
d²
sup 191.N, (w, b, {(x' y' x', "'',"))" /)))2)) S Ōar
1
Ti
NITi
Td2Q-1
be( be(11.1)
Öas
1
rd2Q-1
from Lemmas 24 and 31.
Then, as 71 su,1E(-1,1] [12m 19,, No (w, b, {(X', y" x², .")'') < achieved.
Vrdaq- PuRQQIICii(Ae (Coog d2
By an argument similar to the one in the proof of Lemma 18, we obtain that with high probability, (21 (w, b, {(X', y w', x0",,"")) ) 2) E log for 2 UP (af...., x,,,,²'),, To
obtain the supremum bound over E [-1, 1], suffices to utilize Lemma 29 again.
In the following sections, we assume that these two events (I) and (II) occur.
C Construction of Attention Matrix
We construct attention matrix I with good approximation property: Section 3.2.2 for the proof outline.
In this section, a set of basis functions H is defined as
H
x
Hep, (zy)) S Pi +...... + S P.P. ....,
We let Bp |H] and introduce numbering H {hi,..., (hg Note that
r r (rr)).
Bp
C.1 Approximation Error of Two-layer Neural Network
As sketched in Section 3.2.2, we need to show that there is a set of two-layer neural networks which can approximate basis functions well. The goal of this section is to show the following proposition. Proposition 33. Assume Ti 11 2) and N,T, +1( Let (w be neurons obtained by line 4 of Algorithmm I with
2 . m - (log d) for constant C.: Assume b is re-initialized in Line 5 of Algorithm 1. Then, with high probability over the distribution of training data and the initialization of model parameters, there exist vectors a ...... a Br R such that for each n E [Br].
(w (1) x + (b) - kn(2)) =0 { m
Σ a
sup iE (
N₂
29

---

[mmm)
0 rp + 1
Σ 2 (w)) (1) x' + bj) - nn (1'4)
N₂
= m holds. Moreover, suPne[Br]]|a ||a"|| 0 ((=))
This proposition states that there exists two-layer neural network ΣØ=1 ago ((wo + +bb) which
approximate the basis function h,, and its second-layer magnitude ||a"|2 only scales with Note that the condition of I and N,7i comes from Corollary 32, which implies that the noise terms of the one-step gradient is well controlled.
C.1.1 Alignment S and Efficient Approximation
From now suppose that 9r, .N. (w, b) satisfies the conditions in Corollary 32 (to simplify the notation, drop the dependency on {(X*, y", "'',", , )) An important previous result [DLS22] is that, if gr,,N (w. b) aligns with the r-dimensional subspace S, then it can be used to approximate polynomials on S efficiently.
Lemma 34 (Lemma 21 in [DLS22))) There exists of symmetric r-dimensional lensors (Th (eskS such that
(n())) = Σ
TT @K
OSK<P
and | |F S r -
Definition 35 (f -laainnnmnn We call that gr..N, (w. b) has f(r, alliimmment with S, if the
following holds:
For any S --symmetrie-ddddmnnniinnnn
r-dimensional tensor T.
Ew-Uuit(gst gI,,NN (w, b) b)th)
and the hidden constant in ">" needs to be uniformly lower bounded.
Lemma 36 (Corollary 4 in [DLS22], adapted). Suppose that gr, (w.b) has f(r, dl-aligument with S. Let T be 0 < k < P-symmetric r-dimensional tensor. Then, there exists vrr(w, b) such that
E Uww-Unif(Sa-1) (wt (w. (gT.. w. b). xx) = (T, x 1ir): Em WT w. b) / f(r, dr,
| (w. b) S f(r.d)+
IN. (w, b)]l
Intuitively speaking, f(r, -llalignment ensures that the MLP gradient Ir,,N, (w. b) has suffi- ciently large component on S. If so, the lemma states that, infinite-width neural network Ew (v (grg,, b).,)).)"" whose hidden layer weight is gr,,N, (w, b) and the activation is 2 + zt, can express the functions (T,
@k) with small output layer Ur(w, b).
Let us show that our ITY,N (T,,,(W,bb) b) has sufficient alignment b satisfies (B.5).
Lemma 37 (Tensor expectation lower bound). Let be a k < P-symmetric r-dimensional tensor and = Wirr Wi:r [|2 (2Q-2
Then
Q
E Ew--Uni((aa1)) (TW
B T 20 20 |
d(0+1)
holds for all 0 Sis S k. Here T0w T.
Proof. Let u ~ A(00, Ia), z x(d) and 2.- Therefore
/rrr|| u1 20 2Q-22
un
We can decompose as
E T = Biz
2)
E 1
8k)2
30

---

holds. On the other hand, let æ ~ S"-1 and 2' ~ (r).. Then we can decompose č as ū = (22)222 a. Thus
T 2 [=((2) 2)
B (k
holds. It implies that
T( 8k [[(=')(1Q22)
E (xx
E[ 2)k]
Similarly,
[[(=')(4Q-))) 2)(k
|
[ 9k 1>
E1| T (k=
B[z D]
satisfied. Now from Corollary 13 in [DI.S22],
T ((
|| Er'
holds. Therefore we obtain
E (22)(4Q0-2))
B[T x @k }2]
E| E[z(2Q-214
(( (1Q=)) 2)1
E S(k 1))
e[2(40-22) 2001
E (2 ((Q (10-2)4
16 11 2)(
El| =( 111]]
B[z qQ 2//
[ 2](k
-2)
E (w Sk- ) ||
d322-13
Corollary 38. Let T be k < P-symmetric r-dimensional tensor and let m(w, b) (b) Beg (2)
QU(Q-1)! Benxr [=2Q] |w1:r|12 100.:rl1:0--u..... wrr Moreover, assume that (B.5) holds. Then,
Buw--Un((3e-))(( m(w, bx)))) 1); R [[][((( b) b)(k=>))( [
holds for all Sisk.
Proof. This is obtained by calibrating the coefficients from the previous lemma.
Lemma 39. Let T be a 0 < k S. P-symmetric r-dimensional tensor. Assume that (B.5) holds. Furthermore, suppose S da, Ti = 0(d2+1,0 and NiT; = 22(d22+1, r). Then,
E Unit (T NT.N, (w, 10 2 11122-1j 1k 2k
T
holds for all 0 S k < P. Therefore, IT,,N (w, b) has d-24d-220 -22+1--22222 2 alligmment with S.
Proof. Decompose NN, (w, b) m(w, b) + r(w, b) as defined in Corollary 32. From the proof of Lemma 11 in [DLS22], we can show that
T Jr,NN, (w, b) b){
Bww.Unif(3d-)
>B 2
T. m b))k22
T Unit(8 (0(2)))))) b)
(C.1)
for some 8(w, b) satisfying
Un ((
Emp- www.Uningg-y///[mmmm 2) -)0=) Umw-viirgs m,llrr 6) [1.
31

---

Here, from Corollaries 32 and 38,
E (u ()) b)
1r
(2-1)) E Un- m(
+
dP-1 Tj
NITi
holds. Now, from the condition of N1 and Ti, and S Vd, we know
x²
d220
+
+
S 1.
d Ti
NITi
d2Q+1
Therefore we can achieve (n---nit((3 (6ww, (1) b)a] I Uni (T, m(w, b)(rr) 9. Apply- ing this and Corollary 38 = to (C.1) yields the result.
C.1.2 Approximation Power of Infinite-width Neural Network
No we can show that some "infinite-widie neural network can approximate the functions h H. Lemma 40 (Lemma 9 in [DLS22], adapted). Let a Unif({ 11, 11,1)) and b Unif([ log d. log d)) Then, for any k 2 0 there exists un (a, b) such that for |a|| S log d,
[a,o[vx(a, (ax b)] R² sup lox b) Ōa(1).
a.b
Lemma 41 (Lemma 12 and Corollary 5 in [DLS22]] adapted). There exists on (a, w, by, b2) such that if
San (x)) = En, bi 16 (a, w, bi (2) (w, b1) + b2)]
where
a Unif(+7) Unif(gd-1).bn Unif([-1,,) 1J), Unif([- d, log d)),
it holds that
(fo. (x/) hn(2)))) 1
sup *E|Na],Ti -ISITT
N₂ 1
sup (fon hr, (x'))) N₂
E E, [on w. b1 b2) Ö(r)), sup I0n w. b2) Ô(rP)
b.
for each n E [Br].
Proof. We first construct dhy (a, w. bi, ba) for basis tensors defined in Lemma 34, such that if
fore (x) = Ea, by 1dTz (a, w. bi 62) (2m
(w, b1), x) + b2)],
it holds that
(Ferpy (ar)) - (TT. (at1) )² 5
1
(C.2)
7e(No),7:+111551
N2¹
sup T 1515T
(Fomp (x0) - (TH, (ox (A))
1
(C.3)
N2
(a (w.bs)
LA(w)-B(b,) where us and vhrg are constructed in Lemma 36 and the events A(w) and B(b1) are defined as
Let ΦाØ (a, w. b1,
(2m)
A(w) = sup 2 (w, (e),b1))
2m Kgn., b1), x')1 S log b1E[-1,1]
sup
b(E(1111
32

---

for E [Ne). Ti +1 T3
21/ 1917911,,M (w, b2)ll $1
sup b
B(b1) (0. 1 C# (2 P +
n M(001000 (z))] 2 Cz..
Then,
forg (x)
22k (a, by) b2) (w. b1)
-,,,w,biik by
(z), + bz)
b1), x)
(2m)) 7 4 (w, b1) (2m)
-Rab 2 [B(())) Fe
I Laww) Ear da (a, by) (2 (1, x)) b2)]
where a' ~ Unif((+13). Here, if æ E UVen i......... ..... 2x, x²).
Mapp (w. b1) (2m))
[vx(a, ( (2m a' '9.. (w, b1), x:) + be]]
Ew
=FE ehrr (wo, by NoT, (w, b1), x²) + Eu bhry (w, b1)(1 (9r) (wn, bı). xx)²
Note that the absolute value of second term above can be upper bounded by Eu - [A(w))] 1/2 (wTT (wr, bi) bi) (w, bi), typ] 1/2 Recall that A(w) occurs with high probability from Corollary 32, and Ew (VTr (w, (1 (w, )2] 1/2 is the second term be ((1/ /N2).
polynomial in d. Hence
Moreover, as IP P(((b)) )] 1/2 from Lemma 20, we have
E2, 2 II( Ew Mur (wD, by ()) (w. by).). xx)
=Ba, [2 (T x }] Es, [2(1 IB x 0k)] (TH x @k).
Therefore, (C.2) and (C.3) holds. Noting that 1n vrd20-T /C1 /Cic////looo d) C2 We obtain Eg w. by by lobn (a, w. b1, b1)
4
E. [UA (a, b2) ] by Blb())) v (w, bi b1)
(2m)) j2kk 2k
En lun (a, ((d) 12
(2m 22k22
θ M P-t
and
sup
Φ (a, w. bz, b2)l
a, b,
2
1 (2)) (B(h)
whrp (w,b1)l
2m high
2
(1(20-1),27)
11 (wn, by)l TT) T
(2m))
2k
2ō k P_tk
from Lemmas 36 and 39. To show the original claim, it suffices to set Φn (a, w, b2) D ΦTv (a, w. b2), which yields the upper bounds on the magnitude of on because D p. SP and Ee tky Put S.P.
33

---

C.1.3 Proof of Proposition 33
We discretize Lemma establish Proposition 33.
Proof. [Proof of Proposition 33] First, in the same discretization manner the proof of Lemma 13
in [DLS22)] we can show that there exists a a Br such that for each n. E [Bp],
2
m Lajo ajo ( 1=1
(2 r j.j In (w) (,, xt) + +bg) kn(())
sup
(N
0
m
N₂
2 ( 0 x² tbb) hn kn((z²)
Σ cajo
(2
sup T1 T111S1TT2
1
-
N2
and supne[Ba] ||a||3 = 0 = with high probability ( b, denotes the biases at the first and second initialization, respectively). Recall that wy (1) 2n j. (w) (0) b,0))
2mp E f(X', y'',,",, W p()), b()) w f(X', y' x²; ( r (()). It remains to bound the effect of the second term. By Lemma 18, we know that the norm of second term is o( m Similar to the argument in the proof of Lemma 18, the inner product between the second term and x', x (i E [N2]. Ti + < < T2) is also Ō 0 ( with high probability. From the Cauchy- Schwarz inequality and the Lipschitz continuity of o,
1 (w, Σ dgj ( (2 2 i.
(0) .b,0))) x) + bb)
Σ a 0
(1) [w) ( 20 + +b;))
=ō = = =0 =0(1
P
m
We obtain the assertion.
C.2
Concentration of Correlation
Next, we evaluate the discrepancy between the empirical correlation 1 No 2==1
Σ No
v,(((r,)) and its
expectation for each / E H. which decides the approximation error in (b) of (3.1).
Recall that h E H can be written in the form as
h(x) = II
Hep, (2.) (Q S Pl + + P.P. 2 O......r. 2 2).
Vp.!
Also note that
N₂ ,),( - [xh(w)()))
N₂
N2
N2.
5Th( +
Σ
a. ((2
Ew (1x, ()),))))((x)).
N2 1=1
N₂ j=1
where (2)2) EP
He,(z). First we give an upper bound for
1 N₂ Eye o. ((z,, B))
En (<x,)))((z)))
34

---

Lemma 42. Let D < 2P be the degree of (x( (a)).) when we regard it as an r-variable polynomial with respect to ,1y)..;Ty)) For all 0 there exist absolute constants M and C depending only D such that
P
oz
Eg lo (1,,
21
N2
N/
2/A
<2exp
min ISKSD
CDM2
holds, where CD and M are constants.
Proof.
Let F(N1,...,NN) N₂ Σ a. (1),,))))(z,))) It is degree- D polynomial for Fi. F1 .... NN,.r which are standard Gaussian variables. Also,
[ æN)]lE== VN2 (I.,
satisfied. Then, applying Lemma 15 to F' yields the desired result.
Then, our aim is to bound (o. 3 ((),))))(())]l {[]
Lemma 43. Let h(x)= = II vo: Hcp. (zi) and P + .... For P
|E (a. (1x, B (1))2 S R&C,
holds where Ck a constant depending only on For other k. [1
Proof. First,
Σ c
(He
i!
holds. To calculate W
(He, ((x note that from Lemma 8,
(
@1 +.......a
i!
Hep (I1) Vpn!
Hcp, (Ir) Vpr!
3 B.- Heg (x1) ... He (1r)
(C.4)
21 ......h
qr
9,20
Also, from the discussion in Appendix A.3, Ex 03 8"r He,((x, G3 )((x))) is equal the coeffi- On
de."
cient of (C.4) with respect to Tा...., He,, (x1) (x,) Therefore, B|V (He B))h(=))] +
if and only if + = k. Then, | [H (() =
2
(k- PJ! (He ((x, 3)){(())
F'
9 Here, we can observe that the (fI....,ji jk) entry of the tensor is equal to
pl (()()))(())))
(k-P)
81!.... .......8 sp!
Ck-P (s) - p1)! (sy 1e--Pil pr)!
B:1-Pi 35--Pr
Vpn! I., 20.......-P.2..
Vpi....
where 8; Ek- =i (i [r]). There are 81!-8.1 and thus
index sets (Ji,...,Je jR) having the same entry,
C B (He (1,,3))((()))) (A: P)!
Ad (
81!...8p! 8p
Σ
C-P pz! p s....... sp! (s) pı)! ..... (s, -pr)!
sk
35

---

(89 P
Σ
1
k!
(91 + pi)! (q, + pr)! 2
k- P pi! ..p! (a) + pp)!.... (4r + p.)
91!... qp!
41 41 +4,
P,9,20
(19
33-)2
(k P}! k!
(q) + Pi)!.. (gr + pr)
Σ
k- qı! q, (k- P)!
prlqi pil...pplqi 9p
q: P.g, 20
(89 83.0)2
(C.5)
We can verify that there exists constant C2 depending only on P such that < Cs holds uniformly; for example, one can naively obtain
k! (o (t- p ting
(k-P)
k!
( (41 + p1)1 .................... (gr + +pr)!
S kl(n + P1)! .. (ar + pr) S klk..
(k- P)! p1! 1r ..!!!!
Here we used (01 + P1) + .... + (gr + Pr) = k.
Moreover, Em t......ogg = -g. P.
q.!....!! (8 ... 32. = 113000- P11. holds. Hence (C.5) is
(k-
upper bounded by 2
k- Ox R.C.
Corollary 44. With high probability,
vyh(o,)) Eglor (elo =
2
VN₂
holds.
Proof. Plugging the results above and = O ANz (log d) into Lemma 42 yields
1
c.. - Egg lo. ((xx [x))
x
(log d)
N₂
VN²
with high probability. Moreover, we can easily verify method used in the proof of Lemma 26.
1 No
Σee sy((z;))
by the vNo
C.3 Proof of Proposition 2
We prove the following Proposition 2 using the preparations above.
...,,'y,, ,'y,) and y (vi...., VYN,) T where T₁ + SLST + T2. With high probability over the data distribution and random initialization the parameters, there exists such that
Proposition 2. Let X = (æf,...,
T3P
2PP
1(X',3' 00; ((),rr.b) -T
N₂
holds for all t E [Ti + 1...., T2). where ((1) is obtained by line 4 of Algorithm and b is in-iiiiialized in Line of Algorithm where the conditions on Ti. m, m,7, N,,T, are identical to satisfied.
that Theorem 1. Moreover, ||]]| = O(r2P/m
Proof. It suffices show the proposition for each Ti +1sts + T2, and then take the union bound. We drop the superscript for simplicity.
Note that, from Lemma 8 and E. (m((()))) (x=] = In- for h W E H. we can expand f.( as fx(x) ERe
Br Ex~N(0 If (x)) (x)) Now let A = aa Br
B Bp By where
36

---

{a") J is constructed in Proposition 33. Then, by setting = AAT
ow T
x + b1)
AA (W X b1
y
N2
(ow (1)
bm)
Bp
N2
Σ (2) +e((z))
Elfahnll +
My hn (23,) [8l
+
>
Uj (z,)
N2
N2
j=1
(x)
Bp
1
Σ (xr) + ex()))
B[schn]+ +
My hn (1x) - [[(sam]]]
+
Σ (z))(
N₂
N₂
n=1
j=1
hn (x)
holds, where (x)) j=1 a3 x bby - (n) satisfits [<,( = /
(1)
from Proposition 33.
To bound the error, we note that the following holds for each n:
As (x) is a polynomial whose degree is at most P, and Eg hn)) = O(1), (A., 2 P/2
(log d/2 holds with high probability from Lemma 15.
From lemma 8, [B[ (a)),, (a))] (a=)] = O(1) holds.
From Corollary 17, |y5l S (log d) P/2 holds with high probability.
Hence, for each n E [B]] we obtain
N₂ hn (2z)) -E[f E[falin]
N2
(x:) +
(x))
E B[fahin] + 0(1)
Σ
(x,)
N₂
N₂
0(1)
Öly + No
j=1
j=1
Ōl xt/
0 p + N2
m
hn (x) E[fainn
pp m
P
=ō
+
N₂
m
N₂
Combining this with Bp Pr and m P) we arrive at the upper bound.
Finally for [[]]] [1, we have
Br
Br Σ n=1
|T <
Σ n=1
(a")
|| = -Õ(r2P) /m).
D
Generalization Error Analysis and Proof of Theorem
Note that after one-step gradient descent, (1m 11 Ōa,((1) is ensured with high probability from Lemma 18 and Corollary 32.
37

---

D.1 Rademacher Complexity Bound
Let W = W(1) and suppose bI.... .......,. bm are biases obtained in line of Algorithm 1. Let
((w) (1) ) x + b1)
FN.G = (X. yx)) + To(W X b1
YITT sG)
N
(o((w))) bm)
be the set of transformers where the norm of the attention matrix is constrained, and let
(ad, -(x'),(00
sup JEFN,o
Σ f''f(x'' . y'
be its Rademacher complexity, where e' ~ Unif((#1f). Here contexts X E RddN and y E IR of length N. We can evaluate the Rademacher complexity as follows. Proposition 45.
Gm
Radr(FN,c) = :d.r
VT
holds, when 1byl (1) and | (1) all j E [m].
Proof. First, it holds that
Radr(FN,c
T Σ
f(( y² x²))
sup
JEFN.G
T
t=1
((w) ( b1) : o ((w()) T Tax + bm)
To(W()) x +bb
E
sup ST
T
N
t=1
(1 (1) X b1 1yy ((w((() (") b1)
T
E
sup
|
T
||]
SG
N
(o((m))) T nn + bm)
a ((u (1) Txx + b1)
G
(W X + b1 "
B
T
N
L ((wm (1) T + 3m))
(ww (1)
( Txx + +b1)
T
(W X
b1 bl
E
T
N
o ((w (2)) Txx bm) FF
((u ()) T b1)
(b) G
OWW() + b1
TE x y.m
T
N
[o((wm) (1) T + bm) F
27
G
Ex a
N
( (1) ay + TTa
(2)o(((((()) x + bz)
N
k
==1
38

---

(6)
G
N1o.at)) (V.(ar) ) Ta x + bx (2)(o(()))))) ( + bi)2
k11=1
(d)
G
for (2f.(z)
+ 26t En ((2((w) ( 2')²
Txx') + 266)]
where the annotated equalityinnequaliies are obtained by (a) Cauchy-Schwand inequality, (b)
1r |Il G.111 E. ||A.11.. (c) convexity of f()) 2², and (d)(a(a B))² I
(a + 33) s 20² 2,3
Let us further evaluate the last line: from the assumption and Lemma 7, know that || ||. 10x. and (x)) are all (1): in particular, we easily obtain that E=-~((o.ra) ) T xx)² (1).
Consequently, the last line is G VT\ Exi= E (2(())))) )))
+ (d,r(1).
and remains to evaluate ((,,La)... ( [<xx)) z0)3]
On this term, we obtain
B. [1xx)) (2) ( x²) - Bm, [S. (ar)") ((10g xx4 1/2
=(1 (::
Lemma 7).
Hence we can conclude Rad7 (FN,,(F,,)))) = =0
Gm VT
Corollary 46. Let
((ww ) Txx + bi)
To(W(1)) X bl
>/ I1 so}
FN.G
(X,
L ((wmm)
(2) ) To + bml )
uf (X, w _
Then, Rady = Ōu.
Gm vT
holds when |by| = Öd,((1) and | 11 Õar(1) for all
j E [m].
Proof. Note that
T *[x*, " a0"
E Exygieeeee
sup SEFN.
[=1
1
T
+
1 T
< (x9),))))
sup JEFNNG
é e'f yy,x"")))))) x²)) + +B
T
tm1
This follows from sup/eeFw.e T ET sponding to r = O. The second term in the last line can be bounded as
e( y x²) 2 0, as FN.G contains zero function corre-
T
(
T
27
y
S
Σ
c'y'
1=1
0
1
El ((())))
VT
VT
by Lemma
39

---

D.2 Prompt Length-free Generalization Bound
The ICL risk for prompt length N was defined as
RN( Ex gf... MI.N. æ; W.T.b) -yl.
where the length of XINN and fI.N is fixed to N. In this section, we show that the ICL risk "converges" when the context length is sufficiently large.
Proposition 47. Assume that ||wy|| = (1,r(1) and byر = (1) each j E [m]. Then,
[Rx(F) Rs (f)) (1
holds.
Proof.
Note that
[Rx(f) - Ras(f)
<B[IS(X I.N, DI.N, a; WI J(XIM I:M, MI.M 2; W. r, b)l]
a [w)) + b1)
OW IN + b1 DINN
xW + b1 b Mui
=E
N
M
o w T x + 6m)
o ( x + b1)
(e)
0 (W X IN + 1 N
(W X I.M + b1 M M
F|
M
0 (w,, x + bm
27 1/2
[o X
61 b1))yn
( X1:M5 + b1 b1 131.M
<|
N
M
1/2
( b1)
E
5 ( T x + bm)
where (a) holds because a Tb = oo (ab S ||||]||bb ||5 |||||||||||||||
Let
us
bound
ooW t o(WX )an:
[27
which is to
ET
Ex.N
(+ E. oww x. by)y M En oww] xx b,)
'I] For simplic-
ity, we drop the subscript j and obtain
27
1
Ex y
ow a; b)ye
ow w + b)y/
N
M
1 N
<2B
5 (ww + b)y
-Elo(w x + b]yb
N
i=1
27
+2E
aww x; + +byy - -E[ow= + b)u]
M
1=1
V
V[(o
T
x + b)y]]
+ M
2
(++b)y)]
2
2
+
Blo(w
x b) (z) + +)
N
M
2
2
+
x)² (z)) + +x33]] (2
N
M
40

---

Here (a) is obtained in the same vein as the derivation of (d) in the proof of Proposition 45. Moreover,
in the same way as that proof, we can show that (2 T x)² + (2f))2f.(2)) + +252)] = (1,,
Consequently we get IE
<wW +b1
o( X.MM IM +b1 L. Jgriar
||]
( + 2/ Mr aa, (m))
Next, we observe that
o (w) x + b1)
E
S 2 Σ
[[(w] xx² + b3] = (m))
o (w T m x bm)]
==1
Putting all things together, we arrive at
- - +1M
D.3 Proof of Theorem
Finally we are ready to prove our main theorem.
Proof. [Proof of Theorem 1] Let î be the attention matrix constructed in Proposition 2 and let P be the minimizer of the ridge regression problem (line in Algorithm 1). By the equivalence between optimization with L2 regularization and norm-constrained optimization, there exists A2 such that
T |r < ||] = ((122)())))) /m),
1 TT
Σ - y'= y"; xt; w r bi]
T2 1=T++1
Ti+T2
Σ
(y) - -f=(xx y x² ( r b))²
T2 (=T++1
T1+T2
(y/ - f(x' x'''""''" ( , T.)))
T2 t=Ti+1
Then, from Proposition 2,
1
Ti+T
Σ T1+1
1v' - f(x' y x (1)) r", - -
73P
2P
T3
N₂
holds.
We evaluate RN, (f) T where f f(xt, y x²: (1)) r b) using the Rademacher complexity.
(f)
Ti Ti-T2
Σ
10d - f(x y x (L),
T2
T1+1
1 T++T:
ly' f(x' ((XX''''a x w()) w(,)rr,,b)) .b))
RN,(F)
Σ
Te T=
P3P
72P
+
N₂
1 TitT
sup JEËx,IFI
RN,(f)
Σ
1v - f(x yd x w
T2
t=T+1
41

---

holds, where F is defined in Corollary 46. We can evaluate the expectation value of the second term of the last line as, using et l.d.
Unif(#1}.
T1+T.
RN)
1
Σ
1 - y2 x', 2/))
sup SEFN.IRI
T2 1=1++
(0) 2E 20X0)00000000000000
1 T1-12
cl 5(xx, y² 1111
sup xEFx,IFIF
T2
[=T1+1
(6)
T: +T=
2 Radı, (F (FN..1111o)) Eyie
Te T+++
27 1/2
Ti+Is
< Radys (Fo (FN.IIT(ie) + y.e
Σ T2 T=
d'y2
T4P
< Radr, (F (NIT)T +
E [v
VT2
T2
Here, (a) holds from the standard symmetrization argument (see eq. (4.17) in [Wai19] for example), (b) is derived using eq. (1) in [Mau16], noting that (Ir |z - y| is /2-Lipschitz continuous, and (c) follows from the bound of ||T]| in Proposition 2, together with Lemma 7 and Proposition 45. Note that
sup JEFx.ITIF
RN)
1 T2
Σ +
ly - f(x y x 3')l
(D.1)
nonnegative because f(X, y. v. y) is included in F NN. Fll Hence from Markov's inequality,
(D.1) is Ō
paP
with probability at least 0.995. Hence we obtain
T2
p2P
Rx
m
N2
T2
Now we have done the upper bound for RN2 (f). For Ry- (/), we can use Proposition 47. Note that all the desired events occur with high probability, except for the one described above, which occurs with probability 0.995. This completes the proof of Theorem 1.
E Derivation of Simplified Self-attention Module
We derive equation (2.6), following the same line of argument as [ZFB23]. As we assumed that e
W PV WKQ are in the form as (2.4), the output of the attention layer is written as
K
E
E
0
01xm
fAUD (E; PV w
E
E
(E.1)
01xm
V
N
Note that we adopt the right-bottom entry (Fxem (E; PV wKO output of the query. By (E.1), Hence we have
as a prediction for
m+1,N+1
(JAun WPV WKQ))
m+1,N+1
42

---

o (ww) T b1)
E
K
*
01xm
*
o(w mn x + bm)
0
=
+01xm
E
N
o [w) T x + d1)
K
:
E
( w T x + bm)
m
-0
UN
0]
N
a (w x + b1)
o(W + b1
K
T ( m 2 + bm
yNl
N
ow x + b1)
WX + b1
N
o ow T + bm)
Setting vk'
yields equation (2.6).
Details of Experiments
F1 Detailed Experimental Settings
For the experiments in Section 4, test loss was averaged over 128 independent tasks, with 2048 and 128 independent queries generated for cach task in the experiment of Figure 1 and 2, respectively. During the validation of the experiment of Figure 1, the coefficients {ci in the single-index model were fixed to be (c2, ca) 2 21/2, 2 31/2) to reduce the variance in the baseline methods.
GPT-2 model. The experiments in Section were based on the architecture used in [GTLV22]] whose backbone is the GPT-2 model [RWC+ 19]. Given the prompt ((x, vt)) N+1 the embedding E E (dx(2N+2) was constructed as E [1,,,1..... 9I...., NN, ÿNJ where ÿ: [Mr, 0, .,.... OOT. This embedding was mapped to E E 12 (2N through the read-in layer. The transformed embedding É was then passed through the 12-layer GPT-2 backbone with 8 heads, with dropout disabled. Finally, the output was converted into 2N + 2-dimensional vector through the read-out layer, where the (2k- - 1) entry (k [N + 1]) corresponds to the prediction (= ví ..... x 1 ax)) for ys, that is, the answer for the query æt reading the context of length h: - 1. We used the Adam optimizer [KB15] with a learning rate of 0.0001. The training loss at each step was set as B ΣE EN+
1
N+1 (v)- Úu( ... ./ where B = 8 is the minibatch size.
Baseline algorithms. Baseline algorithms used in the comparison experiment are as follows:
Two-layer neural network f(x) = Σラ a,o(w] x) where a = ReLU and m = 256. The initialization scale followed the mean-field regime, with a; i.dd Unif([#1/m}) and wy iid
~
11). The architecture was trained using the Adam optimizer, with a learning rate of 0.1 for the first layer and 0.1/m /mm for the second layer, along with weight dacay rate of 0.0001. Early stopping was applied to prevent overfitting.
Kernel ridge regression using the Gaussian RBF kernel. The ridge regression parameter was set to À 0.01.
Again, note that these baseline algorithms were trained from scratch for each validation task, whereas the parameters of the pretrained GPT-2 model remained unchanged across all validation tasks.
'https://githnn httpcchhww..o.ttst/teeectttttttt http:tttttipii-cootttttt-- cot.--oooottttttnnnnggg. learningg MIT License.
43

---

F.2 Experiment on Our Architecture and Algorithm
In addition to the GPT-2 experiments, we also con- ducted comparison between our Algorithm on our simplified architecture and baseline algorithms. Data
100
generated as yf = H ()) where B L.d. 3n E Rd 3 1 0]
prediction risk
: |||3]| = 1} with d 32 and = 2. We pretrained our transformer ar- chitecture (2.6) with m 8, 000, using Algorithm with Ti 10 N1
10 771 10 (recall that
10
(egdd-Cyy C. should be large) and T- 3, 000, N₂ = 512, A2 0.1. We used mini-batch SGD to find approximate optimizer I' of the ridge regression. For validation, we altered the context length from 2 2 and compared its test error with one- step GD (training the first layer via one-step gradient descent, and then conducting ridge regression on the sec- ond layer) on two-layer NN of width 8, 000 and kernel ridge regression with RBF kernel working directly on the test prompt. In Figure 3 we see that our simplified transformer also outperforms neural network and kernel method, especially in short context regime.
? m rd220- 1 (log d)
keme!
two-layer
pretrained TF
10 in-context sample size N'
Figure In-context generalization error kernel ridge regression, neural network + one- step gradient descent, and Algorithm on our Transformer (2.6).
44

---

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:
They are presented accurately.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in 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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We mentioned the limitation of our result and the future challenge in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in 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 the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
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]
45

---

We presented full-proofs for our results in the appendix, and sketched the proof in Section 3.2.
Justification:
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 or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
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 of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: See Appendix F.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer 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 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 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 it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) If the contribution is new model (e.g., a large language model), then there should either be 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 deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to 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?
46

---

Answer:
[No]
This is a theory paper where the toy experiments are not central to our contribution.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. ec/
cesssonoolion iiyy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. cc/ httpc//ggiieee/oodaooe
www.ssssiiiicyyyy for more details.
The authors should provide instructions 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. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: Yes]
Justification: See Section 4 and Appendix F.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
See Section 4 and Appendix F.
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).
47

---

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: [No]
Justification: We only have toy experiments and is not a central part of this paper.
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 NeurlPS Code of Ethics https http://neurppppp http:///ww...tttti.cc.ieeeeeeee.
didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the
Answer: [Yes]
Justification: Yes, and we checked the code.
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: Our results are purely theoretical working on a small simplified model, thus there seems no 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.
48

---

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 negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We don't have any models or datasets to release.
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 not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses 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: [Yes]
Justification:
We included the information in Appendix F.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or 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 (c.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 package should be provided. For popular datasets, paperswi thcode. com com/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.
49

---

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 don't release any assets.
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: Our paper does not include such subjects.
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 include such subjects.
Guidelines:
The answer NA means that the paper does not 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.
50


Reviews:
[{"review_id": "cD3hjVgiyD", "reviewer": "Reviewer_TSPc", "summary": "This work shows that a single layer transformer can learn low dimensional structure in pretraining, and use it to learn a single index function in-context with a number of samples that scales in the size of the low-dimensional structure. Algorithm: The MLP layer is trained for one step of GD and the attention weights are trained for the rest of the time.", "questions": "There are three interesting things in this model: \n1. The non-linearity\n2. The single index function\n3. The low rank structure\nBut I am not sure how they individually play a role. The low rank structure I think serves the purpose of making the context length smaller, and isnt necessary, but what about the others?\n\nCan you include some examples of function classes satisfying Assumption 1 with low P? For instance, what if my ground truth $\\sigma_*$ is ReLU?\n\nCan this be empirically verified? An experiment that shows that w aligns with S under true GD would be nice.\n\nIs there some easy interpretation of Gamma? Without the MLP, for linear functions, it has been shown to be the inverse covariance of the data. With the MLP it appears to play a similar role.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper shows that low dimensional structure can be leveraged in pre-training. It is generally well written.", "weaknesses": "1. The MLP is before the attention layer instead of after. (My initial impression was that the attention would do something like a projection onto the low dimensional subspace and the MLP would learn the non-linearity somehow). Can the authors say something about this? \n2. There is no soft-max. \n3. This is not true GD/SGD (and I think this is not obvious from the introduction). \n    1. The MLP is only trained for one iteration.\n    2.  The function is initialized to compute 0, and following that the $w_i$ are exactly “erased” by the regularization.\n    3. The biases are then reset to be uniform on a unit interval.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback. We address the technical concerns below. \n\n---\n\n**Nonstandard architecture and training procedure**\n\nWe agree that Algorithm 1 deviates from the most standard training procedure in practice. We make the following remarks on our architecture and optimization algorithm. \n\n- **MLP before attention.** This simplified theoretical setting reflects the scenario where lower layers of the trained transformer act as data representations on top of which the upper layers perform ICL.\nFor example, see Figure 1 in [[Guo et al. 2023]](https://openreview.net/pdf?id=ikwEDva1JZ): they considered a Transformer architecture with alternating MLP and attention layers and suggested that the earlier layers grasp the features of the true function, while the subsequent layers act as linear attention.\nOur architecture, with an MLP layer for feature learning followed by a linear attention module, idealizes such process occurring in more complicated Transformer.\nThe same architecture has been analyzed in prior works on in-context learning of nonlinear features such as [[Kim and Suzuki 2024]](https://proceedings.mlr.press/v235/kim24af.html). \n\n- **Linear attention.** In the ICL theory literature, the Softmax attention is often simplified to linear attention to facilitate the optimization analysis. Such simplification can capture several prominent aspects of practical transformer training dynamics -- see [[Ahn et al. 2023]](https://openreview.net/pdf?id=0uI5415ry7) and references therein. Also, in our setting we already observe interesting in-context behavior using a linear attention. Extending our analysis to Softmax attention is an interesting future direction. \n\n- **Layer-wise training and one-step GD.** \nIn the neural network theory literature, layer-wise training and reinitialized bias units are fairly standard algorithmic modifications when proving end-to-end sample complexity guarantees, as seen in many prior works on feature learning [[Damian et al. 2022]](https://proceedings.mlr.press/v178/damian22a/damian22a.pdf) [[Abbe et al. 2023]](https://proceedings.mlr.press/v195/abbe23a/abbe23a.pdf). Without such modifications, stronger target assumptions are generally needed, such as matching link function [[Ben Arous et al. 2020]](https://jmlr.csail.mit.edu/papers/volume22/20-1288/20-1288.pdf).  \nIn addition, training the representation for one gradient step (followed by SGD on the top layer) is a well-studied optimization procedure for two-layer neural networks (with the hope being that the \"early phase\" of representation learning already extracts useful features) -- see [[Ba et al. 2022]](https://openreview.net/pdf?id=akddwRG6EGi) [[Barak et al. 2022]](https://openreview.net/pdf?id=8XWP2ewX-im); but to our knowledge, we are the first to apply this dynamics to nonlinear transformers in the ICL setting and establish rigorous statistical guarantees.  \n\n---\n\n**Comments on the problem setting** \n\nThe single-index setting is arguable the simplest nonlinear extension of linear regression in which a nonlinear link function is applied to the response. Moreover, for this function class, the statistical efficiency of simple algorithms (that can be implemented on the in-context examples) has been extensively studied. We therefore believe that it is a natural first step for the theoretical understanding of nonlinear ICL.  \nAs for the low-rank structure, our motivation is to introduce \"prior\" information of the function class which cannot be exploited by algorithms that directly learn from the test prompt, but can be adapted to during pretraining (see paragraph starting line 40 and 63). This enables ICL to outperform baseline algorithms such as linear/kernel regression. \n\n---\n\n**\"Can you include some examples of function classes satisfying Assumption 1?\"**\n\nOur Assumption 1 requires the link function to be a bounded-degree polynomial, such as a Hermite polynomial of degree $P$. This assumption is analogous to those appearing in the feature learning literature, see [[Abbe et al. 2023]](https://proceedings.mlr.press/v195/abbe23a/abbe23a.pdf) and [[Ba et al. 2023]](https://openreview.net/pdf?id=HlIAoCHDWW). \n\n---\n\n**\"An experiment that shows that w aligns with S under true GD would be nice\"**\n\nFollowing your suggestion, we have included numerical experiments demonstrating the alignment phenomenon in SGD with multiple steps. Please refer to Figure 2 in the pdf document in our general response. \n\n---\n\n**\"Is there some easy interpretation of Gamma?\"** \n\nIntuitively speaking, $\\Gamma$ plays the role of approximating (nonlinear) basis functions that express the true function. As shown in Section 4.2, $\\Gamma$ is decomposed as $ADA$, where $A$ is analogous to the second-layer of a two-layer MLP which can approximate basis functions $\\{g_i\\}$. This is one crucial difference from the linear case, where the attention matrix only needs to learn the data covariance.  \nWe will provide more explanations on this in the revision. \n\n---\n\nWe would be happy to clarify any further concerns/questions in the discussion period.", "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": "DVyVRsXwRx", "reviewer": "Reviewer_3HLy", "summary": "* The authors theoretically study the capacity for (nonlinear) transformer\n  models to perform in-context learning with a low-complexity function class.\n* The set-up involves pretraining transformers on sequences encoding\n  input-output examples from a reduced-rank Gaussian single-index model.\n  * The use of Gaussian single index models gives the transformer a chance to\n    implement a more intricate (but still analytically tractable) in-context\n    learning algorithm than in-context linear regression for linear data.\n  * The restriction of the in-context target functions to a low-rank subspace\n    gives the transformer an opportunity to internalise (during pretraining)\n    a 'prior' that it can rely on during in-context learning for a sample\n    efficiency advantage.\n* Within this setting the authors prove their main theorem, characterising\n  the in-context sample efficiency of a transformer trained under certain\n  conditions.\n  * Importantly, the efficiency scales with the dimension of the\n    low-dimensional model subspace rather than the ambient dimension.\n  * This suggests that the pretrained transformer model is capable not only\n    of performing ICL of Gaussian single-index models, but also of encoding\n    the 'prior' and gaining a corresponding sample efficiency advantage from\n    doing so.", "questions": "To begin with I will flag that before reading this paper I was not familiar\nwith Gaussian single-index models. Based on the information in the paper and\nmy own light additional reading, I think I have assembled enough\nunderstanding to ask sensible questions, but I could be mistaken, and I\napologise in advance if my questions are ill-formed.\n\n0. To this point, do the authors know of an accessible reference work that covers the topic of Gaussian single-index models? In case other potential readers are, like myself, not familiar with the topic, it could be useful to include a reference to a statistics textbook or similar.\n\n1. A basic question: I don't think I understand line 52: \"In single-index\n   model, outputs only depend on the direction of $\\beta_t$ in the\n   $d$-dimensional input space.\" Is this just foreshadowing the constraint\n   that $\\beta_t$ is restricted to the unit sphere in the $r$-dimensional\n   subspace? (Perhaps the constraint on the magnitude of $\\beta$ should be\n   mentioned in the introduction.)\n\n2. I understand that the first goal addressed by this work is to create a\n   setting where transformers exhibit nonlinear in-context learning behaviour\n   (and I consider this a strength of the work). However, the shift from\n   regression to single-index models brings with it additional statistical\n   and architectural complexity. Therefore, with respect to your second\n   stated goal of demonstrating adaptivity of ICL to low-dimensional\n   structure in the training function class, I am wondering if it would be\n   possible to get away with a simpler demonstration:\n\n    1. In particular, what would happen if one considered the problem setting\n       of in-context linear regression where the regression tasks are confined\n       to a low-dimensional subspace (or a unit sphere on such a subspace)? Do\n       you think that such a setting would offer a similar chance for the\n       transformer to show its ability to internalise informative 'prior'\n       information, but without the added complexity of considering\n       single-index models?\n\n    2. If I am not mistaken, this simpler low-rank regression setting is\n       *not* a special case of your setting, since it would require the link\n       function to be always the identity function which violates your\n       assumptions about the coefficients in the Hermite expansion. Is that\n       correct?\n\n    3. Nevertheless, would you conjecture that a similar result to your main\n       theorem---a generalisation bound that scales with the dimension of the\n       low-dimensional subspace rather than that of the ambient\n       dimension---would hold in such a setting?\n\nFinally, some questions about future work.\n\n3. Your theoretical work does not offer an empirical component. However, your\n   analysis appears to make some clear predictions about what a transformer\n   trained in this setting would learn, in terms of its ICL performance for\n   this kind of data and even the internal mechanisms with which it might\n   learn to perform ICL (such as using the MLP layers to encode the\n   low-dimensional subspace).\n   \n    1. Do you think these performance predictions would be borne out by\n       empirical experiments to this effect?\n\n    2. Do you think that trained transformers would implement ICL in this\n       setting using the mechanisms you have proposed?\n\n    3. Are there any differences between your theoretical training algorithm\n       and modern transformer training procedures that you expect might cause\n       these predictions not to hold?", "limitations": "The authors have adequately addressed the limitations of their work. Assumptions are clearly stated and documented. My only concern in this regard is covered in the \"Weaknesses\" section of this review.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "I thank the authors for their contributions in this important area of\nresearch. In my opinion the work shows at least the following strengths.\n\n1. The work is well motivated by an important research goal within the field:\n  advancing our theoretical understanding of more complex forms of ICL\n  exhibited by transformers. This helps to bridge the gap between the\n  nuanced ICL capabilities of LLMs and the simpler abilities of in-context\n  linear regression transformers, for example.\n\n2. The setting proposed by the authors, of learning Gaussian single-index\n  models, appears to provide a neat example of a non-linear in-context\n  learning capability making non-trivial use of a transformer's MLP layers,\n  while still being analytically tractable and therefore amenable to\n  informative comparison with baselines from the statistics and ML\n  literature, as the authors have demonstrated.\n\n3. The main theoretical result by the paper clearly demonstrates that\n  pretrained transformers are capable of learning ICL for this nonlinear\n  function class and that they are able to take advantage of the special\n  low-dimensional structure of the function class in-context to improve their\n  sample complexity over an ICL method that considers arbitrary directions.\n\n4. The paper is reasonably well-organised and the writing is clear. Though for\n  me the technical content is outside of my background and  zone, with\n  careful study I was able to understand the setting and theorem statement\n  thanks to the detailed technical description and accompanying discussion\n  provided by the authors.", "weaknesses": "I lack the background to evaluate the soundness of the main theorem, but the\nhigh-level structure of the main theorem statement seems plausible to me and\nso I am open to taking the authors at their word.\n\nInstead, the concerns I have with the paper are about the framing of the\ncontributions. In particular, my concern focuses on the second research\nquestion, \"Can a pretrained transformer adapt to certain structures of the\ntarget function class, and how does such adaptivity contribute to the\nstatistical efficiency of ICL?\" (Lines 45/46), to which the authors answer,\nin summary, that, yes, there is a 'separation' (improved sample complexity)\nbetween these transformers and \"algorithms that directly learn the\nsingle-index function from the test prompt\".\n\nMy concerns with this framing are as follows:\n\n1. **These baselines are potentially misleading.** The authors compare the\n   sample complexity of a transformer model that has access to more\n   information (through pretraining) against that of a baseline that lacks\n   this information. It is not surprising that prior information can reduce\n   sample complexity and this is why the transformer's performance beats this\n   baseline.\n     \n     I think the appropriate framing would be more along the lines of saying\n     that the transformer is able to (in some sense) 'match' other algorithms\n     that have this same access to prior information, if these baselines are\n     available in the literature.\n\n     Note: I grant that the authors have shown that transformers can take\n     advantage of this additional information, and this is a meaningful\n     contribution. The specific concern is with the misleading comparison\n     that might suggest transformers have some radical sample efficiency\n     properties, when the improvement over these baselines has clearly come\n     from access to additional information that the baselines don't have.\n\n2. **This does not appear to be a completely novel observation.** There are at least\n   empirical examples of transformers showing improved sample efficiency\n   through making use of prior information gleaned from pretraining. For\n   example, this is implied by the task diversity experiments of Raventos et\n   al. (cited by the authors). In the low task-diversity regime, the model\n   appears to memorise a short list of task vectors and when asked to perform\n   ICL of these particular tasks it would obviously outperform a ridge\n   regression model entertaining a larger hypothesis class. Therefore it\n   appears that in some related work this second question has already been\n   addressed.\n\n     *Note:* These empirical observations are different from the present\n     work's theoretical characterisation of the sample efficiency advantage\n     of the pretrained transformer and I think this does not fundamentally\n     detract from the authors' claims to novelty. However, I think it would\n     be appropriate to introduce this question in the context of these\n     previous attempts to answer it.\n\nI would invite the authors to consider changing the framing of their paper to\nmore clearly qualify the nature of their contribution with respect to this\nsecond research question.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.\n\n---\n\n**Potentially misleading baseline and prior works**\n\nThank you for the insightful comments. Indeed there are empirical examples of transformers achieving better statistical efficiency by making use of prior information of the function class -- we do not claim that we are introducing a new phenomenon. We will discuss [Raventos et al. 2023] in the Introduction and clarify the overlap in the high-level message. Our motivation is to provide rigorous end-to-end learning guarantee that captures this adaptivity in an idealized theoretical setting.  \nWe also agree that the comparison of sample complexity (pretrained transformers vs. algorithms on in-context test examples) is not straightforward, since transformers benefit from additional data in the pretraining phase. From the theoretical point of view, our goal is to highlight the statistical benefit of this prior information, since many previous theoretical works on ICL proved that transformers implement a simple \"prior-agnostic\" algorithm, such as one gradient descent step on the least squares objective for linear regression.  \n\n---\n\n**On Gaussian single-index models (Questions 0, 1, and 2)**\n\n* A single-index model is defined by an unknown univariate\n\"link\" function applied to an unknown one-dimensional projection of the input. Efficiently learning this function requires some extent of *adaptivity* to the low-dimensional structure, and hence this function class has been extensively studied in the neural network theory literature, to demonstrate the benefit of representation learning.  Here we list some works covering this topic: [[Ba et al. 2022]](https://openreview.net/pdf?id=akddwRG6EGi) [[Bietti et al. 2022]](https://openreview.net/pdf?id=wt7cd9m2cz2) [[Damian et al. 2023]](https://openreview.net/pdf?id=73XPopmbXH). We however realize that this problem setting has not been introduced in the context of transformer theory. \nWe will revise the Introduction and Problem Setting section to provide more background and intuition. \n\n* Regarding the question 1, there we aimed to simply explain that in single-index models, the function value of $f_*$ changes only along the direction parallel to $\\beta$ in the input space $\\mathbb{R}^d$, because $f_*$ only depends on the inner product $\\langle\\beta,x\\rangle$.\nWe will make that explanation clearer.\n\n* We select this problem setting because it simultaneously meets two desiderata: $(i)$ a nonlinear function class, and $(ii)$ adaptivity to target \"prior\". On one hand, single-index model is arguably the simplest nonlinear extension of linear regression in which a nonlinear link function is applied to the response. On the other, this function class allows us to specify low-dimensional structure which the learning algorithm should adapt to. Moreover, for this function class, the statistical efficiency of simple algorithms (that can be implemented on the in-context examples) has been extensively studied. \n\n* If we are not concerned with point $(i)$ above, it is possible that a linear function class (with low-rank prior) can exhibit similar phenomenology, where a transformer with linear MLP layer can adapt to the class prior.\nMore generally, we indeed conjecture a generalization bound that does not depend on the ambient dimensionality of the problem, but some notion of \"intrinsic dimensionality\" of the function class. \n\n---\n\n**Empirical evidence and future work (Question 3)**\n\n* We conducted experiments that supports our theoretical claim on the statistical efficiency of ICL -- please see the .pdf document attached in our general response. \nSpecifically, in Figure 1 we indeed observe a separation in the test loss between ICL and \"prior-agnostic\" algorithms such as kernel ridge regression, as the theory predicts. \n\n* We believe that the mechanism (employed in our theoretical analysis) of *feature extraction \\& adaptivity to low-dimensionality* by MLP layer + *in-context function approximation* by attention layer could also appear in modern transformer training. However, since our statistical guarantees are derived for a particular architecture and gradient-based optimization procedure, it is likely that the quantitative scalings (e.g., dependence on context length, task diversity) have limited predictive power in realistic settings. \n\n---\n\nWe would be happy to clarify any further concerns/questions in the discussion period.", "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": "vwpc1moPxx", "reviewer": "Reviewer_gjHz", "summary": "This work studies in-context learning on nonlinear functions, which are Hermite polynomials, using Transformers with linear attention and nonlinear MLP. When the nonlinear target function depends on a low-rank dimension, the ICL risk trained by gradient descent is proven to be a function of the rank. The proof supports the theory. I am willing to update my review after the rebuttal.", "questions": "1. I think Algorithm 1 is proposed to handle the non-convexity of the problem, right?", "limitations": "There is no potential negative societal impact of this work.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. ICL on nonlinear functions is a significant problem to solve. \n2. The proof is solid and reasonable to me.\n3. Characterizing the ICL performance as a function of the internal complexity of the nonlinearity of the target function is a novel and nice result.", "weaknesses": "1. The theoretical result seems not to be tight. When $r$ is very close to $d$, the sample complexity is much worse than using two-layer neural networks. However, the Transformer considered in this paper already includes a nonlinear and trainable MLP layer. Then, I think the derived bound is not tight. In other words, why is the bound worse than two-layer NN?\n\n2. No experiments are provided to justify the result, especially for the comparison with kernel methods and two-layer NN.\n\n3. The presentation of the proof sketch can be improved. It is better to emphasize that finding $\\Gamma$ is solving a convex problem (mentioned in line 206) between lines 262 and 264. I think the convexity is the key to the reason why the ICL risk of the constructed $\\Gamma$ is attainable by gradient descent. \n\n4. Some references on ICL optimization and generalization are missing. Li et. al., ICML 2024. \"How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?\" (This work considers softmax attention and nonlinear MLP, which could be mentioned in line 93); Gatmiry et al., ICML 2024. \"Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?\" \nThere is also a concurrent work: Yang et al., 2024. \"In-Context Learning with Representations: Contextual Generalization of Trained Transformers\", https://openreview.net/pdf?id=fShWHkLX3o It will be interesting to compare this concurrent work, although it is not required.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful for your helpful comments. We address the concerns and answer the questions below. \n\n---\n\n**\"The theoretical result seems not to be tight.\"**\n\nNote that although we have an MLP block in our architecture, it is not trainable during the inference time. Indeed a sample complexity of $O(d^{\\Theta(k)})$ for MLPs can be attained when we train the MLP parameters in response to each task (e.g., see [[Bietti et al. 2022]](https://openreview.net/pdf?id=wt7cd9m2cz2)). \nHowever, when considering ICL via Transformer, we need to fix the MLP layer and use the same parameters for all tasks, where the true function and direction $\\beta$ can vary.Thus it is not surprising that our sample complexity is worse than NN+GD (trained on the test prompt) in the $r\\simeq d$ regime. \nInstead, our focus is the regime where $r\\ll d$, for which we show that pretrained transformer can adapt to this low-dimensional structure (target prior) and outperform baseline algorithms that only have access to the in-context examples.\n\n---\n\n**\"No experiments are provided to justify the result.\"**\n\nFollowing your suggestion, we conducted experiments that supports our theoretical claim on the statistical efficiency of ICL -- please see the .pdf document attached in our general response. Specifically, in Figure 1 we observe that when $r\\ll d$, there is indeed a separation in the test loss between pretrained transformers and baseline algorithms that directly learn from the test prompt, including kernel ridge regression and two-layer neural network + gradient descent. \n\n---\n\n\n**\"The presentation of the proof sketch can be improved.\"**\n\nThank you for the suggestion. As you pointed out, the convexity of the ridge regression problem for $\\Gamma$ plays an important role because the optimal solution $\\Gamma^*$ can be efficiently solved; this allows us to state that the approximation and generalization error of $\\Gamma^*$ is not comparable to that of the constructed $\\Gamma$. \nWe will clarify this point in the proof sketch section.\n\n---\n\n**\"references on ICL optimization and generalization are missing\"**\n\nThank you for your suggestion.  We will add citations and comparison with these concurrent works. \n\n---\n\n**\"Algorithm 1 is proposed to handle the non-convexity of the problem?\"**\n\nYes, this simplification using one-step gradient descent circumvents the difficulty of analyzing the entire optimization trajectory with non-convex landscape. Such a strategy is commonly-used in the neural network theory literature -- training the representation for one gradient step (followed by SGD on the top layer) is a well-studied optimization procedure for two-layer neural networks (with the hope being that the \"early phase\" of representation learning already extracts useful features) -- see [[Ba et al. 2022]](https://openreview.net/pdf?id=akddwRG6EGi) [[Barak et al. 2022]](https://openreview.net/pdf?id=8XWP2ewX-im); but to our knowledge, we are the first to apply this dynamics to nonlinear transformers in the ICL setting and establish rigorous statistical guarantees.  \n\n---\nWe would be happy to clarify any further concerns/questions in the discussion period.", "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": "o4vpUzzKW8", "reviewer": "Reviewer_UaEF", "summary": "This paper investigates the ICL abilities of transformers for learning nonlinear function classes. The authors study transformers with nonlinear MLP layers and demonstrate that such transformers can efficiently learn single-index target functions in-context with a prompt length dependent on the function's low-dimensional structure. The proposed method could significantly reduce the statistical complexity compared to direct learning methods.", "questions": "Please see part 3 in the last section.", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The topic of this work is both interesting and important.\n\n2. The analysis provided is rigorous and sound. I think the low-rank result is very encouraging that the risk bound depend on r not d, which could be very helpful for explaining many success of LLMs.\n\n3. The presentation is clear and easy to read.", "weaknesses": "1. Although the ICL capabilities of transformers are indeed powerful and of great interest, the model and pre-training algorithm presented in this paper are not very realistic. While it is common to use simplified models for analyzing LLMs, the authors' pre-training algorithm involves setups (such as gradient descent and (re-)initialization) that deviate from standard LLM training approaches.\n\n2. Given the unrealistic setup, the paper would benefit greatly from numerical experiments demonstrating the effectiveness of the proposed algorithm in learning nonlinear functions. If obtaining these results is challenging, conducting other experiments that could provide intuition or insights into understanding ICL in standard LLMs would also be helpful.\n\n3. While the work seems very relevant to understanding ICLs, the paper lacks real evidence supporting this connection. If I am misunderstanding the relevance, I would appreciate it if the authors could provide more detailed explanations on how this work contributes to understanding the mechanisms of ICL and LLMs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback. We address the technical concerns below. \n\n---\n\n**\"pre-training algorithm involves setups that deviate from standard LLM training approaches\"**\n\nWe agree that Algorithm 1 deviates from the most standard training procedure in practice. We make the following remarks on our architecture and optimization algorithm. \n\n- **MLP before attention.** This simplified theoretical setting reflects the scenario where lower layers of the trained transformer act as data representations on top of which the upper layers perform ICL.\nSee Figure 1 in [[Guo et al. 2023]](https://openreview.net/pdf?id=ikwEDva1JZ): they considered a Transformer architecture with alternating MLP and attention layers and suggested that the earlier layers grasp the features of the true function, while the subsequent layers act as linear attention.Our architecture, with an MLP layer for feature learning followed by a linear attention module, idealizes such process occurring in more complicated Transformer.  \nThe same architecture has been analyzed in prior works on in-context learning of nonlinear features such as [[Kim and Suzuki 2024]](https://proceedings.mlr.press/v235/kim24af.html). \n\n- **Linear attention.** In the ICL theory literature, the Softmax attention is often simplified to linear attention to facilitate the optimization analysis. Such simplification can capture several prominent aspects of practical transformer training dynamics -- see [[Ahn et al. 2023]](https://openreview.net/pdf?id=0uI5415ry7) and references therein. Also, in our setting we already observe interesting in-context behavior using a linear attention. Extending our analysis to Softmax attention is an interesting future direction. \n\n- **Layer-wise training and one-step GD.** \nIn the neural network theory literature, layer-wise training and reinitialized bias units are fairly standard algorithmic modifications when proving end-to-end sample complexity guarantees, as seen in many prior works on feature learning [[Damian et al. 2022]](https://proceedings.mlr.press/v178/damian22a/damian22a.pdf) [[Abbe et al. 2023]](https://proceedings.mlr.press/v195/abbe23a/abbe23a.pdf). Without such modifications, stronger target assumptions are generally needed, such as matching link function [[Ben Arous et al. 2020]](https://jmlr.csail.mit.edu/papers/volume22/20-1288/20-1288.pdf).  \nIn addition, training the representation for one gradient step (followed by SGD on the top layer) is a well-studied optimization procedure for two-layer neural networks (with the hope being that the \"early phase\" of representation learning already extracts useful features) -- see [[Ba et al. 2022]](https://openreview.net/pdf?id=akddwRG6EGi) [[Barak et al. 2022]](https://openreview.net/pdf?id=8XWP2ewX-im); but to our knowledge, we are the first to apply this dynamics to nonlinear transformers in the ICL setting and establish rigorous statistical guarantees.  \n\n---\n\n**\"numerical experiments demonstrating the effectiveness of the proposed algorithm\"**\n\nFollowing your suggestion, we conducted experiments that supports our theoretical claim on the statistical efficiency of ICL -- please see the .pdf document attached in our general response. Specifically, in Figure 1 we indeed observe that when $r\\ll d$, there is a separation in the test loss between pretrained transformers and baseline algorithms that directly learn from the test prompt, including kernel ridge regression and two-layer neural network + gradient descent. \n\n---\n\n**\"how this work contributes to understanding the mechanisms of ICL and LLMs\"**\n\nWe believe that the the following high-level messages from our theoretical results are relevant: \n\n* Our analysis highlights how pretrained transformers can leverage prior information (low-dimensional structure in particular) of the function class and outperform simple baseline algorithms. Such a phenomenon is empirically known (e.g., on linear functions, it had been observed that Transformers adapt to the \"sparsity of problem distribution\" [[Garg et al. 2022]](https://openreview.net/pdf?id=flNZJ2eOet)), but existing theoretical analysis with end-to-end learning guarantees typically showed that ICL implements a simple *prior-agnostic* algorithm, such as one gradient descent step on the least squares objective for linear regression.  \nIn addition,  our finding that transformer can adapt to a low-rank target prior suggests, at a high level, a generalization error scaling that does not depend on the ambient dimensionality of the problem, but some notion of \"intrinsic dimensionality\" of the function class. \n\n* Our analysis makes use of a mechanism for ICL in shallow transformers: *feature extraction \\& adaptivity to low-dimensionality* by MLP layer + *in-context function approximation* by attention layer.  We believe that similar mechanism also appears in modern pretrained transformers, which can provide insight to mechanistic interpretability research.\n\n---\n\nWe would be happy to clarify any further concerns/questions in the discussion period.", "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:
Dear Reviewers and Area Chair,

We appreciate your time and effort to provide detailed comments on our paper. In response to the suggestions and comments on the experimental verification, we conducted the following experiments (the results of which are reported in the attached .pdf).

---

First, we compare the statistical efficiency of our studied transformer model trained via Algorithm 1, against two baseline methods: kernel ridge regression, and two-layer neural networks trained by gradient descent.   

We set $d=64,m=8000$, and the true function is $\mathrm{He}_2(\langle\beta,x\rangle)$, where $\beta$ is randomly drawn from $r=2$-dimensional subspace in $\mathbb{R}^{64}$. At test time, context length is altered from $2^6$ to $2^{10}$, and we compared our Transformer, a two-layer neural network (MLP) with width $m=8000$, and kernel ridge regression using radial basis function kernel.
Hyperparameters for our Algorithm 1 are set as $T_1=300,000,N_1=10,000,\eta_1=1,000,T_2=20,000,N_2=1,500,$ and $\lambda_2=0.01$.        

We plotted the test error for each context length in Figure 1 in the PDF, where we observe that Transformer indeed achieves lower test error.

---

Second, we demonstrate the alignment between the weights of the MLP layer $w$ and the target $r$-dimensional subspace, when the transformer is trained by standard SGD with multiple steps.  

The target function and $d,r,m$ are the same as the first experiment. Here we trained MLP layer $w$, bias $b$ and $\Gamma$ in our architecture *simultaneously* via single-pass SGD with weight decay in $T_1=300,000$ steps, with context length $N_1=5,000$.   
The parameters are initialized as follows (here we employ a *mean-field* scaling to encourage feature learning). 
$w_j\sim \mathrm{Unif}(\mathbb{S}^{d-1})$, $\eta=0.1m,\lambda=10^{-6}$.  $b_j\sim \mathrm{Unif}[-1,1]$, $\eta=0.1m,\lambda=10^{-6}$.  $\Gamma= 0.01I_m/m$, $\eta=0.00001/m,\lambda=10^{-3}$.   

---

We hope that these experiments can address the reviewers' concern on the lack of empirical support.    

Please refer to our detailed response to each reviewer where we respond to the individual comments.

Meta-review: The auhtors investigate the ability for a single-layer softmax attention transformer to in-context learn a low-dimensional target function by training on in-context examples.  The authors leverage a somewhat unnatural variant of GD which involves a single step of (non-convex) GD followed by a re-initialization and ridge regression step, an approach that has appeared prior in the deep learning theory literature.  Using this they show that the sample complexity depends only on the intrinsic dimension of the target function rather than the ambient dimension. 

There was some support among the reviewers, with some stronger and some weaker. I'm inclined to agree with the positive reviewers: I think this work provides a nice contribution to our understanding of the sample complexity of canonical learning problems translated to the in-context learning setting.  I recommend acceptance.
Decision: Accept (poster)