Title: Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context
Venue: NeurIPS 2024
Accepted: yes
Keywords: Transformer

in-context learning

feature learning

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

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

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 single-index target functions f ∗ ( x ) = σ ∗ ( ⟨ x , β ⟩ ) , where the index features β ∈ 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 information exponent of the link functions σ ∗ ) 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.

1 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 Contributions

1.1.1 Function Class: Gaussian Single-index Models

To address the above questions, we study the in-context learning of arguably the simplest nonlinear extension of linear regression where a nonlinearity is applied to the response. Specifically, we consider the following class of single-index target functions, where the t -th pretraining task is constructed as

where σ 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 rankr 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 d -independent in-context sample complexity in an idealized theoretical setting for a shallow transformer optimized via gradient-based pretraining.

10 1

10 2

in-context sample size

N

*

10

2

10

1

10 0

prediction risk

d

=16

d

=32

two-layer NN

pretrained TF

1.1.2 Main Result: Learning Single-index Models In-Context

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)

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 MLPlayer extract useful features and adapt to the low-dimensionality of the function class, whereas the attention layer performs in-context function approximation on top of the learned features.

1.2 Related Works

Theory of in-context learning. Many recent works studied the in-context learning ability of transformers trained by gradient descent. [ZFB23, ACDS23, MHM23, WZC + 23, ZWB24] studied

the training of linear transformer models to learn linear target functions in-context by implementing one gradient descent step in the forward pass. Similar optimization results have been established for looped linear transformers [GSR + 24], transformers with SoftMax attention [HCL23, NDL24, CSWY24a, CSWY24b] or nonlinear MLP layer [KS24, LWL + 24]. Our problem setting resembles [KS24], where a nonlinear MLP block is used to extract features, followed by a linear attention layer; the main difference is that we establish end-to-end guarantees on the optimization and sample complexity for learning a concrete nonlinear function class, whereas [KS24] focused on convergence of optimization. [CCS23] showed that transformers can learn nonlinear functions in-context via a functional gradient mechanism, but no statistical and optimization guarantees were given. If we do not take gradient-based optimization into account, the function class that can be implemented incontext by transformers has been characterized in many prior works [BCW + 23, GHM + 23, ZZYW23, JLLVR24, SGS + 24, KNS24]. These results typically aim to encode specific algorithms (LASSO, gradient descent, etc.) in the forward pass or directly analyze the Bayes-optimal estimator.

Gradient-based learning of low-dimensional functions. The complexity of learning lowdimensional functions with neural network has been extensively studied in the feature learning theory literature. Typical target functions include single-index models [BAGJ21, BES + 22, BBSS22, MHPG + 23, DNGL23, BES + 23] and multi-index models [DLS22, AAM22, AAM23, BBPV23]. While a shallow neural network can efficiently approximate such low-dimensional functions, the efficiency of gradient-based training is governed by properties of the nonlinearity σ ∗ . In the single-index setting, prior works established a sufficient sample size n ≳ d Θ( K ) , where K ∈ N is the information exponent for algorithms utilizing correlational information [BAGJ21, BBSS22, DNGL23, MHWSE23], or the generative exponent for algorithms that employ suitable label transformations [DPVLB24, LOSW24, ADK + 24, JMS24]. Moreover, n ≍ d samples are information theoretically necessary without additional structural assumptions; this entails that estimators that only access the test prompt inevitably pay a sample size that scales with the ambient dimensionality. As we will see, pretrained transformers can avoid this 'curse of dimensionality' by exploiting the low-dimensionality of the task distribution. Low-dimensional structure of the function class similar to our setting has been assumed to study the efficiency of transfer learning [DLS22] and multi-task learning [CHS + 23].

2 Problem Setting

Notations. ∥ · ∥ denotes the ℓ 2 norm for vectors and the ℓ 2 → ℓ 2 operator norm for matrices. For a vector w , we use w a : b for a ≤ b to denote the vector [ w a , w a +1 , . . . , w b ] ⊤ . 1 N denotes the all-one vector of size N . The indicator function of A is denoted by I A . Let N be a nonnegative integer; then [ N ] denotes the set { n ∈ Z | 1 ≤ n ≤ N } . For a nonnegative integer i , the i -th Hermite polynomial is defined as He i ( z ) = ( -1) i e z 2 2 d i d z i e -z 2 2 . 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 . ˜ O ( · ) , ˜ Ω( · ) 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 .

2.1 Data Generating Process

2.1.1 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 , where x i , x ∈ R d and y i ∈ R . The labeled examples 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 as

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 t 1 , . . . , x t M , y t M , x t , y t ) } T t =1 , which is composed of prompts { ( x t 1 , y t 1 , . . . , x t M , y t M , x t ) } T t =1 and responses { y t } T t =1 for queries { x t } T t =1 , where y t = f t ∗ ( x t ) + ς t and ς t i . i . d . ∼ D ς .

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 as

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

2.1.2 Gaussian single-index models

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 as

where ς 1 , . . . , ς N i . i . d . ∼ Unif( {-τ, τ } ) , and the distribution D f ∗ of the true function f ∗ is specified as:

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

Link Function. σ ∗ ( z ) = ∑ P i = Q c i i ! He i ( z ) , where 2 ≤ Q ≤ P . We draw the Hermite coefficients { c i } P i = Q from any distribution satisfying

̸

̸

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 rankr 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, we may 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].

2.2 Student Model: Transformer with Nonlinear MLP Layer

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 d e × 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 as

where ρ is the temperature, and W K , W Q ∈ R d k × d e , W V ∈ R d v × d e and W P ∈ R d e × d v are the key, query, value, and projection matrix, respectively. Following prior theoretical works [ZFB23, ACDS23, MHM23, KNS24], we remove the SoftMax and instead analyze the linear attention with ρ = N ; it has been argued that such simplification can reproduce phenomena in practical transformer training [ACS + 23]. We further simplify the original self-attention module (2.3) by merging W P W V as W PV ∈ R d e × d e and ( W K ) ⊤ W Q as W KQ ∈ R d e × d e , and consider the following parameterization also introduced in [ZFB23, ACDS23, WZC + 23],

where v ∈ R and K ∈ R ( d e -1) × ( d e -1) . Then, the simplified attention module is written as ˜ f Attn ( E ; W PV , W KQ ) = E + W PV 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:

where w 1 , . . . , w m ∈ R d and b 1 , . . . , b m ∈ R are trainable parameters and σ : R → R is a nonlinear activation function; we use σ ( z ) = ReLU( z ) = max { z, 0 } . This is to say, we define the embedding as the 'hidden representation' σ ( w ⊤ x + b ) of a widthm two-layer neural network. For concise notation we write W = ( w 1 , . . . , w m ) ⊤ , b = ( b 1 , . . . , b m ) ⊤ .

Remark 2. We make the following remarks on the considered architecture.

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 ˜ f Attn and MLP embedding (2.5), we can express the model prediction y as

where Γ = v K ⊤ and σ ( WX + b 1 ⊤ 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 .

2.3 Pretraining: Empirical Risk Minimization via Gradient Descent

We pretrain the transformer (2.6) by the gradient-based learning algorithm specified in Algorithm 1, which is inspired by the layer-wise training procedure studied in the neural network theory literature [AAM22, DLS22, BES + 22]. In particular, we update the trainable parameters in a sequential manner.

Algorithm 1: Gradient-based training of transformer with MLP layer

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 ]) .

2 Stage I: Gradient descent for MLP layer

5 Initialize b j ∼ Unif([ -log d, log d ]) .

6 Stage II: Empirical risk minimization for attention layer

7 Draw data { ( x t 1 , y t 1 , . . . , x t N 2 , y t N 2 , x t , y t ) } T 1 + T 2 t = T 1 +1 with prompt length N 2 .

8

∗

Γ

argmin

Γ

1

T

