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

We thank the reviewer for the helpful feedback. We address the technical concerns below.

"pre-training algorithm involves setups that deviate from standard LLM training approaches"

We agree that Algorithm 1 deviates from the most standard training procedure in practice. We make the following remarks on our architecture and optimization algorithm.

• MLP before attention. This simplified theoretical setting reflects the scenario where lower layers of the trained transformer act as data representations on top of which the upper layers perform ICL. See Figure 1 in [Guo et al. 2023]: they considered a Transformer architecture with alternating MLP and attention layers and suggested that the earlier layers grasp the features of the true function, while the subsequent layers act as linear attention.Our architecture, with an MLP layer for feature learning followed by a linear attention module, idealizes such process occurring in more complicated Transformer.
  The same architecture has been analyzed in prior works on in-context learning of nonlinear features such as [Kim and Suzuki 2024].

• Linear attention. In the ICL theory literature, the Softmax attention is often simplified to linear attention to facilitate the optimization analysis. Such simplification can capture several prominent aspects of practical transformer training dynamics -- see [Ahn et al. 2023] and references therein. Also, in our setting we already observe interesting in-context behavior using a linear attention. Extending our analysis to Softmax attention is an interesting future direction.

• Layer-wise training and one-step GD. In the neural network theory literature, layer-wise training and reinitialized bias units are fairly standard algorithmic modifications when proving end-to-end sample complexity guarantees, as seen in many prior works on feature learning [Damian et al. 2022] [Abbe et al. 2023]. Without such modifications, stronger target assumptions are generally needed, such as matching link function [Ben Arous et al. 2020].
  In addition, training the representation for one gradient step (followed by SGD on the top layer) is a well-studied optimization procedure for two-layer neural networks (with the hope being that the "early phase" of representation learning already extracts useful features) -- see [Ba et al. 2022] [Barak et al. 2022]; but to our knowledge, we are the first to apply this dynamics to nonlinear transformers in the ICL setting and establish rigorous statistical guarantees.

"numerical experiments demonstrating the effectiveness of the proposed algorithm"

Following your suggestion, we conducted experiments that supports our theoretical claim on the statistical efficiency of ICL -- please see the .pdf document attached in our general response. Specifically, in Figure 1 we indeed observe that when $r\ll d$, there is a separation in the test loss between pretrained transformers and baseline algorithms that directly learn from the test prompt, including kernel ridge regression and two-layer neural network + gradient descent.

"how this work contributes to understanding the mechanisms of ICL and LLMs"

We believe that the the following high-level messages from our theoretical results are relevant:

• Our analysis highlights how pretrained transformers can leverage prior information (low-dimensional structure in particular) of the function class and outperform simple baseline algorithms. Such a phenomenon is empirically known (e.g., on linear functions, it had been observed that Transformers adapt to the "sparsity of problem distribution" [Garg et al. 2022]), but existing theoretical analysis with end-to-end learning guarantees typically showed that ICL implements a simple prior-agnostic algorithm, such as one gradient descent step on the least squares objective for linear regression.
  In addition, our finding that transformer can adapt to a low-rank target prior suggests, at a high level, a generalization error scaling that does not depend on the ambient dimensionality of the problem, but some notion of "intrinsic dimensionality" of the function class.

• Our analysis makes use of a mechanism for ICL in shallow transformers: feature extraction & adaptivity to low-dimensionality by MLP layer + in-context function approximation by attention layer. We believe that similar mechanism also appears in modern pretrained transformers, which can provide insight to mechanistic interpretability research.

We would be happy to clarify any further concerns/questions in the discussion period.

We thank the reviewer for the updated evaluation. We would like to address the reviewer's outstanding concern with the following remarks:

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

• As mentioned in the Introduction, the single-index target function has been extensively studied in the deep learning theory literature; however, to our knowledge, we are the first to theoretically investigate the learning of this function class in the more challenging in-context setting.
  It is worth noting that even in the "easier" setting of standard neural network training (i.e., not considering ICL), some of the most influential papers on the theory of representation learning have considered algorithmic modifications similar to ours, such as layer-wise training, one-step feature learning, and reinitialized bias units [Abbe et al. 2022] [Damian et al. 2022] [Ba et al. 2022]. These works are not dismissed by the community because they "deviate from standard training approaches"; rather, they serve as an important first step toward understanding the statistical and computational efficiency of neural networks beyond simple linear problems.
  We hope the reviewer can recognize the nontrivial theoretical contribution of our analysis in this context.

We are thankful for your helpful comments. We address the concerns and answer the questions below.

"The theoretical result seems not to be tight."

