Nonlinear transformers can perform inference-time feature learning
Authors: Naoki Nishikawa
Yujin Song
Kazusato Oko
Denny Wu
Taiji Suzuki
Venue: ICML 2025 poster (2025)
Keywords: transformers
in-context learning
feature learning
single-index models
Primary area: theory->learning_theory

Abstract: Pretrained transformers have demonstrated the ability to implement various algorithms at inference time without parameter updates. While theoretical works have established this capability through constructions and approximation guarantees, the optimization and statistical efficiency aspects remain understudied. In this work, we investigate how transformers learn features in-context -- a key mechanism underlying their inference-time adaptivity. We focus on the in-context learning of single-index models $y=\sigma_*(\langle \\boldsymbol{x},\\boldsymbol{\beta}\rangle)$, which are low-dimensional nonlinear functions parameterized by feature vector $\\boldsymbol\beta$. We prove that transformers pretrained by gradient-based optimization can perform *inference-time feature learning*, i.e., extract information of the target features $\\boldsymbol{\beta}$ solely from test prompts (despite $\\boldsymbol
{\beta}$ varying across different prompts), hence achieving an in-context statistical efficiency that surpasses any non-adaptive (fixed-basis) algorithms such as kernel methods. Moreover, we show that the inference-time sample complexity surpasses the Correlational Statistical Query (CSQ) lower bound, owing to nonlinear label transformations naturally induced by the Softmax self-attention mechanism.

Nonlinear transformers can perform inference-time feature learning
Naoki Nishikawa 12
Yujin Song 12
Kazusato Oko 23 Denny Wu 45
Taiji Suzuki"
Abstract
contain input-output pairs exhibiting certain functional rela- tionships. For instance, substantial body of research (Garg et al., 2022; Von Oswald et al., 2023; Ahn et al., 2023; Mahankali et al., 2023a; Zhang et al., 2023; Gatmiry et al., 2024) has demonstrated that linear transformers can perform linear regression in-context. As for nonlinear transform- ers, which possess greater expressive power to implement more complex algorithms (Bai et al., 2023; Cheng et al., 2023; Guo al., 2023; Kim al., 2024), existing optimiza- tion analyses generally do not provide end-to-end statistical guarantees for in-context learning of nontrivial nonlinear function classes (Kim & Suzuki, 2024a; Yang et al., 2024; Bu et al., 2024). A notable exception Oko et al. (2024b), which characterizes the optimization and sample complexity of in-context learning for Gaussian single-index models us- ing shallow transformer. Specifically, they showed that for degree- P single-index function class of dimension the required in-context sample size of the pretrained transformer scales as , which is comparable to the performance of kernel methods on an r-dimensional subspace.
Pretrained transformers have demonstrated the ability to implement various algorithms at infer- ence time without parameter updates. While the- oretical works have established this capability through constructions and approximation guaran- tees, the optimization and statistical efficiency as- pects remain understudied. In this work, we inves- tigate how transformers learn features in-context- key mechanism underlying their inference- time adaptivity. We focus on the in-context learning of single-index models y = a. ((x, 3)), which are low-dimensional nonlinear functions parameter- ized by feature vector B. We prove that transform- ers pretrained by gradient-basee optimization perform inference-tine feature learning, i.e., ex- tract information of the target features 3 solely from test prompts (despite 33 varying across differ- ent prompts), hence achieving an in-context sta- tistical efficiency that surpasses any non-adanee -ddapitt ffxeddbbais algorithms such as kernel methods. Moreover, we show that the inference-time sam- ple complexity surpasses the Correlational Statis- tical Query (CSQ) lower bound, owing to nonlin- ear label transformations naturally induced by the Softmax self-attention mechanism.
Inference-Time Feature Learning. Our starting point is the observation that for learning low-dimensional tar- get functions (e.g., single-index models) on isotropic data, kernel methods and, more generally, non-adaptive (fixed- feature) estimators are statistically suboptimal. Numerous prior studies have shown that neural networks can outper- form kernel methods due to feature learning (Ghorbani et al., 2019; Ba et al., 2022: Damian et al., 2022: Abbe al., 2022), where model parameters adapt the low-dimensional struc- ture of the learning problem during gradient-bucc based training. In the ICL. analysis of Oko et al. (2024b), transformers im- plement a kernel-like inference algorithm on a nonlinear feature map ()) fixed across tasks, constructing the predictor E, (-) by fitting the coefficients a, in-context. In contrast, we expect that better sample complexity can be achieved by algorithm that adaptively selects the feature map at inference time, i.c., varying {-)mm based on the in in-context examples.
Introduction
Large language models (LLMs) are capable of In-Context Learning (ICL) (Brown et al., 2020), where they construct variable inference-time algorithms based on examples pro- vided in the test prompt, without updating the model param- eters. Recent theoretical studies have shown that pretrained transformers (Vaswani et al., 2017) can implement regres- sion algorithms in their forward passes when test prompts
'Equal contribution 'The University of Tokyo, Tokyo, Japan "RIKEN AIP. Tokyo, Japan Unversitty of California. Berkeley "New York University 'Flatiron Institute. Correspondence to: saakkkawwwanki229@@@e www.....tooooo.acc....oooo Yujin Song ressoar.resccrch gmail.com>.
Motivated by the importance of adaptivity in the statistical efficiency of neural networks (Bach, 2017; Suzuki, 2019), in this work we theoretically study inference-time feature learning the ability to adaptively extract latent features of the ground truth from test prompts without any parameter updates. We ask the following question.
Naoki Nishikawa
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

Nonlincar transformers can perform inference-time feature learning
Q1: Can pretrained transformers implement in-context
expansion of o, (see Definition 2). This sample complex- ity improves upon kernel methods, and the exponential dependence on the IE aligns with the CSQ lower bound (Damian et al., 2022: Abbe et al., 2023).
algorithms capable of feature learning,
and outperform non-adaptive methods such as kernel
models?
We study the Gaussian single-index model setting where the class of target functions is given by f. : Rd + E. fxx)= a_(1B, 2²)) Here, a : IR- R is a nonlinear link func- tion, and f. depends only on the projection along 3 E Rd, referred to as the feature vector. The transformer model is provided with test prompts (in-context examples) in the following format:
SQ algorithms. For learners that access full statistical queries (i.e., in the form of B[o(z,)) the complexity of single-index learning is governed by the generative exponent ge(a.) S ie(o.) defined as the lowest possible information exponent after arbitrary L label transforma- tion (Damian et al., 2024); in particular, ge(a=) < 2 for any polynomial link function. Recent works have shown that gradient-based training on transformed labels can achieve a sample complexity of n 2 d((sc((-)), which improves upon the CSQ rate (Dandi et al., 2024; Lee et al., 2024; Arnaboldi et al., 2024; Joshi et al., 2024).
T1 y1. NN UN.
x query
examples
where the labels satisfy y = o. x;)) with feature vector 3 varying across different prompts. In this setup, inference- time feature learning refers to the model extracting /3 by examining in-context examples and using it estimate the all without any parameter updates.
The ICL mechanism in Oko et al. (2024b) operates as a ker- nel regression algorithm, which entails that the in-context sample complexity depends on deg(o.). We expect this dependency to improve with inference-time feature learning. Indeed, drawing from prior works showing that linear trans- formers can implement single-step GD in-context (Von Os- wald et al., 2023; Ahn et al., 2023; Mahankali et al., 2023a; Zhang et al., 2023) which corresponds to a CSQ-like al- gorithm - it is natural to conjecture that a sufficient sample size in our single-index setting scales as (((((())) That being said, when considering the effect of nonlinear soft- max attention, transformers can also introduce nonlinear label transformations at inference time, hence potentially achieving statistical efficiency that depends only on ge(a.) and surpasses the limitations of CSQ-based methods. This motivates us to investigate the following question regarding the ICL sample complexity.
label
This single-index setting captures scenarios where the ground truth exhibits low-dimensional structure, and has been extensively studied in deep learning theory to demon- strate the adaptivity of neural networks (Ba et al., 2022; Bietti et al., 2022: Berthier et al., 2023; Mahankali et al., 2023b). Furthermore, following Oko et al. (2024b) we as- sume that the feature vectors are drawn from some r I d dimensional subspace, which implies that if pretraining iden- tifies the support of the target function class, the in-context sample complexity should only scale with the subspace dimensionality r instead of the ambient dimensionality d.
Sample Complexity of Feature Learning. The sample complexity of learning single-index models on isotropic Gaussian data in Rd has been extensively studied in prior works. For kernel methods, the dimension dependence can be sharply characterized, whereas for adaptive procedures such as gradient-basee based feature learning, the statistical com- plexity can be characterized by (variants of) statistical query (SQ) lower bounds (Kearns, 1998; Reyzin, 2020). We sum- marize these prior results as follow.
Q2: Can transformers with nonlinear self-attention implement inference-tin algorithms that surpass the statistical barrier imposed by CSQ lower bounds?
1.1 Our Contributions
We establish end-to-end optimization and statistical guaran- tees for the in-context learning of the Gaussian single-index function class using a transformer with softmax attention. Our main result is the following.
Kernel methods. If a. is polynomial of degree deg(or,), then kernel methods using fixed feature maps require n 2 (deg(o.) samples to learn J. with low estimation error (Ghorbani et al., 2021; Donhauser et al., 2021).
Theorem (informal). Consider the learning of single- index polynomials o. ((3, x)) where the feature vector f is drawn from an r-dimensional subspace. If we optimize a single-layer transformer with softmax attention via gradient- based pretraining (Algorithm 1) using Tpt Ω ((o(((o)))) tasks, where for each task the prompt length satisfies Npt = () then the model achieves oa(1) in- context prediction risk the number of in-context examples satisfies Nuet = ((r D)). where ge(a.) $ 2 is the generative exponent of a..
CSQ algorithms. For learners that make correlational queries the target (i.e., in the form of B[(()))) ())).
such as one-pass SGD the squared/correlation loss (Ben Arous et al., 2021; Damian et al., 2023) or single- step GD (Damian et al., 2022: Ba et al., 2023), prior works have established a sufficient sample complexity of n 2 (iie(())), where ie(a.) S deg(o,) is the informa- tion exponent defined as the lowest degree in the Hermite
2

---

Nonlincar transformers can perform inference-time feature learning
We make the following remarks on the main theorem.
head mechanisms emerge via associative memory obtained through gradient descent.
While the pretraining sample complexity scales with the ambient dimension d, the required in-context sample size depends only the support dimension of the target func- tions r Σ d. This aligns with the findings of Oko et al. (2024b), which show that transformers can adapt to low- dimensional structure of the target function class through pretraining, thereby enhancing inference efficiency. We note that while Oko et al. (2024b) restricts the analysis to the regime Z Vd, our results applies arbitrary r S d.
Learning Functions with Sparse Feature. Other than single-index models, various function classes with sparse features have been studied, and neural networks' adaptivity to such features via gradient descent have been explored in standard (non-ICL) settings. These function classes include multi-index models (Damian al., 2022; Abbe et al., 2022; 2023; Bietti et al., 2023a; Mousavi-Hosseini Hosseini et al., 2023a), additive models (Oko et al., 2024a; Ren al., 2025), parity functions (Barak et al., 2022; Suzuki et al., 2023; Glasgow, 2024). Kim & Suzuki (2024b) showed that chain-of-thought reasoning of transformers helps solve parity functions effi- ciently, albeit not in an ICL setting.
The inference-time sample complexity of ()(g)(a)) is independent of both the degree and the information expo- nent of the link function d.; in other words. pretrained transformers can implement SQ algorithms in-context with statistical efficiency that surpasses both kernel meth- ods and CSQ learners. Consequently, our ICL complexity improves upon Oko et al. (2024b) in all regimes.
2 Preliminaries
Notations. ||m|| denotes the & norm of a vector v. For matrix A, we denote its lz operator norm and Frobenius norm as ||A||2 and |||||, respectively. Let (e) = (-1)'e 22 d' dar e - be the degree-i (probabilist's) Hermite polynomial. Given the Hermite expansion of a function f. i.e., f(2) ( Ei [i He)()), we define H(f, i) = c; as its degree-i coefficient. gd-1 denotes the unit sphere in Rd. ((A) is the indicator function. diag(ar, ..... (n)) represents the n n diagonal matrix whose (i, i)-th element is az.
Our analysis explicitly shows that softmax self-attention can extract the inner product (3, æt), where 3 varies across tasks, by computing the correlation between nonlinear transformations of both the input x and the label. This mechanism is analogous to the tensor partial trace algo- rithm studied in Damian et al. (2024); see Section 5 for detailed explanation.
On the technical level. we demonstrate that: (i) the Softmax attention can be pretrained to compute xz.z) for each prompt, and (ii) the nonlinearity in the attention reduces the information exponent of the link function o,, which is shown through careful analysis of the nonlinear correlation computed by the attention scores. Mechanism (i) enables inference-time feature learning, whereas mechanism (ii) leads to an information exponent-free sample complexity.
2.1 Preliminaries on In-Context Learning (ICL)
We consider the setting of learning functional relationships in-context as introduced in Akyürek et al. (2022); Garg et al. (2022). The input consists of N labeled example pairs (æi, 3R) E R& x R for = I,.... N, referred to as context or in-context examples, along with a query æ E Rd The entire input sequence (æ1 11. EN: UN. 2²) referred to a prompt, which we sometimes shorten to (X, y, x) where uN] The
1.2 Related Work
In-context Learning of Functions. ICL abilities of linear transformers have been extensively studied in terms of both expressivity and optimization. They are shown to imple- ment linear regression algorithms, including one gradient descent step (Von Oswald al., 2023; Ahn al., 2023; Mahankali et al., 2023a; Zhang et al., 2023), multi-step gra- dient descent (Ahn et al., 2023; Gatmiry et al., 2024), and sparse regression (Bai et al., 2023). (Zhang et al., 2024) showed that linear transformer can adapt to the prior mean of the coefficient vector. We note several related works on ICL of nonlinear transformers that were not discussed in the introduction: Huang et al. (2023); Chen et al. (2024) analyzed in-context linear regression of softmax transform- ers, including optimization analysis, while Li et al. (2024) investigated a classification problem (which is also essen- tially linear). Nichani et al. (2024) examined the ability of optimized softmax transformers to learn causal structures in-context, Bietti et al. (2023b) demonstrated that induction
X = [x] learner is tasked to predict the corresponding output E R for the query æ. The inputs æ,,..., N,,, x are i.i.d. sam- ples drawn from a specific distribution Da. and the outputs V1..... yy are generated as
æN]] and = [yn
y: = fa(s)) +
where f. : IRd - R is a function associated with each prompt, and SN ~ De are i.i.d. noise. The target function f. varies across prompts and is drawn from D,, for each prompt. Our goal is to obtain a model f(X, y, x; 0) with parameters 0, which takes the prompt as input and predicts the query output fax)) without updating any pa- rameters during inference time.
We divide learning into two stages: pretraining and infer- ence. In the pretraining stage, the model is optimized training set (((z), vf...., whx, Vhy x²), )))) which
3

---