2

T

t

1

+

=

T

T

1

2

+1

(

y

-

f

(

X

,

y

,

x

;

W

,

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

3.1 Main Theorem

Our main theorem characterizes the pretraining and in-context sample complexity of the transformer with MLP layer (2.6) optimized by layer-wise gradient-based pretraining outlined in Algorithm 1.

Theorem 1. Given Assumption 1 and 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 2 Q +1 r ) , γ ≍ 1 m 3 2 r 1 2 d Q , η 1 ≍ m 3 2 rd 2 Q -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 as

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

∑

t

t

t

t

(1)

Γ

,

b

))

2

+

λ

2

2

∥

Γ

∥

2

F

.

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.

Comparison against baseline methods. Below we summarize the statistical complexity of commonly-used estimators that only have access to N ∗ in-context examples, but not the pretraining data. Note that the in-context sample complexity all depends on the ambient dimensionality d ≫ r .

Kernel models. Recall that our target function is a degreeP polynomial, and hence kernel ridge regression requires N ∗ ≳ d P in-context examples [GMMM21, DWY21].

CSQ learners. The correlational statistical query (CSQ) lower bound suggests that an algorithm making use of correlational information requires N ∗ ≳ d Θ( Q ) samples to learn a single-index model with information exponent Q [DLS22, AAM22]. This sample complexity can be achieved by online SGD training of shallow neural network [BAGJ21, DNGL23].

Information theoretic limit. Since the single-index model (1.1) contains d unknown parameters, we can infer an information theoretic lower bound of N ∗ ≳ d samples to estimate this function. This complexity can be achieved (up to polylog factors) by tailored SQ algorithms [CM20, DPVLB24] or modified gradient-base training of neural network [LOSW24, ADK + 24, JMS24].

3.2 Proof Sketch of Main Theorem

We provide a sketch of derivation for Theorem 1. The essential mechanism is outlined as follows: after training the 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.

