Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context

Authors: kazusato  oko
yujin  song
taiji  suzuki
denny  wu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 04:49:43

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.

Main paper: 
 introduction: 
 Pretrained transformers [VSP + 17] possess the remarkable ability of in-context learning (ICL) [BMR + 20], whereby 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 transformer can implement a learning algorithm, such as gradient descent on the in-context examples, in its forward pass [HvMM + 19, DSD + 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 on ICL focus on learning the function class of linear models using linear transformers trained via empirical risk minimization. In this setting, it can be shown that minima of the pretraining loss implements one (preconditioned) gradient descent step on the least squares objective computed on the test prompt [ZFB23, ACDS23, MHM23] . This implies that the forward pass of the pretrained transformer can learn linear targets with n ≳ d in-context examples, 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, RPCG23, ZWB24, LLZV + 24].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-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 a weighted-sum estimator to ridge regression as 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 a 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 ContributionsTo 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 single-index target functions, where the t-th pretraining task is constructed asEQUATIONwhere σ t * : R → R is the unknown link function, and β t ∈ R d 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 (x 1 , y 1 , . . . , x N , y N ) and predict the output corresponding to the query x (See Section 2 for details). This problem setting is based on the following considerations.• 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 a link function with degree P and information exponent Q (defined as the index of the smallest non-zero coefficient in the Hermite expansion), we know that kernel methods require at least n ≳ d P samples [GMMM21, DWY21], whereas twolayer neural network trained by gradient descent can achieve better sample complexity n ≳ d Θ(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 highlights the adaptivity of neural network to low-dimensional functions, i.e., single-index model σ * (⟨x i , β⟩) that depends on one direction (index features)β in R d [AAM22, BES + 22, BMZ23].In such setting, the complexity of gradient descent is dominated by the search for direction β ∈ R d , 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 low-dimensional, since the index features for each task β t are drawn from a rank-r subspace. This low-dimensionality of function class cannot be exploited by any algorithms that directly estimate the target function from test prompt, but as we will see, transformers can adapt to this additional structure via gradient-based 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 a GPT-2 model [RWC + 19] (with the same configurations as the incontext linear regression setting in [GTLV22]) to learn the Gaussian single-index task (1.1) with degree-3 link function, and compare its in-context sample complexity against baseline algorithms (see Section 4 for details).In Figure 1 we observe that the pretrained transformer achieves low prediction risk using fewer in-context examples than two baseline algorithms: kernel ridge regression, and neural network trained by gradient descent. Moreover, we observe that unlike the baseline methods,Transformer achieves an in-context sample complexity (almost) independent of the ambient dimensionality d.The goal of this work is to rigorously establish a dindependent in-context sample complexity in an idealized theoretical setting for a shallow transformer optimized via gradient-based pretraining.We characterize the sample complexity of learning (1.1) in-context, using a nonlinear transformer optimized by gradient descent. Each single-index task is specified by an unknown index feature vector β ∈ R d drawn from some r-dimensional subspace, and a link function σ * with degree P and information exponent Q ≤ 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 low-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 : (x 1 , y 1 , . . . , x N , y N , x) → y be a 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-context prediction risk E|f (x) -f * (x)|-τ = o d (1), where τ is 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 ≳ d Θ(Q) , N T ≫ T ≳ d Θ(Q)pretraining , N * ≳ r Θ(P ) inference , m ≳ r Θ(P ) approximation , 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 a single-index model with information exponent Q using a two-layer neural network. On the other hand, the sample complexity for the in-context 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 n ≳ 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 see a separation between pretrained transformer 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. • Our analysis of pretraining reveals the following mechanism analogous to [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. . For a set S, Unif(S) denotes the uniform distribution over S. We denote the unit sphere {x∈ R d | ∥x∥ = 1} by S d-1 . Õ(•), Ω(•) represent O(•)and Ω(•) notations where polylogarithmic terms are hidden. We write a ≲ b when there exists a constant c such that a ≤ cb. If both a ≲ b and b ≲ a holds, we write a ≍ b. 
in-context learning: 
 We first introduce the basic setting of ICL [BMR + 20] of simple function classes as investigated in [GTLV22, ASA + 22]. In each task, the learner is given a sequence of inputs and outputs (x 1 , y 1 , . . . , x N , y N , x) referred to as prompt, wherex i , x ∈ R d and y i ∈ R. The labeled ex- amples X = (x 1 • • • x N ) ∈ R d×N , y = (y 1 • • • y N ) ⊤ ∈ R N are called context, and x is the query. Given input distribution x 1 , . . . , x N , x i.i.d.∼ D x , the output y i is expressed asy i = f * (x i ) + ς i , i ∈ [N ],where f * is the true function describing the input-output relation andς i i.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 f * . In the pretraining phase, we optimize the model parameters given training data from T distinct tasks {(x t 1 , y We say a model learns these functional relations in-context, when the model can predict the output f * (x) corresponding to query x by solely examining the context (X, y), without updating model parameters for each task. Given the pretrained model f (X, y, x; θ) with parameter θ which predicts the label of query x from context (X, y), we define the expected ICL risk asEQUATIONwhere y = f * (x) + ς and the expectation is taken over the in-context data:x 1 , . . . , x N * , x ∼ D x , f * ∼ D f * , ς 1 , . . . , ς N * , ς ∼ D ς .Note that we take the expectation with respect to contexts(X 1:N * , y 1:N * ) ∈ R d×N * × R N * of length N *, in order to examine the behavior of ICL at a specific context length.We consider the situation where the true input-output relation is expressed by single-index models, i.e., functions that only depend on the direction of the index vector β in the input space. An efficient learning algorithm should adapt to this feature and identify the relevant subspace from highdimensional 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 τ ≥ 0 be the noise level. The prompt (x 1 , y 1 , . . . , x N , y N , x) is generated asx 1 , x 2 , . . . , x N , x i.i.d. ∼ D x = N (0, I d ), y i = f * (x i ) + ς i , f * (x i ) = σ * (⟨x i , β⟩), where ς 1 , . . . , ς N i.i.d.∼ Unif({-τ, τ }), and the distribution D f * of the true function f * is specified as:1. Index Features. Let S be an r ≤ d-dimensional linear subspace of R d . We draw β uniformly from the unit sphere S(S) in S, i.e., from S(S):= {β | β ∈ S, ∥β∥ = 1}. 2. Link Function. σ * (z) = P i=Q ci i! He i (z), where 2 ≤ Q ≤ P . We draw the Hermite coefficients {c i } P i=Q from any distribution satisfyingE[c 2 Q ] = Θ d,r (1) ̸ = 0, P i=Q c 2 i ≤ R 2 c (a.s.) and (c Q , . . . , c P ) ̸ = (0, . . . , 0) (a.s.).(2.2)Throughout the paper, we assume that P ≪ d, r and r ≪ d; specifically, we take P = Θ d,r (1) and r ≲ d 1/2 . Note that the condition r ≪ d 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-index 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, i.e, min{i | c i ̸ = 0} in this case; see [BAGJ21, DH18] ) at least Q. This heterogeneity reflects the situation where the difficulty of learning the input-output relation varies across tasks. Note that we allow for different distributions of the Hermite coefficients {c i }: for example, wemay set (c Q , . . . , c P ) ∼ Unif (c Q , . . . c P )| P i=Q c 2 i i! = 1 (manifold of coefficients satisfying E x [f * (x)] = 1), Unif({0, 1} P -Q+1 \ (0, . . . , 0)), or Unif({(1, . . . , 0), . . . , (0, . . . , 1)}).• The condition Q ≥ 2 (i.e., the Hermite expansion of σ * does not contain constant and linear terms) ensures that the gradient update detects the entire r-dimensional subspace instead of the trivial rank-1 component. For generic polynomial σ * , this assumption can be satisfied by a simple preprocessing step that subtracts the low-degree components, as done in [DLS22] .We consider a transformer composed of a single-layer self-attention module preceded by an embedding module using a nonlinear multi-layer perceptron (MLP). Let E ∈ R de×d N be an embedding matrix constructed from prompt (x 1 , y 1 , . . . , x N , y N , x). A single-layer SoftMax self-attention module [VSP + 17] is given asf Attn (E; W P , W V , W K , W Q ) = E + W P W V E • softmax (W K E) ⊤ W Q E ρ , (2.3)where ρ is the temperature, and W K , W Q ∈ R d k ×de , W V ∈EQUATIONwhere v ∈ R and K ∈ R (de-1)×(de-1) . Then, the simplified attention module is written as fAttn (E;W P V , W KQ ) = E + W P V E E ⊤ W KQ E N, and we take the right-bottom entry as the prediction of y corresponding to query x. 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 is not sufficient to learn nonlinear single-index models. Instead, we set d e = m + 1, d N = N + 1 for m ∈ N and construct E using an MLP layer: Remark 2. We make the following remarks on the considered architecture.EQUATION• The MLP embedding layer before attention has been adopted in recent works [GHM + 23, KS24, KNS24]. This setting can be interpreted as an idealized version of the mechanism that lower layers of the transformer construct useful representation, on top of which upper attention layers implement the in-context learning algorithm.• Our architecture is also inspired by recent theoretical analyses of gradient-based 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 fAttn and MLP embedding (2.5), we can express the model prediction y asEQUATIONwhere Γ = vK ⊤ and σ(W X + b1 ⊤ N ) denotes the m × N matrix whose (i, j)-th entry is σ(w ⊤ i x j + b i ); see Appendix E for full derivation. We refer to Γ as the attention matrix.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, we update the trainable parameters in a sequential manner.Input :Learning rate η 1 , weight decay λ 1 , λ 2 , prompt length N 1 , N 2 , number of tasks T 1 , T 2 , attention matrix initialization scale γ.1 Initialize w (0) j ∼ Unif(S d-1 ) (j ∈ [m]); b (0) j ∼ Unif([-1, 1]) (j ∈ [m]); Γ (0) j,j ∼ Unif({±γ}) (j ∈ [m]) and Γ (0) i,j = 0 (i ̸ = j ∈ [m]j ← w (0) j -η 1 ∇ w (0) j 1 T1 T1 t=1 (y t -f (X t , y t , x t ; W (0) , Γ (0) , b (0) )) 2 + λ 1 w (0) j . 5 Initialize b j ∼ Unif([-Γ * ← argmin Γ 1 T2 T1+T2 t=T1+1 (y t -f (X t , y t , x t ; W (1) , Γ, b)) 2 + λ2 2 ∥Γ∥ 2 F . Output :Prediction function x → f (X, y, x; W (1) , Γ * , b).• In Stage I we optimize the parameters of the MLP (embedding) layer, which is a 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 + 22, 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 a reliable estimate of the target subspace S, which enables the subsequent attention layer to implement a sample-efficient in-context learning algorithm. • In 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 Γ performs regression on the polynomial basis (defined by the MLP embedding), which can be 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 
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 r ≲ d 1 2 . We pretrain the transformer (2.6) using Algorithm 1 with m = Ω(r P ),T 1 = Ω(d Q+1 r Q ), N 1 T 1 = Ω(d 2Q+1 r), γ ≍ 1 m 3 2 r 1 2 d Q ,η 1 ≍ m 3 2 rd 2Q-1 2 • (log d) -Cη for constant C η .Then, for appropriately chosen regularization parameters λ 1 , λ 2 > 0, 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 1 can be upper bounded asR N * (f ) -τ = Õ r 3P m + r 4P 1 T 2 + 1 N 2 + 1 N * .Theorem 1 suggests that to achieve low in-context prediction risk, it is sufficient to set T 1 , N 1 = Ω(d Θ(Q) ), and m, T 2 , N 2 , N * = Ω(r Θ(P ) ). Observe that the pretraining complexity T 1 , N 1 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 1 . Remark 3. We make the following remarks.• The sample complexity highlights different roles of the two sources of low dimensionality in our setting. The low dimensionality of single-index f * entails that the pretraining cost scales as N ≳ d Θ(Q) , which is consistent with prior analyses on gradient-based feature learning [BAGJ21, DNGL23] . On the other hand, the low dimensionality of function class (i.e., β lies in r-dimensional subspace) leads to an in-context sample complexity that scales as N * ≳ r Θ(P ) , which, roughly speaking, is the rate achieved by polynomial regression or kernel models on r-dimensional data.• The multiplicative scaling between N 1 and T 1 in the sample complexity suggests that one can tradeoff between the two quantities, that is, pretraining on more diverse tasks (larger T 1 ) can reduce the required pretraining context length N 1 .• Similar low-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 span of all target functions, and the second-layer approximates the nonlinearity. 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. We provide a sketch of derivation for Theorem 1. The essential mechanism is outlined as follows:after training the MLP embedding via one gradient descent step, the MLP parameters {w(1) j } align with the common subspace of the target functions S. Subsequently, the (linear) attention module estimates the input-output relation f * (which varies across tasks) on this r-dimensional subspace. We explain these two ingredients in the ensuing sections.We first show that the first gradient descent step on 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 havew (1) j ≃ η 1 • 2 T1 T1 t=1 y t ∇ w (0) j f (X t , y t , x t ; W (0) , Γ (0) , 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 T 1 and N 1 . For the population gradient, we make use of the Hermite expansion of σ bj (z) = σ(z + b j ) and show that its leading term is proportional to w(0) 1:r 0 d-r, which is contained in the target subspace S; see Appendix B for details.Next we construct the attention matrix Γ which satisfies the following approximation property: Proposition 2 (Informal). There exists Γ such that with high probability,f (X t , y t , x t ; W (1) , Γ, b) -y t -τ = Õ r 3P /m + r 2P /N 2 for all t ∈ {T 1 + 1, . . . , T 2 }. Moreover, we have ∥ Γ∥ F = Õ r 2P /m .We provide an intuitive explanation of this construction. Recall that the target function can be written in its Hermite expansion f * (x) = P i=Q ci i! He i (⟨x, β⟩). We build an orthonormal basis for f * as follows. Let {β 1 , . . . , β r } be an orthonormal basis of S. Then, any f * can be expressed as a linear combination of functions inH = { r j=1 He pj (⟨β j , •⟩) | Q ≤ p 1 + • • • + p r ≤ P, p 1 ≥ 0, . . . , p r ≥ 0}, where E x∼N (0,I d ) [h(x)h ′ (x)] = I h=h ′ holds for h, h ′ ∈ H.We write H = {h 1 , . . . , h B P } with B P = |H| = Θ(r P ), and observe that a two-layer neural network can be constructed to approximate each h n . Specifically, there exist vectors a 1 , . . . ,a B P ∈ R m such that m j=1 a n j σ ⟨w (1) j , x⟩ + b j ≃ h n (x) for each n ∈ [B P ].Consequently, we can build the desired attention matrix using coefficients a 1 , . . . , a B P . Let A = a 1 • • • a B P ∈ R m×B P and Γ = AA ⊤ , the attention module can recover the true function as1 N2 Γσ(W (1) X t + b1 ⊤ N2 )y t , σ(W (1) x + b) = B P n=1 1 N2 N2 i=1 m j=1 a n j σ w (1) j , x i + b j y i m j=1 a n j σ w (1) j , x + b j (a) ≃ B P n=1 1 N2 N2 i=1 h n (x i )y i h n (x) (b) ≃ B P n=1 E[h n (x)f * (x)]h n (x) = f * (x). (3.1)Roughly speaking, the self-attention architecture computes the correlation between the target function and basis element h i 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 Γ at this stage scale only with r up to polylogarithmic terms; this is because W (1) already identifies the low-dimensional target subspace S.Consequently, the in-context sample size N * only scales with the target dimensionality r ≪ d.Finally, we transfer the learning guarantee from the constructed Γ to the (regularized) empirical risk minimization solution Γ * . By the equivalence between optimization with L 2 regularization and norm-constrained optimization, there exists λ 2 such that ∥Γ * ∥ F ≤ ∥ Γ∥ F and the empirical ICL loss by Γ * is no larger than that of Γ. Hence, we can bound the generalization error by Γ * using a standard Rademacher complexity bound for norm-constrained transformers provided in Appendix D.1. 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 in Appendix D.2.4 Synthetic Experiments 
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 models: for each task t, the context {(x t i , y t i )} N i=1 is generated as x t i i.i.d.∼ N (0, I d ) andy t i = P i=Q c t i i! He i (⟨x t i , β t ⟩), where (c t Q , . . . c t P ) i.i.d. ∼ Unif (c Q , . . . c P )| P i=Q c 2 i i! = 1 and β t i.i.d. ∼ Unif β|β = [β 1 , . . . , β r , 0, . . . , 0] ⊤ |, ∥β∥ = 1 . See Appendix F for further details.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 r. For each problem setting the model is pretrained for 100, 000 steps using the data of degree P = 4 and information exponent Q = 2 (see Appendix F for details). In Figure 2 In Figure 1 we compare the in-context sample complexity of the GPT-2 model pretrained by data of Q = 3 and P = 2 against two baseline algorithms that directly learn f * on the test prompt: (i) kernel ridge regression with the Gaussian RBF kernel k(x, x ′ ) = exp -∥x -x ′ ∥ 2 /σ 2 , and (ii) two-layer neural network with ReLU activationf NN (x) = 1 m m i=1 a i σ(⟨x, w i ⟩) trained by the Adam optimizer [KB15].We observe that for r = 8, d = 16, 32, the pretrained transformer outperforms both KRR and two-layer NN; moreover, 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 pretrained 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 pretraining 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 a larger sample complexity. We outline a few limitations and possible future directions.• The in-context sample complexity we derived r Θ(P ) corresponds to that of polynomial regression or 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 r Θ(P ) ). An interesting question is whether transformers can implement a more efficient in-context algorithm that matches the complexity of gradient-based feature learning in r dimensions. This can be achieved if the pretrained model learns features in-context. • Our pretraining complexity is based on one GD step analysis similar to [DLS22, BES + 22] which 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 we modify the pretraining procedure or training objective, as done in [DTA + 24, LOSW24, ADK + 24, JMS24].[GTLV22] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a 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 preprint arXiv:2310.05249, 2023.[HvMM + 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:1906.01820, 2019.[JLLVR24] Hong Jun Jeon, Jason D Lee, Qi Lei, and Benjamin Van Roy. An information-theoretic analysis of in-context learning. arXiv preprint arXiv:2401.15530, 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 P. 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 arXiv:2408.12186, 2024.[KS24] Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv:2402.01258, 2024.[LLZV + 24] Yue M Lu, Mary I Letey, Jacob A Zavatone-Veth, Anindita Maiti, and Cengiz Pehlevan. Asymptotic theory of in-context learning by linear attention. arXiv preprint arXiv:2405.11751, 2024.[LOSW24] Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with sgd near the information-theoretic limit. arXiv preprint arXiv:2406.01581, 2024.[  
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [No] Justification: We only have toy experiments and it 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 of compute required for each of the individual experimental runs as well as 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 or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
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 is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: We don't have any models or datasets to release.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We included the information in Appendix F.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. 
 
 Appendix: We consider the high-dimensional 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 S = {(x 1 , . . . , x r , 0, . . . , 0) ⊤ | x 1 , . . . , x r ∈ R} and β ∼ Unif(S(S)) = Unif({(β 1 , . . . , β r , 0, . . . , 0) ⊤ | β 2 1 + • • • + β 2 r = 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 Definition 3. We say that an event A occurs with high probability when there exists a sufficiently large constant C * which does not depend on d, r and1 -P r[A] ≤ O(d -C * )holds.We implicitly assume that we can redefine the constant C * to be sufficiently large as needed. The lemma below is a basic example and will be used in the proofs: Lemma 4. Let x ∼ N (0, 1). Then, |x| ≲ √ log d holds with high probability.Proof. From the tail bound of Gaussian, P(|x| ≥ t) ≤ 2 exp (-t 2 /2) holds. Hence we getP(|x| ≥ √ 2C * log d) ≤ O(d -C * ).Note that C * can be redefined by changing the hidden constant in |x| ≲ √ log d.If A 1 , . . . , A M occurs with high probability whereM = O(poly(d)), then A 1 ∩ • • • ∩ A M alsooccurs with high probability (by redefining C * ). In particular, throughout this paper we assume that m, N 1 , N 2 , T 1 , T 2 = O(poly(d)), which allows us to take such unions.In this paper, a k-tensor is a multidimensional array which has k indices: for example, matrices are 2-tensors. Let A be a k-tensor. A i1,...,i k denotes (i 1 , . . . , i k )-th entry of A. Let A be a k-tensor and B be a l-tensor where k ≥ l. A(B) denotes a k -l tensor whose (i 1 , . . . , i k-l )-th entry is A(B) i1,...,i k-l = j1,...,j l A i1,...,i k-l ,j1,...,j l B j1,...,j l , and is defined only when sizes are compatible. If k = l, we sometimes write A(B) as A • B or ⟨A, B⟩. Let v ∈ R d be a vector and k be a positive integer. Then,v ⊗k ∈ R d×•••×d denotes a k-tensor whose (i 1 , . . . , i k )-th entry is v i1 • • • v i k . Let f (x) : R d → R be a d-variable differentiable function. A k-tensor ∇ k f (x) is defined as ∇ k f (x) i1,...,i k = ∂ ∂x i1 • • • ∂ ∂x i k f (x).The following properties can be verified easily. Lemma 5. For tensors A and B, ∥A(B)∥ 2 F ≤ ∥A∥ 2 F ∥B∥ 2 F holds. Lemma 6. For a vector v, ∥v ⊗k ∥ 2 F = (∥v∥ 2 2 ) k holds.We frequently use (probablists') Hermite polynomials, which is defined asHe i (z) = (-1) i e z 2 2 d i dz i e -z 2 2 ,where i is a nonnegative integer. We often make use of the orthogonality property: E z∼N (0,1) [He i (z)He j (z)] = i!δ i,j . The Hermite expansion for σ : R → R is defined asσ(z) = i≥0 ai i! He i (z) where a i = E z∼N (0,1) [σ(z)He i (z)]. Similarly, the multivariate Hermite expansion for f : R d → R is defined as f (z) = i1≥0,...,i d ≥0 ai 1 ,...,i d (i1)!•••(i d )! He i1 (z 1 ) • • • He i d (z d ), where a i1,...,i d = E z1,...,z d ∼N (0,1) [f (z)He i1 (z 1 ) • • • He i d (z d )]. The coefficient a i1,...,i d can also be obtained by a i1,...,i d = E z1,...,z d ∼N (0,1)∂ i 1 ∂z i 1 1 • • • ∂ i d ∂z i d d f (z) .Consider f * in our problem setting (Assumption 1). We can bound E[f * (x) 2 ] and E[f * (x) 4 ] as follows.Lemma 7. Under Assumption 1, E x,f * [f * (x) 2 ] = Θ d,r (1) and E x,f * [f * (x) 4 ] = O d,r (1) holds.Proof. From (2.2) and P = Θ d,r (1),EQUATION. We can bound E x [f * (x) 4 ] = E z∼N (0,E z∼N (0,1),{ci}      P i=Q c i i! He i (z)   4    ≤ Q≤i,j,k,l≤P E[|c i c j c k c l |] i!j!k!l! E[|He i (z)He j (z)He k (z)He l (z)|] ≤ Q≤i,j,k,l≤P R 4 c E z [He i (z) 4 ] 1/4 E z [He j (z) 4 ] 1/4 E z [He k (z) 4 ] 1/4 E z [He l (z) 4 ] 1/4 ≤P 4 R 4 c ζ 4 P .The lemma below is useful to find a basis of the set of single-index functions.Lemma 8. Suppose β ∈ S(S). Then,He p (⟨x, β⟩) = p1+•••+pr=p p1≥0,...,pr≥0 p! p 1 ! • • • p r ! • β p1 1 • • • β pr r • He p1 (x 1 ) • • • He pr (x r )holds.Proof. Note that E x1,...,xr∼N (0,1)∂ i 1 ∂x i 1 1 • • • ∂ ir ∂x ir rHe p (⟨x, β⟩) is nonzero only wheni 1 + • • • + i r = p, because ∂ i 1 ∂x i 1 1 • • • ∂ ir ∂x ir r He p (⟨x, β⟩) = p(p -1) • • • (p -(i 1 + • • • + i r ) + 1)β i1 1 • • • β ir r He p-(i1+•••+ir) (⟨x, β⟩) and E x1,...,xr∼N (0,1) [He q (⟨x, β⟩)] = E z∼N (0,1) [He q (z)] = 0 if q > 0. When i 1 + • • • + i r = p, then E x1,...,xr∼N (0,1) ∂ i 1 ∂x i 1 1 • • • ∂ ir ∂x ir r He p (⟨x, β⟩) = p!β i1 1 • • • β ir rholds. Thus, from the multivariate Hermite expansion we obtain the claim.Corollary 9. Let f * (x) = P i=Q ci i! He i (⟨x, β⟩). Then, E x [∇ k f * (x)] = c k β ⊗k if Q ≤ k ≤ Pand otherwise it is the zero tensor.Lemmas on random vectors. Let χ(d) be the distribution ofz 2 1 + • • • + z 2 d where z i i.i.d.∼ N (0, 1). Note that the Gaussian vector v ∼ N (0, I d ) can be decomposed as v = zw where z ∼ χ(d) and w ∼ Unif(S d-1 ), because the norm and the direction of the Gaussian vector are independent of each other. Lemma 10 (Example 2.11 in [Wai19]). For z ∼ χ(d) and 0 < t < 1,P[|z 2 -d| ≥ dt] ≤ 2 exp(-dt 2 /8) holds. Lemma 11. Let C > 0. For w ∼ Unif(S d-1 ), ∥w 1:r ∥ 2 ≤ 4Crd log d holds with probability at least 1 -2 exp(-d/32) -2rd -C (then with high probability).Proof. Let z 2 ∼ χ 2 (d) be independent from w. We observe that the distribution of z 2 ∥w 1:r ∥ 2 is the χ 2 (r)-distribution. First, from Lemma 10, with probability at least 1 -2 exp(-d/32), z 2 ≥ d/2 holds. Moreover, we can say that (z ′ ) 2 ∼ χ 2 (r) satisfies (z ′ ) 2 ≤ 2Cr log d with high probability for sufficiently large constant: let us decompose(z ′ ) 2 = v 2 1 + • • • + v 2 r where v i ∼ N (0, 1); from Lemma 4, v 2i ≤ 2C log d holds with probability at least 1 -2d -C , and thus (z ′ ) 2 ≤ 2Cr log d holds with probability at least 1 -2rd -C . Taking the union bound yields the result.Lemma 12. For k ≥ 1 and r ≤ d,E w∼S d-1 [∥w 1:r ∥ 2k ] = O d,r r d k holds. Proof. Since ∥w 1:r ∥ 2 is a polynomial of w, Lemma 24 in [DLS22] yields E w∼S d-1 [∥w 1:r ∥ 2k ] ≲ E w∼S d-1 [∥w 1:r ∥ 4 ] k/2 . Here, E w∼Unif(S d-1 ) [∥w 1:r ∥ 4 ] = E v i ∼N (0,1) [(v 2 1 +•••+v 2 r ) 2 ] E z∼χ 4 (d) [z] = 3r+r(r-1) d(d+2) = Θ(r 2 /d 2 ) yields the result.Lemma 13 (Sub Gaussian vector). Let x be a d-dimensional random vector and suppose that ⟨u, x⟩ is zero-mean σ 2 -sub Gaussian for any u ∈ R d such that ∥u∥ = 1. Then, ∥x∥ = Õ(σ √ d) holds with high probabillity. Lemma 14. Let X be a σ 2 -sub Gaussian random variable and Y be a bounded random variable satisfying|Y | ≤ R almost surely. Then, XY -E[XY ] is O(σ 2 R 2 )-sub Gaussian. Proof. XY is σ 2 R 2 -Gaussian by examining the tail probability. Hence, XY -E[XY ] is O(σ 2 R 2 )-sub Gaussian from Lemma 2.6.8 in [Ver18].Polynomial concentration. We frequently use the following bounds on an output data y (which is a polynomial of x). Lemma 15 (Theorem 1.5 in [GSS21] ). Let x ∼ N (0, I n ) and f be a degree-D n-variable polynomial function. Then, for all t ≥ 0 it holds thatP(|f (x) -E x [f (x)]| ≥ t) ≤ 2 exp - 1 C D M 2 min 1≤k≤D t ∥E x [∇ k f (x)]∥ F 2/k, where M is an absolute constant and C D is a constant depending only on D (if ∥E x [∇ k f (x)]∥ F = 0 then we ignore the term with k).Lemma 16. Let x ∼ N (0, I d ) and y = f * (x) + ς = P i=Q ci i! He i (⟨x, β⟩) + ς where c i and β are drawn according to Assumption 1, and ς ∼ Unif{±τ }. Then, for any fixed {c i } and β on the support of their distributions, it holds thatP(|y| ≥ t + τ ) ≤ 2 exp - 1 C P M 2 min 2≤k≤P t R c 2/k, where M is an absolute constant and C P is a constant depending only on P .Proof. Note that from Corollary 9, E x [f * (x)] = 0, E x [∇f (x)] = 0 and E x [∇ i f * (x)] = c i (β) ⊗i ; by (2.2), Lemma 6 and ∥β∥ = 1, ∥E x [∇ i f * (x)]∥ F ≤ R cThis 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 ∈ [m],wEQUATIONholds where σ bj (z) := σ(z + b j ). Here, x t i and y t i is the input and output of i-th context at t-th task, respectively. First we show that the final term is sufficiently small, if we set γ = |Γ (0) j,j | sufficiently small. Lemma 18. With high probability over the distribution of {(X t , y t , x t , y t )} T1 t=1 and {w(0) j } m j=1 , 2η 1 1 T 1 T1 t=1 f (X t , y t , x t ; W (0) , Γ (0) , b (0) )∇ w (0) j f (X t , y t , x t ; W (0) , Γ (0) , b (0) ) = Õ η 1 γ 2 m √ d . Proof. Recall that f (X t , y t , x t ; W (0) , Γ (0) , b (0) ) = m j=1 Γ (0) j,j 1 N 1 N1 i=1 y t i σ bj (w (0) j ⊤ x t i ) σ bj (w (0) j ⊤ x t ), ∇ w (0) j f (X t , y t , x t ; W (0) , Γ (0) , b (0) ) = Γ (0) j,j 1 N 1 N1 i=1 y t i σ ′ bj (w (0) j ⊤ x t i )x t i σ bj (w (0) j ⊤ x t ) + Γ (0) j,j 1 N 1 N1 i=1 y t i σ bj (w (0) j ⊤ x t i ) σ ′ bj (w (0) j ⊤ x t )x t .Fix i, t, j and consider the inner product between w (0) j ∼ Unif(S d-1 ) and x t i ∼ N (0, I d ). From the rotational invariance, without loss of generality, we can assume that w (0) j = [1, 0, . . . , 0] ⊤ andEQUATIONj , x t i ∼ N (0, 1). From Lemma 4, w We can take the union bound with respect to all the i, t, j, retaining the high probability bound (see Appendix A). Therefore, noting that |b j | ≤ 1, we obtain|f (X t , y t , x t ; W (0) , Γ (0) , b (0) )| ≲ mγ(log d) P/2+1 and ∥∇ w (0) j f (X t , y t , x t ; W (0) , Γ (0) , b (0) )∥ ≲ γ(log d) P/2+1/2 √ d.Thus we 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 f . Letg T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) := T1 i=1 1 T 1 y t σ b (w ⊤ x t ) 1 N 1 N1 i=1 y t i σ ′ b (w ⊤ x t i )x t i + T1 i=1 1 T 1 y t σ ′ b (w ⊤ x t )x t 1 N 1 N1 i=1 y t i σ b (w ⊤ x t i ), (B.2) and g(w, b) = E[g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )],where the expectation is taken with respect to the data {(X t , y t , x t , y t )} T1 t=1 . Note that w(1)EQUATIONj , {(X t , y t , x t , y t )} T1 t=1 ) + Õ η 1 γ 2 m √ d from Lemma 18.In Appendix B.1, we analyze the explicit form of g(w, 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 g T1,N1 around g.First, note thatg(w, b) = E yσ b (w ⊤ x)y i σ ′ b (w ⊤ x i )x i + E yσ ′ b (w ⊤ x)xy i σ b (w ⊤ x i ) = E f * E {xi},x,{ςi},ς (f * (x) + ς)σ b (w ⊤ x)(f * (x i ) + ς i )σ ′ b (w ⊤ x i )x i + E f * E {xi},x,{ςi},ς (f * (x) + ς)σ ′ b (w ⊤ x)x(f * (x i ) + ς i )σ b (w ⊤ x i ) = E f * E x,ς (f * (x) + ς)σ b (w ⊤ x) E {xi},{ςi} (f * (x i ) + ς i )σ ′ b (w ⊤ x i )x i + E f * E x,ς (f * (x) + ς)σ ′ b (w ⊤ x)x E {xi},{ςi} (f * (x i ) + ς i )σ b (w ⊤ x i ) = 2E f * [E x [f * (x)σ ′ b (w ⊤ x)x]E x [f * (x)σ b (w ⊤ x)]].Here E f * denotes the expectation with respect to the distribution of the true function f * (i.e., the distribution of {c E x [f * (x)σ ′ b (w ⊤ x)x] = P -1 i=Q-1 a i+1 (b)E x [∇ i+1 f * (x)](w ⊗i ) i! + w P i=Q a i+2 (b)E x [∇ i f * (x)](w ⊗i ) i! , (B.3) E x [f * (x)σ b (w ⊤ x)] = P i=Q a i (b)E x [∇ i f * (x)](w ⊗i ) i! . Proof. Consider a function g(w ⊤ x) = i≥0 ãi i! He i (w ⊤ x). Here E x [f * (x)g(w ⊤ x)] can be calculated as E x [f * (x)g(w ⊤ x)] = i≥0 ãi i! E x [f * (x)He i (w ⊤ x)] = i≥0 i1+•••+i d =i i1≥0,...,i d ≥0 ãi w i1 1 • • • w i d d i 1 ! • • • i d ! E x [f * (x)He i1 (x 1 ) • • • He i d (x d )] (∵ Lemma 8) = i≥0 i1+•••+i d =i i1≥0,...,i d ≥0 ãi w i1 1 • • • w i d d i 1 ! • • • i d ! E x ∂ i1 ∂x i1 1 • • • ∂ i d ∂x i d r f * (x) = i≥0 ãi E x [∇ i f * (x)](w ⊗i ) i! .Where the final line is obtained noting that the tensor∇ i f * (x) has i! i1!•••i d ! entries whose value is ∂ i 1 ∂x i 1 1 • • • ∂ i d ∂x i d r f * (x).Note that E x [∇ i f * (x)] = E x ∇ i P j=Q cj j! He j (⟨x, β⟩) is a nonzero tensor only when Q ≤ i ≤ P , from Lemma 8. We can also show (B.3) asE x [f * (x)σ ′ b (w ⊤ x)x] =E x [∇f * (x)σ ′ b (w ⊤ x)] + E x [f * (x)σ ′′ b (w ⊤ x)]w (∵ Stein's lemma) = i≥0 a i+1 (b)E x [∇ i+1 f * (x)](w ⊗i ) i! + w i≥0 a i+2 (b)E x [∇ i f * (x)](w ⊗i ) i! = P -1 i=Q-1 a i+1 (b)E x [∇ i+1 f * (x)](w ⊗i ) i! + w P i=Q a i+2 (b)E x [∇ i f * (x)](w ⊗i ) i!Using Lemma 19, g(w, b) can be expanded asg(w, b) = E f * 2 P -1 i=Q-1 a i+1 (b)E x [∇ i+1 f * (x)](w ⊗i ) i! + w P i=Q a i+2 (b)E x [∇ i f * (x)](w ⊗i ) i! P i=Q a i (b)E x [∇ i f * (x)](w ⊗i ) i! =: E f * 2a Q (b) (Q -1)! E x [∇ Q f * (x)](w ⊗(Q-1) ) a Q (b)E x [∇ Q f * (x)](w ⊗Q ) Q! + s(w, b) = 2a Q (b) 2 Q!(Q -1)! E f * [E x [∇ Q f * (x)](w ⊗(Q-1) ) • E x [∇ Q f * (x)](w ⊗Q )] + s(w, b), (B.4)where s(w, b):=2E f * P -1 i=Q-1 a i+1 (b)E x [∇ i+1 f * (x)](w ⊗i ) i! + w P i=Q a i+2 (b)E x [∇ i f * (x)](w ⊗i ) i! • P i=Q a i (b)E x [∇ i f * (x)](w ⊗i ) i! - a Q (b) 2 (Q -1)! E x [∇ Q f * (x)](w ⊗(Q-1) ) E x [∇ Q f * (x)](w ⊗Q ) Q!expresses terms in the asymptotic expansion which are not the leading term. Note that, in this explicit formulation, a i (b) = E z∼N (0,1) [σ b (z)He i (z)] appears. It will be helpful to show that this coefficient is well-controlled. From continuity we know that there existsC L such that E z∼N (0,1) [σ b (z)He Q (z)] ≥ C L with probability 1/2.In the following, we first calculate the main term and then upper bound the residual terms.Let us calculate the main term in (B.4) explicitly. Lemma 21.2aQ (b) 2 Q!(Q -1)! E f * [E x [∇ Q f * (x)](w ⊗(Q-1) ) • E x [∇ Q f * (x)](w ⊗Q )] = 2a Q (b) 2 E c Q [c 2 Q ] Q!(Q -1)! (2Q -1)!! E z∼χr [z 2Q ] ∥w 1:r ∥ 2Q-2 w 1:r 0 d-r holds.Proof. From Corollary 9, it holds that ⊗Q . Then, we obtainE x [∇ Q f * (x)] = c Q (β)E f * [E x [∇ Q f * (x)](w ⊗(Q-1) ) • E x [∇ Q f * (x)](w ⊗Q )] =E c Q [c 2 Q ]E β [(β ⊗Q w ⊗(Q-1) ) • (β ⊗Q • w ⊗Q )] =E c Q [c 2 Q ]E β [β ⊗2Q ](w ⊗(2Q-1) ). Next, let us show E β [β ⊗2Q ](w ⊗(2Q-1) ) = (2Q -1)!! E z∼χr [z 2Q ] ∥w 1:r ∥ 2Q-2 w 1:r 0 d-r .First, note that by letting β ′ ∼ N (0, Σ β ) whereΣ β = I r 0 r 0 ⊤ r O d-r, and z ∼ χ r independent from β, β ′ ∼ βz holds (recall that we assumed that β ∼ Unif({(β 1 , . . . , β r , 0, . . . , 0) ⊤ | β 2 1 + • • • + β 2 r = 1}) -see Appendix A). Then, E β [β ⊗2Q ] = 1 E z∼χr [z 2Q ] E v ′ 1:r ∼N (0,Ir),v ′ r+1:d =0 [v ′ ⊗2Q ].It suffices to show that E v∼N (0,Ir) [v ⊗2Q ](z ⊗(2Q-1) ) = (2Q -1)!!z∥z∥ 2Q-2 for z ∈ R r . As an example, we show that the first entry of E v∼N (0,Ir) [v ⊗2Q ](z ⊗(2Q-1) ) equals to (2Q -1)!!z 1 ∥z∥ 2Q-2 : the other components can be calculated similarly. First, we expand(E v∼N (0,Ir) [v ⊗2Q ](z ⊗(2Q-1) )) 1 as (E v∼N (0,Ir) [v ⊗2Q ](z ⊗(2Q-1) )) 1 = i1,...,i 2Q-1 ∈[r] E[v 1 v i1 • • • v i 2Q-1 ]z i1 • • • z i 2Q-1 . Note that E[v 1 v i1 • • • v i 2Q-1 ] ̸ = 0 if and only if the degree of v 1 v i1 • • • v i 2Q-1 with respect to v j is even, for all j ∈ [r].Then, we obtain(E v∼N (0,Ir) [v ⊗2Q ](z ⊗(2Q-1) )) 1 = q1+•••+qr=Q q1≥1,q2,...,qr≥0 c q1,...,qr E[v 2q1 1 v 2q2 2 • • • v 2qr r ]z 2q1-1 1 z 2q2 2 • • • z 2qr r =z 1 q1+•••+qr=Q q1≥1,q2,...,qr≥0 c q1,...,qr E[v 2q1 1 v 2q2 2 • • • v 2qr r ]z 2q1-2 1 z 2q2 2 • • • z 2qr r ,where c q1,...,qr = (2Q-1)! (2q1-1)!(2q2)!•••(2qr)! is the number of ways to assign (2q 1 -1) "1"s, (2q 2 ) "2"s, ..., among i 1 to i 2Q-1 .Note that E v∼N (0,1) [v 2q ] = (2q -1)!!. Therefore, we obtain(E v∼N (0,Ir) [v ⊗2Q ](z ⊗(2Q-1) )) 1 =z 1 q1+•••+qr=Q q1≥1,q2,...,qr≥0 (2Q -1)! (2q 1 -1)!(2q 2 )! • • • (2q r )! (2q 1 -1)!! • • • (2q r -1)!!z 2q1-2 1 z 2q2 2 • • • z 2qr r =z 1 q1+•••+qr=Q q1≥1,q2,...,qr≥0 (Q -1)! (q 1 -1)!(q 2 )! • • • (q r )! (2Q -1)!!(z 2 1 ) q1-1 (z 2 2 ) q2 • • • (z 2 r ) qr =(2Q -1)!!z 1 (z 2 1 + • • • + z 2 r ) Q-1 (∵ multinomial theorem)which concludes the assertion.By Lemma 21, we observe that the main term of the expected one-step gradient aligns to the rdimensional subspace S. In other words, the MLP layer captures the information of S via pretraining, which will be useful in the later stage.First, similarly to the Lemma 21, we can obtain the explicit formulation of the entire g(w, b). Lemma 22.g(w, b)= Q-1≤i≤P -1, Q≤j≤P, i+j is odd 2a i+1 (b)a j (b)E[c i+1 c j ] i!j! (i + j)!! E z∼χr [z i+j+1 ] ∥w 1:r ∥ i+j-1 w 1:r 0 d-r + w Q≤i≤P, Q≤j≤P, i+j is even 2a i+2 (b)a j (b)E[c i c j ] i!j! (i + j -1)!! E z∼χr [z i+j ] ∥w 1:r ∥ i+j (B.6) holds.Proof. Note thatE β [β ⊗n ](w ⊗(n-1) ) =    (n-1)!! Ez∼χ r [z n ] ∥w 1:r ∥ n-2 w 1:r 0 d-r (n is even) 0 (n is odd) and E β [β ⊗n ](w ⊗n ) = (n-1)!! Ez∼χ r [z n ] ∥w 1:r ∥ n (n is even) 0 (n is odd) ,which can be obtained similarly to the proof of Lemma 21.Let us upper bound the non-leading terms; we need to show that the moment of first r components in residual term s(w, b) are sufficiently small. Proof. Note that ∥s(w, b) 1:r ∥ 2 is a polynomial in w; from Lemma 24 in [DLS22] , it suffices to show the case k = 1.s(w, b) 1:r = Q-1≤i≤P -1, Q≤j≤P, i+j is odd, i,j̸ =(Q-1,Q) 2a i+1 (b)a j (b)E[c i+1 c j ] i!j! (i + j)!! E z∼χr [z i+j+1 ] ∥w 1:r ∥ i+j-1 w 1:r + Q≤i≤P, Q≤j≤P, i+j is even 2a i+2 (b)a j (b)E[c i c j ] i!j! (i + j -1)!! E z∼χr [z i+j ] ∥w 1:r ∥ i+j w 1:r . By Minkowski's inequality E[∥x + y∥ 4 ] 1/4 ≤ E[∥x∥ 4 ] 1/4 + E[∥y∥ 4 ] 1/4, it suffices to show that (B.7) holds for each term. Note thatE[c i+1 c j ] ≤ R 2 c and E[c i c j ] ≤ R 2 c from Assumption 1, |a i (b)| ≤ C H from Lemma 20, and E w [∥w 1:r ∥ 4k ] = O((r/d) 2k) from Lemma 12. Moreover, it is known that E z∼χr [z 2l ] = Θ(r l ). Putting these things together yields the assertion.It is also needed in later stages to give a high probability upper bound on the ∥g(w, b)∥; Lemma 24.sup b∈[-1,1] ∥g(w, b)∥ = Õd,r 1 rd 2Q-1holds with high probability over w ∼ Unif(S d-1 ).Proof. 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 g(w, b), instead of g(w, b) 1:r ).Now we control the deviation of g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) from g(w, b). Lemma 25. Under Assumption 1,sup f * ,b∈[-1,1] ∥E x [f * (x)σ ′ b (w ⊤ x)x]∥ ≲ r d log d Q-1 , (B.8) sup f * ,b∈[-1,1] ∥E x [f * (x)σ ′ b (w ⊤ x)x] 1:r ∥ ≲ r d log d Q-1 , (B.9) sup f * ,b∈[-1,1] |E x [f * (x)σ b (w ⊤ x)]| ≲ r d log d Q (B.10)holds with high probability over w ∼ Unif(S d-1 ), respectively. Here sup f * denotes the supremum over the support of f * whose distribution is specified in Assumption 1.Proof. We bound the leading terma Q (b)Ex[∇ Q f * (x)](w ⊗(Q-1) ) (Q-1)! in (B.3) of Lemma 19 to show (B.8). From Lemmas 5 and 6, we havea Q (b)E x [∇ Q f * (x)](w ⊗(Q-1) ) (Q -1)! 2 = a Q (b)c Q (Q -1)! (β) ⊗Q (w ⊗(Q-1) ) = a Q (b)c Q (Q -1)! (β) ⊗Q (w ⊗(Q-1) 1:r ) 2 ≤ a Q (b)c Q (Q -1)! 2 ∥β∥ 2Q ∥w 1:r ∥ 2(Q-1) = a Q (b)c Q (Q -1)! 2 ∥w 1:r ∥ 2(Q-1) .Moreover, Lemma 20 tells us that |a Q (b)| ≤ C H always holds when b ∈ [-1, 1], and Assumption 1 ensures that |c Q | ≤ R c . Therefore, from Lemma 11, we can show thata Q (b)Ex[∇ Q f * (x)](w ⊗(Q-1) ) (Q-1)! 2 ≤ C H Rc (Q-1)! 2 4Cr d log d Q-1 with probability at least 1 -2 exp(-d/32) -2rd -C . (B.9) and (B.10) can be obtained similarly.The following lemma is parallel to Lemma 19 in [DLS22] . Then, with high probability over the distribution of (X t , y t , x t , y t ), it holds that1 N 1 N1 i=1 y t i σ ′ b (w ⊤ x t i )x t i -E x [f t * (x)σ ′ b (w ⊤ x)x] ≤ Õ d N 1 , 1 N 1 N1 i=1 y t i σ ′ b (w ⊤ x t i )x t i -E x [f t * (x)σ ′ b (w ⊤ x)x] 1:r ≤ Õ r N 1 , 1 N 1 N1 i=1 y t i σ b (w ⊤ x t i ) -E x [f t * (x)σ b (w ⊤ x)] ≤ Õ 1 N 1 .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, |f t * (x)| ≤ R where R ≲ (log d) P/2 holds with high probability for all t ∈ [T 1 ]. Conditioned on this,1 N 1 N1 i=1 f t * (x t i )σ ′ b (w ⊤ x t i )x t i -E x [f t * (x)σ ′ b (w ⊤ x)x] ≤ 1 N 1 N1 i=1 f t * (x t i )I f tfor sufficiently large C; hence this term can be ignored. Finally, ς i σ ′ b (w ⊤ x t i )⟨x t i , u⟩ is τ 2 -sub Gaussian for any ∥u∥ = 1 and then ∥ 1N1 N1 i=1 ςσ ′ b (w ⊤ x t i )x t i ∥ = Õ( d/N 1 ) with high probability.By similar procedure, we can obtain a bound of ∥g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) -g(w, b)∥.Lemma 27. Fix any b ∈ [-1, 1] and w ∈ S d-1 satisfying the conditions (B.8)-(B.10). Then, with high probability over the distribution of {(X t , y t , x t , y t )} T1 t=1 ,∥g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) -g(w, b)∥ = Õ  r Q-1 d Q-1 d T 1 + d 2 N 1 T 1   , ∥(g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) -g(w, b)) 1:r ∥ = Õ  r Q-1 d Q-1 r T 1 + r 2 N 1 T 1   , holds.Proof. In this proof, the term "with high probability" indicates the probability with respect to the distribution of {(X t , y t , x t , y t )} T1 t=1 . Letz 1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) = T1 t=1 1 T 1 y t σ b (w ⊤ x t ) 1 N 1 N1 i=1 y t i σ ′ b (w ⊤ x t i )x t i and z 2 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) = T1 t=1 1 T 1 y t σ ′ b (w ⊤ x t )x t 1 N 1 N1 i=1 y t i σ b (w ⊤ x t i ).Note that g T1,N1 = z 1 + z 2 holds.Let us consider z 1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ). Since we assume (B.8)-(B.10) hold, by taking union bound for Lemma 26 with respect to t ∈ [T 1 ], with high probability∥ 1 N1 N1 i=1 y t i σ ′ b (w ⊤ x t i )x t i ∥ ≤ R = Õ d N1 + r Q-1 d Q-1for each t. Moreover, from Corollary 17, |y t | ≤ R ′ = Õ(1) for each t.Here, we notice that the procedure in the proof of Lemma 26 can be applied, yielding∥z 1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) -E[z 1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )]∥ = Õ RR ′ d/T 1 = Õ r Q-1 d Q-1 + d N 1 d T 1with high probability. Similarly we can obtain• ∥z 2 -E[z 2 ]∥ = Õ r Q d Q + 1 N1 d T1 , • ∥(z 1 -E[z 1 ]) 1:r ∥ = Õ r Q-1 d Q-1 + r N1 r T1, and• ∥(z 2 -E[z 2 ]) 1:r ∥ = Õ r Q d Q + 1 N1 r T1 ,which concludes the proof. We need to obtain a bound uniformly over b ∈ [-1, 1]. We first introduce the following definition.Definition 28. Fix any w ∈ S d-1 and X = Lemma 29. Fix any w ∈ S d-1 . Then, with high probability over the distribution of {(X t , y t , x t , y t )} T1 t=1 , the following holds:T1 t=1 {x t 1 , . . . , x t N1 , x t } ⊂ R d . Then, define a finite disjoint partition B(w, X) = {B 1 , . . . , B N (w,X) } of [-1, 1], i.e., ∪ i B i = [-1, 1] and B i ∩ B j = ∅ (i ̸ = j)• g(w, b) is L 1 -Lipschitz continuous with respect to b in the interval [-1, 1].• Let X = T1 t=1 {x t 1 , . . . , x t N1 , x t }. Then, g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) is L 2 -Lipschitz continuous with respect to b in each interval B i ∈ B(w, X).Here Next, let us see g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ). Recall that g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ):= T1 i=1 1 T 1 y t σ b (w ⊤ x t ) 1 N 1 N1 i=1 y t i σ ′ b (w ⊤ x t i )x t i + T1 i=1 1 T 1 y t σ ′ b (w ⊤ x t )x t 1 N 1 N1 i=1 y t i σ b (w ⊤ x t i ).From -Lipschitz, in each interval in B(w, X), g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) is L 2 = Õ( √ d)-Lipschitz.Corollary 30. With high probability over the distribution of w and {(X t , y t , x t , y t )} T1 t=1 ,sup b∈[-1,1] ∥g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) -g(w, b)∥ = Õ  r Q-1 d Q-1 d T 1 + d 2 N 1 T 1   , sup b∈[-1,1] ∥(g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) -g(w, b)) 1:r ∥ = Õ  r Q-1 d Q-1 r T 1 + r 2 N 1 T 1   (B.11) holds.Proof. Let X = T1 t=1 {x t 1 , . . . , x t N1 , x t } and consider B(w, X) = {B 1 , . . . , B N (w,X) }. Then, let B be the union of Õ r Q-1 d Q-1 r T1 + r 2 N1T1 / √ d -coverings of B i ∈ B(w, X). Since |B| is polynomial in d,|∥g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) -g(w, b)∥ -∥g T1,N1 (w, π(b), {(X t , y t , x t , y t )} T1 t=1 ) -g(w, π(b))∥| ≤ Õ  r Q-1 d Q-1 r T 1 + r 2 N 1 T 1   holds.Lemma 31. With high probability over the distribution of {(X t , y t , x t , y t )} T1 t=1 ,• It holds that sup b∈[-1,1] ∥g(w, b) -g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )∥ = Õ  r Q-1 d Q-1 d T 1 + d 2 N 1 T 1   (B.12)with high probability over w ∼ Unif(S d-1 ).• It holds thatsup b∈[-1,1] E w∼Unif(S d-1 ) [∥g(w, b) -g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )∥ j ] 1/j = Õ  r Q-1 d Q-1 r T 1 + r 2 N 1 T 1   (B.13) for j ∈ [4P ].Proof. If an event A occurs with probability at least 1 -d -C over the simultaneous distribution of (w, {(X t , y t , x t , y t )} T1 t=1 ), then we can say that• with probability at least 1-d C/2 over {(X t , y t , x t , y t )} T1 t=1 , P w {A | {(X t , y t , x t , y t )} T1 t=1 )} ≥ 1 -d -C/2 holds.(B.12) follows from this remark and Corollary 30.For (B.13), with high probability over the distribution of {(X t , y t , x t , y t )} T1 t=1 ,P w {(B.11) | {(X t , y t , x t , y t )} T1 t=1 } ≥ 1 -d -Cis satisfied for some C. Also, from Lemma 10 and Corollary 17, with high probability over the distribution of {(X t , y t , x t , y t )} T1 t=1 , we have∥x t i ∥, ∥x t ∥ = O( √ d) and |y t | = Õ(1). From the definition of g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) (B.2) and expan- sion of g(w, b) (B.6), sup w ∥g(w, b) -g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )∥ ≤ sup w ∥g(w, b)∥ + sup w ∥g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )∥ is upper bounded by Õ(d).Therefore, with high probabilitysup b∈[-1,1] E w∼Unif(S d-1 ) [∥g(w, b) -g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )∥ j ] 1/j ≤E w∼Unif(S d-1 )   sup b∈[-1,1] ∥g(w, b) -g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )∥ j   1/j = Õ       1 -d -C   r Q-1 d Q-1 r T 1 + r 2 N 1 T 1   j + d -C d j    1/j     .By setting sufficiently large C, we arrive at the claim.Summary. We combine the obtained results and show that ∥g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )∥ and the moment of the residual terms are bounded with high probability.Corollary 32. With high probability over the distribution of {(X t , y t , x t , y t )} T1 t=1 ,(I) g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 ) = m(w, b) + r(w, b, {(X t , y t , x t , y t )} T1 t=1 ) holds where m(w, b) := 2a Q (b) 2 E c Q [c 2 Q ] Q!(Q -1)! (2Q -1)!! E z∼χr [z 2Q ] ∥w 1:r ∥ 2Q-2 w 1:r 0 d-r ,andsup b∈[-1,1] E w∼Unif(S d-1 ) r(w, b, {(X t , y t , x t , y t )} T1 t=1 ) 1:r j 1/j = Õd,r   r Q-1 d Q-1 r T 1 + r 2 N 1 T 1 + r d 2Q+1   . (II) If η 1 γ ≤ √ rd 2Q-1 /C 1 (log d) C2 for sufficiently large C 1 and C 2 , T 1 ≳ r Q d Q+1 , and N 1 T 1 ≳ rd 2Q+1 , then sup b∈[-1,1] ∥2η 1 γg T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )∥ ≤ 1holds with high probability over w ∼ Unif(S d-1 ). Moreover,sup b∈[-1,1]2η 1 γg T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 , z) ≤ log d holds with high probability over w ∼ Unif(S d-1 ) and the training data for the second stage {(X t , y t , x t , y t )} T2 t=T1+1 , where z ∈ T2 t=T1+1 {x t 1 , . . . , x t N2 , x t }. Proof. Combining Lemmas 23 and 31 yields (I). For (II), by the condition for T 1 and N 1 , we obtainsup b∈[-1,1] ∥g T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 )∥ ≤ Õd,r   r Q-1 d Q-1 d T 1 + d 2 N 1 T 1 + 1 rd 2Q-1   = Õd,r 1 rd 2Q-1from Lemmas 24 and 31.Then, asη 1 γ ≲ √ rd 2Q-1 /C 1 (log d) C2 , sup b∈[-1,1] ∥2η 1 γg T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1)∥ ≤ 1 is achieved. By an argument similar to the one in the proof of Lemma 18, we obtain that with high probability,2η 1 γg T1,N1 (w, b, {(X t , y t , x t , y t )} T1 t=1 , z) ≲ log d for z ∈ T2 t=T1+1 {x t 1 , . . . , x t N2 , x t }.To obtain the supremum bound over b ∈ [-1, 1], it suffices to utilize Lemma 29 again.In the following sections, we assume that these two events (I) and (II) occur.We construct an attention matrix Γ with a good approximation property: see Section 3.2.2 for the proof outline. In this section, a set of basis functions H is defined asH =    x → r j=1 1 √ j! He pj (x j ) Q ≤ p 1 + • • • + p r ≤ P, p 1 ≥ 0, . . . , p r ≥ 0    .We let B P := |H| and introduce a numbering H = {h 1 , . . . , h B P }. Note thatB P = P i=Q r + i -1 i = Θ(r P ).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 T 1 = Ω(d Q+1 r Q ) and N 1 T 1 = Ω(d 2Q+1 r). Let {w (1) j } j be neurons obtained by line 4 of Algorithm 1 with γ ≍ 1 m 3 2 r 1 2 d Q ,η 1 ≍ m 3 2 rd 2Q-1 2 • (log d) -Cη 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 1 , . . . ,a B P ∈ R m such that for each n ∈ [B P ], sup i∈[N2],T1+1≤t≤T2   m j=1 a n j σ(w (1) j ⊤ x t i + b j ) -h n (x t i )   2 = Õ r P m + 1 N 2 , sup T1+1≤t≤T2   m j=1 a n j σ(w (1) j ⊤ x t + b j ) -h n (x t )   2 = Õ r P m + 1 N 2 holds. Moreover, sup n∈[B P ] ∥a n ∥ 2 = Õ r P m .This proposition states that there exists a two-layer neural network m j=1 a n j σ ⟨w(1) j , •⟩ + b j which can approximate the basis function h n , and its second-layer magnitude ∥a n ∥ 2 only scales with r. Note that the condition of T 1 and N 1 T 1 comes from Corollary 32, which implies that the noise terms of the one-step gradient is well controlled.From now on, we suppose that g T1,N1 (w, b) satisfies the conditions in Corollary 32 (to simplify the notation, we drop the dependency on {(X t , y t , x t , y t )} T1 t=1). An important previous result [DLS22] is that, if g T1,N1 (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 a set of symmetric r-dimensional tensors{T n k } 0≤k≤P,n∈[B P ] such that h n (x) = 0≤k≤P T n k , x ⊗k 1:r and ∥T n k ∥ F ≲ r P -k 4 . Definition 35 (f (r, d)-alignment).We call that g T1,N1 (w, b) has f (r, d)-alignment with S, if the following holds:• For any 0 ≤ k ≤ P -symmetric r-dimensional tensor T , E w∼Unif(S d-1 ) [ T , g T1,N1 (w, b) ⊗k 1:r 2 ] ≳ f (r, d) k ∥T ∥ 2 F ,and the hidden constant in "≳" needs to be uniformly lower bounded.Lemma 36 (Corollary 4 in [DLS22] , adapted). Suppose that g T1,N1 (w, b) has f (r, d)-alignment with S. Let T be a 0 ≤ k ≤ P -symmetric r-dimensional tensor. Then, there exists ψ T (w, b) such thatE w∼Unif(S d-1 ) [ψ T (w, b)⟨g T1,N1 (w, b), x⟩ k ] = T , x ⊗k 1:r , E w∼Unif(S d-1 ) [ψ T (w, b) 2 ] ≲ f (r, d) -k ∥T ∥ 2 F , |ψ T (w, b)| ≲ f (r, d) -k ∥T ∥ F ∥g T1,N1 (w, b)∥ k .Intuitively speaking, f (r, d)-alignment ensures that the MLP gradient g T1,N1 (w, b) has sufficiently large component on S. If so, the lemma states that, an infinite-width neural networkE w [ψ T (w, b)⟨g T1,N1 (w, b), •⟩ k ],whose hidden layer weight is g T1,N1 (w, b) and the activation is z → z k , can express the functions T , x ⊗k 1:r with small output layer ψ T (w, b). Let us show that our g T1,N1 (w, b) has sufficient alignment if b satisfies (B.5). Lemma 37 (Tensor expectation lower bound). Let T be a k ≤ P -symmetric r-dimensional tensor and w = w 1:r ∥w 1:r ∥ 2Q-2 . ThenE w∼Unif(S d-1 ) [T ( w⊗k ) 2 ] ≳ r Qi d (Q+1)i E[∥T ( w⊗k-i )∥ 2 F ]holds for all 0 ≤ i ≤ k. Here T ( w⊗0 ) := T .Proof. Let u ∼ N (0, I d ), z ∼ χ(d) and ū = u 1:r ∥u 1:r ∥ 2Q-2 . We can decompose ū as ū = z 2Q-1 w. ThereforeE[T ( ū⊗k ) 2 ] = E[z (4Q-2)k ]E[T ( w⊗k ) 2 ]holds. On the other hand, let x ∼ S r-1 and z ′ ∼ χ(r). Then we can decompose ū as ū = (z ′ ) 2Q-1 x. ThusE[T ( ū⊗k ) 2 ] = E[(z ′ ) (4Q-2)k ]E[T (x ⊗k ) 2 ]holds. It implies thatE[T ( w⊗k ) 2 ] = E[(z ′ ) (4Q-2)k ] E[z (4Q-2)k ] E[T (x ⊗k ) 2 ].Similarly,E[∥T ( w⊗k-i )∥ 2 F ] = E[(z ′ ) (4Q-2)(k-i) ] E[z (4Q-2)(k-i) ] E[∥T (x ⊗(k-i) )∥ 2 F ]is satisfied. Now from Corollary 13 in [DLS22] ,E[∥T (x ⊗(k-i) )∥ 2 F ] ≲ r i E[T (x ⊗k ) 2 ] holds. Therefore we obtain E[T ( w⊗k ) 2 ] = E[(z ′ ) (4Q-2)k ] E[z (4Q-2)k ] E[T (x ⊗k ) 2 ] ≳ r -i E[(z ′ ) (4Q-2)k ] E[z (4Q-2)k ] E[∥T (x ⊗(k-i) )∥ 2 F ] = r -i E[(z ′ ) (4Q-2)k ] E[z (4Q-2)k ] E[z (4Q-2)(k-i) ] E[(z ′ ) (4Q-2)(k-i) ] E[∥T ( w⊗k-i )∥ 2 F ] = r (2Q-2)i d (2Q-1)i E[∥T ( w⊗k-i )∥ 2 F ].Corollary 38. Let T be a k ≤ P -symmetric r-dimensional tensor and let m(w, b)1:r = 2a Q (b) 2 Q!(Q-1)! Ec Q [c 2 Q ](2Q-1)!! Ez∼χ r [z 2Q ]∥w 1:r ∥ 2Q-2 w 1:r . Moreover, assume that (B.5) holds. Then,E w∼Unif(S d-1 ) [ T , m(w, b) ⊗k 1:r 2 ] ≳ 1 d (2Q-1)i r 2i E[∥T (m(w, b) ⊗k-i 1:r )∥ 2 F ] holds for all 0 ≤ i ≤ k.Proof. This is obtained by calibrating the coefficients from the previous lemma.Lemma 39. Let T be a 0 ≤ k ≤ P -symmetric r-dimensional tensor. Assume that (B.5) holds. Furthermore, suppose r ≲ d 1 2 , T 1 = Ω(d Q+1 r Q ) and N 1 T 1 = Ω(d 2Q+1 r). Then, E w∼Unif(S d-1 ) [ T , g T1,N1 (w, b) ⊗k 1:r 2 ] ≳ 1 d (2Q-1)k r 2k ∥T ∥ 2 F holds for all 0 ≤ k ≤ P . Therefore, g T1,N1 (w, b) has d -2Q+1 r -2 -alignment with S.E w∼Unif(S d-1 ) [ T , g T1,N1 (w, b) ⊗k 1:r 2 ] ≥E w∼Unif(S d-1 ) 1 2 T , m(w, b) ⊗k 1:r 2 -E w∼Unif(S d-1 ) [δ(w, b) 2 ] (C.1)for some δ(w, b) satisfyingE w∼Unif(S d-1 ) [δ(w, b) 2 ] ≲ k i=1 E w∼Unif(S d-1 ) [∥T (m(w, b) ⊗k-i 1:r )∥ 2 F ] E w∼Unif(S d-1 ) [∥r(w, b) 1:r ∥ 4i ].Here, from Corollaries 32 and 38,E w∼Unif(S d-1 ) [δ(w, b) 2 ] ≲ k i=1 d (2Q-1)i r 2i E w∼Unif(S d-1 ) [ T , m(w, b) ⊗k 1:r 2 ] Õd,r r Q-1 d Q-1 r T 1 + r 2 N 1 T 1 + r d 2Q+1 iholds. Now, from the condition of N 1 and T 1 , and r ≲ √ d, we knowd (2Q-1) r 2 r Q-1 d Q-1 r T 1 + r 2 N 1 T 1 + r d 2Q+1 ≲ 1. Therefore we can achieve E w∼Unif(S d-1 ) [δ(w, b) 2 ] ≤ 1 4 E w∼Unif(S d-1 ) [ T , m(w, b) ⊗k 1:r 2 ].Applying this and Corollary 38 (i = k) to (C.1) yields the result.No we can show that some "infinite-width" neural network can approximate the functions h ∈ H. Lemma 40(Lemma 9 in [DLS22], adapted). Let a ∼ Unif({-1, 1}) and b ∼ Unif([-log d, log d]).Then, for any k ≥ 0 there exists v k (a, b) such that for |x| ≤ log d,E a,b [v k (a, b)σ(ax + b)] = x k , sup a,b |v k (a, b)| = Õd (1).Lemma 41 (Lemma 12 and Corollary 5 in [DLS22], adapted). There existsϕ n (a, w, b 1 , b 2 ) such that if f ϕn (x) := E a,w,b1,b2 [ϕ n (a, w, b 1 , b 2 )σ(2η 1 ag T1,N1 (w, b 1 ) + b 2 )]wherea ∼ Unif{±γ}, w ∼ Unif(S d-1 ), b 1 ∼ Unif([-1, 1]), b 2 ∼ Unif([-log d, log d]), it holds that sup i∈[N2],T1+1≤t≤T2 f ϕn (x t i ) -h n (x t i ) 2 ≲ 1 N 2 , sup T1+1≤t≤T2 f ϕn (x t ) -h n (x t ) 2 ≲ 1 N 2 , E a,w,b1,b2 [ϕ n (a, w, b 1 , b 2 ) 2 ] = Õ(r P ), sup a,w,b1,b2 |ϕ n (a, w, b 1 , b 2 )| = Õ(r P ) for each n ∈ [B P ]. Proof. We first construct ϕ T n k (a, w, b 1 , b 2 ) for basis tensors defined in Lemma 34, such that if f ϕ T n k (x) := E a,w,b1,b2 [ϕ T n k (a, w, b 1 , b 2 )σ(⟨2η 1 ag T1,N1 (w, b 1 ), x⟩ + b 2 )], it holds that sup i∈[N2],T1+1≤t≤T2 f ϕ T n k (x t i ) -T n k , (x t i ) ⊗k 1:r 2 ≲ 1 N 2 , (C.2) sup T1+1≤t≤T2 f ϕ T n k (x t ) -T n k , (x t ) ⊗k 1:r 2 ≲ 1 N 2 . (C.3) Let ϕ T n k (a, w, b 1 , b 2 ) = 2v k (a,b2)ψ T n k (w,b1) (2η1) k γ k I A(w) I B(b1)where v k and ψ T n k are constructed in Lemma 36 and the events A(w) and B(b 1 ) are defined asA(w) := sup b1∈[-1,1] 2η 1 γ| g T1,N1 (w, b 1 ), x t i |, sup b1∈[-1,1] 2η 1 γ| g T1,N1 (w, b 1 ), x t | ≤ log d for i ∈ [N 2 ], T 1 + 1 ≤ t ≤ T 2 ∩ sup b1∈[-1,1] ∥2η 1 γg T1,N1 (w, b 1 )∥ ≤ 1 B(b 1 ) := E z∼N (0,1) [σ b1 (z)He i (z)] ≤ C H (2 ≤ i ≤ P + 2) ∩ E z∼N (0,1) [σ b1 (z)He Q (z)] ≥ C L .Then,f ϕ T n k (x) =E a,w,b1,b2 2v k (a, b 2 )ψ T n k (w, b 1 ) (2η 1 ) k γ k I A(w) I B(b1) σ(⟨2η 1 ag T1,N1 (w, b 1 ), x⟩ + b 2 ) =E b1 2 • I B(b1) E w ψ T n k (w, b 1 ) (2η 1 ) k γ k I A(w) E a ′ ,b2 [v k (a, b 2 )σ(⟨2η 1 a ′ γg T1,N1 (w, b 1 ), x⟩ + b 2 )] ,wherea ′ ∼ Unif({±1}). Here, if x ∈ T2 t=T1+1 {x t 1 , . . . , x t N2 , x t }, E w ψ T n k (w, b 1 ) (2η 1 ) k γ k I A(w) E a ′ ∼Unif({±1}),b2 [v k (a, b 2 )σ(⟨2η 1 a ′ γg T1,N1 (w, b 1 ), x⟩ + b 2 )] =E w ψ T n k (w, b 1 )⟨g T1,N1 (w, b 1 ), x⟩ k + E w ψ T n k (w, b 1 )(1 -I A(w) )⟨g T1,N1 (w, b 1 ), x⟩ k .Note that the absolute value of second term above can be upper bounded byE w (1 -I A(w) ) 1/2 E w (ψ T n k (w, b 1 )⟨g T1,N1 (w, b 1 ), x⟩ k ) 2 1/2. Recall that A(w) occurs with high probability from Corollary 32, andE w (ψ T n k (w, b 1 )⟨g T1,N1 (w, b 1 ), x⟩ k ) 2 1/2 is polynomial in d.Hence we can set the second term to be O(1/ √ N 2 ).Moreover, as P[B(b 1 )] = 1/2 from Lemma 20, we haveE b1 2 • I B(b1) E w ψ T n k (w, b 1 )⟨g T1,N1 (w, b 1 ), x⟩ k =E b1 2 • T n k , x ⊗k + E b1 2(1 -I B(b1) ) T n k , x ⊗k = T n k , x ⊗k .Therefore, (C.2) and (C.3) holds. Noting that EQUATIONWe discretize Lemma 41 to establish Proposition 33.Proof. [Proof of Proposition 33] First, in the same discretization manner as the proof of Lemma 13 in [DLS22] , we can show that there exists a 1 , . . . , a B P such that for each n ∈ [B P ],supEQUATIONj , b j denotes the biases at the first and second initialization, respectively). Recall that w(1) j = 2η 1 Γ (0) j,j g T1,N1 (w (0) j , b (0) j ) - 2η 1 1 T1 T1 t=1 f (X t , y t , x t ; W (0) , Γ (0) , b (0) )∇ w (0) j f (X t , y t , x t ; W (0) , Γ (0) , b (0)). It remains to bound the effect of the second term. By Lemma 18, we know that the norm of second term is Õ η 1 γ 2 m √ d . Similar to the argument in the proof of Lemma 18, the inner product between the second term and x t , a n j σ wx t i (i ∈ [N 2 ], T 1 + 1 ≤ t ≤ T 2 ) is also Õ η 1 γ 2 m √ d ,(1)j , x + b j = Õ r P m η 1 γ 2 m √ d √ m = Õ r P m .We obtain the assertion.Next, we evaluate the discrepancy between the empirical correlation 1 N2 N2 j=1 y j h(x j ) and its expectation for each h ∈ H, which decides the approximation error in (b) of (3.1). Recall that h ∈ H can be written in the form ash(x) = r i=1 1 √ p i ! He pi (x i ) (Q ≤ p 1 + • • • + p r ≤ P, p 1 ≥ 0, . . . , p r ≥ 0). Also note that 1 N 2 N2 j=1 y j h(x j ) -E x,ς [yh(x)] = 1 N 2 N2 j=1 ς j h(x j ) + 1 N 2 N2 j=1 σ * (⟨x j , β⟩)h(x j ) -E x [σ * (⟨x, β⟩)h(x)],where σ * (z) = P i=Q ci i! He i (z). First we give an upper bound for| 1 N2 N2 j=1 σ * (⟨x j , β⟩)h(x j ) - E x [σ * (⟨x, β⟩)h(x)]|.Lemma 42. Let D ≤ 2P be the degree of σ * (⟨x, β⟩)h(x), when we regard it as an r-variable polynomial with respect to (x 1 , . . . , x r ). For all t > 0 there exist absolute constants M and C D depending only on D such thatP   1 N 2 N2 j=1 σ * (⟨x j , β⟩)h(x j ) -E x [σ * (⟨x, β⟩)h(x)] ≥ t   ≤2 exp - 1 C D M 2 min 1≤k≤D t √ N 2 ∥E[∇ k x (σ * (⟨x, β⟩)h(x))]∥ F 2/kholds, where C D and M are constants.Proof. Let F (x 1 , . . . , x N ) = 1 N2 N2 j=1 σ * (⟨x j , β⟩)h(x j ).It is a degree-D polynomial for x 1,1 , . . . , x 1,r , . . . , x N2,r , which are standard Gaussian variables. Also,∥E[∇ k F (x 1 , . . . , x N )]∥ F = N 2 -1 ∥E[∇ k (σ * (⟨x, β⟩)h(x))]∥ Fis satisfied. Then, applying Lemma 15 to F yields the desired result.Then, our aim is to bound∥E[∇ k (σ * (⟨x, β⟩)h(x))]∥ F . Lemma 43. Let h(x) = r i=1 1 √ pi He pi (x i ) and P = p 1 + • • • + p r . For Q + P ≤ k ≤ P + P , ∥E[∇ k (σ * (⟨x, β⟩)h(x))]∥ 2 F ≤ R 2 c Ck holds where Ck is a constant depending only on k. For other k, ∥E[∇ k (σ * (⟨x, β⟩)h(x))]∥ F = 0. Proof. First, ∥E[∇ k (σ * (⟨x, β⟩)h(x))]∥ 2 F = P i=Q c i i! E[∇ k (He i (⟨x, β⟩)h(x))] 2 F holds.To calculate E[∇ k (He i (⟨x, β⟩)h(x))], note that from Lemma 8,He i (⟨x, β⟩)h(x) =   q1+•••+qr=i qj ≥0 i! q 1 ! • • • q r ! β q1 1 • • • β qr r He q1 (x 1 ) • • • He qr (x r )   He p1 (x 1 ) √ p 1 ! • • • He pr (x r ) √ p r ! . (C.4) Also, from the discussion in Appendix A.3, E x ∂ s 1 ∂x s 1 1 • • • ∂ sr ∂z srr He i (⟨x, β⟩)h(x) is equal to the coefficient of (C.4) with respect to1 s1!•••sr! He s1 (x 1 ) • • • He sr (x r ). Therefore, E[∇ k (He i (⟨x, β⟩)h(x))] ̸ = O if and only if i + P = k.Then,P i=Q ci i! E[∇ k (He i (⟨x, β⟩)h(x))] 2 F = c k-P (k-P )! E[∇ k (He k-P (⟨x, β⟩)h(x))] 2 F .Here, we can observe that the (j 1 , . . . , j k )-th entry of the tensorc k-P (k-P )! E[∇ k (He k-P (⟨x, β⟩)h(x))] is equal to c k- P s 1 ! • • • s r ! (s 1 -p 1 )! • • • (s r -p r )! • β s1-p1 1 • • • β sr-pr r 1 √ p 1 ! • • • √ p r ! I s1-p1≥0,...,sr-pr≥0wheres i = k k ′ =1 I j k ′ =i (i ∈ [r]). There are k! s1!•••sr! index sets (j 1 , . . . , j k ) having the same entry, and thusc k- P (k -P )! E[∇ k (He k-P (⟨x, β⟩)h(x))] 2 F = s1+•••+sr=k c 2 k- P 1 p 1 ! • • • p r ! k! s 1 ! • • • s r ! s 1 ! • • • s r ! (s 1 -p 1 )! • • • (s r -p r )! 2 • (β s1-p1 1 • • • β sr-pr r ) 2 I s1-p1≥0,...,sr-pr≥0 = q1+•••+qr=k-P ,qj ≥0 c 2 k- P 1 p 1 ! • • • p r ! k! (q 1 + p 1 )! • • • (q r + p r )! (q 1 + p 1 )! • • • (q r + p r )! q 1 ! • • • q r ! 2 • (β q1 1 • • • β qr r ) 2 = q1+•••+qr=k-P ,qj ≥0 c 2 k- P (k -P )! q 1 ! • • • q r ! k! (k -P )! (q 1 + p 1 )! • • • (q r + p r )! p 1 ! • • • p r !q 1 ! • • • q r ! • (β q1 1 • • • β qr r ) 2 . (C.5)We can verify that there exists a constant Ck depending only on P such thatk! (k-P )! (q1+p1)!•••(qr+pr)! p1!•••pr!q1!•••qr!≤ Ck holds uniformly; for example, one can naively obtaink! (k -P )! (q 1 + p 1 )! • • • (q r + p r )! p 1 ! • • • p r !q 1 ! • • • q r ! ≤ k!(q 1 + p 1 )! • • • (q r + p r )! ≤ k!k k .Here we used (q1 + p 1 ) + • • • + (q r + p r ) = k. Moreover, q1+•••+qr=k-P ,qi≥0 (k-P )! q1!•••qr! (β q1 1 • • • β qr r ) 2 = ∥β ⊗k-P ∥ 2 F = 1 holds. Hence (C.5) is upper bounded by c 2 k- P Ck ≤ R 2 c Ck .Corollary 44. With high probability,1 N 2 N2 j=1 y j h(x j ) -E x [σ * (⟨x, β⟩)h(x)] = Õ 1 √ N 2 holds.Proof. Plugging the results above and t = Θ( 1 √ N2 (log d) P ) into Lemma 42 yields1 N 2 N2 j=1 σ * (⟨x j , β⟩)h(x j ) -E x [σ * (⟨x, β⟩)h(x)] ≲ 1 √ N 2 (log d) Pwith high probability. Moreover, we can easily verify 1N2 N2 j=1 ς j h(x j ) = Õ 1 √ N2by the method used in the proof of Lemma 26.We prove the following Proposition 2 using the preparations above. Proposition 2. Let X t = (x t 1 , . . . , x t N2 ) and y t = (y t 1 , . . . , y t N2 ) ⊤ where T 1 + 1 ≤ t ≤ T 1 + T 2 . With high probability over the data distribution and random initialization of the parameters, there exists Γ such thatf (X t , y t , x t ; W (1) , Γ, b) -y t -τ = Õ  r 3P m + r 2P N 2  holds for all t ∈ {T 1 + 1, . . . , T 2 }, where W (1) is obtained by line 4 of Algorithm 1 and b is re-initialized in Line 5 of Algorithm 1, where the conditions on η 1 , λ 1 , m, γ, N 1 , T 1 are identical to that in Theorem 1. Moreover, ∥ Γ∥ F = Õ r 2P /m is satisfied.Proof. It suffices to show the proposition for each T 1 + 1 ≤ t ≤ T 1 + T 2 , and then take the union bound. We drop the superscript t for simplicity. Note that, from Lemma 8 and E x [h(x)h ′ (x)] = I h=h ′ for h, h ′ ∈ H, we can expand f * (x) as f * (x) ={a n } B P n=1 is constructed in Proposition 33. Then, by setting Γ = AA ⊤ ,AA ⊤ σ(W (1) X + b1 ⊤ N )y N 2 ,     σ(w (1) 1 ⊤ x + b 1 ) . . . σ(w (1) m ⊤ x + b m )     -y = B P n=1 (h n (x) + ϵ n (x))   E[f * h n ] +   1 N 2 N2 j=1 y j h n (x j ) -E[f * h n ]   + 1 N 2 N2 j=1 y j ϵ n (x j )   -f * (x) -ς = B P n=1 (h n (x) + ϵ n (x))   E[f * h n ] +   1 N 2 N2 j=1 y j h n (x j ) -E[f * h n ]   + 1 N 2 N2 j=1 y j ϵ n (x j )   - B P n=1 h n (x)E[f * h n ] -ςholds, where ϵ n (x) := m j=1 a n j σ(w(1) j ⊤ x + b j ) -h n (x) satisfits |ϵ n (x)| = Õ r P /m + 1/N 2 from Proposition 33.To bound the error, we note that the following holds for each n:• As h n (x) is a polynomial whose degree is at most P , andE x [∇ k h n (x)] = O(1), |h n (x)| ≲(log d) P/2 holds with high probability from Lemma 15.• From lemma 8, |E[f * (x)h n (x)]| = O(1) holds.• From Corollary 17, |y j | ≲ (log d) P/2 holds with high probability.Hence, for each n ∈ [B P ], we obtainh n (x) Õ(1) + ϵ n (x) Õ( r P m + 1 N 2 ) E[f * h n ] O(1) +   1 N 2 N2 j=1 y j h n (x j ) -E[f * h n ]   Õ( 1 N 2 ) + 1 N 2 N2 j=1 y j ϵ n (x j ) Õ( r P m + 1 N 2 ) -h n (x)E[f * h n ] = Õ  r P m + 1 N 2 + r P m + 1 N 2   .Combining this with B P = Θ(r P ) and m = Ω(r P ), we arrive at the upper bound. Finally for ∥ Γ∥ F , we have∥ Γ∥ F ≤ B P n=1 ∥a n (a n ) ⊤ ∥ F = B P n=1 ∥a n ∥ 2 = Õ r 2P /m .Note that after one-step gradient descent, ∥wj ∥ = Õd,r (1) is ensured with high probability from Lemma 18 and Corollary 32.Let W = W (1) and suppose b 1 , . . . , b m are biases obtained in line 5 of Algorithm 1. LetEQUATION1 ) ⊤ x + b 1 ) . . .(1)m ) ⊤ x + b m )    ∥Γ∥ F ≤ Gbe the set of transformers where the norm of the attention matrix is constrained, and letRad T (F N,G ) = E {X t },{y t },{x t },{ϵ t } sup f ∈F N,G 1 T T t=1 ϵ t f (X t , y t , x t )be its Rademacher complexity, where ϵ t ∼ Unif({±1}). Here contexts X ∈ R d×N and y ∈ R N are of length N . We can evaluate the Rademacher complexity as follows.Proposition 45.Rad T (F N,G ) = Õd,r Gm √ Tholds, when |b j | = Õd,r (1) and ∥w(1)j ∥ = Õd,r (1) for all j ∈ [m].Proof. First, it holds thatEQUATION1 ) ⊤ x t + b 1 ) . . .(1)m ) ⊤ x t + b m )       (a) ≤ E     1 T sup ∥Γ∥ F ≤G ∥Γ∥ F T t=1 ϵ t σ(W (1) X t + b1 ⊤ N )y t N    σ((w (1) 1 ) ⊤ x t + b 1 ) . . . σ((w (1) m ) ⊤ x t + b m )    ⊤ F     = G T E     T t=1 ϵ t σ(W (1) X t + b1 ⊤ N )y t N    σ((w (1) 1 ) ⊤ x t + b 1 ) . . . σ((w (1) m ) ⊤ x t + b m )    ⊤ F     ≤ G T E      T t=1 ϵ t σ(W (1) X t + b1 ⊤ N )y t N    σ((w (1) 1 ) ⊤ x t + b 1 ) . . . σ((w (1) m ) ⊤ x t + b m )    ⊤ 2 F      (b) = G T T E X,y,x      σ(W (1) X + b1 ⊤ N )y N    σ((w (1) 1 ) ⊤ x + b 1 ). . .(1)EQUATIONl ) ⊤ x + b l )   2    (c) ≤ G √ T m k,l=1 E x∼N (0,I d ),x ′ ∼N (0,I d ),f * ,ς (f * (x) + ς) 2 σ((w (1) k ) ⊤ x + b k ) 2 σ((w (1) l ) ⊤ x ′ + b l ) 2 (d) ≤ G √ T m k,l=1 E x,f * ,ς (2f * (x) 2 + 2ς 2 )(2((w (1) k ) ⊤ x) 2 + 2b 2 k ) E x ′ (2((w (1) l ) ⊤ x ′ ) 2 + 2b 2 l ) ,where the annotated equality/inequalities are obtained by (a) Cauchy-Schwarz inequality, (b)E ϵt∼Unif{±1} [∥ t ϵ t A t ∥ 2 F ] = t ∥A t ∥ 2 F , (c) convexity of f (z) = z 2 , and (d)(σ(α + β)) 2 ≤ (α + β) 2 ≤ 2α 2 + 2β 2 .Let us further evaluate the last line: from the assumption and Lemma 7, we know that ∥w k ∥, |b k | and E[f * (x) 2 ] are all Õ(1): in particular, we can easily obtain that E x∼N (0,I d ) ((wEQUATION, and remains to evaluateEQUATIONk ) ⊤ x) 2 . On this term, we obtainE x,f * f * (x) 2 ((w (1) k ) ⊤ x) 2 ≤ E x,f * f * (x) 4 1/2 E x ((w (1) k ) ⊤ x) 4 1/2 = Õ(1) (∵ Lemma 7).EQUATION1 ) ⊤ x + b 1 ) . . .(1)m ) ⊤ x + b m )    ∥Γ∥ F ≤ G ∪ (X, y, x, y) → y .Then, Rad T ( FN,G ) = Õd,rGm √ T holds when |b j | = Õd,r (1) and ∥w(1)j ∥ = Õd,r (1) for all j ∈ [m]. Proof. Note that E {X t },{y t },{x t },{y t },{ϵ t } sup f ∈ FN,G 1 T T t=1 ϵ t f (X t , y t , x t , y t ) ≤E {X t },{y t },{x t },{ϵ t } sup f ∈F N,G 1 T T t=1 ϵ t f (X t , y t , x t ) + E {y t },{ϵ t } 1 T T t=1 ϵ t y t .This follows from sup f ∈F N,G 1 T T t=1 ϵ t f (X t , y t , x t ) ≥ 0, as F N,G contains zero function corresponding to Γ = O. The second term in the last line can be bounded asE {y t },{ϵ t } 1 T T t=1 ϵ t y t ≤ E {y t },{ϵ t }   1 T T t=1 ϵ t y t 2   = 1 √ T E[(y t ) 2 ] = O 1 √ T by Lemma 7.The ICL risk for prompt length N was defined asR N (f ) = E X,y,x,y [|f (X 1:N , y 1:N , x; W , Γ, b) -y|],where the length of X 1:N and y 1: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 ∥w j ∥ = Õd,r (1) and |b j | = Õd,r (1) for each j ∈ [m]. Then,|R N (f ) -R M (f )| = Õ ∥Γ∥ F m 1/N + 1/M holds.Proof. Note that|R N (f ) -R M (f )| ≤E[|f (X 1:N , y 1:N , x; W , Γ, b) -f (X 1:M , y 1:M , x; W , Γ, b)|] =E    Γ σ(W X 1:N + b1 ⊤ N )y 1:N N - σ(W X 1:M + b1 ⊤ M )y 1:M M ,    σ(w ⊤ 1 x + b 1 ) . . . σ(w ⊤ m x + b m )       (a) ≤ ∥Γ∥ F E    σ(W X 1:N + b1 ⊤ N )y 1:N N - σ(W X 1:M + b1 ⊤ M )y 1:M M    σ(w ⊤ 1 x + b 1 ) . . . σ(w ⊤ m x + b m )       ≤∥Γ∥ F E σ(W X 1:N + b1 ⊤ N )y 1:N N - σ(W X 1:M + b1 ⊤ M )y 1:M M 2 1/2 • E       σ(w ⊤ 1 x + b 1 ) . . . σ(w ⊤ m x + b m )    2    1/2, where (a) holds because a⊤ Γb = Γ • (ab ⊤ ) ≤ ∥Γ∥ F ∥ab ⊤ ∥ F = ∥Γ∥ F ∥a∥∥b∥. Let us bound E σ(W X 1:N +b1 ⊤ N )y 1:N N - σ(W X 1:M +b1 ⊤ M )y 1:M M 2 , which is to m j=1 E f * E X,y 1 N N i=1 σ(w ⊤ j x i + b j )y i -1 M M i=1 σ(w ⊤ j x i + b j )y i 2 .For simplicity, we drop the subscript j and obtainE X,y   1 N N i=1 σ(w ⊤ x i + b)y i - 1 M M i=1 σ(w ⊤ x i + b)y i 2   ≤2E   1 N N i=1 σ(w ⊤ x i + b)y i -E σ(w ⊤ x + b)y 2   +2E   1 M M i=1 σ(w ⊤ x i + b)y i -E σ(w ⊤ x + b)y 2   = 2 N V σ(w ⊤ x + b)y + 2 M V σ(w ⊤ x + b)y ≤ 2 N + 2 M E σ(w ⊤ x + b) 2 (f * (x) + ς) 2 (a) ≤ 2 N + 2 M E (2(w ⊤ x) 2 + 2b 2 )(2f * (x) 2 + 2ς 2 ) .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 E (2(w ⊤ x) 2 + 2b 2 )(2f * (x) 2 + 2ς 2 ) = Õd,r (1).E σ(W X 1:N +b1 ⊤ N )y 1:N N - σ(W X 1:M +b1 ⊤ M )y 1:M M 2 = 2 N + 2 M Õd,r (m).Next, we observe thatE       σ(w ⊤ 1 x + b 1 ) . . . σ(w ⊤ m x + b m )    2    ≤ 2 m j=1 E[(w ⊤ j x) 2 + b 2 j ] = Õd,r (m).Putting all things together, we arrive at|R N (f ) -R M (f )| = Õ ∥Γ∥ F m 1/N + 1/M .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 Γ * be the minimizer of the ridge regression problem (line 8 in Algorithm 1). By the equivalence between optimization with L 2 regularization and norm-constrained optimization, there exists λ 2 such that∥Γ * ∥ F ≤ ∥ Γ∥ F = Õ r 2P /m , 1 T 2 T1+T2 t=T1+1 |y t -f (X t , y t , x t ; W (1) , Γ * , b)| 2 ≤ 1 T 2 T1+T2 t=T1+1 (y t -f (X t , y t , x t ; W (1) , Γ * , b)) 2 ≤ 1 T 2 T1+T2 t=T1+1 (y t -f (X t , y t , x t ; W (1) , Γ, b)) 2 .Then, from Proposition 2,1 T 2 T1+T2 t=T1+1 |y t -f (X t , y t , x t ; W (1) , Γ * , b)| -τ = Õ  r 3P m + r 2P N 2   holds. We evaluate R N2 (f ) -τ where f = f (X t , y t , x t ; W (1) , Γ * , b) using the Rademacher complexity. R N2 (f ) -τ = 1 T 2 T1+T2 t=T1+1 |y t -f (X t , y t , x t ; W (1) , Γ * , b)| + R N2 (f ) - 1 T 2 T1+T2 t=T1+1 |y t -f (X t , y t , x t ; W (1) , Γ * , b)| -τ ≤ Õ  r 3P m + r 2P N 2   + sup f ∈ FN,∥ Γ∥ F R N2 (f ) - 1 T 2 T1+T2 t=T1+1 |y t -f (X t , y t , x t , y t )|holds, where F is defined in Corollary 46. We can evaluate the expectation value of the second term of the last line as, using ϵ t i.i.d. ∼ Unif{±1},E {X t },{y t },{x t },{y t }   sup f ∈ FN,∥ Γ∥ F R N2 (f ) - 1 T 2 T1+T2 t=T1+1 |y t -f (X t , y t , x t , y t )|   (a) ≤ 2E {X t },{y t },{x t },{y t },{ϵ t }   sup f ∈ FN,∥ Γ∥ F 1 T 2 T1+T2 t=T1+1 ϵ t |y t -f (X t , y t , x t , y t )|   (b) ≲ Rad T2 ( FN2,∥ Γ∥ F ) + E y,ϵ 1 T 2 T1+T2 t=T1+1 ϵ t y t ≤ Rad T2 ( FN2,∥ Γ∥ F ) + E y,ϵ   1 T 2 T1+T2 t=T1+1 ϵ t y t 2   1/2 ≤ Rad T2 ( FN2,∥ Γ∥ F ) + 1 √ T 2 E y 2 1/2 (c) = Õ  r 4P T 2   .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 (x, y) ⊤ → |x -y| is √ 2-Lipschitz continuous, and (c) follows from the bound of ∥ Γ∥ F in Proposition 2, together with Lemma 7 and Proposition 45. Note thatsup f ∈ FN,∥ Γ∥ F R N2 (f ) - 1 T 2 T1+T2 t=T1+1 |y t -f (X t , y t , x t , y t )| (D.1) is nonnegative because f (X, y, x, y) = y is included in FN,∥ Γ∥ F . Hence from Markov's inequality, (D.1) is Õ r 4P T2with probability at least 0.995. Hence we obtainR N2 (f ) -τ = Õ  r 3P m + r 2P N 2 + r 4P T 2   .Now we have done the upper bound for R N2 (f ). For R N * (f ), 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.We derive equation (2.6), following the same line of argument as [ZFB23]. As we assumed that W P V W KQ are in the form as (2.4), the output of the attention layer is written asfAttn (E; W P V , W KQ ) = E + * * 0 1×m v E • E ⊤ K * 0 1×m * E N . (E.1)Note that we adopt the right-bottom entry fAttn (E;W P V , W KQ ) m+1,N +1as a prediction for output of the query. By (E.1). Hence we havefAttn (E; W P V , W KQ ) m+1,N +1 =[0 1×m v]E • E ⊤ K * 0 1×m *     σ(w ⊤ 1 x + b 1 ) . . . σ(w ⊤ m x + b m ) 0     N =v[y 1 • • • y N 0] • E ⊤     K    σ(w ⊤ 1 x + b 1 ) . . . σ(w ⊤ m x + b m )    0     N = v[y 1 • • • y N ] • σ(W X + b1 ⊤ N ) ⊤ K    σ(w ⊤ 1 x + b 1 ) . . . σ(w ⊤ m x + b m )    N = vK ⊤ σ(W X + b1 ⊤ N )y N ,    σ(w ⊤ 1 x + b 1 ) . . . σ(w ⊤ m x + b m )    .Setting Γ = vK ⊤ yields equation (2.6).For the experiments in Section 4, test loss was averaged over 128 independent tasks, with 2048 and 128 independent queries generated for each task in the experiment of Figure 1 and 2, respectively. During the validation of the experiment of Figure 1 , the coefficients {c i } in the single-index model were fixed to be (c 2 , c 3 ) = ( √ 2 • 2!/2, √ 2 • 3!/2) to reduce the variance in the baseline methods.GPT-2 model. The experiments in Section 4 were based on the architecture used in [GTLV22]1 , whose backbone is the GPT-2 model [RWC + 19]. Given the prompt {(x i , y i )} N +1 i=1 , the embedding E ∈ R d×(2N +2) was constructed as E = [x 1 , ỹ1 , . . . , x N , ỹN ] where ỹi = [y i , 0, . . . , 0] ⊤ . This embedding was mapped to Ẽ ∈ R 256×(2N +2) 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 a 2N + 2-dimensional vector through the read-out layer, where the (2k -1)-th entry (k ∈ [N + 1]) corresponds to the prediction ŷk (x t 1 , y t 1 . . . , x t k-1 , y t k-1 , x t k ) for y k , that is, the answer for the query x t k reading the context of length k -1. We used the Adam optimizer [KB15] with a learning rate of 0.0001. The training loss at each step was set as1 B B t=1 N +1 k=1 y t k -ŷk (x t 1 , y t 1 . . . , x t k-1 , y t k-1 , x t k )2 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) = m j=1 a j σ(w ⊤ j x) where σ = ReLU and m = 256. The initialization scale followed the mean-field regime, with a j i.i.d.∼ Unif({±1/m}) and w j i.i.d.Unif(S d-1 ). The architecture was trained using the Adam optimizer, with a learning rate of 0.1 for the first layer and 0.1/m for the second layer, along with a 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. In addition to the GPT-2 experiments, we also conducted comparison between our Algorithm 1 on our simplified architecture and baseline algorithms. Data is generated as y t i = He 2 (⟨x t i , β t ⟩), where β t i.i.d. ∼ Unif{β ∈ R d | β = [β 1:r , 0] ⊤ , ∥β∥ = 1} with d = 32 and r = 2. We pretrained our transformer architecture (2.6) with m = 8, 000, using Algorithm 1 with T 1 = 10 5 , N 1 = 10 4 , η 1 = 10 5 (recall that η 1 ≍ m 3 2 rd 2Q-1 2 • (log d) -Cη should be large) and T 2 = 3, 000, N 2 = 512, λ 2 = 0.1. We used mini-batch SGD to find an approximate optimizer Γ * of the ridge regression. For validation, we altered the context length from 2 4.5 to 2 9 and compared its test error with onestep GD (training the first layer via one-step gradient descent, and then conducting ridge regression on the second layer) on a 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.• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It 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 a 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 of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.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 circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).38th Conference on Neural Information Processing Systems (NeurIPS 2024).B P n=1 E x∼N (0,I d ) [f * (x)h n (x)]h n (x). Now let A = a 1 • • • a B P ∈ R m×B P , wherehttps://github.com/dtsip/in-context-learning, MIT License. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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? 
2. There is no soft-max. 
3. This is not true GD/SGD (and I think this is not obvious from the introduction). 
    1. The MLP is only trained for one iteration.
    2.  The function is initialized to compute 0, and following that the $w_i$ are exactly “erased” by the regularization.
    3. The biases are then reset to be uniform on a unit interval.
questions: questions: There are three interesting things in this model: 
1. The non-linearity
2. The single index function
3. The low rank structure
But 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?

Can you include some examples of function classes satisfying Assumption 1 with low P? For instance, what if my ground truth $\sigma_*$ is ReLU?

Can this be empirically verified? An experiment that shows that w aligns with S under true GD would be nice.

Is 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: limitations: Yes
paper_summary: paper_summary: * The authors theoretically study the capacity for (nonlinear) transformer
  models to perform in-context learning with a low-complexity function class.
* The set-up involves pretraining transformers on sequences encoding
  input-output examples from a reduced-rank Gaussian single-index model.
  * The use of Gaussian single index models gives the transformer a chance to
    implement a more intricate (but still analytically tractable) in-context
    learning algorithm than in-context linear regression for linear data.
  * The restriction of the in-context target functions to a low-rank subspace
    gives the transformer an opportunity to internalise (during pretraining)
    a 'prior' that it can rely on during in-context learning for a sample
    efficiency advantage.
* Within this setting the authors prove their main theorem, characterising
  the in-context sample efficiency of a transformer trained under certain
  conditions.
  * Importantly, the efficiency scales with the dimension of the
    low-dimensional model subspace rather than the ambient dimension.
  * This suggests that the pretrained transformer model is capable not only
    of performing ICL of Gaussian single-index models, but also of encoding
    the 'prior' and gaining a corresponding sample efficiency advantage from
    doing so.
strength_weakness: strength_weakness: I lack the background to evaluate the soundness of the main theorem, but the
high-level structure of the main theorem statement seems plausible to me and
so I am open to taking the authors at their word.

Instead, the concerns I have with the paper are about the framing of the
contributions. In particular, my concern focuses on the second research
question, "Can a pretrained transformer adapt to certain structures of the
target function class, and how does such adaptivity contribute to the
statistical efficiency of ICL?" (Lines 45/46), to which the authors answer,
in summary, that, yes, there is a 'separation' (improved sample complexity)
between these transformers and "algorithms that directly learn the
single-index function from the test prompt".

My concerns with this framing are as follows:

1. **These baselines are potentially misleading.** The authors compare the
   sample complexity of a transformer model that has access to more
   information (through pretraining) against that of a baseline that lacks
   this information. It is not surprising that prior information can reduce
   sample complexity and this is why the transformer's performance beats this
   baseline.
     
     I think the appropriate framing would be more along the lines of saying
     that the transformer is able to (in some sense) 'match' other algorithms
     that have this same access to prior information, if these baselines are
     available in the literature.

     Note: I grant that the authors have shown that transformers can take
     advantage of this additional information, and this is a meaningful
     contribution. The specific concern is with the misleading comparison
     that might suggest transformers have some radical sample efficiency
     properties, when the improvement over these baselines has clearly come
     from access to additional information that the baselines don't have.

2. **This does not appear to be a completely novel observation.** There are at least
   empirical examples of transformers showing improved sample efficiency
   through making use of prior information gleaned from pretraining. For
   example, this is implied by the task diversity experiments of Raventos et
   al. (cited by the authors). In the low task-diversity regime, the model
   appears to memorise a short list of task vectors and when asked to perform
   ICL of these particular tasks it would obviously outperform a ridge
   regression model entertaining a larger hypothesis class. Therefore it
   appears that in some related work this second question has already been
   addressed.

     *Note:* These empirical observations are different from the present
     work's theoretical characterisation of the sample efficiency advantage
     of the pretrained transformer and I think this does not fundamentally
     detract from the authors' claims to novelty. However, I think it would
     be appropriate to introduce this question in the context of these
     previous attempts to answer it.

I would invite the authors to consider changing the framing of their paper to
more clearly qualify the nature of their contribution with respect to this
second research question.
questions: questions: To begin with I will flag that before reading this paper I was not familiar
with Gaussian single-index models. Based on the information in the paper and
my own light additional reading, I think I have assembled enough
understanding to ask sensible questions, but I could be mistaken, and I
apologise in advance if my questions are ill-formed.

0. 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.

1. A basic question: I don't think I understand line 52: "In single-index
   model, outputs only depend on the direction of $\beta_t$ in the
   $d$-dimensional input space." Is this just foreshadowing the constraint
   that $\beta_t$ is restricted to the unit sphere in the $r$-dimensional
   subspace? (Perhaps the constraint on the magnitude of $\beta$ should be
   mentioned in the introduction.)

2. I understand that the first goal addressed by this work is to create a
   setting where transformers exhibit nonlinear in-context learning behaviour
   (and I consider this a strength of the work). However, the shift from
   regression to single-index models brings with it additional statistical
   and architectural complexity. Therefore, with respect to your second
   stated goal of demonstrating adaptivity of ICL to low-dimensional
   structure in the training function class, I am wondering if it would be
   possible to get away with a simpler demonstration:

    1. In particular, what would happen if one considered the problem setting
       of in-context linear regression where the regression tasks are confined
       to a low-dimensional subspace (or a unit sphere on such a subspace)? Do
       you think that such a setting would offer a similar chance for the
       transformer to show its ability to internalise informative 'prior'
       information, but without the added complexity of considering
       single-index models?

    2. If I am not mistaken, this simpler low-rank regression setting is
       *not* a special case of your setting, since it would require the link
       function to be always the identity function which violates your
       assumptions about the coefficients in the Hermite expansion. Is that
       correct?

    3. Nevertheless, would you conjecture that a similar result to your main
       theorem---a generalisation bound that scales with the dimension of the
       low-dimensional subspace rather than that of the ambient
       dimension---would hold in such a setting?

Finally, some questions about future work.

3. Your theoretical work does not offer an empirical component. However, your
   analysis appears to make some clear predictions about what a transformer
   trained in this setting would learn, in terms of its ICL performance for
   this kind of data and even the internal mechanisms with which it might
   learn to perform ICL (such as using the MLP layers to encode the
   low-dimensional subspace).
   
    1. Do you think these performance predictions would be borne out by
       empirical experiments to this effect?

    2. Do you think that trained transformers would implement ICL in this
       setting using the mechanisms you have proposed?

    3. Are there any differences between your theoretical training algorithm
       and modern transformer training procedures that you expect might cause
       these predictions not to hold?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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?

2. No experiments are provided to justify the result, especially for the comparison with kernel methods and two-layer NN.

3. 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. 

4. 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?" 
There 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.
questions: questions: 1. I think Algorithm 1 is proposed to handle the non-convexity of the problem, right?
limitations: limitations: There is no potential negative societal impact of this work.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. 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.

3. 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.
questions: questions: Please see part 3 in the last section.
limitations: limitations: N/A

Comments: title: Summary of our responses and revisions
Thanks again
Thanks and follow-up questions about baselines
comment: Dear Area Chair,

As the rebuttal period draws to a close, we would like to briefly summarize our responses and revisions to our manuscript to address the reviewers' concerns.
- **Experimental verification.** Several reviewers noted that the manuscript could be strengthened by an experiment that supports our theoretical findings. Our [general response](https://openreview.net/forum?id=uHcG5Y6fdB&noteId=Ygau6T3C2K) presented experiments that demonstrate a separation in statistical efficiency.
- **Motivation of the studied setting.** Following the reviewers' suggestions, we have expanded our discussion on the background on single-index models. This *nonlinear* function class is common in the NN theory literature (see our [reply to reviewer 3HLy](https://openreview.net/forum?id=uHcG5Y6fdB&noteId=K3Xao71geS)), but to our knowledge, we are the first to provide end-to-end learning guarantees (with gradient-based training) in the more challenging *in-context* scenario, going beyond prior works on *linear* ICL. This idealized setting allows us to rigorously study $(i)$ how pretrained transformers adapt to the prior information of the function class, and $(ii)$ how the MLP block can act as a feature extractor.
- **Training procedure.** We agree with the reviewers that Algorithm 1 deviates from the most standard training paradigm.  We have added further explanations that this layer-wise training (with one-step GD) is commonly used in the theoretical analysis of representation learning [[Damian et al. 2022]](https://proceedings.mlr.press/v178/damian22a/damian22a.pdf) [[Ba et al. 2022]](https://proceedings.neurips.cc/paper_files/paper/2022/file/f7e7fabd73b3df96c54a320862afcb78-Paper-Conference.pdf) to handle the non-convexity, but we are the first to adapt such a procedure to the ICL setting (see our [reply to reviewer UaEF](https://openreview.net/forum?id=uHcG5Y6fdB&noteId=6WcM6yn1y9)).
---
We appreciate the reviewers' feedback which helped us improve the manuscript, and we hope that our nontrivial theoretical contributions can be recognized (in light of our rebuttal and revisions).
Thank you for your response. The empirical verification of alignment is nice. I also appreciate the use of the low rank structure as a means to learn something in pretraining. I will raise my score to a 5. I am still somewhat unsatisfied with the analysis of gradient descent. I understand that this is proof technique is somewhat common in the analysis of single hidden layer networks, but I believe this should be described as its own algorithm, since the main difficulty with analyzing gradient descent *is* dealing with the weights of the first layer.
We thank the reviewer for the updated evaluation. We would like to address the reviewer's outstanding concern with the following remarks: 

* As a theory paper, our primary goal is to provide the *first* end-to-end learning guarantee (that accounts for the model width, number of tasks, prompt length, etc.) when learning a *nonlinear* class of functions in context. The single-index function class is significantly more challenging than the linear regression class examined in prior works. Our statistical guarantees entail a separation in statistical efficiency between pretrained transformers and algorithms that directly learn from the test prompt, while also demonstrating the adaptivity to "prior knowledge" of the target function and the benefits of nonlinear MLP layer. We believe that rigorously establishing these points, even in an idealized setting, represents an important contribution.  
We would also like to point out that even in the simpler linear setting, existing end-to-end analyses typically assume specific model parameterization and training procedures (e.g., gradient flow, infinite tasks) [[Zhang et al. 2023]](https://www.jmlr.org/papers/volume25/23-1042/23-1042.pdf) [[Ahn et al. 2023]](https://openreview.net/pdf?id=LziniAXEI9). We argue that these results do not offer more "*intuition or insights into understanding ICL in standard LLMs*" than our work does.  

* As mentioned in the Introduction, the single-index target function has been extensively studied in the deep learning theory literature; however, to our knowledge, we are the first to theoretically investigate the learning of this function class in the more challenging *in-context* setting.  
It is worth noting that even in the "easier" setting of standard neural network training (i.e., not considering ICL), some of the most influential papers on the theory of representation learning have considered algorithmic modifications similar to ours, such as layer-wise training, one-step feature learning, and reinitialized bias units [[Abbe et al. 2022]](https://proceedings.mlr.press/v178/abbe22a/abbe22a.pdf) [[Damian et al. 2022]](https://proceedings.mlr.press/v178/damian22a/damian22a.pdf) [[Ba et al. 2022]](https://proceedings.neurips.cc/paper_files/paper/2022/file/f7e7fabd73b3df96c54a320862afcb78-Paper-Conference.pdf). 
These works are not dismissed by the community because they "*deviate from standard training approaches*"; rather, they serve as an important first step toward understanding the statistical and computational efficiency of neural networks beyond simple linear problems.  
We hope the reviewer can recognize the nontrivial theoretical contribution of our analysis in this context.
Thanks for the rebuttal. I think my concerns are partially addressed. While I appreciate the effort of conducting relevant numerical experiments, I feel like it only provides marginal intuition and understanding. I also understand that it is not easy to provide a thorough numerical experiment in such a short rebuttal period. 

Therefore, my concerns are only partially addressed. I will maintain my original rating, but I will take your effort into consideration in the later stage.
Dear Reviewer,

As the discussion period is ending in two days, we wanted to follow up to check whether our response has adequately addressed your concerns for an improved score. If there are any remaining questions or concerns, we would be happy to provide further clarification.

Thank you for your time and consideration.
Thank you for your reply. I found these comparisons informative and they have helped me better understand and appreciate the contribution of this paper.
We thank the reviewer for the positive evaluation and insightful followup questions. 

**Sample efficiency for single-index model estimators assuming knowledge of the low-rank subspace**

If the relevant $r$-dimensional subspace is known to the learner, then information theoretically $n\simeq r$ samples would be sufficient and necessary (because intuitively speaking we are estimating $r$ parameters). However, this information theoretic limit does not entail the existence of efficient (polynomial time) algorithms.  
As for existing estimators, it is known that doing kernel ridge regression (KRR) on the low-rank subspace requires $n\gtrsim r^{P}$ samples, where $P$ is the degree of the nonlinear link function $\sigma_*$; this contrasts the $n\gtrsim d^P$ rate when the subspace is not known. The most sample-efficient estimator that we are aware of is the partial-trace procedure studied in [[Damian et al. 2024]](https://arxiv.org/pdf/2403.05529), which achieves a sample complexity of $n\gtrsim r^{k_*/2}$ (assuming knowledge of the $r$-dimensional subspace) where $k_*\le P$ is the *generative exponent* depending on structure of the link function; this tailored algorithm makes use of the tail of Gaussian data, and the obtained complexity matches the SQ lower bound under polynomial compute.   
Our generalization error bound for ICL entails a sufficient sample size of $n\gtrsim r^{\Theta(P)}$, which is parallel to the sample complexity achieved by KRR on the low-rank subspace. We however note that the exponent of $r$ in Theorem 1 does not exactly match the KRR rate due to additional multiplicative factors; closing this gap is an interesting future direction. 

**Empirical validation of theoretical predictions**

At a *qualitative* level, our theoretical analysis suggests a statistical efficiency separation between pretrained transformers and baseline algorithms on the in-context examples when $r\ll d$, which is supported by Figure 1.  
Obtaining a more *quantitative* comparison, such as evaluating the tightness of our generalization bound in terms of polynomial dependence on $d,r$, would be more challenging computationally, as it requires the training of an ensemble of transformer models, each with training data of different dimensionality and context length (analogous to Figure 3 in [[Ba et al. 2023]](https://openreview.net/pdf?id=HlIAoCHDWW)). Following your suggestion, we are currently working on such an experiment with varying dimensionality to verify that the sample complexity of the in-context phase only depends on the subspace dimensionality $r$ rather than the ambient dimensionality $d$.
Thank you for your response. I think the proposed revisions (especially adding more background about single-index models and better contextualisation with prior work including the work of Raventos).

I also welcome the proof-of-concept experiments that you have added, which I think improve the validity of this paper's contribution, assuming the experiments would be documented and discussed in detail in a future revision of the paper.

At the moment I maintain my recommendation that the paper should be accepted.

I have two further questions, both regarding baselines of one kind or another:

1. Regarding theoretical in-context sample complexity of the pre-trained transformer in this setting: Is it known what is the sample efficiency for other single-index model estimators *assuming knowledge of the low-rank subspace* like the pre-trained transformer? I understand why you have emphasised the comparison between your result and known sample efficiency for estimators without this prior knowledge. I think if you could *also* show that the transformer performs in some sense *as well as* other estimators with the same knowledge, that would strengthen your contribution further.

2. Regarding the empirical demonstration of the separation, how does the performance of each estimator you studied compare to theoretical predictions? I am assuming the transformers used for these experiments (described as 'our transformer' in your top-level rebuttal) matches fairly closely the subject of your theoretical analysis, so a direct comparison may be possible in some form. At the moment, the high-level story of 'sample efficiency separation' is corroborated by the experiments, but I'm asking if it is possible to show that the experiments also corroborate additional details of your theoretical results.
Dear Reviewer,

As the discussion period is ending in three days, we wanted to follow up to check whether our response has adequately addressed your concerns for an improved score. If there are any remaining questions or concerns, we would be happy to provide further clarification.

Thank you for your time and consideration.
Thank you for your response. I will raise the score to positive. The most attractive thing to me in this work is the framework for studying nonlinear functions for ICL, which characterizes the impacts of nonlinearity.
We thank the reviewer for the followup question. We make the following clarifications: 

- In the experiment, the MLP parameters of the transformer model are pretrained using Algorithm 1, and *fixed* during the in-context phase (across all tasks). Whereas for the two-layer MLP baseline, the parameters are directly trained on the in-context examples. 
- We set the MLP layer of the transformer and the first layer of the baseline two-layer network to be of the same size ($m\times d$). The transformer model also has ($m\times m$) attention parameters to be optimized.  
We note that the experiment aims to compare the sample complexity of ICL against baseline algorithms (which is the reviewer's concern in *Weakness 1*), not the parameter efficiency. A side remark is that for algorithms that directly learn the target function from in-context examples (e.g., kernel regression and two-layer NN + GD), the information theoretically optimal sample complexity is $n\gtrsim d$, meaning that when $r\ll d$, these baseline algorithms cannot match the in-context efficiency of transformers (which is independent of $d$) even with *infinite* compute or parameters. 

We would be happy to clarify any further concerns/questions in the discussion period.
Thank you for the response. My remaining question focuses on Weakness 1. I can get the point that you mainly study the case when $r\ll d$. That's OK. I just want to know more details of the comparison. For example, 1, in experiments, do you still fix the MLP layer or let it be trainable? 2. How many trainable parameters are there for the Transformer you use and the two-layer neural networks you use in the experiment? I would like to see whether the comparison makes sense.