Note that although we have an MLP block in our architecture, it is not trainable during the inference time. Indeed a sample complexity of $O(d^{\Theta(k)})$ for MLPs can be attained when we train the MLP parameters in response to each task (e.g., see [Bietti et al. 2022]). However, when considering ICL via Transformer, we need to fix the MLP layer and use the same parameters for all tasks, where the true function and direction $\beta$ can vary.Thus it is not surprising that our sample complexity is worse than NN+GD (trained on the test prompt) in the $r\simeq d$ regime. Instead, our focus is the regime where $r\ll d$, for which we show that pretrained transformer can adapt to this low-dimensional structure (target prior) and outperform baseline algorithms that only have access to the in-context examples.

"No experiments are provided to justify the result."

Following your suggestion, we conducted experiments that supports our theoretical claim on the statistical efficiency of ICL -- please see the .pdf document attached in our general response. Specifically, in Figure 1 we observe that when $r\ll d$, there is indeed a separation in the test loss between pretrained transformers and baseline algorithms that directly learn from the test prompt, including kernel ridge regression and two-layer neural network + gradient descent.

"The presentation of the proof sketch can be improved."

Thank you for the suggestion. As you pointed out, the convexity of the ridge regression problem for $\Gamma$ plays an important role because the optimal solution $\Gamma^*$ can be efficiently solved; this allows us to state that the approximation and generalization error of $\Gamma^*$ is not comparable to that of the constructed $\Gamma$. We will clarify this point in the proof sketch section.

"references on ICL optimization and generalization are missing"

Thank you for your suggestion. We will add citations and comparison with these concurrent works.

"Algorithm 1 is proposed to handle the non-convexity of the problem?"

Yes, this simplification using one-step gradient descent circumvents the difficulty of analyzing the entire optimization trajectory with non-convex landscape. Such a strategy is commonly-used in the neural network theory literature -- training the representation for one gradient step (followed by SGD on the top layer) is a well-studied optimization procedure for two-layer neural networks (with the hope being that the "early phase" of representation learning already extracts useful features) -- see [Ba et al. 2022] [Barak et al. 2022]; but to our knowledge, we are the first to apply this dynamics to nonlinear transformers in the ICL setting and establish rigorous statistical guarantees.

We would be happy to clarify any further concerns/questions in the discussion period.

We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.

Potentially misleading baseline and prior works

Thank you for the insightful comments. Indeed there are empirical examples of transformers achieving better statistical efficiency by making use of prior information of the function class -- we do not claim that we are introducing a new phenomenon. We will discuss [Raventos et al. 2023] in the Introduction and clarify the overlap in the high-level message. Our motivation is to provide rigorous end-to-end learning guarantee that captures this adaptivity in an idealized theoretical setting.
We also agree that the comparison of sample complexity (pretrained transformers vs. algorithms on in-context test examples) is not straightforward, since transformers benefit from additional data in the pretraining phase. From the theoretical point of view, our goal is to highlight the statistical benefit of this prior information, since many previous theoretical works on ICL proved that transformers implement a simple "prior-agnostic" algorithm, such as one gradient descent step on the least squares objective for linear regression.

On Gaussian single-index models (Questions 0, 1, and 2)

• A single-index model is defined by an unknown univariate "link" function applied to an unknown one-dimensional projection of the input. Efficiently learning this function requires some extent of adaptivity to the low-dimensional structure, and hence this function class has been extensively studied in the neural network theory literature, to demonstrate the benefit of representation learning. Here we list some works covering this topic: [Ba et al. 2022] [Bietti et al. 2022] [Damian et al. 2023]. We however realize that this problem setting has not been introduced in the context of transformer theory. We will revise the Introduction and Problem Setting section to provide more background and intuition.

• Regarding the question 1, there we aimed to simply explain that in single-index models, the function value of $f_*$ changes only along the direction parallel to $\beta$ in the input space $\mathbb{R}^d$, because $f_*$ only depends on the inner product $\langle\beta,x\rangle$. We will make that explanation clearer.

• We select this problem setting because it simultaneously meets two desiderata: $(i)$ a nonlinear function class, and $(ii)$ adaptivity to target "prior". On one hand, single-index model is arguably the simplest nonlinear extension of linear regression in which a nonlinear link function is applied to the response. On the other, this function class allows us to specify low-dimensional structure which the learning algorithm should adapt to. Moreover, for this function class, the statistical efficiency of simple algorithms (that can be implemented on the in-context examples) has been extensively studied.