3.2.1 Training the MLP Layer

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 have

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 σ b j ( 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.

3.2.2 Attention Matrix with Good Approximation Property

Next we construct the attention matrix ¯ Γ which satisfies the following approximation property:

Proposition 2 (Informal) . There exists ¯ Γ such that with high probability,

for all t ∈ { T 1 +1 , . . . , T 2 } . Moreover, we have ∥ ¯ Γ ∥ F = ˜ O ( r 2 P /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 c i 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 in H = { ∏ r j =1 He p j ( ⟨ β 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 as

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 .

3.2.3 Generalization Error Analysis

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

4.1 Experimental Setting

We pretrain a GPT-2 model [RWC + 19] to learn the Gaussian single-index function class (1.1). Specifically, we consider the 12-layer architecture (with 22.3M parameters) used in [GTLV22] for in-context linear regression. The pretraining data is generated from random single-index 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 ) and y 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.

4.2 Empirical Findings

Ambient dimension-free sample complexity. In Figure 2 we examine how the in-context sample complexity of the GPT-2 model depends on the ambient dimensionality d and the function class dimensionality 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(a) we observe that for fixed r = 8 , varying the ambient dimensionality d = 16 , 32 , 64 leads to negligible change in the model performance for the in-context phase. In contrast, Figure 2(b) illustrates that for fixed d , the required sample size N ∗ scales with the dimensionality of the function class r = 2 , 4 , 8 . This confirms our theoretical finding that transformers can adapt to low-dimensional structure of the distribution of target functions via gradient-based pretraining.

0

50

100

150

200

250

in-context sample size

N

*

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

prediction risk

d=16,r=8

d=32,r=8

d=64,r=8

(a) ICL risk - varying

d

, fixed

r

.

0

50

100

150

200

250

in-context sample size

N

*

0.0

0.2

0.4

0.6

0.8

1.0

1.2

prediction risk

d=64,r=2

d=64,r=4

d=64,r=8

(b) ICL risk - fixed

d

, varying

r

.

Superiority over baseline algorithms. 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 activation f 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.

5 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].

References

[AAM22] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The mergedstaircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory , pages 4782-4887. PMLR, 2022.

[AAM23] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. In Conference on Learning Theory , pages 2552-2623. PMLR, 2023.

[ACDS23] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems , 36, 2023.

[ACS + 23] Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv preprint arXiv:2310.01082 , 2023.

[ADK + 24] Luca Arnaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iuvant: Data repetition allows sgd to learn high-dimensional multi-index functions. arXiv preprint arXiv:2405.15459 , 2024.

[ASA + 22] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661 , 2022.

[BAGJ21] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research , 22(1):4788-4838, 2021.

[BBPV23] Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivien. On learning Gaussian multiindex models with gradient flow. arXiv preprint arXiv:2310.19793 , 2023.

[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems , 35, 2022.

[BCW + 23] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in Neural Information Processing Systems , 36, 2023.

[BEG + 22] Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems , 35, 2022.

[BES + 22] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems , 35, 2022.

[BES + 23] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu. Learning in the presence of low-dimensional structure: A spiked random matrix perspective. Advances in Neural Information Processing Systems , 36, 2023.

[BKM + 19] Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions in high-dimensional generalized linear models. Proceedings of the National Academy of Sciences , 116(12):5451-5460, 2019.

[BL19] Yu Bai and Jason D Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. arXiv preprint arXiv:1910.01619 , 2019.

[BMR + 20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems , 33, 2020.

[BMZ23] Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.00055 , 2023.

[CCS23] Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.06528 , 2023.

[CHS + 23] Liam Collins, Hamed Hassani, Mahdi Soltanolkotabi, Aryan Mokhtari, and Sanjay Shakkottai. Provable multi-task representation learning by two-layer relu neural networks. arXiv preprint arXiv:2307.06887 , 2023.

[CM20] Sitan Chen and Raghu Meka. Learning polynomials in few relevant dimensions. In Conference on Learning Theory , pages 1161-1227. PMLR, 2020.

[CSWY24a] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.19442 , 2024.

[CSWY24b] Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. arXiv preprint arXiv:2409.10559 , 2024.

[DH18] Rishabh Dudeja and Daniel Hsu. Learning single-index models in gaussian space. In Conference On Learning Theory , pages 1887-1930. PMLR, 2018.

[DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory , pages 5413-5452. PMLR, 2022.

[DNGL23] Alex Damian, Eshaan Nichani, Rong Ge, and Jason D. Lee. Smoothing the landscape boosts the signal for SGD: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems , 36, 2023.

[DPVLB24] Alex Damian, Loucas Pillaud-Vivien, Jason D Lee, and Joan Bruna. The computational complexity of learning gaussian single-index models. arXiv preprint arXiv:2403.05529 , 2024.

[DSD + 22] Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers. arXiv preprint arXiv:2212.10559 , 2022.

[DTA + 24] Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborová, and Florent Krzakala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. arXiv preprint arXiv:2402.03220 , 2024.

[DWY21] Konstantin Donhauser, Mingqi Wu, and Fanny Yang. How rotational invariance of common kernels prevents generalization in high dimensions. In International Conference on Machine Learning , pages 2804-2814. PMLR, 2021.

[GHM + 23] Tianyu Guo, Wei Hu, Song Mei, Huan Wang, Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? a case study on learning with representations. arXiv preprint arXiv:2310.10616 , 2023.

[GMMM21] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Linearized two-layers neural networks in high dimension. The Annals of Statistics , 49(2):1029-1054, 2021.

[GSR + 24] Khashayar Gatmiry, Nikunj Saunshi, Sashank J Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi-step gradient descent for in-context learning? In International Conference on Machine Learning , 2024.

[GSS21] Friedrich Götze, Holger Sambale, and Arthur Sinulis. Concentration inequalities for polynomials in α -sub-exponential random variables. Electronic Journal of Probability , 26:1 - 22, 2021.

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

[LWL + 24] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. How do nonlinear transformers learn and generalize in in-context learning? In International Conference on Machine Learning , 2024.

[Mau16] Andreas Maurer. A vector-contraction inequality for rademacher complexities. In Algorithmic Learning Theory: 27th International Conference, ALT 2016, Bari, Italy, October 19-21, 2016, Proceedings 27 , pages 3-17. Springer, 2016.

[MHM23] Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576 , 2023.

[MHPG + 23] Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with SGD. In International Conference on Learning Representations , 2023.

[MHWSE23] Alireza Mousavi-Hosseini, Denny Wu, Taiji Suzuki, and Murat A Erdogdu. Gradientbased feature learning under structured data. Advances in Neural Information Processing Systems , 36, 2023.

[MM18] Marco Mondelli and Andrea Montanari. Fundamental limits of weak recovery with applications to phase retrieval. In Conference On Learning Theory , pages 1445-1450. PMLR, 2018.

[MZD + 23] Arvind Mahankali, Haochen Zhang, Kefan Dong, Margalit Glasgow, and Tengyu Ma. Beyond ntk with vanilla gradient descent: A mean-field analysis of neural networks with polynomial width, samples, and time. Advances in Neural Information Processing Systems , 36, 2023.

[NDL24] Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735 , 2024.

[OSSW24] Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Learning sum of diverse features: computational hardness and efficient gradient-based training for ridge combinations. arXiv preprint arXiv:2406.11828 , 2024.

[RPCG23] Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in Neural Information Processing Systems , 36, 2023.

[RWC + 19] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

[SGS + 24] Michael E Sander, Raja Giryes, Taiji Suzuki, Mathieu Blondel, and Gabriel Peyré. How do transformers perform in-context autoregressive learning? arXiv preprint arXiv:2402.05787 , 2024.

[Ver18] Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science . Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.

[VONR + 23] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174. PMLR, 2023.

[VSP + 17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017.

[Wai19] Martin J Wainwright. High-Dimensional Statistics: A Non-Asymptotic Viewpoint . Cambridge University Press, 2019.

[WZC + 23] Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv:2310.08391 , 2023.

[ZFB23] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927 , 2023.

[ZWB24] Ruiqi Zhang, Jingfeng Wu, and Peter L Bartlett. In-context learning of a linear transformer block: Benefits of the mlp component and one-step gd initialization. arXiv preprint arXiv:2402.14951 , 2024.

[ZZYW23] Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305.19420 , 2023.

Table of Contents

A Preliminaries

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 and

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 get P ( | x | ≥ √ 2 C ∗ 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 where M = O (poly( d )) , then A 1 ∩ · · · ∩ A M also occurs 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.

A.2 Tensor Notations

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 i 1 ,...,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

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 i 1 · · · v i k .

Let f ( x ) : R d → R be a d -variable differentiable function. A k -tensor ∇ k f ( x ) is defined as

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.

A.3 Hermite Polynomials

We frequently use (probablists') Hermite polynomials, which is defined as

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 a i 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 ) = ∑ i 1 ≥ 0 ,...,i d ≥ 0 a i 1 ,...,i d ( i 1 )! ··· ( i d )! He i 1 ( z 1 ) · · · He i d ( z d ) , where a i 1 ,...,i d = E z 1 ,...,z d ∼N (0 , 1) [ f ( z )He i 1 ( z 1 ) · · · He i d ( z d )] . The coefficient a i 1 ,...,i d can also be obtained by a i 1 ,...,i d = E z 1 ,...,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) ,

We can bound E x [ f ∗ ( x ) 4 ] = E z ∼N (0 , 1) , { c i } [ ( ∑ P i = Q c i i ! He i ( z ) ) 4 ] naively as follows: let ζ 4 P = max Q ≤ i ≤ P E z ∼ N (0 , 1) [He i ( z ) 4 ] . ζ P is an O (1) quantity depending only on P and Q . Then,

The lemma below is useful to find a basis of the set of single-index functions.

Lemma 8. Suppose β ∈ S ( S ) . Then,

holds.

Proof. Note that E x 1 ,...,x r ∼N (0 , 1) [ ∂ i 1 ∂x i 1 1 · · · ∂ i r ∂x i r r He p ( ⟨ x , β ⟩ ) ] is nonzero only when i 1 + · · · + i r = p , because ∂ i 1 ∂x i 1 1 · · · ∂ i r ∂x i r r He p ( ⟨ x , β ⟩ ) = p ( p -1) · · · ( p -( i 1 + · · · + i r ) + 1) β i 1 1 · · · β i r r He p -( i 1 + ··· + i r ) ( ⟨ x , β ⟩ ) and E x 1 ,...,x r ∼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 x 1 ,...,x r ∼N (0 , 1) [ ∂ i 1 ∂x i 1 1 · · · ∂ i r ∂x i r r He p ( ⟨ x , β ⟩ ) ] = p ! β i 1 1 · · · β i r r holds. Thus, from the multivariate Hermite expansion we obtain the claim.

Corollary 9. Let f ∗ ( x ) = ∑ P i = Q c i i ! He i ( ⟨ x , β ⟩ ) . Then, E x [ ∇ k f ∗ ( x )] = c k β ⊗ k if Q ≤ k ≤ P and otherwise it is the zero tensor.

A.4 Other Auxiliary Lemmas

Lemmas on random vectors. Let χ ( d ) be the distribution of √ z 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 = z w 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 ,

holds.

Lemma 11. Let C > 0 . For w ∼ Unif( S d -1 ) , ∥ w 1: r ∥ 2 ≤ 4 Cr d log d holds with probability at least 1 -2 exp( -d/ 32) -2 rd -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 ≤ 2 Cr 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 2 i ≤ 2 C log d holds with probability at least 1 -2 d -C , and thus ( z ′ ) 2 ≤ 2 Cr log d holds with probability at least 1 -2 rd -C . Taking the union bound yields the result.

Lemma 12. For k ≥ 1 and r ≤ d ,

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 ∥ 2 k ] ≲ 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 ] = 3 r + 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 ∥ = ˜ O ( σ √ 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 degreeD n -variable polynomial function. Then, for all t ≥ 0 it holds that

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 c i 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 that

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 c holds. Lemma 15 yields the result.

Corollary 17. Fix any { c i } and β on the support of their distributions. Then, | f ∗ ( x ) | , | y | ≲ (log d ) P/ 2 holds with high probability over the distribution of x and ς .

B Proofs for the MLP Layer

This section analyzes the behavior of MLP weight w after one gradient descent step (line 4 in Algorithm 1), as sketched in Section 3.2.1.

We set λ 1 = η -1 1 . From line 4 of Algorithm 1, for each j ∈ [ m ] ,

holds where σ b j ( 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.

Proof. Recall that

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] ⊤ and x t i ∼ N (0 , I d ) . Therefore, 〈 w (0) j , x t i 〉 ∼ N (0 , 1) . From Lemma 4, 〈 w (0) j , x t i 〉 2 ≲ log d holds with high probability. Moreover, from Lemma 10 and Corollary 17, we know that | y t i | , | y t | ≲ (log d ) P/ 2

and ∥ x t i ∥ , ∥ x t ∥ ≲ √ d holds. 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

and

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

and g ( w , b ) = E [ g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 )] , where the expectation is taken with respect to the data { ( X t , y t , x t , y t ) } T 1 t =1 . Note that

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 T 1 ,N 1 around g .

B.1 Calculation of Population Gradient

B.1.1 Hermite Expansion

First, note that

Here E f ∗ denotes the expectation with respect to the distribution of the true function f ∗ (i.e., the distribution of { c i } P i = Q and β ) specified in Assumption 1. Now let σ b ( z ) = ∑ i ≥ 0 a i ( b ) i ! He i ( z ) be the Hermite expansion of student activation σ b ( z ) = ReLU( z + b ) .

Lemma 19. It holds that

Proof. Consider a function g ( w ⊤ x ) = ∑ i ≥ 0 ˜ a i i ! He i ( w ⊤ x ) . Here E x [ f ∗ ( x ) g ( w ⊤ x )] can be calculated as

Where the final line is obtained noting that the tensor ∇ i f ∗ ( x ) has i ! i 1 ! ··· 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 c j j ! He j ( ⟨ x , β ⟩ ) ] is a nonzero tensor only when Q ≤ i ≤ P , from Lemma 8. We can also show (B.3) as

Using Lemma 19, g ( w , b ) can be expanded as

where

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.

Lemma 20. There are constants C H and C L depending only on Q and P satisfying the following condition. Let σ b ( z ) = σ ( z + b ) and b ∼ Unif([ -1 , 1]) . Then,

holds with probability 1/2. Especially, ∣ ∣ E z ∼N (0 , 1) [ σ b ( z )He i ( z )] ∣ ∣ ≤ C H (2 ≤ i ≤ P +2) holds with probability 1.

Proof. From Lemma 15 in [BES + 23], we know that for i ≥ 2 , ∣ ∣ E z ∼N (0 , 1) [ σ b ( z )He i ( z )] ∣ ∣ = e -b 2 / 2 √ 2 π | He i -2 ( b ) | holds. First, as b ∼ Unif([ -1 , 1]) , by setting sufficiently large C H , | He i ( b ) | ≤ C H holds for all 0 ≤ i ≤ P . Therefore, as ∣ ∣ E z ∼N (0 , 1) [ σ b ( z )He i ( z )] ∣ ∣ ≤ | He i -2 ( b ) | , ∣ ∣ E z ∼N (0 , 1) [ σ b ( z )He i ( z )] ∣ ∣ ≤ C H holds for all 2 ≤ i ≤ P +2 with probability one.

From continuity we know that there exists C 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.

B.1.2 Calculation of the Main Term

Let us calculate the main term in (B.4) explicitly. Lemma 21.

holds.

Proof. From Corollary 9, it holds that E x [ ∇ Q f ∗ ( x )] = c Q ( β ) ⊗ Q . Then, we obtain

Next, let us show

First, note that by letting β ′ ∼ N (0 , Σ β ) where

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,

It suffices to show that E v ∼N (0 ,I r ) [ v ⊗ 2 Q ]( z ⊗ (2 Q -1) ) = (2 Q -1)!! z ∥ z ∥ 2 Q -2 for z ∈ R r . As an example, we show that the first entry of E v ∼N (0 ,I r ) [ v ⊗ 2 Q ]( z ⊗ (2 Q -1) ) equals to (2 Q -1)!! z 1 ∥ z ∥ 2 Q -2 : the other components can be calculated similarly. First, we expand ( E v ∼N (0 ,I r ) [ v ⊗ 2 Q ]( z ⊗ (2 Q -1) )) 1 as

̸

Note that E [ v 1 v i 1 · · · v i 2 Q -1 ] = 0 if and only if the degree of v 1 v i 1 · · · v i 2 Q -1 with respect to v j is even, for all j ∈ [ r ] . Then, we obtain

where c q 1 ,...,q r = (2 Q -1)! (2 q 1 -1)!(2 q 2 )! ··· (2 q r )! is the number of ways to assign (2 q 1 -1) '1's, (2 q 2 ) '2's, ..., among i 1 to i 2 Q -1 .

Note that E v ∼N (0 , 1) [ v 2 q ] = (2 q -1)!! . Therefore, we obtain

which concludes the assertion.

By Lemma 21, we observe that the main term of the expected one-step gradient aligns to the r -dimensional subspace S . In other words, the MLP layer captures the information of S via pretraining, which will be useful in the later stage.

B.1.3 Bounding Residual Terms

First, similarly to the Lemma 21, we can obtain the explicit formulation of the entire g ( w , b ) . Lemma 22.

holds.

Proof. Note that

and

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.

,

Lemma 23. Let j = 4 k where k ∈ [ P ] . then,

holds uniformly over all b ∈ [ -1 , 1] .

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 .

̸

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 that E [ 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 ∥ 4 k ] = O (( r/d ) 2 k ) from Lemma 12. Moreover, it is known that E z ∼ χ r [ z 2 l ] = Θ( 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.

holds 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 ).

B.2 Concentration of Empirical Gradient

Now we control the deviation of g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 ) from g ( w , b ) . Lemma 25. Under Assumption 1,

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 term a Q ( b ) E x [ ∇ Q f ∗ ( x )]( w ⊗ ( Q -1) ) ( Q -1)! in (B.3) of Lemma 19 to show (B.8). From Lemmas 5 and 6, we have

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 that ∥ ∥ ∥ a Q ( b ) E x [ ∇ Q f ∗ ( x )]( w ⊗ ( Q -1) ) ( Q -1)! ∥ ∥ ∥ 2 ≤ ( C H R c ( Q -1)! ) 2 ( 4 Cr d log d ) Q -1 with probability at least 1 -2 exp( -d/ 32) -2 rd -C . (B.9) and (B.10) can be obtained similarly.

The following lemma is parallel to Lemma 19 in [DLS22].

Lemma 26. Fix any t ∈ [ T 1 ] , b ∈ [ -1 , 1] and w ∈ S d -1 satisfying the conditions (B.8) to (B.10) . Then, with high probability over the distribution of ( X t , y t , x t , y t ) , it holds that

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,

Note that f t ∗ ( x t i ) I f t ∗ ( x t i ) ≤ R σ ′ b ( w ⊤ x t i ) ⟨ x t i , u ⟩ -E x [ f t ∗ ( x ) I f t ∗ ( x ) ≤ R σ ′ b ( w ⊤ x ) ⟨ x , u ⟩ ] is R 2 -sub Gaussian for any ∥ u ∥ = 1 from Lemma 14. Hence the first term can be upper bounded by ˜ O ( R √ d N 1 ) with high probability from Lemma 13. For the second term, we observe that

for 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 ∥ 1 N 1 ∑ N 1 i =1 ςσ ′ b ( w ⊤ x t i ) x t i ∥ = ˜ O ( √ d/N 1 ) with high probability.

By similar procedure, we can obtain a bound of ∥ g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 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 ) } T 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 ) } T 1 t =1 . Let