Nonlinear transformers can perform inference-time feature learning
includes pp distinct tasks from D/.- In the inference stage, we present prompt (æ1, 3/..-..; : NNoss : UNsss : 2)) with context length Nest to the pretrained model, which
Note that, by definition, ge(az) S ie(a.) I deg(o.). It
is known that for polynomial link functions, there exists monomial transformation o. + o that reduces the infor- mation exponent to 1 or 2.
aims to predict the label fz(x). To evaluate the model's l's average performance at inference, we use test ICL, error defined as
Lemma 3 (Lee et al. (2024), Proposition 6). It holds that
{L 1 (if o is not even)
RICL
= ED,, DD. (INNa M1:NNost x; 0) yl]- Nestt
ge(a.) 11 2 (if a. is even)
(2.1)
Based on this setup, we emphasize that there are two distinct notions of "sample complexity" ICL:
Moreover ge(o.) = minz21 ie(o!) holds.
The pretraining task (sample) complexity refers to the task (sample) size Tpt (Nps) used by the pretraining algo- rithm to optimize the parameters 0.
Sample complexity of learning f.. The sample com- plexity of various learning algorithms are dictated by the quantities Definition 2. Linear/kernel models on fixed basis require deog(oo)) samples to learn f. (Ghor- bani et al., 2021; Donhauser SGD (Ben Arous et al., 2021: Damian et al., 2023) or single-step gradient descent (Damian et al., 2022; Ba et al., 2023; Dandi et al., 2023) on two-layer neural net- work دEー1 ( 2) + bb)), which can adaptively select feature maps (a 1) + by)) j= achieves a better sam- such dependence on the
The inference sample complexity refers to the context length Niet required at inference time to ensure small RICL characterizing the statistical complexity of the constructed in-context learning algorithm.
al., 2021). Online
We consider the regime where Npr: Tpt and Noett are poly- nomial in the dimensionality d. and track the dimension dependence in the sample complexity for both the pretrain- ing and inference phase.
ple complexity of ((d information exponent also appears in correlational statistical query (CSQ) lower bounds (Damian et al., 2022). Recent works have shown that data reuse (Araaboldi et al., 2024;
2.2 Single-index Models
For the class of target functions f., we consider Gaussian single-index models where the function value depends only on a single direction ,3 in the d-dimensional input space. This setting reflects the scenario where learners have to cap- ture the unknown feature vector ,3 from ligg-dimmesional input in Rd, and hence has been extensively studied in fea- ture learning theory literature (Bai & Lee, 2019; Ba al., 2022; Bietti et al., 2022; Mousavi-Hosseini et al., 2023a; Mahankali et al., 2023b; Berthier et al., 2024).
Lee et al., 2024) or loss modification (Joshi et al., 2024) can introduce nonlinear label transformations such that the "effective" information exponent is reduced to ge(a.). Consequently, the sample complexity can be improved to ((dvo(.)--VV) which matches the SQ lower bound (2d)) (for polynomial o.) up to polylogarithmic factors.
ICL sample complexity. For the ICL setting, (Oko et al., 2024b) studied the single-index function class where B is uniformly drawn from the unit sphere in an r(S vd)- dimensional linear subspace of Rd. The authors derived pretraining task and sample complexity of ((((tie(o.)) and inference-time sample complexity of (( Note that this inference complexity matches that of kernel methods on the r-dimensional subspace. Our focus (see Q1 and Q2 in the introduction) is to investigate whether transformers can be pretrained to achieve inference com- plexity that surpasses kernel methods or CSQ algorithms.
Definition 2. Let x ~ (0, Ia)- (Gaussian) single-index models refer to functions of the form
: + R, fx( = (33,
where a. : IR- + I is the link function and B3 is the fea- ture vector. We assume o. is a polynomial and satisfies [x= O. We define three key properties: 1. deg(o.) denotes the (polynomial) degree of o..
2. The information exponent (Dudeja & Hsu. 2018; Ben Arous et al., 2021) is defined as
3 Problem Setting
3.1 Task Setup
ie(a.) = minfi (((a,, i) #0.
We define the task distribution as follows.
3. The generative exponent (Damian et al., 2024) is
Assumption 4. Let Oa(1) be the noise level. A prompt
ge(a.) == min minfi H(h o oz, i) # 0},
(æ1, V1, EN UN. xx) is generated as:
KEL
where 2 denotes the set of all transformations from R to R for Py = 10,
x Ii...., "N...,,N,22
2N, x ~ Dz = No
y. = f(xx)) Gr, G Unif((-T,7)).
where a.# is the forw-forward by a...
fa (ax)) = 0.((B
4

---

Nonlinear transformers can perform inference-time feature learning
For the Irue function fa, we make the following assumptions:
=a =aTo (WFFAttnE)) Auln(E)
(3.1)
a, is a polynomial that remains across tasks. We as- sume the normalization conditions [===N(0..)0==0)]]
where o is applied entry-wise.
0 and (===)(o 1.
We further introduce some simplifications: we merge pa- rameters as WKQ = (W We E [R(++))((++) and W W W WV Rm)((+1), and write
S. denotes r-dimensional linear subspace Rf for
S d. For each prompt, [3 is drawn uniformly from the unit sphere Supp(B) == (3 / E Sr, ||31| = =1 in Sr-
WKQ
r
0dx1
WFV
[Omxd fj
In Assumption 4, feature vector B is newly drawn for each task. We allow for scenarios where B3 has limited (r-dimensional) support in Rd,, and we will show that the inference- time sample complexity adapts to this support. Note that unlike (Oko et al., 2024b) which established simi- lar adaptivity, we do not assume low-dimensionality (r o(d)); instead result holds for any 1 S S.d.
Remark 5.
01x
1
where r E IRddd and e E R==1,, Similar simplifications (zeroing out specific sub-matrices) have been adopted in re-
cent theoretical works on the expressivity and optimization of transformers (Ahn et al., 2023; Mahankali et al., 2023a; Wu et al., 2023; Kim & Suzuki, 2024a).
Consequently, frF can be rewritten
3.2 Student Model and Pretraining Algorithm
Student Model. We train a single-layer trans- former (Vaswani et al., 2017) with SoftMax attention. For a prompt (I,, M1..... NN, UN, we construct the embedding E E ()x
FTF (X, y, vi b,
m ajo
N i=1 y,eltile Tw/p
Li=1 eus/peR ! Tælp + b1)
1/
(3.2)
as
E
F1
EN
x
E ER (R(++))(N+))
with the new parameters (T, U, b, a): see Appendix A for the derivation. Hereafter, we refer to I as the attention matrix and consider (v, b, a) as the ML. parameters.
y1
UN
1
We first convert E to Attn(E) through the Softmax atten- tion layer defined as
Algorithm I, we specify our gradient- based pretraining algorithm to optimize the pa- rameters (I, v. a). The algorithm uses Tpt pretraining tasks {(Xi.N,.. i:N 1 Vi.N,,, ""))) Tea and optimizes the (fg regularized) squared loss
Pretraining Algorithm.
Attn( E)
V E Softmax(Mask(pp> 1. (W* K E) T)))),
where p 0 is the temperature, and W Wº, wV E I(d+))((++1) are attention parameters. The Softmax func- tion is applied to each column vector, while the Mask function replaces all entries in the final row of p p-1 (WK E) T "WEE E ((N++))(N++)) with- -00
1 T.
(XTe(XiE,,)))) (= 1+1
Σ
u,; y8/²
27.
Remark 6. Observe that the first N rows of p-1 (W E) @Q E represent the correlations between x: and for E [N]. while the last row corresponds to the corre- lation between x and itself. Consequently, the last row is (lypically) approximately vd times larger in magnitude than the others. Thus, apply masking avoid the Softmax attention from focusing on the uninformative final row.
for i = 1, here To 0 and Ti T2 = Tpt. The training procedure is divided into two stages.
In Stage I, we perform a single gradient descent step on the attention matrix I using the éz-regularized squared loss L, (T, v, b, a) + A |[T] 1/ computed over T, independent pretraining tasks. This optimization problem is nonconvex with respect to r. This single gradient descent update originates from theoretical studies on feature learning Ba et al. (2022); Damian et al. (2022), where it is shown that the first gradient update captures low-dimensional features of the true function. More recent results have highlighted that one gradient step in transformers can encode information for implementing algorithms such as kernel methods on low-dimensional subspaces (Oko et al., 2024b) or induction heads (Bietti et al., 2023b).
Then, we apply position-wii wise multi-layer perceptron (MLP) with activation function a-) and parameters WF E Rmx((+1), b E RT and a E R", where is the network width. Throughout this paper, we set a (=) max (2, 0) (ReLU activation). We use the (N + 1)-th element of the MLP output as the prediction for the query output. Overall, the model's prediction fIF is given by
frF(X, y, x; W W W. WF b)
In Stage II, we optimize the MLP parameters. Specifically, we randomize the weight v and bias b, and perform ridge
=MLP Attn N
5

---

Nonlinear transformers can perform inference-time feature learning
Algorithm 1 Gradient-based training of transformer
4.2 Implications of the Main Theorem
Input Learning rate 71, 72, regularization rate A1. A2.
Our main theorem suggests the statistical capabilities and optimization properties of transformers (see Table 4.1).
initialization scale a, temperature p
Initialize (()) Ta/vd, v00) ~ (nni(((+)))) b())
Om, (0) ol
(i) inference-Time Feature Learning. First, consider the standard case where r = d: the inference-time sample com- plexity is (( Since ge(a.) S deg(a.) always holds, our derived inference-time sample complexity surpasses the lower bound (((dcg(r))) for non ron-adaptive algorithms based on fixed basis such as kernel methods. This implies that pretrained transformers can implement an algorithm capable of inference-tine feature learning, that is, they can adapt to the feature vector (3 of single-index functions, which varies across test tasks. Beyond merely establishing the inference- time sample complexity, we verify in Section 5 that the attention mechanism indeed captures the feature vector B.
Stage I: Gradient descent on Attention Matrix
r + r(0) 7 V v(0), a()) +
Stage I:: Optimization of MLP Layer
Initialize bj Unif 1, 11) --- =v()) a" + arg mina L2(T", v² b", a) + A/ aal
Output: Prediction frr ( y, x; r b" a²)
regression the output weights a. Note that the only role of the MLP layer is to fit the polynomial link function Since our primary focus is the optimization of attention matrix, we adopt this simpler algorithm for MLP which admits a closed-form solution due to the convex objective.
Moreover, when B restricted to an r S d-dimensional sub- space in R2 (see Assumption 4), the inference-ting sample complexity depends only on r. In other words, transformers can significantly reduce the inference-time sample complex- ity when the support of 33 is low-dimensional, i.e., r < d.
Main Result: Transformer Performs Inference-Time Feature Learning
(ii) Beyond CSQ. Our result further shows that the inference-time sample complexity can surpass the sample complexity lower bound d=((=-//21 for CSQ algorithms. For instance, the sample complexity upper bounds for one-step gradient descent (Ba et al., 2023) and one-pass SGD (Ben Arous et al., 2021) on two-layer neural networks all depends exponentially on the information exponent. In contrast, in our analysis, the single-layer transformer im- plements an in-context algorithm that is statistically more efficient than CSQ learners, due to the nonlinear attention.
4.1
Main Theorem
Now we present our main theorem, which establishes an optimization guarantee for Algorithm and a bound on the inference-time sample complexity.
Theorem 1. Let fTF(X, y, x; b", be a transformer pretrained via Algorithm 1 with MLP width m = (()2ge(a.)) and initialization scale (0m-- log for constant C.. Then, there exist hyperparameters >1, M,, and 7/2 such that the following hold with probability least 0.99 over the training data and random initialization:
(iii) Optimization Guarantee. While recent works have revealed that nonlinear transformers can implement rich al- gorithms by constructing specific parameter configurations, our result provides an optimization guarantee for pretrain- ing along with the pretraining task/sample complexity of ((i((r)) for Tpt and Npt (see in Theorem 1). We estab- lish this by proving that single gradient descent step can construct attention matrix that approximates E|33 J. which enables inference- time feature learning, and that the MLP layer can successfully fit the link function d..
1. (Optimization) Suppose that the pretraining task size and the pretraining context length satisfy Ti, Npt and T2 (((33go(o.)/2 (i.e., Tpt d(0)die(o.)+2 (/))(.))/2)) Then, Algorithm I yields parameters (T", b², a") such that the empirical loss
1 Ti+T. Σ y - fTr(XXXXX/ V'''', I., I", v² ",,bb,)))))))) a²)l = oa(1).
To T1+1
Remark 7. For simplicity of analysis, we assume that the link function a. is fixed across tasks. Task- link functions may be accommodated by adding an additional attention layer fit the link function adaptively from the test prompt, done in (Oko et al., 2024b). An alternative could be to fit the MLP parameter a using each test sample: such test-time training constitutes a relatively lightweight optimization, where it only involves training the MLP layer, since the inner attention layers have already extracted the relevant feature information (B.10 x)).
Moreover, we have ||a"|| = Ö(133o(0.))
2. (Inference-time sample complexity) At inference time, if the (in-context Jtest prompt length satisfies
Nrest 52(r3gelo.))2 ./2)
then the ICL error (2.1) RICL od(1).
6

---

Nonlincar transformers can perform inference-time feature learning
the Softmax attention module leads both (i) inference- time feature learning capability, and (ii) statistical effi- ciency that surpasses CSQ algorithms. Observe that the term ett/pe=; Tm/p appears in the architecture (3.2), where ex Tw/p generates a series of polynomials in x via Her- mite expansion, which enables detection of the feature di- rection /3 from the nonlinearly transformed signal eli/p, yielding (t). On the other hand, the factor eat/p reduces the information exponent of a. to its generative exponent ge(a.) ) through the exponential transformation e/P, yield- ing (ii). Hereafter, we outline our analysis step by step,
Regression on test prompt
Kerne
CSQ
SQ
(dde((o))
(
(gg((_))
In-context learning
(Oko et al., 2024b)
This work
Pretraining: d((ic((-))
Pretraining: (iii(()))
Inference: (((deg(a)) (n) Inference: , 1) p2)
Table 1. Comparison of (inference-time) sample complexity. Here, r denotes the dimensionality of the support from which the task vector B is drawn (See Assumption 4). [*]] is only appli- cable when r S vd, whereas [*2] holds for any r S d.
5.1 Gradient Descent on Attention Matrix
Assuming that fTF is sufficiently small at initialization, we can approximate the population (expected) gradient by the population correlational gradient ElyVrfrF]. We show that the magnitude of this correlation is (dd-(ie(o.)-1))3 A standard matrix concentration argument then yields the re- quired pretraining task and sample complexity of (i(io()))), ensuring that the correlational signal is not hidden within the noise due to finite sample size.
4.3 Comparison Against Prior Work
Most recent studies on single-index models have focused on the standard supervised regression setting rather than in- context learning, with the exception of (Oko et al., 2024b) which investigates the ICI. complexity of single-index learn- ing under a low-dimensional assumption. Here we highlight the distinction from the work; see Table 4.1 for the compari- son of relevant sample complexities.
Furthermore, we demonstrate that the single-step gradient on I is approximately proportional to E| This implies that the attention matrix acquires information about the r- dimensional support of /3 via pretraining, leading to an inference-time sample complexity independent of d.
(i) Improved ICL sample complexity. In (Oko et al., 2024b), the derived inference-time sample complexity is ( (( which corresponds to the complexity of non-adaptive methods (for instance, kernel methods) on an r-dimensional subspace. contrast, we establish an improved sample complexity of (((()))))))
5.2 Feature Extraction by Softmax Attention
Next we show that the pretrained attention module can extract the feature information (B, directly from ob- serving the test prompt. Recall that the model output is given by fTF = Σド= (3 X, y, x; I*) + by), where
(ii) Removal of low-dimensional assumption. (Oko et al., 2024b) assumed that the dimensionality of pre- training task distribution is much smaller than the ambient dimensionality d, whereas our analysis do not rely on this assumption.
g( y, T) = N -N " / r=/e N exi/ng") Fa/e is the attention output. We establish the following proposition, which demonstrates that g contains the feature information: Proposition 8 (Informal). Suppose Niest π (1330((0))2) For I' optimized in Algorithm 1, we have
(iii) Technical innovations. To derive the improved inference-time sample complexity, we demonstrate that the nonlinear transformation applicd to the output label {ur} re- duces the information exponent of the link function a, to its generative exponent. This requires a careful tracking of the corresponding Hermite coefficients through SoftMax at- tention block, which constitutes one of the most technically challenging parts of the analysis. This goes beyond (Oko et al., 2024a) where the authors showed that a linear atten- tion computes correlations on the raw label {yr} and a fixed kernel basis {o(x.)} (constructed by a fixed MLP block). See the next section for the technical overview.
( (xx ge(a-) VT
xx) C. C
with high probability over the prompt, where Ca and Co are ((1/poly log d) quantities independent of the data.
We see that the nonlinearity introduced by the Softmax attention plays a key role in deriving the proposition above.
5 Proof Overview
Warm-Up: Linear Attention Module. First consider the counterpart of with a linear attention module Rii (X, y. x; I) = N-' Vi(y x Tx). If the con- text length N is sufficiently large, the second term con- Txx Moreover, it is known that if (2)00(0 * O, or equivalently, ie(a.) 1,
We provide a proof sketch for our main Theorem 1, where the main goal is to show that the self-attention mechanism provably extracts the feature vector (3 for each prompt. The key idea is that the exponential transformation in
centrates around Elyyz TTx.
7

---

Nonlinear transformers can perform inference-time feature learning
then Elya: ] 3 is satisfied. Overall, the attention mod- ule successfully computes 3T Tax 3 in-context, if T is (nearly) spanned by B. However, if ie(a.) > 1, then Elyx"]= O and the direction 3 is not identified. Thus we must exploit the nonlinearity to handle high-information- exponent scenarios.
Finally, we bound the discrepancy between the training and test errors via the Rademacher complexity of the class of transformers, leveraging norm bound on ||a*|| derived from that of ||a']] See Appendices D and E for details.
6 Synthetic Experiment
Exploiting Softmax. To consider the nonlinear attention, first show that the exponential transformation reduce the information exponent of to its generative exponent. Specifically, we show the following:
We conduct numerical experiments on synthetic data to compare the in-context learning algorithm implemented by nonlinear transformers against -on-adaptive kemel methods. We train a 6-layer GPT-2 model (Radford et al., 2019) to learn the Gaussian single-index task. See Appendix for detailed settings of the experiments.
Lemma 9 (Informal). The information exponent of (x) and o. exp(o.) is equal to ge(a.), where o. is clipped version of / p (see Appendix B).
1.4
Based this property, we show that Softmax attention can compute (x, ,) in-context for any polynomial link o.;
1.2
consider the denominator N Σ p Tæ/e of g- For N sufficiently large, the quantity approximates
Generalization error
1.0
0.8
Eg, (o( [expo a.
o..oo..(.... ((3, (1) ((Tx, x1) /p]]
0.6
Now we consider Hermite expansions of expox ((B. /p)) and exp (TTx, x1) /p), and apply the property Ex, ((w, He, ((w'//(wwa))))) (()))))) & I[i 5) (w, w').
0.4-
d=4
d=8
0.2
d=
12
Fortunately, exp (TTx, x1) /p) contains all degrees of Hermite components, and from the lemma above, ((3,2p))) has a nonzero coefficient for Hege(o.)- Thus, the attention module computes (z)()))( if is nearly spanned by B. Crucially, since the correlational signal strength depends only the generative exponent ge(a=), the required context length N does not scale with ic(a.).
d=
0.0
20
40
60
80
100
120
Context length
1.4
d=4
1.2
d=8
d=12
Generalizattinn error
1.0
d=16
We also show that the numerator of g retains this feature information. The remaining technical challenge is to ensure that the information is not lost due to the correlation between the denominator and numerator: we tackle this through careful order analysis. See Appendix B for details.
0.8
0.6
0.4
0.2-
5.3 Fitting the MLP Layer and Test Error Analysis
To complete the optimization analysis, we show that the MLP layer can fit the unknown link function via ridge re- gression (line 5 in Algorithm 1). First, we construct an ML.P parameter a' such that
0.0
20
40
60
80
100
120
Context length
Figure 6.1. Generalization versus prompt length for ker- nel ridge regression (Top) and pretrained GPT-2 model (Bottom).
T + T2
(TF(X ,' "",,,,"",""" ,',,,""
T
6.1 Comparison with Non-Adaptive Algorithm
+
For each test task f, we generate data as æi.,... 2 N x NO Ia), f Uni (i.e., r with of o, ((B" x/)) = Hez( x for i E [N]. We compare the performance of two approaches:
is sufficiently small. Let ((a') be the quantity above: thanks to the equivalence between the fz-regularized con- vex optimization and its norm-constrained counterpart, we can show that there exists a ridge parameter X2 such that E(a") < E(aa) and ||a"|| < ||a'|| hold, where a" is the parameter obtained by solving ridge regression.
In-context learning using a 6-layers GPT-2 model, con- figured as in (Garg et al., 2022; Oko et al., 2024b). We
8

---

Nonlinear transformers can perform inference-time feature learning
pretrain the GPT-2 model using on synthetic datasets {Xi i:Npe" Vi.N.. x', y" Jtm and evaluate its estimation error Tiem Dien (y² - (farrz(Xi.... Vi.Now x²) across different values of Noest and d.
that the Softmax attention can extract feature information by leveraging nonlinear transformations applied to labels.
We outline several future directions. First, the pretraining procedure and its current sample complexity of ( may potentially be improved to dBige(o.)) by going beyond one gradient step (Dandi et al., 2024; Arnaboldi et al., 2024; Lee et al., 2024). Second, extending our analysis to broader function classes - such as non- oolyyooolaial σ., multi-index targets, or anisotropic covariates (Ghorbani et al., 2020; Re- finetti al., 2021; Mousavi-Hosseini et al., 2023b) is an- other interesting direction. Finally, while our study demon- strates both "memorization" of the r-dimensional support of the target function class and inference- time feature learning, it remains an open question whether a tradeoff or certain interaction exists between these two mechanisms.
Kernel method based on the Gaussian RBF kernel k(æ,, x²) = exp(
==11²/o²}- For each task, the kernel is trained from scratch using the test prompt (X i Vi...... as training data.
We set d = 8, 12, 16 and measure the test error at Noest [2 1J (k- = 2........ 7). the performance of kernel models deteriorates rapidly due to the sample complexity lower bound ddog(a.) = d², whereas the pretrained transformer exhibits a much more favorable dimension scaling in the in-context sample complexity.
..........
In Figure 6.1 we observe that
6.2 Sample Complexity Scaling
Acknowledgements
To further validate our theoretical claim that pretrained trans- formers can achieve test-time sample complexity beyond the CSQ lower bound, we conduct sample complexity probing experiment under the same setting as in Section 6.1. In Figure 6.2, we vary the input dimensionality and measure the minimal number of samples required to reach specified test error thresholds. The estimated inference-time sample = d=.1, outperform-
NN was partially supported by JST ACT-X (JPMJAX24CK) and JST BOOST (JPMJBS2418). YS was partially sup- ported by JST CREST (JPMJCR2115). TS was partially supported by JSPS KAKENHI (24K02905) and IST CREST (JPMICR2015), This research supported by the National Research Foundation, Singapore and the Ministry of Dig- ital Development and Information under the AI Visiting Professorship Programme (award number AIVP-2024-004), Any opinions, findings and conclusions or recommenda- tions expressed in this material are those of the author(s) and do not reflect the views of National Research Founda- tion. Singapore and the Ministry of Digital Development and Information.
complexity scales approximately ing both the kernel lower bound of d and the CSQ lower bound of 2 d=.5 for the link function (.).) = Hea which has information exponent 3.
First N to achieve threshold
10a
Impact Statement
6x 10
This paper presents work whose goal is to advance the field of Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
4 10
slope-1.16
th05
slope=1.14
th05
3 10
slope=1.13
threshold=0.1
slope-1.15
threshold-O.11
10 Input dimensionality
References
Abbe, F.,, Adsera, F. B., and Misiakiewicz, T. The merged- staircase property: necessary and nearly sufficient con- dition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pp. 4782-4887. PMLR, 2022.
Figure 6.2. Sample complexity of pretrained GPT-2 on the degree- 3 single-index task y Hes ((B, x²). We plot the minimal test- time sample size Nust required to achieve test error thresholds (0.05, 0.075, 1, 0.125) against input dimensionality d log-log scale. The slope of each plot is estimated using least squares.
Abbe, E., Adsera, E. B., and Misiakiewicz, T. SGD learning on neural networks: leap complexity and saddle-to-sadde dynamics. Conference on Learning Theory, pp. 2552- 2623. PMLR, 2023.
7 Conclusion and Future Direction
We studied in-context learning of single-index polynomi- als and demonstrated that pretrained transformers achieve inference-time sample complexity that surpasses both kernel methods and CSQ lower bounds. Our analysis highlights
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn to implement preconditioned gradient de-

---

Nonlinear transformers can perform inference-time feature learning
scent for in-context learning. Advances in Neural Infor- mation Processing Systems, 36, 2023.
Bietti, A., Bruna, J., and Pillaud-Vivien L. On learning Gaussian multi-index models with gradient flow. arXiv preprint 1xXiv:2310.19799 2023a.
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. What learning algorithm is in-context learn- ing? investigations with linear models. arXiv preprint arXiv:2211, 15661, 2022.
Bietti, A., Cabannes, V., Bouchacourt, D., Jegou, H., and Bottou, L. Birth of transformer: A memory view- point. In Thirly-seventh Conference on Neural Infor- mation Processing Systems, 2023b. URL https:// // openreview net forum? (orrmmid==X22EBLLSSS
Amaaoldi, L., Dandi, Y., Krzakala, F. Pesce, L., and Stephan, L. Repetita iuvant: Data repetition allows sgd to learn high-dimensional multi-index functions. arXiv preprint arXiv:2405. 15459, 2024.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33, 2020.
Ba, J., Erdogdu, M. A., Suzuki, T., Wang, Z., Wu, D., and Yang, G. High-dimensional asymptotics of feature leamn ing: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35, 2022.
Bu, D., Huang, W., Han, A., Nitanda, A., Suzuki, T., Zhang, Q., and Wong, H.-S. Provably transformers harness multi- concept word semantics for efficient in-context learning. arXiv preprint arXiv:241 1.02 199, 2024.
Ba, J., Erdogdu, M. A., Suzuki, T., Wang, Z., and Wu, D. Learning in the presence of low-dimensional structure: A spiked random matrix perspective. Advances in Neurul Information Processing Systems, 36, 2023.
Chen, S., Sheen, H., Wang. T., and Yang. Z. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint 240i::2402./94422 2024.
Bach, F. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research, 18(19):1-53, 2017.
Cheng, X., Chen, Y., and Sra, Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:2312.06522 2023.
Bai, Y. and Lee, J. D. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. arXiv preprint 1x10.01910001111 2019.
Damian, A., Lee, J., and Soltanolkotabi, M. Neural net- works can learn representations with gradient descent. In Conference on Learning Theory, pp. 5413-5452. PMLR, 2022.
Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. Trans- formers statisticians: Provable in-context learning with in-context algorithm selection. Advances in Neural Infor- mation Processing Systems, 36, 2023.
Damian, A., Nichani, E., Ge, R., and Lee, J. D. Smoothing the landscape boosts the signal for SGD: Optimal sample complexity for learning single index models. Advances in Neural Information Processing Systems, 36, 2023.
Barak, B., Edelman, B., Goel, S., Kakade, S., Malach, E., and Zhang, C. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35, 2022.
Damian, A., Pillaud- Vivien, L., Lee, J. J., and Bruna, The computational complexity of learning gaussian single- index models. arXiv preprint arXiv:2403.055 403.05599, 2024.
Ben Arous, G., Gheissari, R., and Jagannath, A. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research, 22(1):4788-488 2021.
Dandi, Y., Krzakala, F., Loureiro, B., Pesce, L., and Stephan, L. How two-layer neural networks learn, one (giant) step at a time. arXiv preprint arXiv:2.305. 18270, 2023.
Berthier, R., Montanari, A., and Zhou, K. Learning time- scales in two-layers neural networks. arXiv preprint arXiv:2303.0000 00055 2023.
Dandi, Y., Troiani, E., Arnaboldi, I.., Pesce, I.., 7deeorooáá L., and Krzakala, F. 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.
Berthier, R., Montanari, A., and Zhou, K. Learning time- scales in two-layers neural networks. Foundations of Computational Mathematics, pp. 1-84, 2024.
Bietti, A., Bruna, J., Sanford, C., and Song, M. J. Learn- ing single-index models with shallow neural networks. Advances in Neural Information Processing Systems, 35, 2022.
Donhauser, K., Wu, M., and Yang, F. How rotational invari- ance of common kernels prevents generalization in high dimensions. International Conference on Machine Learning, pp. 2804-2814 PMLR, 2021.
10

---

Nonlinear transformers can perform inference-time feature learning
Dudeja, R. and Hsu, D. Learning single-index models in gaussian space. In Conference On Learning Theory, pp- 1887-1930. PMLR, 2018.
Kim, and Suzuki, T. Transformers learn nonlinear fea- tures in context: Nonconvex mean-field dynamics the attention landscape. arXiv preprint arXiv:2402.01258. 2024a.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. What can transformers learn in-context? a case study of sim- ple function classes. Advances in Neural Information Processing Systems, 35, 2022.
Kim, J. and Suzuki, T. Transformers provably solve parity efficiently with chain of thought, 2024b. URL https: / //arxiv.orz org/ abs 2210 08633.
Gatmiry, K., Saunshi, N., Reddi, S. J., Jegelka, S., and Kumar, S. Can looped transformers learn to implement multi-step gradient descent for in-context learning? In International Conference on Machine Learning, 2024.
Kim, J., Nakamaki, T., and Suzuki, T. Transformers min- imax optimal nonparametric in-context learners. arXiv preprint arXiv:2408. /2/86, 2024.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Ghorbani, B., Mei, S., Misiakiewicz, T., and Monta- nari, A. Limitations of lazy training of two-layers neural network. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Process- ing Systems, volume 32. Curran Associates, Inc., 2019. URL
Lee, J. D., Oko, K., Suzuki, T., and Wu, D. Neural network learns low-dimensional polynomials with sgd near the information-theoretii limit. arXiv preprint arXiv:2406.01 01.4400601188 581, 2024.
cc cc//apper_iille//ppprr///////// (assfabbce4aaffc0088853348868883 pdf.
Li, H., Wang, M., Lu, S., Cui, X., and Chen, P.-Y. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. CoRR, abs/2402.15607, 2024. URL https: //oii//dii///eo 100/10.
Ghorbani, B., Mei, S., Misiakiewicz, T., and Montanari, A. When do neural networks outperform kernel methods? Advances in Neural Information Processing Systems, 33: 14820-14830, 2020.
48550 arw.iv.viiiv../.vv. v.. 2402 15607.
Mahankali, A., Hashimoto, T. B., and Ma, T. One step of gradient descent is provably the optimal in-context learner with one layer of linear self- attention. arXiv preprint arXiv:2307.03577 2023a.
Ghorbani, B., Mei, S., Misiakiewicz, T., and Montanari, A. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2):1029-1054 2021.
Mahankali. A., Zhang, H., Dong, K., Glasgow, M., and Ma. T. Beyond ntk with vanilla gradient descent: A mean- field analysis of neural networks with polynomial width, samples, and time. Advances in Neural Information Pro- cessing Systems, 36, 2023b.
Glasgow, M. SGD finds then tunes features in two-layer neural networks with near-optimal sample complexity: A case study in the XOR problem. The Twelfth International Conference on Learning Representations, 2024. URL https:t////. httpass:///p/eeeeeeeeeeee http:/eeettrrooeeeeeee id=Hg0J1xx2B16
Mousavi-Hosseini, A., Park, S., Girotti, M., Mitliagkas, I., and Erdogdu, M. A. Neural networks efficiently learn low- dimensional representations with SGD. In International Conference on Learning Representations, 2023a.
Guo, T., Hu, W., Mei, S., Wang, H., Xiong, C., Savarese, S., and Bai, Y. How do transformers learn in in-context beyond simple functions? a case study on learning with representations. arXiv preprint 1.ti::2310.16111 2023.
Mousavi-Hossenie A., Wu, D., Suzuki, T., and Erdogdu, M. A. Gradient-based feature learning under structured data. Advances in Neural Information Processing Sys- tems, 36, 2023b.
Huang, Y., Cheng, Y., and Liang, Y. In-context convergence of transformers. arXiv preprint arXiv:2310.05244 2023.
Nichani, E., Damian, A., and Lee, J. D. How transform- ers learn causal structure with gradient descent. arXiv preprint 140.v22402.14735, 2024.
Joshi, N., Misiakiewicz, T., and Srebro, N. On the com- plexity of learning sparse functions with statistical and gradient queries. arXiv preprint arXiv:2407.05622, 2024.
Oko, K., Song. Y.. Suzuki, T., and Wu, D. Learning sum diverse features: computational hardness and efficient gradient- based training for ridge combinations. arXiv preprint arXiv:2406. /828, 2024a.
Kearns, M. Efficient noise-tolerant learning from statistical queries. Journal of the ACM (JACM), 45(6):93-1006, 1998.

---

Nonlinear transformers can perform inference-time feature learning
Oko, K., Song, Y., Suzuki, T., and Wu, D. Pretrained trans- former efficiently learns low-dimensional target func- tions in-context, 2024b. URL https://arxiv.orr/ abs/2411. 02544.
Yang. T., Huang, Y., Liang, Y., and Chi, Y. In-context learning with representations: Contextual generalization of trained transformers, arXiv preprint arXiv:2408. 10144, 2024.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners. 2019.
Zhang, R., Frei, S., and Bartlett, L. Trained trans-
formers learn linear models in-context, arXiv preprint arXiv:2306. 099.7, 2023.
Refinetti, M., Goldt, S., Krzakala, F,, and Zdeborová, L. Classifying high-diman dimenssona gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference on Machine Learming, pp. 8936- 8947. PMLR, 2021.
Zhang, R., Wu, J., and Bartlett, P. In-context learn- ing of a linear transformer block: Benefits of the mlp component and one-step gd initialization, arXiv preprint wwxh:2242.11955,, 2024.
Ren, Y., Nichani, E., Wu, D., and Lee, J. D. Emergence and scaling laws in sgd learning of shallow neural networks. arXiv preprint arXiv:2504. 19983, 2025.
Reyzin, L.
Statistical queries and statistical algo- rithms: Foundations and applications. arXiv preprint arXiv:2004.00557 2020.
Suzuki, T. Adaptivity of deep reLU network for learning in besov and mixed smooth besov spaces: optimal rate and curse of dimensionality. In International Conference on Learning Representations, 2019. URL https:// openreview. net. forum? eetieeee1ecscccc..
Suzuki, T., Wu, D., Oko, K., and Nitanda, A. Feature leamn ing via mean-field langevin dynamics: classifying sparse parities and beyond. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M.. and Levine, S. Advances in Neural Information Processing Systems, volume 36, Curran Associates, Inc., 2023.
pp. 34536-34556.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I. At- tention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp- 35151-35174, PMLR, 2023.
Wainwright, M. J. High--Dienssooal Statistics: A Non- Asymptotic Viewpoint. Cambridge University Press, 2019.
Wu, J., Zou, D., Chen, Z., Braverman, V., Gu, Q., and Bartlett, P. L. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint 2tti:2330008.000 2023.
12

---

Nonlinear transformers can perform inference-time feature learning
Preliminaries
A.1 High Probability Events
We frequently use the term "with high probability" or "w.h.p.": Definition 10. We state that an event A occurs with high probability, when
P[A] 2 d Cubp
holds for sufficiently large constant Cwhp- Here, Cwhp depends on neither d nor v.
An illustrative example is Gaussian tail bound: for 2 ~ NO, 1) and t 2 0, it holds that
P|z > 1) S exp -1//2]-
Plugging whp log d into above, we can see that S 2C log holds with probability at least 1 - d Cwhp, In this situation, we say that "2 I O(vlog d) with high probability" this means that we can redefine Cwhp to be sufficiently large, by changing the hidden constant in O (Vlog d).
Throughout this paper, whenever we take the intersection of high-probabill events A Ams m = O(polyd) is always ensured. Thus, (by redefining Cwhp.) we can state that A n ....... n Am also occurs with high probability.
Moreover, we consider the high-dimensional setting, i.e., our result holds for all d 2 Cnd where Cha is constant independent of both d and r.
Unless otherwise specified, the notations O, O, Ô, θ, are considered with respect to d. Note that we allow r to scale with d, while treating deg(o.) as an (11) quantity.
A.2 Estimator Constructed by Single-Layer Transformer
We now show that the estimator constructed by the transformer (3.1) is explicitly written as (3.2), given our parameter configuration
WKQ=
r
Oax1
WFV
=OOmxd v).
Oixd
1
From (3.1) and the configuration above, the estimator fre is written as, recalling that
111
UN
fTF = FV
ESofttax(Mastie -1 EE W KQ E.,++))) + b) WKQ
x Tx/ + 91/p
= TTo
FV ESoftmax
+
x Txx yn
-00
x Tx/p + v1/p
To
pyy
PyN
v) Softmax
+
x. Tx: / WN/D
m
Tw/p
aj 1/ eu//p
+ bj
Ta/p
1=1
Hence we obtain (3.2).
B In-Context Feature Extraction via Softmax Self-Attention Module
We show that the softmax self-attention has the ability to extract the feature vector 3 of the single-index function a. o_((2, B)) (this section precedes the optimization analysis, as the techniques developed here are later utilized in the optimization theory). Crucially, we show that the nonlinear transformation applied z via softmax attention module reduces the effective
13

---

Nonlinear transformers can perform inference-time feature learning
information exponent. Consequently, the generative exponent (Damian al., 2024) of the link function o, governs the inference-time complexity, which surpasses the CSQ limit (Damian al., 2022). Let
Σ i= yeehi Tm/e Ei= eui/pe Tw/p
TTm/p
g(X, yx
(B.1)
be the output of the attention layer. The goal of this section is to show the following Proposition 11, stating that g can extract the inner product between æ and [3 which varies across tasks.
Proposition 11. Let = (log Cr d) and K = (llgg C* d) where Cp and Cx are any constants satisfying Cp 2 deg 0. 1/2
1 and Cs 2 max (2 matrix G written as
G for + +,, 1,eegg( 4} (e;(o.) defined Definition 14). Suppose NYF
G-U DU NN
(B.2)
where D = diag 1,..., 1,0, 1.0, -....,)) U is an orthogonal matrix, and N is matrix satisfying |N Ie O(IVa).
d-r
Now, N = (((3gg(0.)) then with high probability over length-N prompts (X VLNN, x), it holds that
g(X.y0 xY
ge(o.)
P1 + P2
(Ga,3) F
+ -ggooo/ log 2dog(a 1od ++2 dd).
(B.3)
Here, Pi and P2 are independent of the data and satisfies P1 = oa(1) and P2 O2 (log d) Crz)) where Cp, is a constant dependent d. but is uniformly upper bounded using deg(ox) and odd(oo) = E0-Noo.1)oo.oo (2))), where o ox,(oa(.)) is polynomial constructed by extracting only the odd parts a..
/T U DU Vr. which matches the condition of Proposition 11. Since U is spanned by the basis of the support of in this case, (Ga,, B) in (B.3) satisfies (Ga, B3 ~ (x, (). This implies that, if the context length inference time is ((r22(a_)), then the pretrained attention layer can extract the feature vector B which varies across different prompts and compute (x, B). This highlights the inference-time feature learning ability of the pretrained attention layer. Moreover, the inference-tine sample complexity depends on neither deg(a.) nor ie(a.), which surpasses kernel and CSQ methods. The fuctor log introduced just for compatibility with later stages.
Remark 12. Later we will show that I after pretraining satisfies I
22oglo_)) d in the residual term is
Preparations. We use the following property, which states that there exists polynomial transformation which reduces the information exponent of to its generative exponent. Recall that H(f, i) denotes the degree-i Hermite coefficient of (that is, f(2) = Σizo d
H(f.i) Hez(z)).
Lemma 13 (LLee et al., 2024), Proposition 6). The following holds:
Ifge((..) = 1, then there exists i I C₁ (deg(o.), odd(o))) such that [H(o', 1)] 2 Di De(deg(o)),odd((,.)
Ifge(a.) = 2, then there exists i I Cz(deg(a.) (o.)) such that [H(a),2) 2)] 2 D2(deg(a.)).
Here, (ide((o,.) odd(a.)) odd((a) and (deg(ar,) od)(a.)) are constants which only depend on deg(o.) and odd(o.) = Ez (0_ oda (z) where (.)) is a polynomial constructed by extracting only the odd parts of a.. Similarly, Cz(deg(o.)) )) and D2(deg(o.)) are constants depending only on deg(o.).
Definition 14, For i 2 1, let es(a.) be the minimal j 2 such that H(o2, i) + 0 holds. If such j does not exist, then define e((o.) = co. (Especially from Lemma 13, (g.().)(o.)) can be uniformly upper bounded using the degree and the odd part of o..)
Another key property is that y is bounded with high probability.
Lemma 15 (Oko et al., 2024b), Corollary 17). If ||]|| = 1 and x ~ NO, Ia). then (o,((3, ((3,))) S (log (//2 holds with high probability.
14

---

Nonlinear transformers can perform inference-time feature learning
(=(=)
if
(2(2)
This motivates us to introduce new function (2
log
Note that o. (13, x))
0 (otherwise)
(.((,,)))// holds with high probability from Lemma 15 and the condition of p.
Here we show the following lemma appeared in Section 5, stating that the nonlinear transformation exp(-) can reduce the information exponent.
Lemma 9. For p 2 1, if ep(o.) xc, then H(exp(o.).)) @ (
(log (ppple)) and H(ō. exp(o.).
@ [og (1oga)dorptos)) holds. Here, the hidden constants in θ notations may depend on o, and p. If ef( = oc, then H(exp(o.)), i) can taken to of O(d where a sufficiently large constant.
This immediately implies that the information exponent of explo,) and o. exp(o.) is equal to ge(o.).
Proof. We consider H(exp(o...) p) for ep(ow) oo. The other cases can be derived similarly. From the Taylor expansion, we obtain
Cpep(a.)
H(exp(6.).P) cp(o.)! H(o, H(o, (olax), p)) +
H
P))
Σ
H(o.'p)
i!
d!
(o. H1
(1)
For (D)
H(o,)) .p) = N(e )x(o)) (2)]
ppr(o.) E E[Hep(=)) (m)) (x(r))))) penlo.)
1
(e (x)( (1log))
holds. Here, from Hölder's inequality, we know
( (x)((o())) 2 p(log d)-
(,) p(log d) 1) (B.4)
Moreover, B[[(o,, (z) 2 p(log d) -1] O(d-') holds due the high probability property. Consequently, (I) is of ((log d) copelo))
For (II), note that H : .p) is dominated by Nे E(Ho (=)] from the similar argument as above. Moreover, from Hölder's inequality that
1 ilpp
1
E
(2,o,)),
He
(2)]
V (2)
Here, as (o.) = (e.ss,,, is of (lpllogd))C((ppo..3)3 Finally, let bound (III). From the definition of or, get
E Elo 2
(2)]
is the constant depending only on σ., and p. we can state that (II)
1
< i! E[HHe (1)
H(..'' p)
(log d) -
Moreover, from the bound residual terms of the Taylor expansion,
it (log d) (log d) -ceel(oo--)))-1
1
12Cpeolo.)11
is satisfied.
Overall, we conclude that H(exp(o.). = [log d) (poplor)
15

---

Nonlinear transformers can perform inference-time feature learning
Lemma 9 does not deal with the case where 0. Here we treat this case:
Lemma 16. H ox 0) = - ou(1) and H(exp F., 0) = + oa(1) holds.
Proof. Note that p(o./p,0 0) and H((o./p)².0) 1 from Assumption 4. For oe exp Σizu t and expo. = Eiz0 1 स o conducting the same argument as Lemma 9 yields the assertion.
Let us derive bounds for some important quantities.
Lemma 17. Suppose a, x. N(0, La) and 3 Supp(3) = {B [ B3 Ga = ( Vlog d) holds with high probability.
B E S,, || = 1}. Then,
x
='G'' Gæ' pRo/t
= o(1) and
Proof. First we can confirm that
x Gæ' psVTT x - DUdl/pp/i+ x Na' N = o(1)
holds with high probability: from the rotational invariance, y U and y Un' standard Gaussian vectors. Then,
the first term is o(1) with high probability from Lemma 29 and pk (loog d). To show that the second term is o(1), can again use Lemma 29 for the case = d and the fact || 22 = O(d 1d21)
Next we show
T Ga 3 DU B Næ ((llo
For the first term, note that ||B U DU S 1 holds. Then, (for any fixed (3) B U DUx the Gaussian with variance most one. Thus, Gaussian tail bound indicates that B U DU= S Vlog with high probability. The second term be also derived using the fact that |N O()) holds with high probability from Lemma 28.
Lemma 18. With high probability over the prompt (X y, x)
exp
Y1 p
exp
x. Gæ
pn//r
)
Emishi
exp (o.((B,0))) + $1 ))epp exp
x Ga pr
- (N-3)
and
N
V/
exp
exp
x Gæ /F
Exiic
(((()xz))) 6..(1. m1))
12 exp
x Ga
p
1=1
=Ō
(-=
holds.
Proof.
Let us show the first assertion: the second one is derived the same manner.
From
Lemma
17, N-1 EN pxcx(x/ Gz/ Gx/pk//)
N-1 1 Σm ((3(
+
x (if æ S c)
///p)xxp(clippx] Gx/ FF/pN/F, 1)) with high probability where clip(æ,c)
Conditioned on
0 (if [z]
c)
this, with high probability,
N
+ exp(vx/D)exp((r) Ga/pevT) Emiik ((3, +1 (1/e)xxp(a; Gæ pw)
N
F. ((B....))+G.i// Ga/pp/F:1) Emics
[ (1.1 clpppe Ga/pe
(1)
16

---

Nonlinear transformers can perform inference-time feature learning
+ Ee ( Ga/pe/FF1)
Buiic ( Ga/pavs
(11)
holds. (I) is of Ō(N- with high probability from Hoeffding's inequality, noting that exp(o. (13, x;),),))) + // G//)exppclippp; G Bx/ppyr; 1)) is O(1)-bounded. (II) is upper-bounded as, using Hölder's inequality,
Ga/pRVT) 2 1) Fi (exp(o.((,,))) m1) + /1/p) 1/4 Ee [(1 exp(2r Ga/ (F)/) 1/4 B[T((>]] Ga pkap/pprr) 2 1)911/2
Ez, [exp(o.((3 x1)) (1/p)(1 exp(x
<lb [11- [xp(a] G pm ()) ,))"}//+ O(d-c))
where C is a constant which can be taken to be sufficiently large from the definition of high probability events. We conclude that (II) is negligible by showing that Es,; [(1 exp(æ; Ga/pni /r)))) = ((1) with high probability over æ. It suffices to show that Ex, [[exp(>] Gz/ Gz/pk/F)) (F)] = Es, [exp(1x] Ga/pn pF)F) = (1) with high probability for / E [4]. Here, it
holds that
Ex, [exp(la] Ga/
= -E -V
[exp(Lx] U DUz/PeyTJexp(laas Nz/pxVT))
<Bar [exp(2lx] DUx 1/2 E Em [exp(2lx; Næ/ N pr)rr)
D (UUx/px/\r)
(B.5)
Note that 22 px /Fl I and /2/N2/ I holds with high probability over æ ~ N(O, Ia): for the former one, we can use Lemma 28 as Ux ~ (O, La). For the latter one, we again apply Lemma 28 (for the case where r = d), together with || = (()/Va). In this case, (B.5) can be upper-bounder by
E "a
where A2 which is known to be (1). This completes the proof.
Proof of Proposition 11. Now we prove Proposition 11. We first naively evaluate the denominator and numerator of (B.1) separately, and deal with their correlation afterwards.
Proof. Plugging the parameter values into (B.1) yields
N-1 1 D (us/ (ww/p)eex(uw//))exppa Gat
(B.1) N-1 Di= exp(vr/p)exp(ar) Ga/pevT)
First, let evaluate the denominator N EN' distribution (x)....,,N)) NN)), this equal to the clipped version N " exp(o. ((3, xz)) + p/p))xp((x Gz/pnyT).
i=1 exp(ys/ p)exp(xx _ Ga /t//////r). Note that, with high probability over the
Its expectation (assume fixed now) is calculated as, using exp(A=) = Σiz0 (P((2/2) X (HHe
E (exp(C1/p)expff ((3 a. ) Ga/ppyT))
|G prer
x Ga
= (exp(T/p) -xp(-r/p)) Ez,
H(exp of, He x
((,,x1)))) exxp
Gaz ||
120
(+)
and
(*) = =Eg;
i!
H(exp o., ()He: x11)
Σ
j!
1
prvrr
Gx
exp
| 2 x2²
He, ((21, Ga/ (ee]((,,Ga///
200
120
17

---

Nonlinear transformers can perform inference-time feature learning
|Garr]]
Garl
H(expoo,
pKVVF)
exp
(19, (L9,Ga//l
Ga/l/////l (Lemma 34)
| 2p²k2r
+ (eexpoo,gelo...
Ga
Re(o.)
G2 ?
=H(exp 6x, 6x, e)exp
ge(a,))
PKVTT
exp
2p2k2r/
(13, Ga/ G
H(cxpoo,,) 1)
|Ga|
|G 2
i!
PRVTT
exp
2p²
(13, Ga/
|G Gxxl))
(B.6)
(2ge(a.)++
We can upper bound the final term E: indeed, as
G pkvr
< with high probability,
H(expoo..)
||Ga|]] pF
Gall 2p2
(13, Gæ Ge/l|Ga/l
exp
(2e((_)++1
H(expo., (B,, Gar/pEVT)))
Σ
(2ge(o))+1
holds. Moreover, from Lemma 33, we know that
H(cxpo,,
(H(expo.,1))
(a..war)llexpd e e².
i!
i!
220
Then, get
H(expo
((3, G///G///)//)/
S
Σ
(B,Ga//PVFF) 1 (B, Gz/px 80(o))+
i!
(=g
12ge(a.)+1
with high probability.
Overall, letting hip = //2 (cxp(T/p) + cxp(-T/p)) +a(1) (recall that we assumed T = 0()), with high probability we obtain from Lemma 18 that
N
cxp(vx//)exp(a: Ga/ppn/r)
= =,,H(eexoo.. 0)exp
2p2
Gx
+ kp
H(expo.,
ge(a,,)
exp
xp2x2r ((3, /R/ (olo)) ( + -
Ga
Next we turn to the numerator: the expectation of N Σi (xxp) p)xpppp pR/pp//F) /Tiis is
[(o. (ttp:/m))t/////ppp//////////pppp ((3, (1 (xxp(rr G G pP)
Gx
2 Ga
= k Eo
전 H((6 @iloeexpoo. o., (B3, (exp
Ga
120
Ga
ai Ga
+ kg Ex:
H(exp o., (()) x1) exp
|G
120
where Kg = 2 (exp(r/p) exp(-T/p)) = 0x(r2/0²)) Hence, by the same procedure as the analysis of the denominator, we obtain
N N-1
Σ / ()cxp()s/ (xx p/xppp////exxp(x Gz/pp/7
18

---

Nonlinear transformers can perform inference-time feature learning
|G
H(o. expox,
|G ?
Hoo exp o., 0)exp
2p²
+ kp
ge(o.)!
exp
2p²x22/ |x 2p2
(13, G2/ pK) )s) ((r))
G
He. ge(o.)!
+ K,Hl(eep o. 0)exp
((3, Ga/ PR))a/ + +5'
2p k²²r)
exp
where E x (B. px/pp//i 80(1))+1 with high probability. Let us derive the claim of the proposition by an- alyzing the correlation between the numerator and the denominator. Let Ao H(o exp 0), Age(o.) H(o, expo,, ge(o), (((eo)) Bo
H), Bgo(o) g((.)).. Note that Age(o.) and
B Bge(o.) are of (logd) (( and then
E E S (B, GI/ (x//pNVT) s0(o.)+1
(((oggt(a../22t) d. (px)
(:: Lemma 17)
E 00 log (0. d. (px) -go(o.) ) n 1/4) (: ge(a.) s 2)
E ((((pk) go(())) x k 1111-3e((o)/44 (:: w(log 80 1/241/2
holds from Lemma 17. Moreover, set N = (((28c(o.)/2) to ensure that N-1/2 term also of o((p)) x 1-17-3066 -3g((o))/4)) In this case, E, E. N are all of o((px) (((p))_(((_))ppppp (elr-)) 3go(e.)/4x xxlla)/xx1122) 1/2) from the condition 12122 = -(A go(o.)) Then we get, letting 2 = (B. Gx) (pxvr).
8(X, v( x; I)
k BB + Ap + Kg gola goo(a (20(())) - kpA2o(a.) 200(0)) + (h.o.t.)
=p
Bu + k B Ro(o.) )20(0_)28((..) 280(0_) ) + (h.o.t.)
K,Bo Ao + KB Bgol, ( =20((0)) ApAgelo.) )=0((.)) (h.o.t.)
=p
kp B1 +
B Beolor) 280(o.) + (h.o.t.))
Bo
K,Bo + Kp + KKBBBboa Bgo(o
kp kpAgolo,) + (h.o..)) (1- Brolos) (n)(a)) + (h.o.t.))
Bo
Bo
Br + Ko
ko koBoAgola.) + k K,Bo Bo + +kpAoBello.) 2Ro(O.) + (h.o..))
=p
Bo
kpBB
W/ Bo + Ao
BoAooo,,..)
Ai gelo..) _g0(o,) 2g000)) + (h.o...)
=p
kg Bo
B.
Where (h.o.t.) means the terms of (((p)) A Age(o, )r 3/ 1/2). Finally, let care about the correlation between the numerator and the denominator, that is, verify that k Ag(o)) An Bgo(o.) does not vanish. Here recall that Ao = oa(1)) and Bu 1 + oa(1) from Lemma 16, and that Age(a.) and Bgo(a.) are of the same order Hence, k, Bogeo = (Age(o))) and Byoec)) = k Agoo(o) = O (Apn(o.)) dominates these terms. Then, taking
holds. Then we can conclude that
W Bo + Ap
P1
Bo
and
Pa ((pp) 3o(o.) K2 Agoggo(r))
-ppA go((o)) (= ((((pk) Be(o.)
-B((o)) -o(o.))Agooa.))) A2o(a.)))
yields the assertion. Finally, the order of p-(h.o.t) is o(Par 3gd (/ Using 1² 1/2 2ilog the proof.
d) completes
We are ready to show the following Proposition 8 outlined Section 5.
19

---

Nonlinear transformers can perform inference-time feature learning
p Proposition 8 (Formal), Suppose N, and k. are satisfying the condition in Proposition 11. Let
rEg NN
KVF
for some matrix N satisfying || 0(1/Va). Then,
(x, .)
ge(a.)
,(x, Pi + P3
+ +oa(P2r 3go(a.)/4
1/ log- 2 2deg(o.))+ d)
/rr
holds with high probability over the distribution of the prompl (X; y, X), where Pi = oa(1) and P2 = Oz ((log Cps))
Proof. From the assumption on the support of B (Assumption 4), there exists a d x d orthogonal matrix U such that ar, UB- Unif og,..., a,,0,........................... ..... 0)]a] +... a² = 1}. Here note that EIUSSTUT) = r ddag l..... O......
holds: then, I satisfies the decomposition (B.2) with G = rrEg133T] + N.
ge(o-) in (B.3). Note that (rEg|33TTar,
Let us simplify
UT diag(1...... .....,,..,,,.. 1. 0,.,.... 0YU.
U = (), Moreover, (Na,3) = o(r-1/2
1/2 log 2deg(a.)+1 d) holds with high probability from Lemma 30.
Overall,
(x, 3) Vr
8o(o.) + o(r- 1/2 1/2g""" deg d) + od (pr
g (x) P2
log 12 -2deg(o.)=2 d)
is satisfied. Then, the case where ge(a.) = 1 is derived immediately. For the case where ge(o.) = 2, noting that (m.3) 1/2 log d) with high probability, we can use (p q)² p² = ((pp + q) for the second term.
Vr
C Optimizing the Attention Matrix
In this section we develop the optimization guarantee for the attention matrix P: see Section for the proof sketch. We show that one gradient descent step contains information about E[3B"] to satisfy the conditions in Proposition 8. If JTF is sufficiently small the initialization, we can approximate the full gradient using the correlational (X","" yy" r"; r, (0), b(O), a(0)) r-r(o)) Hence, we first analyze the model gradient Vrfir(X', ,,,",;;''" v(0), b(0), a()), utilizing the technique developed in Appendix B.
1 gradient Ti
Lemma 19. Let
Σ vt/ (ilppeeinned aloVax Ea Σ en aa///ad aloVd
E
[_(x) (ff
p_(x)
Also, let A; = H(O., ex i), B, = H(exp for 2 0 where F (2) =
Tog
(0 (otherwiss))
Then, using ko = 1/ (exp(T/p) + exp(-e////)) = (1), K2 = 13 (exp(TT/p) - -ex((T//)) = O2(r2/p²) and 2 = (3.z) ovd it holds that
&1 =k exp ( ||0111
Ait1 zi
kpp
||a:|/
A
2p²d/
i!
B +
pvd
exp
2p2d
i!
x
120
$20
+ K, exp
2
|
B x²
i!
B
+
pvd
p
exp
2p²d
|| 1|2
B
i!
2
x
$2:
12:
20

---

Nonlinear transformers can perform inference-time feature learning
+ +0(d
and
||0112
Bi+1 x¹
&2 =k exp
exp
B
2ppd
$20
kp
||0||
Bi 2¹
x+0(d122-111)
exp
a
x+
pvd
2p2d
i!
i2.
with high probability over the data distribution
Proof. Let us analyze &2 first. Note that & expectation of this quantity is calculated
E
(0 =/p/// holds with high probability. The
Em
o
a/paa [e //
=kp Reg
Bi He; (B3 x21 e= m/ a/OV 10
i!
120
(1) kpp TTn
Bi+1 Ho,((B, 201) (3)) w/oaa oo
kip
x
Exn
B
He: x21) =0 w/pvd
pvd
i!
120
$20
|| 2pdd
Bi+1 i!
kp B +
||0|1
B Ö(Vd/N)
kp exp
x
exp
pvd
2pdd
i!
120
120
Here we used Stein's 's lemma to derive (a), and the argument to derive (B.6) for (b). Finally, the deviation from the expectation is upper-bounded using the exactly same technique as the proof of Lemma 18: the only difference is that we now treat d-dimensional vectors, meaning that the bound coming from Hoeffding's inequality is of Ó (Vd/N) (Lemma 31). Similarly, we can analyze E as
Erar (o. ((3, + Gi/p)e (19.2.
= kkp
[o.((3, x ((3.m))) ((( a/o x1
+ Kn
(18,m,))) " a//////
||< 12
A1+1 /
kp
||0||
A y
=k exp
2p
B +
PVd
exp
2ppd
x
$20
i20
+ Kg exp
||20|(2)
2p²d
B 2t
B
pvd
k/
exp
||
2pdd
B zt
x + δ ((d//N).
120
Lemma 20. Let r(0)= = Ia/Vd and v(0). b(0) and a0)) are determined in Algorithm 1. Then, assuming Npt (((2))00..)) with high probability over the data (X (Xi.No V1:No1 æ) (simply written as (X, y, x)) ), it holds that
Vr/Tp(X, y, (0). (()) a aa)) (000
'For a vector v, v = Ō(...) means |||| = Ō
21

---

Nonlinear transformers can perform inference-time feature learning
=0 (Po + Piz
)-( 1000))--11 )...)))) 33x a.hm nax
where Im = K [m] | e0(0) TH, Po...., (ie(a_)-2 O(1), Pee(o))-1 O(0 and
||||| S
O (dc(icc log
C. d) (where Cn can be taken to be sufficiently large).
Proof. For simplicity we define N = Npe. We can calculate the model gradient as
VrfiF(X, y, x; T, v b(0), (0)
yyss(Xy.x.y)) i= yichh/o 2/ova
y o'
where
s(X, y, x.
cuin w////////va EE (0 ( EE 1 a/pvdz.) Lyi/ (J:lp) e a/ xx) N
(Ein otk/pee! 2/pvdj2
First, we can see from the same procedure as the Proposition 11 (with d) that the content of o'(+) above dominated by Pi(1 + oa(l)) with high probability, where P1 in Proposition 11 is positive, which means that 1 if and only if oj 1. Then,
N We =lovd
o' 1// N o a/ova
Vrfrr(X, v00, 6(0). a a(0)) alm s( y. y)
holds with high probability. We proceed calculate s(X, y, ar, v) above. From the proof of Proposition 11 and Lemma 19, have
(x y. N. s(X, y. v. (m1 (X y, 2. E2 y, 2)
y. y. EX, v. x. -))
(C.1)
where
||0||
A
|
B₁ 2' Σ 120
का X y, x. v) = kp exp
+ Kr exp
+ +
2p2d
i!
2p2d
i!
120
||a/|/
B
72(X, U. x v) kp exp
i! 2²
+ 1(NN
2p
$20
||01|2
Ait1 x²
kp
||
4
61 kk exp
2d
i!
B +
pVd
exp
2p2d
i!
x
$20
$200
||a0|/2
B 2i
||
B
+ K2 exp
2ppd
i!
B +
pvd
p
exp
2p2d
i!
2
x
220
1200
((11/22-112)
and
|
Bit1
$2 =kp exp
B
2p²d
120
22

---

Nonlinear transformers can perform inference-time feature learning
kp pvd
exp
| 2p2d
Σ
Bi i!
x +
-+0(d-22--122)
F20
(recall z = (1,0c) pvd
Before proceeding the further analysis, we note the intuition for the goal of this calculation. We can see that terms proportional to will be canceled out in (C.1). Therefore, we expect that (C.1) can be expanded using 2 as Exzo C+2 B. Thus, Vr/Tr(A X, y, I, v0)), b(0), a()) alm Σk20 Ca Bax holds. Then, to calculate the expected correlational gradient Ely Vr frFl, it suffices to calculate Efy-2 Bx = (10.(((3)))) x/p/d)* Bx Bx]]=
E[o (3" d/ 33 + E[o, ((3, x) ((( x/pv k- -33 /pvd].
Now, if k < ie(o.) 1, the first term ([ a) ( // 33 is zero because a degree-k polynomial does not have Hermite component of degree 2 ie(a') = ie(o.) - 1. If ks = ie(o.) then the expectation is nonzero and proportional to E| 1, as desired. We can also show that this is the leading term in the entire gradient. Therefore, we focus on checking _ie(o.)-1 - term in (C.1).
Let proceed with the strategy above. Set N (((r2do(ox)+2) to ensure that all the noise terms in Ti T2, E2 log C d) for sufficiently large C. Then, we can simplify s ((X, y, x. v) as s( ,, (1e(
are of (()())))))))))))/1///1/ log -C d), where
r X, u. O r( r(X, y, 2,,
(o(a.)- 1-1
1o(0.) 1-1 Σ
[e(a.) [-
A
B
Bi+1
kp
2
k p
kp
Σ
i!
(a1)
(02)
1e(v.)-1
(e(o.)-1
le(o.)
B₁
A4t1 zt
B 2b i!
kp
Σ
kp
+ Kn
i!
(b1)
(bb)
(kpB))-P(
2w + 3w² +... (-1)i0(0.) (1
+ 0(210(0.))
(6)
where
B1 +
280
Bg 2²
+
Bagg 1-1 200(0.)-1.
1111B
Here we used the fact that T2/ exp
2pdd
|||1|
Bo 1
Bi
+
B2 x2 +
+ (le(o.)-1)1B,
Bia -in(r))-1) ((2in(7))).
Bo
2Bo
There are a lot of terms of degree 3 _ie(o.)-1 above, but we can show that the terms
Bie(o.) (kpAo+k,Bo)) ko (ie(a.) 1)!
K Bso(a.) ie(a.)-1 (1))))(.kBB (o) Bo kp (ie(a.)- - 1)! (ie(o.) (ooc)--1)
Aic(o.)
jo(o.)-1. (.,.)--.k,B)))
(k)
(C.2)
made from the first term of (a1) (and (br )), the last term of (az) (and (b2)), and the first term of (c), dominates over other terms, i... other terms having degree 2io(c.)-1 are of oa((C.2)) considering the coefficients.
(i 2 1) holds where e, is defined in Definition 14. Moreover, from Lemma 16 it holds that Ag = 2 and Bn = @(1). Note that from the definition of ea and the information exponent, we can see that Eio(o.) = 1 and then Alc(o.), ,, Bie(s.) 0(p-1) holds. Thus, noting that Auto.) kp = (1) and K,- = O(p-2), indeed the term Bok (Ie(o.)-I)) (hip Bb) -2 O(p dominates over other terms.
First we calculate the order of (C.2). From Lemma 9, A,, B, = θ polto->
Next we show the superiority of this term over other 2in(a.)-1 terms (in terms of the degree of p in coefficients): all the terms are made by multiplying a term in (a1) (or (bi)), term in (az) (or (b2)), and a term in (c). Let us see the p factor in
23

---

Nonlinear transformers can perform inference-time feature learning
the coefficients: that again from the definition of the information exponent, eq(a,) (and then A,, B, O(p 2)) if
1 S i ie(g.). Hence, be comparable to p=1 should pick the zeroth-order term from (a1) and (b1). Here the only optimal way is to pick out the zi<(=.)-1 term in (a2) and (b2), which is nothing but (C.2), Otherwise, we should pick out also the zeroth-order term in (02) and (b2). In this case, we should use gi((r))1 term in (c). but this is suboptimal because B./ Bo w is already (p)-))
Getting back to (C.1), we now showed that
s sX y. (Po P1= Po((.))-120000 10(0)))11 +...)8 33+ O(r-ddddd''0 (le(a, (telo.)+1)/2 (e log log- d)
where Pi((o.)-1 (pp=)) which completes the proof.
Lemma 21. Let ((0) = Ia/Va and v(0). b(0) and a()) are determined in Algorithm 1. Then, assuming Npt 2(12d (2))))) with high probability over the data (X, y, x), it holds that
Ti E (frr(X,, y', ",,;; (()) (0066(0)..)) a()) - )²)
V 2T
r=r(0)
=0 (ie(a))) ( (ie((a)))) | 1 O(amr- d-io(o)) log Cn d) + Ola2m²
with high probability over the pretraining data distribution and random initialization, where On is sufficiently large constant.
Proof. Now we have
T Σ
T Σ G G
V
- -²²²² .(0)
2T
Ti
Ti
(-r(0)
where f = (X' y : r(0), (0). b(0), a))) and G Vrfrr(X' y', x r, v(0). b(0), a(0) We first upuer-bound the first term 7. E fOG. From Lemma 15, with high probability yf ((1) holds for all and t. Therefore,
Nyt y.e Tm/p
|fTe(X', ,y x v(0), ((0) a()) s am
i=
(am)
Npt
p000 Tw/p Two
holds with high probability. Also, from Lemma 20, have
,, [o =aLm + +....... .... Pio(a.) (0(00--11111 1 )3x² + aL nxx
where Po...., Pe(o.))-2
O(p-1) and ||||| - 0(r 1d d (ie(o.)-1)2 1)/2 log Co -Ca d). Moreover, ||>||| (VVd) holds with high probability from Lemma 28, meaning | = Oaam d). Overall,
Σ /' Ō(a2m²
(C.3)
ensured with high probability, Next let us calculate
= aLm Ti EVG t=1
T
Σ
PI(3)) ... Pielo.) ) (2 +.... 1.....))))8
(x') (x')T + (n
T
T
....
24

---

Nonlinear transformers can perform inference-time feature learning
Note that the norm of m' is ((1) with high probability. Then, alen Et= m (x')) - ElaLm mi (x (x') = Ti
O(amd/ /T) holds with high probability from Corollary O(r-'d d d log Cn d log doo d) holds (additional log dogo/.))//2 d comes from the order of ys see Lemma 15. dil(a.)- Hereafter we absorb this by retaking C.). Thus, if d then
32.
Moreover, ||n |
1 Ti
V''
aLmB/ma '] O(amr 'd io(a))/2 log log- d)
(C.4)
satisfied. remains to calculate the expectation Elyz Bx (IVo) ((3, x)) (B, x)* Bx " analyze E[mæ ] From Stein's 's lemma, we have
(ovd) (oVd)~"Epall lo (13, (3, x3 Bx
=(ova) (oVad)~EEppa [Em ((B, (000).)))) m)) (B, 0)) [33 + (ov " (13, (3, a² k- 113
and from Lemma 34 and the definition of the information exponent, Ex oo ( (B, x)) (3, (B,x) is zero if & < ie(a.) =
ie(o.) 1 and nonzero if l ie(a=) 1. If k ie(a.) - 1, then (pvd) "E EE oo ((S, (B, (Bxx)"188T x (pva) (i(e) [8 is satisfied from Lemma 34. We can observe that other terms are also proportional to E but dominated by this term. From Lm m with high probability and Pie(o.)-1 = O we have
a almBmma 1
((amdd((eee)))-) p0 (le((, 133
(C.5)
Putting (C.3), (C.4) and (C.5) together completes the proof.
Setting the hyperparameters for Lemma 21 immediately yields the following conclusion of the optimization of I: Proposition 22. By setting the hyperparameters as am
(ie(e.)+11/2 log Ca d), Npt, (o)2)
7/1 = O(a - 10ma-mmmm> (aa-lii((o.))1)) log -Cus d)r 1/2 and d = 7h'" for constants Cy, and Ca. after one gredient descent step (Line Algorithm 1) it holds that
1/2 log l= d (rEg(33T] ( +N)
with high probability over the data distribution, where || Oa(1/Va) holds, where Cx can be taken to be sufficiently large.
This proposition can be immediately connected to Proposition 8.
D Optimizing the MLP Layer
In this section, give an optimization guarantee for the MLP layer. We begin with constructing the output layer parameter Rm yielding sufficient approximation guarantee.
Lemma 23 (Construction of Continuous Output Layer). Let l = 1 or 2 and assume ge(a.) L. Suppose that there exists
g(x) such that
g(x) = P; + P2 (13, x))²
11/ + n(x),
where Pt P2(1), 0a ((og d) Cr.) ), and [n(x)) oa(P 31/4 log log- deglo. d) holds (these conditions are indeed identical to the conclusion of Proposition 8). Then, there exists n(U, b) such that
Eur (-))))(-))) (v) g(x) + b)] (1B, x3)] = oa(1)
with high probability over æ ~ N(O, Ia). Moreover sup., (TV, b)] = Ö(r') holds.
Proof. Note that 1 = 2, then o, is even: this means that there exists polynomial s. such that (2(2) = sx(2") for = 2, Here let (2) Σ
deg(s.) st Now from Lemma 9 in Damian al. (2022), there exists (v, 6) such that
sup,,> (v, b) O(1) and
B Und (v, + b)] **
25

---

Nonlinear transformers can perform inference-time feature learning
for any |=| S 1. If we define
deg(s.) Σ
Th (v. ba 1lo0-log= 1lo-22)) a alogg2 d
log d.
k=0
where a P2r-1/2 then
Bnn-thiif(t1,,bb(laooood do ( (((()))(((())) + +b]
deg(s.)
( ba log 3 d) log (o((2( b)
sg 9 a,a log?
a log"
k=0
deg(s.)
deg(s.)
8% log 2k
(B)p)Uni((()))bb1-1)) (1, log d b)]
Σ
8k2 =0(2),
k=
k=
if 12 lol -2 2d| < 1. Note that sup,,6 (1, b)l = O(aa log log dee(( d) holds.
Furthermore, let define events E_ (v. b) and E.+ (v,
(_ b) = (v (v, = -v= i-p + log d))
[P] a log" d)),
and then define
1
x(u.b) 2a log d
(I(E (v, (B)
(D.1)
Now holds that
E Un--Unl(((1)ddd----)))) b)o(v g(xx) + bb]
Eenips LallPP-lllga log d. (1
n(x) + b)]
+ } B E Po d, Po
Pi) a x\)² n(x) - b)]
Rep Un)
log² | [r' (u. () (B, + b]] oa(1) a. ((),x)) + oa(1)
holds (note that (G, x))
log" d with high probability from the tail bound for the standard Gaussian).
Here, we used the Lipschitz continuity of the ReLU function and the condition [n())] os(a log deg( d) derive (*).
Lemma 24 (Construction of Discrete Output Layer). Under the condition of Lemma 23, there exists a" E R such that
m Σ
ajolo g(x) +b)) xx((3xx)) - oa(1)
(D.2)
with high probability over NOO. Ia). Moreover, || m 3/2 + 31 holds with high probability.
Proof. Let a, m ( by) where æ is constructed in Lemma 23. As g(x) O(1) and sup,, ((u Ö(r²). from Hoeffding's inequality we have
m m Σ
Tl, g(x) + b,) Ev,b () g(x) + b)] Ö(r'm -1/2)
with high probability and then we obtain (D.2).
26

---

Nonlinear transformers can perform inference-time feature learning
Let us upper bound ||a||2 m- m プラ= 77 (T,,, b,)2. Again from Hoeffding's inequality, we obtain
m
T(B b,) Eu,b/T(U, 6)3]
m
and then ramains to upper-bound [v, b) Naively sup,,o T(v, b)2 Ö holds, but from the definition (D.1).
(u, b) is nonzero with probability O(r-1/2), meaning (e, ((131/2). Overall, ||a| O m -3/2 +
/ m 1) is satisfied with high probability.
Proof of the Optimization Part of Theorem Now we are ready to show the optimization part (part 1) in Theorem 1.
Proof of Part 1, Theorem We note that the condition in Lemma 23 is satisfied from Proposition 8, and the condition in Proposition 8 is ensured by Proposition 22. Let a' be the output parameter constructed in Lemma 24: from the equivalence between fz-regularized and norm-constrained optimization algorithms, there exists X2 such that
Ti +
T
2
Ti +TB
1
Σ
y/ - y' , r," v² ,,,"",',, b" "','',a']]
Σ
(v' - -TF(TX , ,",,",
T2 1=++1
T2
1=1+
T Ti+Ta
<
Σ
(y² S-r(( v. x' r v², a'))
T2 t=
$ (T o1)) (:: Lemma 24, m (((2200)))))
-3/4 second term 3ge((-) 4m 1/2 dominant.
and it holds that ||a"| S ||a'l Ö(180(a.) 3/4 3gelo.))44 11/2). Note that as we assumed m 52 (g((o))))) the
E Inference-Timee Time Estimation Error Analysis
We are ready to analyze the ICL loss at the inference time defined by
RCL Newe (T", b", a²) Ep..D, De (I
(X I:No : VI.N 2,," 00, b" a") -
In particular, our primary interest is in deriving the inference-time time sample complexity, i.e., the number of in-context examples Noes in the test prompt needed to ensure that the ICL loss is o1..
E.1 ICL Error at Ntest Npt
We first treat the special case where Nost Npt holds: in this case, as we have an optimization guarantee at context length Npt. it is sufficient to bound generalization gap via a standard Ramemacher complexity analysis.
Recall that g is defined as the attention output, i.e.,
ΣN e Tw:/p Σ
g (X, v.
e Te.jp
In this section, fix r, v and b. We define Fa describe the set of transformers where the norm of the output layer of the MLP is constrained:
FA
+ (X, (X.
Σ
ajolo, g(X v. x; v;
by)
||a|| S A
1=1
Moreover, let
Rady = (
sup Σ JeFA T
1
f((x
1=1
be its Rademacher complexity, where et Unif (1+])) Here, we fix the context length to be N = Oppolyd), i.e, suppose X E RdxN and y E N Then, we obtain the upper bound of Radr(Fx) as follows.
27

---

Nonlinear transformers can perform inference-time feature learning
Lemma 25. When ||||| B and v/ +1, then it holds that
A(B + vm)
(ady(FA)= =
VT
Proof. From the definition of FA, we have
T
Radr (Fa) 18 ( Me
sup
aoo (u g''",, bj)
a: a|
T
1=1
j=1
1 T
v',,," +b)))
sup a: al SA
Dj
T t=1
j=1
From Cauchy-Schward inequality and Jensen's inequality, we have
Radr S [xxx),
sup a:lla A
|||||
T
e'al g(X yy xt; r) + bb))
T g(x, y', x'; I) + bb)
S A (00)61
caluy
T
Next, we obtain the upper bound of the component in the square root. Since et and c" are independent if + t', we have
27
Σ
fa s(x' x0; + bi)
é' (0, g(X' y x';) + by)o (1) x(x", y" ("'a":r) bb))
= -Exxittoeaa
T
1
(1) (X x';
T
Now, a is ReLU activation function, which means that x22)² 22 for IR Therefore, we have
m
T
a ( g(( 3 x²; 1) + b,)²
j=1
m
<
Exxc,,00))a
g 3 x 1) + b,)
<
E, E(x"),(u')(cr)
(o) 9 x)² 63)
T2
B2T + m Σ E May [x y6 ,: r)3]
< T
Finally, we bound the term Elg(X' yy, x a.::r:'' 1)2]. For 1, N, let
evilp Ta/l
x)
ΣN=1 edv eue/pe r.alpp
28

---

Nonlinear transformers can perform inference-time feature learning
N Then, have y, Ex (x, 1)4; and Σ 1=
Pa(X y, v,
Pa
Hence, from Jensen's inequality, we have
N
x'x
Σ Pi y''' "'' x²)))
Σ P: x²)))
Now, we define random variables Z; 2 (i 1,...N as
Z p; y x² (v/) 7i 10=m((X px(X', "0',"") T((())) z (log d) d)deslo.)+t))
Thus, using Hölder's inequality, we have
E (xx')))) 2,1
< E(x pi((' 3" x² (1)))))) ((A)) (log (logd)des(on) d)des(00))))
s [(x)) [P: y² xf ,²){] x (x))),e)) [(v))]] 1/p
(ot)² /2p
(log dddugla.)+) 1/4
s ((x)) [())'']] '/p (1x))= 2 (log (log.)) 1/1
Using Lemma 15 and the definition of high probability events, can set the quantity above to be O(N=3) (note that assume N = pdly Consequently, we have
N
E (w)))) y ,',
r);
B[2.]
i=1
N
< 01)
B[27]
< 0(11+ +E
max
{ (v/)² .P((())())) < (log (oggddeoloss)+)))
i=1....N
< O(llgg d)deglo.)+1).
(E.1)
In conclusion, obtain
2
BTT + mT -0(1)) = =0 - mm)
Radr (Fa) < A
7
VT
Then, we arrive at the following lemma:
Lemma 26. Suppose running Algorithm 1 under the conditions specified in Theorem 1. Then, with probability at least 0.995 over the training data and the random initialization of Algorithm 1, it holds that R RICL, (T" o0" b", aaa) oo((). Npc
Ti+ Proof. Now we know from Part 1 of Theorem 1 that T - lvy STF( (X',y' x², r v","", b" ",,,a))))))))))))))))))"))))"" = oa(1) with high probability. Moreover, we have parameter norm bounds ||a") = 3pp3)((0)))44 11/2) and |bb| I vm. Plugging these values into and B Lemma 25 yields
(100e((o))/4
Radr, (FA) =Ō
VTz
Using standard symmetrization technique (ef: Appendix D.3 in (Oko et al., 2024b)) yields RICI oz(1)+ō with probability at least 0.995. Noting that T2 3(ee)))()))2) we arrive at the assertion.
(~ Aaolo.)/4
/T5
29