• If we are not concerned with point $(i)$ above, it is possible that a linear function class (with low-rank prior) can exhibit similar phenomenology, where a transformer with linear MLP layer can adapt to the class prior. More generally, we indeed conjecture a generalization bound that does not depend on the ambient dimensionality of the problem, but some notion of "intrinsic dimensionality" of the function class.

Empirical evidence and future work (Question 3)

• We conducted experiments that supports our theoretical claim on the statistical efficiency of ICL -- please see the .pdf document attached in our general response. Specifically, in Figure 1 we indeed observe a separation in the test loss between ICL and "prior-agnostic" algorithms such as kernel ridge regression, as the theory predicts.

• We believe that the mechanism (employed in our theoretical analysis) of feature extraction & adaptivity to low-dimensionality by MLP layer + in-context function approximation by attention layer could also appear in modern transformer training. However, since our statistical guarantees are derived for a particular architecture and gradient-based optimization procedure, it is likely that the quantitative scalings (e.g., dependence on context length, task diversity) have limited predictive power in realistic settings.

We would be happy to clarify any further concerns/questions in the discussion period.

We thank the reviewer for the helpful feedback. We address the technical concerns below.

Nonstandard architecture and training procedure

We agree that Algorithm 1 deviates from the most standard training procedure in practice. We make the following remarks on our architecture and optimization algorithm.

• MLP before attention. This simplified theoretical setting reflects the scenario where lower layers of the trained transformer act as data representations on top of which the upper layers perform ICL. For example, see Figure 1 in [Guo et al. 2023]: they considered a Transformer architecture with alternating MLP and attention layers and suggested that the earlier layers grasp the features of the true function, while the subsequent layers act as linear attention. Our architecture, with an MLP layer for feature learning followed by a linear attention module, idealizes such process occurring in more complicated Transformer. The same architecture has been analyzed in prior works on in-context learning of nonlinear features such as [Kim and Suzuki 2024].

• Linear attention. In the ICL theory literature, the Softmax attention is often simplified to linear attention to facilitate the optimization analysis. Such simplification can capture several prominent aspects of practical transformer training dynamics -- see [Ahn et al. 2023] and references therein. Also, in our setting we already observe interesting in-context behavior using a linear attention. Extending our analysis to Softmax attention is an interesting future direction.

• Layer-wise training and one-step GD. In the neural network theory literature, layer-wise training and reinitialized bias units are fairly standard algorithmic modifications when proving end-to-end sample complexity guarantees, as seen in many prior works on feature learning [Damian et al. 2022] [Abbe et al. 2023]. Without such modifications, stronger target assumptions are generally needed, such as matching link function [Ben Arous et al. 2020].
  In addition, training the representation for one gradient step (followed by SGD on the top layer) is a well-studied optimization procedure for two-layer neural networks (with the hope being that the "early phase" of representation learning already extracts useful features) -- see [Ba et al. 2022] [Barak et al. 2022]; but to our knowledge, we are the first to apply this dynamics to nonlinear transformers in the ICL setting and establish rigorous statistical guarantees.

Comments on the problem setting

The single-index setting is arguable the simplest nonlinear extension of linear regression in which a nonlinear link function is applied to the response. Moreover, for this function class, the statistical efficiency of simple algorithms (that can be implemented on the in-context examples) has been extensively studied. We therefore believe that it is a natural first step for the theoretical understanding of nonlinear ICL.
As for the low-rank structure, our motivation is to introduce "prior" information of the function class which cannot be exploited by algorithms that directly learn from the test prompt, but can be adapted to during pretraining (see paragraph starting line 40 and 63). This enables ICL to outperform baseline algorithms such as linear/kernel regression.

"Can you include some examples of function classes satisfying Assumption 1?"

Our Assumption 1 requires the link function to be a bounded-degree polynomial, such as a Hermite polynomial of degree $P$. This assumption is analogous to those appearing in the feature learning literature, see [Abbe et al. 2023] and [Ba et al. 2023].

"An experiment that shows that w aligns with S under true GD would be nice"

Following your suggestion, we have included numerical experiments demonstrating the alignment phenomenon in SGD with multiple steps. Please refer to Figure 2 in the pdf document in our general response.

"Is there some easy interpretation of Gamma?"

Intuitively speaking, $\Gamma$ plays the role of approximating (nonlinear) basis functions that express the true function. As shown in Section 4.2, $\Gamma$ is decomposed as $ADA$, where $A$ is analogous to the second-layer of a two-layer MLP which can approximate basis functions $\{g_i\}$. This is one crucial difference from the linear case, where the attention matrix only needs to learn the data covariance.
We will provide more explanations on this in the revision.

We would be happy to clarify any further concerns/questions in the discussion period.