and

Note that g T 1 ,N 1 = z 1 + z 2 holds.

Let us consider z 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 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 N 1 ∑ N 1 i =1 y t i σ ′ b ( w ⊤ x t i ) x t i ∥ ≤

Here, we notice that the procedure in the proof of Lemma 26 can be applied, yielding

with high probability. Similarly we can obtain

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 = ⋃ T 1 t =1 { x t 1 , . . . , x t N 1 , 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 ) as follows: b and b ′ belong to the same B i if and only if sign( ⟨ w , x ⟩ + b ) = sign( ⟨ w , x ⟩ + b ′ ) for all x ∈ X .

It is clear that |B ( w , X ) | ≲ | X | ≤ ( N 1 +1) T 1 holds because it suffices to divide [ -1 , 1] at the point satisfying ⟨ w , x ⟩ + b = 0 for each x ∈ 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 ) } T 1 t =1 , the following holds:

g ( w , b ) is L 1 -Lipschitz continuous with respect to b in the interval [ -1 , 1] .

Let X = ⋃ T 1 t =1 { x t 1 , . . . , x t N 1 , x t } . Then, g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 ) is L 2 -Lipschitz continuous with respect to b in each interval B i ∈ B ( w , X ) .

Here L 1 = O (1) and L 2 = ˜ O ( √ d ) do not depend on w .