---

Nonlinear transformers can perform inference-time feature learning
E.2 ICL Error at General Mest
Finally, we extend the generalization result to any inference-time context length Ntest, which concludes the proof of Theorem 1.
Lemma 27. For N I polyd,
Ex Bx VUN,,, "","", ",,,",""" a²) ((r
holds.
Proof. Note that
fre(X MIN, ,", a") Σ
ajoov INN vi;,);;"") b7)
j=1
where Egg( (000X M1:N, 2; r- (1)) from (E.1). Then,
[Ir X I:N, VI:,, x; r" 0" b' a²)
aa 11 <
+ b7)²
SE
|aa1|
(INN, < =(((360(o.)123 +
holds.
Finally, we prove the Part 2 of Theorem 1.
Proof of Part 2, Theorem I. Note that
R ICLL "²,b",aa)) ICL (T", v², "0"""" b², a²))"
<E|lFr X 1:N U1:No1 1 v", b", a") FIr( (rr(X1.00 I:Nea Mr:Noa I. 00, b", a")l
holds. From Proposition 11, if Npt: Noett = (((13ge(0.)12)) then
lg( X 1:N yi.Nioa vy r*) g( X IIN VI:N, az; r*] S oalr 3go(a.)/4 log log- d)
holds with high probability (here, C can be made sufficiently large by taking the constant C₁ to be sufficiently large), where the upper bound quantity at right hand side does not depend on the data. Note that
fTF (X 1:N. MINNo. x;; r v² b", a²) fFFF(X 1:NNest V1:N :. r 00, b", a")l
1a5l lo X 1:N VI x;; r*) b7) (V3g((X N:Nua Vi:Noa T; r*) + b7)
< |a Vm- ca(r 3ge(o.)/4 3/4 log d) (t. Cauchy-Schwarz).
Using ||a" Ö (ga(oo )/4 -122 in Part of Theorem now obtain
|frr X1:N,. D1:Np x; D", v², b,,,""))))))"""""" a²) fir( (X1: PI:Nreas x; P TT,"","" ",,,b, ",,,""aa"" a²)l I oa(1)
with high probability. It remains convert this bound into the expectation. Note that the os(1) quantity at right-hand side
does not depend the data: we write this quantity as qa- We state that Pi]fre [X1:N VI:N,. x; r", o" b", a²)
Fre( X N:Ne V1:Noa x;; I 0², b" a*)) < ga] < 1 - d-1 from the definition of the high probability event. We write this event as E. Consequently,
Bl]fTr X1:N ,1:N,,: x; D D" a²) FIT X XI:Naas U1:Nees xz r D² b" a")]]
30

---

Nonlinear transformers can perform inference-time feature learning
Sqa + [ STF( X 1:N VI.N, az r 00, ,,,,,, a²) /TFF( X 1:Noc M1:Neom w. I b. a" ()((( [(E))
Sqa + [ STF X 1I:Np U1:No my r v² b, , a²) - /TF( X N::Nec Niew R. r V. b. a )/311// (d
Sq. + 2 (X 1:N VI.Npa a; r 0¹ b. )²² + 2fTF (X L:Nies & Ni E; r 0. b. a-) )²)//1/2 (d (d=+)//2
Σqa + 0(p 2)
From Lemma 27. Then, from ge(a.) < and d, we obtain | n², b" ,0)aa R (T 0² b" a oa
and arrive the conclusion from Lemma 26.
F
Auxiliary Lemmas
F.1 High Probability Events on Gaussian and Spherical Variables
Lemma 28 ((Wainwright, 2019), Example 2.11). Let
v 2 + ... + x? where ~ (0, Then, for any 0 <
FIl - rt $ 2 eexp( (
holds. Consequently, =² < Orrvlog d) holds with high probability, and if d, d/2 2² S 3d/2 holds with high probability.
Lemma 29. Letr
NO L)) NOO, I,) independent of æ, and r r matrix dependent of neither æ nor | 2 log holds with high probability.
y. Then, (xx Gy) [[
Proof. x be decomposed
2B, where B Uni and = ||||| independent of B. Thus, (x, Gy)
(aG B, v) holds. From Lemma 28 we have G|= B1 O [[| with high probability, and if we fix
satisfying this, then (zG* T B. B.y)² = [[ Vlog d) with high probability again from Lemma 28 (in the case
r == 1). This yields the assertion.
Corollary 30. Let S d, B Unif(S y NN, I+) independent of B. and a v matrix G dependent of neither nor y. Then, (3, Gy) = O[I Glllllloo holds with high probability.
Lemma 31. Let æi...., N ~ NO0 Ia) and let , 2N be i...d. random variables satisfying |2a| S C C(2 might depend æ;) with high probability. Then
N Σ 2,; B[ara]] SÓ(CVa/N)
holds with high probability.
Proof. Define 2 I[[2g] S C)², Note that zι0, is a sub-Gaussian vector, i.e., (2/0,, u) C-sub Gaussian for any (dd1( Then, standard concentration bound for sub-Gaussiaa vectors yields
a - B[sfai] <((oV//N)
with high probability. As N- -i=1 2,,; = N EN z/m; with high probability, it remains show that [[[[
[|2 is sufficiently small. This can be seen by
[ [[ [[
= 2| 2 C)æ1Jll
= E[T([=1 2 |)
< Odd-C.).
where C. can be taken to be sufficiently large from the definition the high-probability event. This completes the proof.
Applying Lemma 31 to each row yields the following:
31

---

Nonlinear transformers can perform inference-time feature learning
Corollary 32. Let $i,...,, aN NO, Fa) and let o1.... un be i.i.d. random vectors satisfying ||ol1| C (vy might depend on ai) with high probability. Then,
Σ vada E(v
S <((Cd/VN)
holds with high probability.
F.2 Hermite Polynomials
Lemma 33. For g(z) Σizo /) (2),,,)
20 i!
holds.
Lemma 34 (Damian et al., 2023), Property 1). For a. B E gd-1
Ex-Noo.r He, a) (Z, Bi) = i! II(i = 5) (a. By
holds.
G Experimental Details
For the experiment in Section 6, used 6-layer GPT-2 model (Radford al., 2019) following the configuration used in (Garg et al., 2022; Oko et al., 2024b). Precisely, the model take the length-((N 1) prompt (X1.N41, V1:N+1, as input. Then, we first make the embedding as
E = Win
æ1, y pad
N1, D N+1 E (2N+2)
pad
(G.1)
where y.aa [y,, 0, .......) O and Win E Rxd a trainable read-in parameter. A 6-layers GPT-2 backbone with attention heads, whose configuration same as that Garg et al. (2022) then convert this embedding into E' Dx where D 256, and we obtain
22 22N +
(w T E'
(G.2)
where out E RDD. is also trainable. Here, we use 221-1 as the prediction of based on 31...... ... Ti-1, 0: and ay Causal attention mask ensure that there is no information leakage. We pretrain the model using the Adam (Kingma & Ba, 2015) optimizer with learning rate 0.0001 on the mean-squared loss calculated over all the positions. We utilize curriculum learning setting which is also adopted in (Garg et al., 2022) to reduce the pretraining cost. we start from = and increased the input dimensionality of by two until the target dimensionality.
The test loss averaged over 256 independent runs, where each run contains 8, 1,9 independent batches.
32


Reviews:
[{"review_id": "0BtUMqE5KU", "reviewer": "Reviewer_V9j6", "summary": "This work studies in-context learning (ICL) of single-index models $\\( y = \\sigma(\\langle x, \\beta \\rangle) \\)$ using nonlinear transformers, focusing on inference-time sample complexity. The authors propose a two-stage training approach: (1) a single gradient descent step on the attention matrix to capture feature structure, and (2) ridge regression on the MLP layer to fit the nonlinear transformation. Unlike *Oko et al. (2024b)*, which assumes a low-dimensional feature structure, this work removes this restriction while achieving better inference-time sample complexity (Table 4.1).  \n\nThe main findings show that nonlinear transformers outperform non-adaptive methods (e.g., kernel methods) by extracting task-specific features at inference time, without parameter updates. Moreover, the softmax attention mechanism plays a crucial role in feature learning, helping transformers surpass the Correlational Statistical Query (CSQ) lower bound, making them more sample-efficient than CSQ-based methods. The study provides rigorous theoretical guarantees on both pretraining complexity and inference complexity , demonstrating that pretrained transformers can adaptively learn features with superior statistical efficiency compared to traditional approaches.\n\n## Update after rebuttal\nThank you for the thoughtful and detailed rebuttal. The clarification of the theoretical contributions—particularly the improved inference-time sample complexity and the removal of low-dimensional assumptions—meaningfully strengthens the paper. The additional experiments further support the main claims. Based on this, I am increasing my score and recommend acceptance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback. We address the technical comments and questions below. \n\n**Distinctions from Oko et al. (2024b)**\n\nBased on your suggestion, we will include a dedicated comparison section in the paper; here we briefly summarize the key differences. \n\n- **Improved inference-time sample complexity** – In Oko et al. (2024b), the derived inference-time sample complexity is $O(r^{\\Theta(deg(\\sigma_*))})$, which corresponds to the sample complexity of kernel methods on an $r$-dimensional subspace. In our paper, we derived an improved sample complexity of $O(r^{\\Theta(gen(\\sigma_*))})$ by showing that nonlinear transformers are capable of inference-time feature learning (beyond the kernel regime).\n- **Absence of low-dimensional assumption** – While Oko et al. (2024b) assumes that the dimensionality $r$ of pretraining task distribution is much smaller than the ambient dimensionality $d$, our analysis does not rely on this assumption. \n- **Theoretical analysis** – To derive the improved sample complexity $O(r^{\\Theta(gen(\\sigma_*))})$, we demonstrate that the nonlinear transformation applied to the output label $\\{y_i\\}$ reduces the information exponent of the link function $\\sigma_*$ to its generative exponent. In particular, it is essential to show that the degree-$gen(\\sigma_*)$ term in the Hermite expansion of the nonlinearly transformed labels (via the softmax transformation) is non-vanishing. This requires a careful tracking of the corresponding Hermite coefficients, which constitutes one of the most technically challenging parts of the analysis. This aspect is not addressed in Oko et al. (2024b), where it is only shown that the linear attention module computes correlations on the raw label $\\{y_i\\}$ and a fixed kernel basis $\\{\\phi(x_i)\\}$ (constructed by a fixed MLP block). \n\n\n**Additional experimental evaluation**\n\nThank you for the suggestions on experiments. We make the following remarks. \n- Our goal is to rigorously study the statistical efficiency of ICL in a well-defined target function class (single-index models). Therefore, similar to prior works that studied ICL for linear regression [Garg et al., 2022; Von Oswald et al., 2023; Ahn et al., 2023; Mahankali et al., 2023a; Zhang et al., 2023], we conducted experiments (on GPT2 models) in a controlled synthetic setting to probe the sample complexity scaling. Pretraining on real-world datasets where the information/generative exponent of the target function is unknown is tangential to our main contribution. \n- To further validate the derived test-time sample complexity rate, we have added additional experiments to validate that the sample size scaling is indeed beyond CSQ (please refer to [[link]](https://osf.io/p87yc?view_only=be342dc212bb49cb8f72309fdc37d376), see also our response to reviewer 7xsp for details). Following your suggestion, we are currently running ablation studies to probe the necessity of the Softmax attention. \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.", "claims_and_evidence": "The paper provides theoretical justification for its claims, particularly in demonstrating improved inference-time sample complexity over kernel methods and surpassing the CSQ lower bound. The main results are well-supported by formal optimization guarantees and statistical complexity bounds. Additionally, Table 4.1 effectively contrasts the findings with prior work (*Oko et al., 2024b*), highlighting the relaxed dimensional assumptions and improved efficiency.  However, some claims could benefit from stronger empirical validation. The synthetic experiments with GPT-2 provide limited practical evidence, as they do not comprehensively test real-world generalization or scalability. Additionally, while the softmax attention mechanism’s role in feature extraction is theoretically well-motivated, a deeper ablation study on its necessity and effectiveness would strengthen the argument. Overall, the core theoretical claims appear well-supported, but further empirical evidence would enhance the credibility of the results.", "methods_and_evaluation_criteria": "Yes, both the studied setting (in-context learning of single-index models with nonlinear transformers) and evaluation criteria (inference-time sample complexity) make sense. The single-index model is a well-established framework in statistical learning and serves as a meaningful testbed for understanding how transformers learn features at inference time. The focus on inference-time sample complexity is particularly relevant, as it directly measures the efficiency of transformers in extracting task-specific information without parameter updates—aligning well with real-world applications of in-context learning.", "theoretical_claims": "The theoretical claims, particularly regarding inference-time sample complexity and softmax attention’s role in feature extraction, are well-structured and follow established techniques. However, I did not verify every step in detail. Key results, such as the gradient descent analysis on the attention matrix and the Hermite expansion argument, appear reasonable but would require closer scrutiny to confirm correctness. No obvious errors stood out, but a deeper review is needed for full validation.", "experimental_designs_or_analyses": "The synthetic experiment using GPT-2 provides a reasonable comparison between in-context learning with transformers and kernel methods, aligning with the paper’s theoretical claims. The choice of a single-index model with a Hermite polynomial link function is appropriate for testing inference-time feature learning.", "relation_to_broader_scientific_literature": "I agree—showing that nonlinear transformers achieve better sample complexity than kernel methods is a strong contribution to the theoretical literature on in-context learning (ICL). This work extends Oko et al. (2024b) by improving sample complexity bounds under more general assumptions, reinforcing that transformers can surpass fixed-basis methods through adaptive feature learning. More broadly, the paper connects to research on neural network adaptivity, highlighting how softmax attention enables efficient inference-time learning. These findings contribute to the study of statistical efficiency in deep learning, showing that transformers can generalize more efficiently than traditional approaches.", "essential_references_not_discussed": "To the best of my knowledge, the paper appropriately cites and discusses relevant prior work.", "other_strengths_and_weaknesses": "Strengths\n\nThe paper tackles a more challenging setting than most theoretical works in the literature, which often rely on simplified transformer models. By incorporating softmax attention and an MLP layer, the study resembles real-world transformer architectures, making its findings more applicable. This more realistic setup allows the authors to establish better sample complexity results than *Oko et al. (2024b)*, demonstrating the advantages of nonlinear transformers for inference-time feature learning.  \n\nWeaknesses\n\n1. Presentation and Organization Need Improvement: \n   - The paper’s structure is difficult to follow, making it harder to grasp the key contributions at first glance.  \n   - Theorem 1 could be stated more clearly to improve readability.  \n   - The distinctions from Oko et al. (2024b) are spread throughout the paper, making them hard to evaluate. A dedicated comparison section in the related work would clarify these contributions.  \n\n2. Limited Experimental Evaluation:\n   - The paper only includes a single synthetic experiment, which is not sufficient to validate the theoretical claims in real-world settings.  \n   - Adding an experiment on real-world data would significantly strengthen the empirical support for the paper’s conclusions.", "other_comments_or_suggestions": "-", "questions_for_authors": "- Distinctions from Oko et al. (2024b): While the paper improves sample complexity results compared to *Oko et al. (2024b)*, the distinctions are spread across different sections. Could you provide a more explicit and structured comparison in the related work section to clarify the key differences?  \n\n- Role of Softmax Attention: The paper argues that softmax attention plays a crucial role in inference-time feature learning. Have you considered ablation studies or comparisons with linear attention to isolate its impact more clearly?  \n\n- Empirical Evaluation: The experimental validation is limited to a single synthetic experiment. Would it be possible to evaluate the method on real-world datasets, and if so, what would be the key challenges?  \n\n- Scalability of Pretraining Requirements: The pretraining task/sample complexity scales with \\( d^{\\Theta(ie(\\sigma))} \\). In practical applications, would this requirement be computationally feasible, especially for large-scale transformers?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fv8oIE8gyV", "reviewer": "Reviewer_7xsp", "summary": "This paper studies the optimization and statistical guarantees for the in-context learning of the Gaussian single-index function class using a Transformer with softmax attention. The derived inference-time sample complexity is tighter than the existing works, which indicates that pre-trained transformers can implement full statistical query (SQ) algorithms. During the proof, a mechanism of ICL is also characterized.\n\n--------------------------\n## Update after rebuttal.\n\nMy concerns are addressed overall. I prefer to keep the current score and support an acceptance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback. We address the technical comments below. \n\n**Additional experiments** \n\nThank you for the suggestions on experiments. We have conducted an additional experiment to probe the test time sample complexity of GPT-2 models for learning single-index target functions, and demonstrate that the sample complexity surpasses the CSQ lower bound.  \n- We pretrained GPT-2 on a degree-3 single-index model  $y=He_3(\\langle\\beta,x\\rangle)$, in the setting where $r=d$, and plotted the relationship between the ambient dimensionality $d$ and the minimal test prompt length $N_\\mathrm{test}$ required to surpass given test error thresholds. \n- In [[link]](https://osf.io/p87yc?view_only=be342dc212bb49cb8f72309fdc37d376), we observe that the estimated inference-time sample complexity is approximately $d^{1.13}\\sim d^{1.16}$, which is lower than the kernel lower bound of $d^3$ and CSQ lower bound of $d^{1.5}$ for $\\mathrm{He}_3$. Note that the CSQ lower bound also suggests a statistical hardness for online SGD – see e.g. [[Damian et al. 2023]](https://openreview.net/forum?id=73XPopmbXH&noteId=eMjWYg6qko).  \n\n\n**Clarification on Algorithm 1**\n\nWe make the following remarks on Algorithm 1. \n- Indeed Algorithm 1 involves a few ingredients that simplify the theoretical analysis, similar to recent works in the feature learning theory literature. In particular, a single gradient descent step (Stage I) has been used to model the early phase of gradient-based training on neural networks [[Ba et al. 2022]](https://openreview.net/forum?id=akddwRG6EGi) [[Damian et al. 2022]](https://proceedings.mlr.press/v178/damian22a/damian22a.pdf) and transformers [Oko et al. 2024b], during which the first-layer parameters align with the feature vector of the single-index target and escape saddle points in the nonconvex landscape.  The layer-wise training strategy – dividing the algorithm into two stages, one gradient descent step on the inner layer and ridge regression for the outer layer – is also common in these prior works, where the lower layer captures the relevant feature, while the top layer (+ reinitialized bias units) approximates the link function. \n- While our theoretical result applies to a specific parameterization and gradient-based learning algorithm, we empirically demonstrated in the additional experiment above that standard end-to-end training of a GPT2 model can achieve an in-context sample complexity that goes beyond the CSQ lower bound. We believe this suggests that mechanisms similar to our theoretical construction may be present in more realistic settings. \n\n\n\n**Distinctions from Oko et al. (2024b)**\n\nWe will include a dedicated comparison section in the paper; here we briefly summarize the key differences. \n\n- Improved inference-time sample complexity – In Oko et al. (2024b), the derived inference-time sample complexity is $O(r^{\\Theta(deg(\\sigma_*))})$, which corresponds to the sample complexity of kernel methods on an $r$-dimensional subspace. In our paper, we derived an improved sample complexity of $O(r^{\\Theta(gen(\\sigma_*))})$ by showing that nonlinear transformers are capable of inference-time feature learning (beyond the kernel regime).\n- Absence of low-dimensional assumption – While Oko et al. (2024b) assumes that the dimensionality $r$ of pretraining task distribution is much smaller than the ambient dimensionality $d$, our analysis does not rely on this assumption. \n- Theoretical analysis – To derive the improved sample complexity $O(r^{\\Theta(gen(\\sigma_*))})$, we demonstrate that the nonlinear transformation applied to the output label $\\{y_i\\}$ reduces the information exponent of the link function $\\sigma_*$ to its generative exponent. In particular, it is essential to show that the degree-$gen(\\sigma_*)$ term in the Hermite expansion of the nonlinearly transformed labels (via the softmax transformation) is non-vanishing. This requires a careful tracking of the corresponding Hermite coefficients, which constitutes one of the most technically challenging parts of the analysis. This aspect is not addressed in Oko et al. (2024b), where it is only shown that the linear attention module computes correlations on the raw label $\\{y_i\\}$ and a fixed kernel basis $\\{\\phi(x_i)\\}$ (constructed by a fixed MLP block). \n \n\n**Practical insights**\n\nOur goal is to rigorously study the statistical efficiency of ICL in a well-defined target function class (single-index models). Therefore, our analysis does not directly yield algorithmic improvements in practical settings. This being said, we believe that the mechanism that we theoretically investigate (in-context feature learning) serves as a foundation of future studies on the capabilities and limitations of ICL compared to other algorithms that act on the test prompt. \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.", "claims_and_evidence": "The analysis is very impressive and solid.", "methods_and_evaluation_criteria": "N/A. This paper is mainly theoretical.", "theoretical_claims": "I checked the proof sketch, which looks reasonable and rigorous.", "experimental_designs_or_analyses": "The experimental design makes sense, although I feel the experiments are not enough. Please check the part of weaknesses.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The paper is well-written, and the theoretical analysis is impressive. I would like to mention some weaknesses here. \n\n1. Experiments are not enough. (a) Figure 6.1 only compares the kernel method and ICL, which correspond to the kernel bound and the SQ lower bound. It is better to also show the result of the CSQ bound, which shall correspond to one-pass SGD on two-layer neural networks from lines 316-318 Right. (b) It is better if you can conduct some quantitative experiments on the order of the sample complexity, even if you use a much simpler model. \n\n2. Algorithm 1 seems quite complicated. Is it designed for the simplicity of theoretical analysis, or is there any practical training strategy that is close to Algorithm 1?", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Can you summarize the technical novelty compared with [Oko et al. 2024b], especially about the difference in your settings (if any) and why you can have a stronger result than theirs?\n\n2. Can you discuss some practical insights from your results? For example, what can people learn from your results to improve the training and testing performance?\n\nOko et al. 2024b. Pretrained transformer efficiently learns low-dimensional target functions in-context.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HmLgO13eF0", "reviewer": "Reviewer_hCez", "summary": "This paper studies the tasking of learning single-index models $y = \\sigma(x \\cdot \\beta)$ using a two-layer single-head \nsoftmax transformer. The authors prove that a pretrained transformer can solve this task in context (different \n$\\beta$ in different prompts). When $\\beta$ is sampled from the unit sphere in an $r$-dimensional subspace of \n$\\mathbb{R}^n$, the authors provide the following bounds on the sample complexity ($\\mathrm{IE}$ and $\\mathrm{GE}$\nstand for the information and generative exponents of the link function $\\sigma$, respectively): \n* For pretraining, the number of prompts is $\\tilde{O}( r^2 d^{\\mathrm{IE}+2} \\vee r^{1.5 \\mathrm{GE}} )$ and \n  the length/number of samples in each prompt is $\\tilde{O}( r^{1.5 \\mathrm{GE}} )$.\n* For inference, the length/number of samples in the prompt is $\\tilde{O}( r^{1.5 \\mathrm{GE}} )$, which matches \n  the SQ lower bound up to a multiplicative constant in the exponent. \n\nThe pretraining algorithm is a rather standard (standard in this area instead of in practice): they first\ntrain the attention layer using one large step of GD and then train the output layer using ridge regression. They show\nthat after the first stage, the attention layer can recover the relevant $r$-dimensional subspace, and after the second \nstage, the model learns to reconstruct the link function using samples.\n\n## update after rebuttal\nI thank the authors for the response. I agree that while the correlation can be easily extracted from the attention in a linear way, leveraging the nonlinearity to improve the sample complexity is non-trivial. Hence, I'll keep my score (4).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback. We will correct the typos and improve the writing of the manuscript. We address the technical points below. \n\n**Properties of single-index target**\n\nWe agree with the reviewer that our current learnability result relies on the Euclidean inner product in the attention layer. We make the following remarks. \n- While the correlation between samples can be easily extracted from the attention, our analysis makes use of the Softmax function to transform the training data in some nonlinear fashion --- such nonlinear transformation is crucial to obtain an improved in-context sample complexity. In contrast, prior works using linear attention layers can only make use of (linear) correlation through the attention block. \n- We speculate that deeper architectures would be able to construct more complicated correlations between data points and further improve the in-context sample complexity (e.g., see recent works on the gradient-based learning of three-layer neural networks [[Nichani et al. 2024]](https://openreview.net/forum?id=fShubymWrc&noteId=efqPkOAyOx) [[Dandi et al. 2025]](https://arxiv.org/abs/2502.13961)), which is an important future direction.\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.", "claims_and_evidence": "This is a theoretical paper and the authors provide complete proofs for their claims. \nThey also include some synthetic experiments to support their claims.", "methods_and_evaluation_criteria": "Yes. (In-context) sample complexity is a standard metric to measure an algorithm's ability to learn single-index models\n(in context).", "theoretical_claims": "I checked the proof strategy but did not check the detailed calculation. The proof looks believable to me, as \nits components are either standard (ridge regression on the output layer) or believable (but non-trivial) extensions of \nexisting argument (e.g., learning the subspace using one large GD step, lower the IE to GE using monomial transformation \nthat comes from expanding the nonlinearity).", "experimental_designs_or_analyses": "This is a theory paper and the synthetic experiments they provide are serviceable.", "relation_to_broader_scientific_literature": "The main contribution of this paper over [Oko et al., 2024b] is that it improves the in-context sample complexity from \n$r^{\\Theta(\\mathrm{deg})}$ to $r^{\\Theta(\\mathrm{ge})}$. The former can be achieved via using kernel methods (in\nthe relevant subspace) while to get the GE bound, the model needs to recover the actual direction $\\beta$. \n\nIn addition, the authors use softmax attentions instead of linear attentions (which is widely used in theory as it is \nusually easier to analyze) and the existence of this nonlinearity leads to a bound beyond the CSQ lower bound \n$r^{\\Theta(\\mathrm{ie})}$ (which is what one will get from using algorithms such as online SGD to learn \nsingle-index models).", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Overall, this is a well-written paper and has strong results. See the previous sections for the strengths. \n\nOne potential issue of this paper is that it relies heavily on certain properties of single-index models. In particular, \nin this setting, the correct way to measure the correlation between two samples (in a given prompt) is exactly the \nstandard Euclidean inner product, so the attention layer only needs to identify the relevant subspace. This is more of \na coincidence instead of something universal.", "other_comments_or_suggestions": "Typos:\n* line 188 right: ... inference-time sample complexity of $\\tilde{O}( r^{\\Theta(\\mathrm{ie}(\\sigma_*))} )$ ....\n  Should be $\\tilde{O}( r^{\\Theta(\\mathrm{deg}(\\sigma_*))} )$.\n* line 229 right: $\\Gamma \\in \\mathbb{R}^m$. Should be $\\Gamma \\in \\mathbb{R}^{d \\times d}$.", "questions_for_authors": "See the \"Other Strengths And Weaknesses\" section.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dA6ouSgEqk", "reviewer": "Reviewer_asP8", "summary": "This paper studies the in-context learning capacities of transformers when the prompt sequences are given by a (possibly low-dimensional) Gaussian single-index model. When the length of the prompt sequences exceed certain (information-theoretic) limit, transformer trained with modified gradient descent method is shown to achieve vanishing in-context prediction error.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the helpful feedback.  We address the technical comments and questions below. \n\n**On Algorithm 1**\n\nWe make the following remarks on Algorithm 1. \n- The strategy of merging attention matrices and zeroing out some submatrices is extensively used in theoretical analysis of both linear and nonlinear transformers [Ahn et al. 2023; Mahankali et al. 2023; Kim and Suzuki 2024a; Oko et al. 2024b] – even with such simplifications, theories can reproduce fruitful insights into attention modules.  We also empirically demonstrated in [[link]](https://osf.io/p87yc?view_only=be342dc212bb49cb8f72309fdc37d376) (see also our response to reviewer 7xsp for details) that pretraining GPT2 model can achieve an test-time sample complexity that goes beyond the CSQ lower bound.  We believe this suggests that mechanisms similar to our theoretical construction may be present in more realistic settings.\n- Our layer-wise training procedure is inspired by recent works in the feature learning theory literature, which simplifies statistical and optimization analyses. In particular, a single gradient descent step (Stage I) has been used to model the early phase of gradient-based training on neural networks [[Ba et al. 2022]](https://openreview.net/forum?id=akddwRG6EGi) [[Damian et al. 2022]](https://proceedings.mlr.press/v178/damian22a/damian22a.pdf) and transformers [Oko et al. 2024b], during which the first-layer parameters align with the feature vector of the single-index target and escape saddle points in the nonconvex landscape.  The layer-wise training strategy – dividing the algorithm into two stages, one gradient descent step on the inner layer and ridge regression for the outer layer – is also common in these prior works, where the lower layer captures the relevant feature, while the top layer (+ reinitialized bias units) approximates the link function. \n\n\n**On the NTK Regime**\n\nWe would like to clarify that our analysis is *not* in the NTK/lazy regime. Quite the contrary, the one-step feature learning paradigm has been proposed to establish statistical separations and superiority over the NTK, due to the adaptivity to low-dimensional structure – see [Ba et al. 2022][Damian et al. 2022] for discussion. \n\n**Beyond Gaussian Single-index Models**\n\nThank you for the suggestion. Extending our analysis beyond Gaussian single-index models is definitely an interesting direction. We make the following remarks. \n- In terms of data distribution, recent works have established universality results for single-index learning beyond Gaussian data, see [[Zweig et al. 2023]](https://openreview.net/forum?id=JkmvrheMe7), by utilizing the orthonormal basis functions of the data distribution or CLT-type arguments. \n- In terms of target functions, a natural extension would be multi-index functions such as XOR/parity, or more interestingly, hierarchical functions that require a deeper model to express. We intend to investigate such extensions as future work.\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.", "claims_and_evidence": "N/A", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "I generally feel that the theoretical results in this paper, to some extent, illustrate the in-context learning capabilities of transformers. The mechanism of how the learned transformer performs feature learning in single-index model is well-explained.\n\nHowever, I have concerns about the “gradient-based training” (Algorithm 1) analyzed in this work, particularly given the setting under consideration:\n\n(1) The K and Q matrices (as well as the V matrix and the matrix from the MLP layer) are merged into a single matrix, with certain coordinates fixed at 0 or 1 (i.e., they are not trained).\n\n(2) In the analysis of optimization, the weight matrices of the attention layer are trained using only a single step of gradient descent.\n\nGiven that the current analysis is already quite complex, it is unclear whether these simplifications are essential to the theoretical claims. In particular, simplification (1) ensures that the output takes the exact form of Equation (3.2), which computes a weighted average over the labels $y_i$ before applying the nonlinear MLP. Additionally, the use of a single gradient descent step suggests the analysis of optimization operates in the NTK regime, which is a clear simplification for theoretical tractability. \n\nOverall, I believe the paper would benefit from clearer explanations and a more detailed justification of this gap.", "experimental_designs_or_analyses": "N/A", "relation_to_broader_scientific_literature": "In-context learning capability of transformers (the backbone of LLMs) is an important topic of both theoretical and practical relevance.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "It is a bit confusing that the dependency on $\\log^{\\mathrm{deg}(\\sigma)}(d)$ is hidden in the complexity bounds.", "questions_for_authors": "As I understand it, the techniques developed in this paper (and prior works) are specific to the Gaussian single-index model. To what extent can they be extended beyond this setting, such as by relaxing the Gaussian assumption?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Single-index models are functions of the form $x \mapsto \sigma(w_* \cdot x)$, where $\sigma$ is typically a nonlinear function characterized by favorable properties in its Hermite polynomial expansion. This paper examines how transformers can learn such functions from in-context examples provided as pairs $(x_i, \sigma(w_* \cdot x_i))$, and provides an analysis of the associated sample complexity. A key finding is that transformers can surpass kernel methods. While single-index models are not representable within a reproducing kernel Hilbert space, transformers exhibit an adaptive capability to learn relevant features, thereby reducing the sample complexity for learning single index models.

This paper differentiates itself from previous studies on in-context learning through its emphasis on the training dynamics of nonlinear attention mechanisms. However, the main theorem establishes a statistical performance bounds without explaining the mechanism of attention layers. In contrast, other studies in the literature have explicitly connected attention mechanisms to optimization procedures, thus explaining phenomena such as out-of-distribution generalization and adaptation to test distributions differing from the training distribution.

The results have been positively received by reviewers, who acknowledge that the paper is well-written and addresses an interesting theoretical topic. Therefore, I recommend acceptance.
Decision: Accept (poster)