Proof. By the explicit form of g ( w , b ) (B.6), we notice that each term is made by taking the product between the term as a j ( b ) a k ( b ) and a vector independent of b whose norm is O (1) . Thus, it suffices to show that a j ( b ) a k ( b ) is O (1) -Lipschitz. This follows from the fact a j = ± e -b 2 / 2 √ 2 π He j -2 ( b ) (see the proof of Lemma 20) and He j -2 ( b )e -b 2 / 2 is O (1) -Lipschitz continuous on the bounded set [ -1 , 1] . Next, let us see g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 ) . Recall that

From Lemma 10 and Corollary 17, we know that | y t i | , | y t | ≲ (log d ) P/ 2 and ∥ x t i ∥ , ∥ x t ∥ ≲ d holds for all i, t with high probability. Assuming this and noting that σ is 1 -Lipschitz, in each interval in B ( w , X ) , g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 ) is L 2 = ˜ O ( √ d ) -Lipschitz.

Corollary 30. With high probability over the distribution of w and { ( X t , y t , x t , y t ) } T 1 t =1 ,

holds.

Proof. Let X = ⋃ T 1 t =1 { x t 1 , . . . , x t N 1 , x t } and consider B ( w , X ) = { B 1 , . . . , B N ( w ,X ) } . Then, let B be the union of ˜ O ((√ r Q -1 d Q -1 √ r T 1 + √ r 2 N 1 T 1 ) / √ d ) -coverings of B i ∈ B ( w , X ) . Since | B | is polynomial in d , Lemma 27 holds with high probability uniformly over b ∈ B . Moreover, from the construction and Lemma 29, for b ∈ [ -1 , 1] \ B , there exists π ( b ) ∈ B such that

holds.

Lemma 31. With high probability over the distribution of { ( X t , y t , x t , y t ) } T 1 t =1 ,

It holds that

with high probability over w ∼ Unif( S d -1 ) .

It holds that

for

j

∈

[4

P

]

.

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 ) } T 1 t =1 ) , then we can say that

with probability at least 1 -d C/ 2 over { ( X t , y t , x t , y t ) } T 1 t =1 , P w { A | { ( X t , y t , x t , y t ) } T 1 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 ) } T 1 t =1 , P w { (B.11) | { ( X t , y t , x t , y t ) } T 1 t =1 } ≥ 1 -d -˜ C is 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 ) } T 1 t =1 , we have ∥ x t i ∥ , ∥ x t ∥ = O ( √ d ) and | y t | = ˜ O (1) . From the definition of g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 ) (B.2) and expansion of g ( w , b ) (B.6), sup w ∥ g ( w , b ) -g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 ) ∥ ≤ sup w ∥ g ( w , b ) ∥ + sup w ∥ g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 ) ∥ is upper bounded by ˜ O ( d ) . Therefore, with high probability

By setting sufficiently large ˜ C , we arrive at the claim.

Summary. We combine the obtained results and show that ∥ g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 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 ) } T 1 t =1 ,

(I) g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 ) = m ( w , b ) + r ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 ) holds

holds with high probability over w ∼ Unif( S d -1 ) . Moreover,

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 ) } T 2 t = T 1 +1 , where z ∈ ⋃ T 2 t = T 1 +1 { x t 1 , . . . , x t N 2 , x t } .

Proof. Combining Lemmas 23 and 31 yields (I). For (II), by the condition for T 1 and N 1 , we obtain

from Lemmas 24 and 31. Then, as η 1 γ ≲ √ rd 2 Q -1 /C 1 (log d ) C 2 , sup b ∈ [ -1 , 1] ∥ 2 η 1 γ g T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 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 T 1 ,N 1 ( w , b, { ( X t , y t , x t , y t ) } T 1 t =1 , z ) 〉∣ ∣ ∣ ≲ log d for z ∈ ⋃ T 2 t = T 1 +1 { x t 1 , . . . , x t N 2 , 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.

C Construction of Attention Matrix

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 as

We let B P := |H| and introduce a numbering H = { h 1 , . . . , h B P } . Note that

C.1 Approximation Error of Two-layer Neural Network

As sketched in Section 3.2.2, we need to show that there is a set of two-layer neural networks which can approximate basis functions well. The goal of this section is to show the following proposition.

Proposition 33. Assume T 1 = ˜ Ω( d Q +1 r Q ) and N 1 T 1 = ˜ Ω( d 2 Q +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 2 Q -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 ] ,

holds. Moreover, sup n ∈ [ B P ] ∥ a n ∥ 2 = ˜ O ( 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.

C.1.1 Alignment to S and Efficient Approximation

From now on, we suppose that g T 1 ,N 1 ( 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 ) } T 1 t =1 ). An important previous result [DLS22] is that, if g T 1 ,N 1 ( 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

and ∥ T n k ∥ F ≲ r P -k 4 .

Definition 35 ( f ( r, d ) -alignment) . We call that g T 1 ,N 1 ( w , b ) has f ( r, d ) -alignment with S , if the following holds:

For any 0 ≤ k ≤ P -symmetric r -dimensional tensor T ,

and the hidden constant in ' ≳ ' needs to be uniformly lower bounded.

Lemma 36 (Corollary 4 in [DLS22], adapted) . Suppose that g T 1 ,N 1 ( 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 that

Intuitively speaking, f ( r, d ) -alignment ensures that the MLP gradient g T 1 ,N 1 ( w , b ) has sufficiently large component on S . If so, the lemma states that, an infinite-width neural network E w [ ψ T ( w , b ) ⟨ g T 1 ,N 1 ( w , b ) , ·⟩ k ] , whose hidden layer weight is g T 1 ,N 1 ( 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 T 1 ,N 1 ( 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 ∥ 2 Q -2 . Then

holds for all 0 ≤ i ≤ k . Here T ( ¯ w ⊗ 0 ) := T .

Proof. Let u ∼ N (0 , I d ) , z ∼ χ ( d ) and ¯ u = u 1: r ∥ u 1: r ∥ 2 Q -2 . We can decompose ¯ u as ¯ u = z 2 Q -1 ¯ w . Therefore

holds. On the other hand, let x ∼ S r -1 and z ′ ∼ χ ( r ) . Then we can decompose ¯ u as ¯ u = ( z ′ ) 2 Q -1 x . Thus

holds. It implies that

Similarly,

is satisfied. Now from Corollary 13 in [DLS22],

holds. Therefore we obtain

Corollary 38. Let T be a k ≤ P -symmetric r -dimensional tensor and let m ( w , b ) 1: r = 2 a Q ( b ) 2 Q !( Q -1)! E c Q [ c 2 Q ](2 Q -1)!! E z ∼ χr [ z 2 Q ] ∥ w 1: r ∥ 2 Q -2 w 1: r . Moreover, assume that (B.5) holds. Then,

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 2 Q +1 r ) . Then,

holds for all 0 ≤ k ≤ P . Therefore, g T 1 ,N 1 ( w , b ) has d -2 Q +1 r -2 -alignment with S .

Proof. Decompose g T 1 ,N 1 ( w , b ) = m ( w , b ) + r ( w , b ) as defined in Corollary 32. From the proof of Lemma 11 in [DLS22], we can show that

for some δ ( w , b ) satisfying

Here, from Corollaries 32 and 38,

holds. Now, from the condition of N 1 and T 1 , and r ≲ √ d , we know

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.

C.1.2 Approximation Power of Infinite-width Neural Network

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 ,

Lemma 41 (Lemma 12 and Corollary 5 in [DLS22], adapted) . There exists ϕ n ( a, w , b 1 , b 2 ) such that if

where

it holds that 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

it holds that

Let ϕ T n k ( a, w , b 1 , b 2 ) = 2 v k ( a,b 2 ) ψ T n k ( w ,b 1 ) (2 η 1 ) k γ k I A ( w ) I B ( b 1 ) where v k and ψ T n k are constructed in Lemma 36 and the events A ( w ) and B ( b 1 ) are defined as

Then,

Note that the absolute value of second term above can be upper bounded by E w [ (1 -I A ( w ) ) ] 1 / 2 E w [ ( ψ T n k ( w , b 1 ) ⟨ g T 1 ,N 1 ( w , b 1 ) , x ⟩ k ) 2 ] 1 / 2 . Recall that A ( w ) occurs with high probability from Corollary 32, and E w [ ( ψ T n k ( w , b 1 ) ⟨ g T 1 ,N 1 ( 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 have

Therefore, (C.2) and (C.3) holds. Noting that η 1 γ ≍ rd 2 Q -1 /C 1 (log d ) C 2 , We obtain

and

from Lemmas 36 and 39. To show the original claim, it suffices to set ϕ n ( a, w , b 1 , b 2 ) = ∑ P k =0 ϕ T n k ( a, w , b 1 , b 2 ) , which yields the upper bounds on the magnitude of ϕ n because ∑ P k =0 r k r P -k 2 ≲ r P and ∑ P k =0 r k r P -k 4 ≲ r P .

C.1.3 Proof of Proposition 33

We 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 ] ,

and sup n ∈ [ B P ] ∥ a n ∥ 2 = ˜ O ( r P m ) with high probability ( b (0) j , b j denotes the biases at the first and second initialization, respectively). Recall that w (1) j = 2 η 1 Γ (0) j,j g T 1 ,N 1 ( w (0) j , b (0) j ) -2 η 1 1 T 1 ∑ T 1 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 ˜ O ( η 1 γ 2 m √ d ) . Similar to the argument in the proof of Lemma 18, the inner product between the second term and x t , x t i ( i ∈ [ N 2 ] , T 1 +1 ≤ t ≤ T 2 ) is also ˜ O ( η 1 γ 2 m √ d ) , with high probability. From the Cauchy-Schwarz inequality and the Lipschitz continuity of σ ,

We obtain the assertion.

C.2 Concentration of Correlation

Next, we evaluate the discrepancy between the empirical correlation 1 N 2 ∑ N 2 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 as

Also note that

where σ ∗ ( z ) = ∑ P i = Q c i i ! He i ( z ) . First we give an upper bound for | 1 N 2 ∑ N 2 j =1 σ ∗ ( ⟨ x j , β ⟩ ) h ( x j ) -E x [ σ ∗ ( ⟨ x , β ⟩ ) h ( x )] | .

Lemma 42. Let D ≤ 2 P 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 that

holds, where C D and M are constants.

Proof. Let F ( x 1 , . . . , x N ) = 1 N 2 ∑ N 2 j =1 σ ∗ ( ⟨ x j , β ⟩ ) h ( x j ) . It is a degreeD polynomial for x 1 , 1 , . . . , x 1 ,r , . . . , x N 2 ,r , which are standard Gaussian variables. Also,

is satisfied. Then, applying Lemma 15 to F

yields the desired result.

Then, our aim is to bound ∥ E [ ∇ k ( σ ∗ ( ⟨ x , β ⟩ ) h ( x ))] ∥ F .

holds where ¯ C k is a constant depending only on k . For other k , ∥ E [ ∇ k ( σ ∗ ( ⟨ x , β ⟩ ) h ( x ))] ∥ F = 0 . Proof. First,

holds. To calculate E [ ∇ k (He i ( ⟨ x , β ⟩ ) h ( x ))] , note that from Lemma 8,

He

i

(

⟨

x

,

β

⟩

)

h

(

x

)

̸

Also, from the discussion in Appendix A.3, E x [ ∂ s 1 ∂x s 1 1 · · · ∂ sr ∂z sr r He i ( ⟨ x , β ⟩ ) h ( x ) ] is equal to the coefficient of (C.4) with respect to 1 s 1 ! ··· s r ! He s 1 ( x 1 ) · · · He s r ( x r ) . Therefore, E [ ∇ k (He i ( ⟨ x , β ⟩ ) h ( x ))] = O if and only if i + ˜ P = k . Then, ∥ ∥ ∥ ∑ P i = Q c i 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 tensor c k -˜ P ( k -˜ P )! E [ ∇ k (He k -˜ P ( ⟨ x , β ⟩ ) h ( x ))] is equal to

where s i = ∑ k k ′ =1 I j k ′ = i ( i ∈ [ r ]) . There are k ! s 1 ! ··· s r ! index sets ( j 1 , . . . , j k ) having the same entry, and thus

We can verify that there exists a constant ¯ C k depending only on P such that k ! ( k -˜ P )! ( ( q 1 + p 1 )! ··· ( q r + p r )! p 1 ! ··· p r ! q 1 ! ··· q r ! ) ≤ ¯ C k holds uniformly; for example, one can naively obtain

Here we used ( q 1 + p 1 ) + · · · +( q r + p r ) = k .

Moreover, ∑ q 1 + ··· + q r = k -˜ P,q i ≥ 0 ( k -˜ P )! q 1 ! ··· q r ! ( β q 1 1 · · · β q r r ) 2 = ∥ β ⊗ k -˜ P ∥ 2 F = 1 holds. Hence (C.5) is upper bounded by c 2 k -˜ P ¯ C k ≤ R 2 c ¯ C k .

Corollary 44. With high probability,

holds.

Proof. Plugging the results above and t = Θ( 1 √ N 2 (log d ) P ) into Lemma 42 yields

with high probability. Moreover, we can easily verify ∣ ∣ ∣ 1 N 2 ∑ N 2 j =1 ς j h ( x j ) ∣ ∣ ∣ = ˜ O ( 1 √ N 2 ) by the method used in the proof of Lemma 26.

C.3 Proof of Proposition 2

We prove the following Proposition 2 using the preparations above.

Proposition 2. Let X t = ( x t 1 , . . . , x t N 2 ) and y t = ( y t 1 , . . . , y t N 2 ) ⊤ 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 that

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 = ˜ O ( r 2 P /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 ) = ∑ 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 , where

{ a n } B P n =1 is constructed in Proposition 33. Then, by setting ¯ Γ = AA ⊤ ,

holds, where ϵ n ( x ) := ∑ m j =1 a n j σ ( w (1) j ⊤ x + b j ) -h n ( x ) satisfits | ϵ n ( x ) | = ˜ O ( √ 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 , and E 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 obtain

Combining this with B P = Θ( r P ) and m = ˜ Ω( r P ) , we arrive at the upper bound. Finally for ∥ ¯ Γ ∥ F , we have

D Generalization Error Analysis and Proof of Theorem 1

Note that after one-step gradient descent, ∥ w (1) j ∥ = ˜ O d,r (1) is ensured with high probability from Lemma 18 and Corollary 32.

D.1 Rademacher Complexity Bound

Let W = W (1) and suppose b 1 , . . . , b m are biases obtained in line 5 of Algorithm 1. Let

be the set of transformers where the norm of the attention matrix is constrained, and let

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.

holds, when | b j | = ˜ O d,r (1) and ∥ w (1) j ∥ = ˜ O d,r (1) for all j ∈ [ m ] .

Proof. First, it holds that

Rad T ( F N,G )

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 ˜ O (1) : in particular, we can easily obtain that E x ∼N (0 , I d ) [ (( w (1) k ) ⊤ x ) 2 ] = ˜ O d,r (1) .

Consequently, the last line is G √ T √ ∑ m k,l =1 E x ∼N (0 , I d ) ,f ∗ [ (2 f ∗ ( x ) 2 )(2(( w (1) k ) ⊤ x ) 2 ) ] + ˜ O d,r (1) , and remains to evaluate E x ∼N (0 , I d ) ,f ∗ [ f ∗ ( x ) 2 (( w (1) k ) ⊤ x ) 2 ] . On this term, we obtain

Hence we can conclude Rad T ( F N,G ) = O ( Gm √ T ) .

Corollary 46. Let

Then, Rad T ( ˜ F N,G ) = ˜ O d,r ( Gm √ T ) holds when | b j | = ˜ O d,r (1) and ∥ w (1) j ∥ = ˜ O d,r (1) for all j ∈ [ m ] .

Proof. Note that

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 as

by Lemma 7.

D.2 Prompt Length-free Generalization Bound

The ICL risk for prompt length N was defined as

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.

. Then,

holds.

Proof. Note that

where (a) holds because a ⊤ Γ b = Γ ◦ ( ab ⊤ ) ≤ ∥ Γ ∥ F ∥ ab ⊤ ∥ F = ∥ Γ ∥ F ∥ a ∥∥ b ∥ . Let us bound E [ ∥ ∥ ∥ ( σ ( WX 1: N + b 1 ⊤ N ) y 1: N N -σ ( WX 1: M + b 1 ⊤ 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 obtain

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 +2 b 2 )(2 f ∗ ( x ) 2 +2 ς 2 ) ] = ˜ O d,r (1) .

Next, we observe that

Putting all things together, we arrive at

D.3 Proof of Theorem 1

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

Then, from Proposition 2,

holds.

We evaluate R N 2 ( f ) -τ where f = f ( X t , y t , x t ; W (1) , Γ ∗ , b ) using the Rademacher complexity.

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 } ,

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 that

is nonnegative because f ( X , y , x , y ) = y is included in ˜ F N, ∥ ¯ Γ ∥ F . Hence from Markov's inequality, (D.1) is ˜ O (√ r 4 P T 2 ) with probability at least 0 . 995 . Hence we obtain

Now we have done the upper bound for R N 2 ( 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.

E Derivation of Simplified Self-attention Module

We derive equation (2.6), following the same line of argument as [ZFB23]. As we assumed that W PV W KQ are in the form as (2.4), the output of the attention layer is written as

Note that we adopt the right-bottom entry ( ˜ f Attn ( E ; W PV , W KQ ) ) m +1 ,N +1 as a prediction for output of the query. By (E.1). Hence we have

Setting Γ = v K ⊤ yields equation (2.6).

F Details of Experiments

F.1 Detailed Experimental Settings

For the experiments in Section 4, test loss was averaged over 128 independent tasks, with 2048 and 128 independent queries generated for 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 × (2 N +2) was constructed as E = [ x 1 , ˜ y 1 , . . . , x N , ˜ y N ] where ˜ y i = [ y i , 0 , . . . , 0] ⊤ . This embedding was mapped to ˜ E ∈ R 256 × (2 N +2) through the read-in layer. The transformed embedding ˜ E was then passed through the 12-layer GPT-2 backbone with 8 heads, with dropout disabled. Finally, the output was converted into a 2 N +2 -dimensional vector through the read-out layer, where the (2 k -1) -th entry ( k ∈ [ N +1]) corresponds to the prediction ˆ y 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 as 1 B ∑ B t =1 ∑ N +1 k =1 ( y t k -ˆ y 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.

F.2 Experiment on Our Architecture and Algorithm 1

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

10 2

in-context sample size

N

*

10

1

10 0

prediction risk

two-layer NN

pretrained TF

the test prompt. In Figure 3 we see that our simplified transformer also outperforms neural network and kernel method, especially in short context regime.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: They are presented accurately.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We mentioned the limitation of our result and the future challenge in Section 5.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We presented full-proofs for our results in the appendix, and sketched the proof in Section 3.2.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: This is a theory paper where the toy experiments are not central to our contribution.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer:[Yes]

Justification: See Section 4 and Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: See Section 4 and Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

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.

8. 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).

9. 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 ?

Answer: [Yes]

Justification: Yes, and we checked the code.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special 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).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our results are purely theoretical working on a small simplified model, thus there seems no societal impact.

Guidelines:

The answer NA means that there 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).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, 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.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We don't release any assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our paper does not include such subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our paper does not include such subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. The topic of this work is both interesting and important.\n\n2. The analysis provided is rigorous and sound. I think the low-rank result is very encouraging that the risk bound depend on r not d, which could be very helpful for explaining many success of LLMs.\n\n3. The presentation is clear and easy to read.\n\nWEAKNESSES:\n1. Although the ICL capabilities of transformers are indeed powerful and of great interest, the model and pre-training algorithm presented in this paper are not very realistic. While it is common to use simplified models for analyzing LLMs, the authors' pre-training algorithm involves setups (such as gradient descent and (re-)initialization) that deviate from standard LLM training approaches.\n\n2. Given the unrealistic setup, the paper would benefit greatly from numerical experiments demonstrating the effectiveness of the proposed algorithm in learning nonlinear functions. If obtaining these results is challenging, conducting other experiments that could provide intuition or insights into understanding ICL in standard LLMs would also be helpful.\n\n3. While the work seems very relevant to understanding ICLs, the paper lacks real evidence supporting this connection. If I am misunderstanding the relevance, I would appreciate it if the authors could provide more detailed explanations on how this work contributes to understanding the mechanisms of ICL and LLMs.\n\nQUESTIONS:\nPlease see part 3 in the last section.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. ICL on nonlinear functions is a significant problem to solve. \n2. The proof is solid and reasonable to me.\n3. Characterizing the ICL performance as a function of the internal complexity of the nonlinearity of the target function is a novel and nice result.\n\nWEAKNESSES:\n1. The theoretical result seems not to be tight. When $r$ is very close to $d$, the sample complexity is much worse than using two-layer neural networks. However, the Transformer considered in this paper already includes a nonlinear and trainable MLP layer. Then, I think the derived bound is not tight. In other words, why is the bound worse than two-layer NN?\n\n2. No experiments are provided to justify the result, especially for the comparison with kernel methods and two-layer NN.\n\n3. The presentation of the proof sketch can be improved. It is better to emphasize that finding $\\Gamma$ is solving a convex problem (mentioned in line 206) between lines 262 and 264. I think the convexity is the key to the reason why the ICL risk of the constructed $\\Gamma$ is attainable by gradient descent. \n\n4. Some references on ICL optimization and generalization are missing. Li et. al., ICML 2024. \"How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?\" (This work considers softmax attention and nonlinear MLP, which could be mentioned in line 93); Gatmiry et al., ICML 2024. \"Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?\" \nThere is also a concurrent work: Yang et al., 2024. \"In-Context Learning with Representations: Contextual Generalization of Trained Transformers\", https://openreview.net/pdf?id=fShWHkLX3o It will be interesting to compare this concurrent work, although it is not required.\n\nQUESTIONS:\n1. I think Algorithm 1 is proposed to handle the non-convexity of the problem, right?\n\nLIMITATIONS:\nThere is no potential negative societal impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n* The authors theoretically study the capacity for (nonlinear) transformer\n  models to perform in-context learning with a low-complexity function class.\n* The set-up involves pretraining transformers on sequences encoding\n  input-output examples from a reduced-rank Gaussian single-index model.\n  * The use of Gaussian single index models gives the transformer a chance to\n    implement a more intricate (but still analytically tractable) in-context\n    learning algorithm than in-context linear regression for linear data.\n  * The restriction of the in-context target functions to a low-rank subspace\n    gives the transformer an opportunity to internalise (during pretraining)\n    a 'prior' that it can rely on during in-context learning for a sample\n    efficiency advantage.\n* Within this setting the authors prove their main theorem, characterising\n  the in-context sample efficiency of a transformer trained under certain\n  conditions.\n  * Importantly, the efficiency scales with the dimension of the\n    low-dimensional model subspace rather than the ambient dimension.\n  * This suggests that the pretrained transformer model is capable not only\n    of performing ICL of Gaussian single-index models, but also of encoding\n    the 'prior' and gaining a corresponding sample efficiency advantage from\n    doing so.\n\nSTRENGTHS:\nI thank the authors for their contributions in this important area of\nresearch. In my opinion the work shows at least the following strengths.\n\n1. The work is well motivated by an important research goal within the field:\n  advancing our theoretical understanding of more complex forms of ICL\n  exhibited by transformers. This helps to bridge the gap between the\n  nuanced ICL capabilities of LLMs and the simpler abilities of in-context\n  linear regression transformers, for example.\n\n2. The setting proposed by the authors, of learning Gaussian single-index\n  models, appears to provide a neat example of a non-linear in-context\n  learning capability making non-trivial use of a transformer's MLP layers,\n  while still being analytically tractable and therefore amenable to\n  informative comparison with baselines from the statistics and ML\n  literature, as the authors have demonstrated.\n\n3. The main theoretical result by the paper clearly demonstrates that\n  pretrained transformers are capable of learning ICL for this nonlinear\n  function class and that they are able to take advantage of the special\n  low-dimensional structure of the function class in-context to improve their\n  sample complexity over an ICL method that considers arbitrary directions.\n\n4. The paper is reasonably well-organised and the writing is clear. Though for\n  me the technical content is outside of my background and  zone, with\n  careful study I was able to understand the setting and theorem statement\n  thanks to the detailed technical description and accompanying discussion\n  provided by the authors.\n\nWEAKNESSES:\nI lack the background to evaluate the soundness of the main theorem, but the\nhigh-level structure of the main theorem statement seems plausible to me and\nso I am open to taking the authors at their word.\n\nInstead, the concerns I have with the paper are about the framing of the\ncontributions. In particular, my concern focuses on the second research\nquestion, \"Can a pretrained transformer adapt to certain structures of the\ntarget function class, and how does such adaptivity contribute to the\nstatistical efficiency of ICL?\" (Lines 45/46), to which the authors answer,\nin summary, that, yes, there is a 'separation' (improved sample complexity)\nbetween these transformers and \"algorithms that directly learn the\nsingle-index function from the test prompt\".\n\nMy concerns with this framing are as follows:\n\n1. **These baselines are potentially misleading.** The authors compare the\n   sample complexity of a transformer model that has access to more\n   information (through pretraining) against that of a baseline that lacks\n   this information. It is not surprising that prior information can reduce\n   sample complexity and this is why the transformer's performance beats this\n   baseline.\n     \n     I think the appropriate framing would be more along the lines of saying\n     that the transformer is able to (in some sense) 'match' other algorithms\n     that have this same access to prior information, if these baselines are\n     available in the literature.\n\n     Note: I grant that the authors have shown that transformers can take\n     advantage of this additional information, and this is a meaningful\n     contribution. The specific concern is with the misleading comparison\n     that might suggest transformers have some radical sample efficiency\n     properties, when the improvement over these baselines has clearly come\n     from access to additional information that the baselines don't have.\n\n2. **This does not appear to be a completely novel observation.** There are at least\n   empirical examples of transformers showing improved sample efficiency\n   through making use of prior information gleaned from pretraining. For\n   example, this is implied by the task diversity experiments of Raventos et\n   al. (cited by the authors). In the low task-diversity regime, the model\n   appears to memorise a short list of task vectors and when asked to perform\n   ICL of these particular tasks it would obviously outperform a ridge\n   regression model entertaining a larger hypothesis class. Therefore it\n   appears that in some related work this second question has already been\n   addressed.\n\n     *Note:* These empirical observations are different from the present\n     work's theoretical characterisation of the sample efficiency advantage\n     of the pretrained transformer and I think this does not fundamentally\n     detract from the authors' claims to novelty. However, I think it would\n     be appropriate to introduce this question in the context of these\n     previous attempts to answer it.\n\nI would invite the authors to consider changing the framing of their paper to\nmore clearly qualify the nature of their contribution with respect to this\nsecond research question.\n\nQUESTIONS:\nTo begin with I will flag that before reading this paper I was not familiar\nwith Gaussian single-index models. Based on the information in the paper and\nmy own light additional reading, I think I have assembled enough\nunderstanding to ask sensible questions, but I could be mistaken, and I\napologise in advance if my questions are ill-formed.\n\n0. To this point, do the authors know of an accessible reference work that covers the topic of Gaussian single-index models? In case other potential readers are, like myself, not familiar with the topic, it could be useful to include a reference to a statistics textbook or similar.\n\n1. A basic question: I don't think I understand line 52: \"In single-index\n   model, outputs only depend on the direction of $\\beta_t$ in the\n   $d$-dimensional input space.\" Is this just foreshadowing the constraint\n   that $\\beta_t$ is restricted to the unit sphere in the $r$-dimensional\n   subspace? (Perhaps the constraint on the magnitude of $\\beta$ should be\n   mentioned in the introduction.)\n\n2. I understand that the first goal addressed by this work is to create a\n   setting where transformers exhibit nonlinear in-context learning behaviour\n   (and I consider this a strength of the work). However, the shift from\n   regression to single-index models brings with it additional statistical\n   and architectural complexity. Therefore, with respect to your second\n   stated goal of demonstrating adaptivity of ICL to low-dimensional\n   structure in the training function class, I am wondering if it would be\n   possible to get away with a simpler demonstration:\n\n    1. In particular, what would happen if one considered the problem setting\n       of in-context linear regression where the regression tasks are confined\n       to a low-dimensional subspace (or a unit sphere on such a subspace)? Do\n       you think that such a setting would offer a similar chance for the\n       transformer to show its ability to internalise informative 'prior'\n       information, but without the added complexity of considering\n       single-index models?\n\n    2. If I am not mistaken, this simpler low-rank regression setting is\n       *not* a special case of your setting, since it would require the link\n       function to be always the identity function which violates your\n       assumptions about the coefficients in the Hermite expansion. Is that\n       correct?\n\n    3. Nevertheless, would you conjecture that a similar result to your main\n       theorem---a generalisation bound that scales with the dimension of the\n       low-dimensional subspace rather than that of the ambient\n       dimension---would hold in such a setting?\n\nFinally, some questions about future work.\n\n3. Your theoretical work does not offer an empirical component. However, your\n   analysis appears to make some clear predictions about what a transformer\n   trained in this setting would learn, in terms of its ICL performance for\n   this kind of data and even the internal mechanisms with which it might\n   learn to perform ICL (such as using the MLP layers to encode the\n   low-dimensional subspace).\n   \n    1. Do you think these performance predictions would be borne out by\n       empirical experiments to this effect?\n\n    2. Do you think that trained transformers would implement ICL in this\n       setting using the mechanisms you have proposed?\n\n    3. Are there any differences between your theoretical training algorithm\n       and modern transformer training procedures that you expect might cause\n       these predictions not to hold?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe paper shows that low dimensional structure can be leveraged in pre-training. It is generally well written.\n\nWEAKNESSES:\n1. The MLP is before the attention layer instead of after. (My initial impression was that the attention would do something like a projection onto the low dimensional subspace and the MLP would learn the non-linearity somehow). Can the authors say something about this? \n2. There is no soft-max. \n3. This is not true GD/SGD (and I think this is not obvious from the introduction). \n    1. The MLP is only trained for one iteration.\n    2.  The function is initialized to compute 0, and following that the $w_i$ are exactly “erased” by the regularization.\n    3. The biases are then reset to be uniform on a unit interval.\n\nQUESTIONS:\nThere are three interesting things in this model: \n1. The non-linearity\n2. The single index function\n3. The low rank structure\nBut I am not sure how they individually play a role. The low rank structure I think serves the purpose of making the context length smaller, and isnt necessary, but what about the others?\n\nCan you include some examples of function classes satisfying Assumption 1 with low P? For instance, what if my ground truth $\\sigma_*$ is ReLU?\n\nCan this be empirically verified? An experiment that shows that w aligns with S under true GD would be nice.\n\nIs there some easy interpretation of Gamma? Without the MLP, for linear functions, it has been shown to be the inverse covariance of the data. With the MLP it appears to play a similar role.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}