On the Training Convergence of Transformers for In-Context Classification

$**References:**$

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

[2] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.

[3] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.15607, 2024.

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

[5] David Samuel. BERTs are Generative In-Context Learners. arXiv, 2406.04823.

[6] Ivan Lee and Nan Jiang and Taylor Berg-Kirkpatrick. Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability. arXiv, 2310.08049

[7] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.

$**Weaknesses:**$ Gaussian data that presents fixed one-token length inputs and outputs. I don't have a problem with Gaussian data, but the framework should be flexible enough to even somewhat resemble real ICL (where the inputs and outputs both can be variable lengths).

$**Reply:**$ In our setting, the prompt length (number of in-context examples) is flexible, and the length of the query and the corresponding output is fixed as one. Most previous papers theoretically studying the ICL of transformers use this setting, e.g. [1-4, 6,7]. Considering the flexible length of queries and corresponding output is an interesting problem.

$**Question:**$ Binary classification is a special case of the multi-class classification. Why write up both?

$**Reply:**$ Because binary classification is a relatively simpler case, its analysis is more concise. We use it as an example to better highlight the theoretical results. Moreover, binary classification, as a special case, has a different structure compared to the case of $c=2$ in the multi-class section. In the multi-class section, when $c=2$, the dimension of the embedding matrix is $(d+2)×(N+1)$, whereas in the binary section, it is $(d+1)×(N+1)$, which is more concise. Therefore, we wrote a separate section dedicated to binary classification.

$**Question:**$ ICL in LLMs is a type of domain adaptation setting, where knowledge about a particular task is scarce in the pre-training corpus. The model needs to be "kindled" with ICL demos to get it to perform better on this task. In contrast, the presented theory first trains the transformers on samples of these tasks, and then requires this data to be "properly distributed" over the space of variations. It is not surprising that the presented results hold in this setting. How do you think the setup needs to change to reflect that realistic ICL in LLMs setting?

$**Reply:**$ We agree that the ICL in LLMs can be a type of domain adaptation setting, and the knowledge about a particular task can be scarce during the per-training of LLMs. However, in this paper, we mainly focus on the theoretical study of the ICL capability of transformers, which is not necessarily a type of domain adaptation setting. For example, many previous papers studying the ICL of transformers [1-4, 6, 7] all considered pre-training transformers on ICL tasks and then testing transformers with the same set of tasks. Moreover, in our setting, the data distribution during pre-training and testing are different. For example, for binary case, during pre-training, $\mu\_{\tau,0}$, $\mu\_{\tau,1}$ and $x\_{\tau, query}$ are sampled according to the specific distribution $P^b\_\Omega(\Lambda)$ and $P_x^b(\mu\_{\tau,0}, \mu\_{\tau,1}, \Lambda)$. However, when testing, the $\mu\_0$ and $\mu\_1$ can be arbitrary two vectors that satisfy Assumption 3.2 and our $x\_{query}$ in testing can be an arbitrary $d$ dimensional vector. Thus, for a particular task with $\mu\_0$, $\mu\_1$ and $x\_{query}$, their corresponding probability during training can be arbitrarily small, which reflects the property -- as you put it -- in realistic ICL that the "knowledge about a particular task is scarce in the pre-training corpus". One interesting change to the setting is pre-training transformers with different types of tasks, like pre-training with both in-context linear regression tasks and in-context classification tasks. We think it is an interesting direction for future research.

$**Question:**$ With respect to the following: probably because our transformer models were only trained with a small prompt length of N = 100. 3 layer GPT2 seems like a small model, why not test with higher N?

$**Reply:**$ In our revised paper, we have conducted experiments on a 3-layer encoder-only transformer with softmax attention and without positional encoding. This setting is close to our theoretical analysis, and as we can see from Figures 1, 2, and 3, the performance of this 3-layer encoder-only transformer is very similar to the single-layer transformer we theoretically studied. As for the original question, we have explained in the original version of the paper right after the sentence you quote: "Similar declined performance when the training prompt length is smaller than the test prompt length has also been observed for in-context linear regression tasks; see e.g. [1]". Similar situations have been widely observed in many places. For example, you can find similar and significant performance degradations in Figure 1 in [1] and Figures 1, 5, 6 in [6]. In our revised paper, we used an encoder-only transformer without positional encoding in the experiment. We note that such performance degradation does not happen and the inference error of this model decreases as the test prompt length $(M)$ increases, proving that some of the insights we obtained from the simplified models also hold for more complex multi-layer non-linear transformers.

$**General response:**$ We need to first clarify that the main focus of this paper is to study the ICL capability of transformers, not the ICL capability of LLMs. We agree that the ICL capability of LLMs is remarkable and studying the ICL ability of LLMs is an interesting and important problem. However, since the theoretical study of the ICL ability of LLMs is much more complex than that of the transformer and is still at a preliminary stage, in this paper, we focus on studying the ICL ability of (simple) transformers, which usually serve as the foundational architectures of most LLMs. In fact, even the theoretical study of the ICL of basic transformers is at an infant stage in the sense that most existing papers, including ours, have to consider some simplified models to make progress. For example, in this paper, we focus on an encoder-only single-layer transformer and, to the best of our knowledge, all existing theoretical studies on the training dynamics of transformers focus only on encoder-only 1-layer transformers [1-4]. Even though we have studied a simplified model in this paper, our newly derived results in the revised paper demonstrate that multi-layer non-linear transformers can exhibit many similar behaviors as the simplified models (see Figures 1, 2, and 3 of the revised paper). Some of the insights we obtained from this simplified model also hold for more complex multi-layer non-linear transformers, which indicates that studying this simplified model can help us have a better understanding of the ICL abilities of real-world-adopting transformers.

$**Weaknesses:**$ Training on hard-coded ICL prompts, when LLMs are trained on next-word prediction (ICL structure is generally not present in the pretraining corpus). This is a major setup difference which makes them incompatible.

$**Reply:**$ Thank you for your question. Yes, most LLMs are based on decoder-only transformers and are trained on next-word prediction. However, there are also some language models such as BERT are based on encoder-only transformers, and many prior papers [5-7] also showed that encoder-only transformers can exhibit remarkable ICL abilities. Moreover, in [6], for many ICL tasks tested in their paper, encoder-only and decoder-only transformers exhibit similar performances. Thus, to simplify the analysis, in this paper, we focus on the encoder-only transformers. Moreover, to the best of our knowledge, all existing theoretical studies on the training dynamics of transformers focus only on encoder-only transformers, e.g., [1-4]. We agree that studying the ICL abilities of decoder-only transformers trained in next-word prediction is also an interesting and important problem. We leave it for future research.

$**Weaknesses:**$ Studying single layer transformers with no non-linear activation functions. This is a good intellectual curiosity but its relevance and usefulness in understanding ICL remains unclear (even classic deep learning theory struggles to present useful insights by studying 2-layer networks). In this paper itself, we see a deviation from expectations when a 3-layer GPT2 architecture with softmax is tested (section 5.2).

$**Reply:**$ We just added and revised our experimental results in the revised paper. We conducted experiments on a 3-layer encoder-only transformer with softmax attention. You can find Figures 1, 2, and 3 in our revised version. From Figure 1, we can see that the real-world multi-layer transformers and the single-layer transformers we study actually exhibit many similarities in performances. For example, from Figure 1, we can see that both models' ICL inference errors decrease as training prompt length ($N$) and test prompt length ($M$) increase, and increase as the number of Gaussian mixtures ($c$) increases. This indicates that some of our insights obtained from studying this simplified model may hold for transformers with more complex structures, and studying this simplified model can help us have a better understanding of the ICL abilities of complex transformers. Moreover, to the best of our knowledge, all existing theoretical studies on the training dynamics of transformers focus only on single-layer transformers, e.g., [1-4]. We agree that studying the ICL abilities of multi-layer transformers is also an interesting and important problem.

Dear Reviewer SwXS,

As a fellow reviewer for this paper, I disagree with your central point about the "real LLMs". As the authors claim and write in the title, abstract, and the main text, this studies Transformers not LLMs. In this regard, it not even needs to study language so it's reasonable the usage of ICL is different from those used in LLMs. I would argue for an acceptance of the work due to its own merit and contributions to the interpretability and theory community beyond its closeness to "real LLMs". Hope you could also reconsider your ratings. Thanks!

Reviewer RKkZ

re: "impressive capability in practice" refers to ICL in LLMs

I would argue that it does not necessarily refer only to LLMs. For example, Garg et. al. (2022)'s work was studying transformers ICL capabilities to fit function classes, such as linear regression, 2-layer ReLU networks, random forests, etc. This works spark a line of research on ICL capabilies of Transformers on abstract mathematical/statistical tasks that are no longer related to language.

re: What is the end goal of this theoretical work

I think this is a common question to ask for every theory work and I'm glad Reviewer SwXS also pointed it out. In my own opinion, studying how Transformers learn in context, and how they achieve such abilities during training and optimization, is crucial in understanding the trustworthiness of how Transformer architecture could be used as a universal computer. For example, Giannou et. al. (2023)'s work has shown that Transformer's variant with looping could be used as programmable computers, in the expressivity sense. Overall, people want to understand the algorithmic abilities of transformers and how training can lead to these expressivity results. Similarly, many real-world LLM works also build on top of the assumption that Transformers are able to optimize context. For example, the line of research says that transformers allow in-context reinforcement learning (Monea et.al. 2024). This also assumes that transformers could implicitly implement some algorithm or heuristics to solve a task. Again, it's scientific to understand how.

Admittedly, the gap between theory and practice sometimes branches away from each other, and the gap could be getting larger, but I still think these studies are meaningful to combat the overhype in the field and help find scientific explanations and solutions for security and AI safety. I would argue that, in order to build truly safe AI systems, we need to understand how they work, and how they gain their abilities. I personally feel this work could be beneficial for that. Let me know what you think.

Reference

Garg, Shivam, Dimitris Tsipras, Percy S. Liang, and Gregory Valiant. "What can transformers learn in-context? a case study of simple function classes." Advances in Neural Information Processing Systems 35 (2023): 30583-30598.

Giannou, Angeliki, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. "Looped transformers as programmable computers." In International Conference on Machine Learning, pp. 11398-11442. PMLR, 2023.

Monea, Giovanni, Antoine Bosselut, Kianté Brantley, and Yoav Artzi. "LLMs Are In-Context Reinforcement Learners." arXiv preprint arXiv:2410.05362 (2024).

Thank you for increasing the score and recognizing our contributions! We would like to add further clarification regarding the comment that the "impressive capability in practice" refers to ICL in LLMs. As Reviewer RKkZ mentioned, ICL capabilities do not necessarily refer to LLMs. For example, Reference [8] empirically showed that transformers have the ICL capabilities to fit function classes. Reference [6] empirically examined and compared the ICL capabilities of different models (CNN, RNN, transformer models, etc.) for various tasks, including linear regression, multiclass classification of Gaussian mixtures, image classification, and language modeling, etc. Both empirical works [6,8] studied the specific ICL abilities of transformers trained with the corresponding tasks, and the primary motivation of our paper is to provide theoretical explanations for those empirical observations of ICL abilities of transformers (in [6, 8] and also in the experimental results of our paper). Our experimental results on single/multi-layer transformers also corroborate our theoretical claims. We hope this paper can provide valuable insights into the theoretical understanding of the ICL mechanisms of transformers. Those insights may be helpful for potential architectural design (as Reviewer SwxS suggested) and building safe AI systems (as Reviewer RKkz indicated).

Thanks again to Reviewers SwxS and RKkz for providing these valuable discussions and feedback.

$**References:**$

[6] Ivan Lee, Nan Jiang, and Taylor Berg-Kirkpatrick. Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability. ICLR 2024

[8] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583–30598, 2022.

$**Weakness:**$ My major concern lies in the analytical novelty of this paper compared to prior work on in-context regression [1]. While this study focuses on the multi-class classification problem, its model and analytical approach appear to share many similarities with [1]. It remains unclear how technically straightforward it is to generalize the results of [1] to multi-classification. Additionally, this paper restricts itself to the linear attention setting, simplifying the analysis and making it somewhat less impactful than [2], which addresses binary classification with strict data assumptions but in the more realistic softmax attention setting. Therefore, a thorough discussion clarifying the technical distinctions and contributions of this work relative to these previous studies would be helpful.

$**Reply:**$ Our technique is different from those used in [1, 2]. In [1], the globally optimal solution (i.e., the parameters of the transformer) has a closed-form expression, and they proved that the 1-layer transformer optimized via gradient flow can converge to this closed-form globally optimal solution. However, in our setting, due to the high non-linearity of our loss function, the global minimizer does not have a closed-form expression. Instead, by analyzing the Taylor expansion near the global minimizer, we prove that the global minimizer consists of a constant plus an error term that is induced by the finite training prompt length ($N$). We further show that the max norm of this error term is bounded, and converges to zero at a rate of $O(1/N)$. Our technical approach to addressing this challenge is new and might be useful in other settings. Moreover, we considered the more practical gradient descent rather than the gradient flow in [1]. In [2], they only studied the binary classification tasks with finite pairwise orthogonal patterns. They generated their data as $x=\mu_j+\kappa v_k$, where $\{\mu\_j\}, {j=1, 2, ..., M\_1}$ are in-domain-relevant patterns and $\{\nu\_k\}, {k=1, 2, ..., M\_2}$ are in-domain-irrelevant patterns, $M\_1\geq M\_2$ and these patterns are all pairwise orthogonal. Thus, the possible distribution of their data is finite and highly limited. In contrast, our work data is drawn according to $P^b(\mu_0,\mu_1,\Lambda)$ or $P^m(\mu, \Lambda)$, and the range and possible distributions of our data are infinite. Thus, we considered more general in-context multi-class classification tasks with infinite patterns while [2] only considered in-context classification tasks with finite patterns, thereby highlighting the distinct contributions and independent interests of our work.

$**Question:**$ For data distribution, why is it essential to preserve the inner product of vectors in the $\Lambda^{-1}$-weighted norm? Is this primarily a technical consideration? It would be helpful if the authors could provide further clarification on the role of data distribution in the analysis.

$**Reply:**$ The primary role of the condition (2) in Assumption 3.1 is to ensure that $\mu\_{\tau, 1}$ and $\mu\_{\tau, 0}$ have the same $\Lambda^{-1}$-weighted norm. Because, if $\mu\_{\tau, 1}$ and $\mu\_{\tau, 0}$ have the different $\Lambda^{-1}$-weighted norms, then, the probability of the ground truth label $y\_{\tau, query}$, $\mathbb{P}(y\_{\tau, query}=1)=\sigma((\mu\_{\tau,1}-\mu\_{\tau,0})^\top \Lambda^{-1} x\_{\tau, query}+ (\mu\_{\tau,1}^\top\Lambda^{-1}\mu\_{\tau,1}-\mu\_{\tau,0}^\top\Lambda^{-1}\mu\_{\tau,0})/2)$, we find it is hard for 1-layer transformer with linear attention to calculate $\mu\_{\tau,1}^\top\Lambda^{-1}\mu\_{\tau,1}-\mu\_{\tau,0}^\top\Lambda^{-1}\mu\_{\tau,0}$ in context. However, we found that a 1-layer transformer with linear attention can approximately calculate $(\mu\_{\tau,1}-\mu\_{\tau,0})^\top \Lambda^{-1} x_{\tau, query}$ in context. Thus, we add the condition (2) in Assumption 3.1. Moreover, the newly added experimental results (Figure 2) in our revised paper also show the necessities of the condition (2). Experimental results in Figure 2 also indicate transformers with more complex structures have better robustness without condition (2). Thus, it is an interesting question whether we can eliminate the need for the condition (2) for more complex transformers. We leave it for future research.

$**Question:**$ The condition (2) in Assumption 3.1 seems unusual to me. Could the authors provide more clarification on this assumption?

$**Reply:**$ The primary role of the condition (2) in Assumption 3.1 is to ensure that $\mu\_{\tau, 1}$ and $\mu\_{\tau, 0}$ have the same $\Lambda^{-1}$-weighted norm. Because, if $\mu\_{\tau, 1}$ and $\mu\_{\tau, 0}$ have the different $\Lambda^{-1}$-weighted norms, then, the probability of the ground truth label $y\_{\tau, query}$, $\mathbb{P}(y\_{\tau, query}=1)=\sigma((\mu\_{\tau,1}-\mu\_{\tau,0})^\top \Lambda^{-1} x\_{\tau, query}+ (\mu\_{\tau,1}^\top\Lambda^{-1}\mu\_{\tau,1}-\mu\_{\tau,0}^\top\Lambda^{-1}\mu\_{\tau,0})/2)$, we find it is hard for 1-layer transformer with linear attention to calculate $\mu\_{\tau,1}^\top\Lambda^{-1}\mu\_{\tau,1}-\mu\_{\tau,0}^\top\Lambda^{-1}\mu\_{\tau,0}$ in context. However, we found that a 1-layer transformer with linear attention can approximately calculate $(\mu\_{\tau,1}-\mu\_{\tau,0})^\top \Lambda^{-1} x_{\tau, query}$ in context. Thus, we add the condition (2) in Assumption 3.1. Moreover, the newly added experimental results (Figure 2) in our revised paper also show the necessities of the condition (2). Experimental results in Figure 2 also indicate transformers with more complex structures have better robustness without condition (2). Thus, it is an interesting question whether we can eliminate the need for the condition (2) for more complex transformers. We leave it for future research.

$**Question:**$ Some papers [1,2,3] have highlighted emergent behaviors in the training dynamics of in-context learning. However, this paper asserts that the transformer will converge to its global minimizer at a linear rate, which appears to contradict those findings. Can the authors discuss this further?

$**Reply:**$ Because [1,2,3] studied the in-context learning of transformers with structures and problems different from ours, the training dynamics of transformers with different structures and for different tasks can be different. Some other papers, such as [5], also proved the linear convergence of transformers for some specific problems. However, all existing theoretical studies on the training dynamics of transformers focus only on single-layer transformers [3-7]. Theoretical understandings of the training dynamics of multi-layer transformers for more complex real-world problems are still unclear and are interesting research directions for future research.

$**References:**$

[4] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.15607, 2024.

[5] Tong Yang, Yu Huang, Yingbin Liang, and Yuejie Chi. In-context learning with representations: Contextual generalization of trained transformers. arXiv preprint arXiv:2408.10147, 2024

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

[7] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.

$**Weakness:**$ This paper claims to be the first to explore the learning dynamics of transformers for in-context classification of Gaussian mixtures and to prove the convergence of training in multi-class classification. However, I find the significance of this assertion unclear, as the paper lacks sufficient detail. Specifically: 1) many prior works have analyzed in-context learning assuming $x$ comes from Gaussian distributions; what additional insights do the results on Gaussian mixtures provide? 2) Why is extending results from binary to multi-class classification considered essential and non-trivial?

$**Reply:**$ 1) Previous studies that assume $x$ is drawn from Gaussian distributions are all focused on the in-context linear regression problem. To the best of our knowledge, we are the first to explore the in-context classification problem under the assumption that $x$ is drawn from Gaussian mixtures. Prior work [4] that studied the in-context classification of transformers assumes the data to be pairwise orthogonal. They generated their data as $x=\mu_j+\kappa v_k$, where $\{\mu\_j\}, {j=1, 2, ..., M\_1}$ are in-domain-relevant patterns and $\{\nu\_k\}, {k=1, 2, ..., M\_2}$ are in-domain-irrelevant patterns, $M\_1\geq M\_2$ and these patterns are all pairwise orthogonal. Thus, the possible distribution of their data is finite and highly limited. In contrast, our work data is drawn according to $P^b(\mu_0,\mu_1,\Lambda)$ or $P^m(\mu, \Lambda)$, and the range and possible distributions of our data are infinite. Hence, we considered a more general situation that our data can have infinite patterns while [4] only considered in-context classification tasks with finite patterns. Thus, we provide additional insights that transformers can perform in-context classification tasks with infinite patterns. 2) Moreover, [4] only considered binary classification. We also provide additional insights that transformers can perform in-context multi-class classification. This is essential because many real-world classification problems are not binary but multi-class. Therefore, explaining how transformers can handle multi-class classification problems in context is an essential question. It is non-trivial because technically, when extending results from binary to multi-class classification, more complicated cross terms in the Taylor expansions of the softmax functions, which are due to the nature of multi-class classification, bring new challenges to the analysis. To address these issues, we derived new bounds on the expected errors of the cross terms in Lemma G.1, G.2, which may be of independent interest to other similar problems.

$**Weakness:**$ The introduction of $\widetilde{L}$ appears to be a key element in proving Theorem 3.1, but its intuition is unclear, and I'm uncertain how it addresses the challenges posed by the non-linear loss function.

$**Reply:**$ In Lemma E.3, we show that as $N\to \infty, L(W)$ will point wisely converge to $\widetilde{L}(W)$. Since we can easily find the global minimizer of $\widetilde{L}(W)$ is $2\Lambda^{-1}$, with the help of $\widetilde{L}(W)$, we can show that as $N\to \infty$, $W^*$, the global minimizer of $L(W)$, will converge to $2\Lambda^{-1}$. Thus, we can denote $W^*=2(\Lambda+G)$. In Lemma E.4, by analyzing the Taylor expansion of the equation $\nabla L(W^*)=0$ at the point $2\Lambda^{-1}$, we address the challenges posed by the non-linear loss function and establish the bound $\|G\|\_{max}=O(N^{-1})$.

$**Weakness:**$ The paper heavily relies on Taylor expansion in its proofs, and I question whether this expansion can accurately approximate the original function. More detail is needed on this aspect.

$**Reply:**$ Yes, we used Taylor expansions in many places in our proofs. However, every time we use the Taylor expansion, we always use the Lagrange form of the remainder to express and bound the approximation error. For example, in the proof of Theorem 3.2, we used the equation

$

    \sigma(a+b)=\sigma(a)+\sigma'(a)b+\frac{\sigma''(\xi(a,b))}{2}b^2, 

$

where $\xi(a,b)$ are real numbers between $a$ and $a+b$. Since $|\sigma''(\xi(a,b))|\leq 1$ and in the proof of Theorem 3.2, we can prove $E[b^2]=o(1/N+1/\sqrt{M})$, we can bound the approximation error smaller than $o(1/N+1/\sqrt{M})$. Similarly, in any other places where we use the Taylor expansion, we always express and bound the approximation error.

$**References:**$

[4] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.15607, 2024.

$**Weakness:**$ The paper primarily focuses on single-layer transformers with simplified linear attention mechanisms. While it provides valuable insights into the convergence properties and error bounds for these models, the findings may not fully extend to more complex multi-layer transformers with softmax or relu attention mechanisms. And the training dynamics of multi-layer transformers could be different to single-layer -- just like multi-layer MLP can be different than linear models.

$**Reply:**$ Thank you for your recognition of our contributions and for raising the question about the extension to multi-layer with softmax or ReLU attention mechanisms. Yes, we agree that the training dynamics and many other properties of multi-layer, non-linear transformers can be different from the single-layer linear transformers we study. However, from the newly added experimental results (Figure 1) in our revised paper, we can see that the real-world multi-layer transformers and the single-layer transformers we studied actually exhibit many similarities in performances. For example, from Figure 1, we can see that both models' ICL inference errors decrease as training prompt length ($N$) and test prompt length ($M$) increase, and increase as the number of Gaussian mixtures ($c$) increases. This indicates that some of our insights obtained from studying this simplified model may still be valuable for transformers with more complex structures, and studying this simplified model can actually help us have a better understanding of the ICL abilities of real-world-adopted transformers. Moreover, the research community is still at the preliminary stage of the theoretical investigations of in-context learning of transformers. To the best of our knowledge, most existing theoretical studies on the convergence behavior focus only on single-layer transformers, e.g., [1-6]. We agree that studying the ICL abilities of multi-layer transformers is also an interesting and important problem. We leave it for future work.

$**Question:**$ Does the same training dynamics results apply to tasks beyond linear regression/classification? Would it be different for other mathematical/statistical tasks, for example, time series, etc.

$**Reply:**$ It is a good question. The training dynamics of transformers with different structures and for different tasks can be different. However, some insights we got from the linear regression/classification may also hold in other mathematical/statistical tasks. For example, we find that for the in-context classification of Gaussian mixtures, the ICL influence errors are affected by the training and testing prompt lengths. We suspect similar behaviors may also hold in other mathematical/statistical tasks. Nevertheless, the training dynamics for other mathematical/statistical tasks is still an interesting open question for future research.

$**References:**$

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

[2] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.

[3] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.15607, 2024.

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

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

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

Your response resolves my concerns and I'm raising my score. I vote strongly for an acceptance of this work, because of its contributions to help the community better understand the training dynamics of Transformers, and in particularly in context learning. I also strongly disagree with Reviewer SwXS and align with the authors of this paper, that this is not an LLM work, and a gap between theoretical understanding of LLMs and empirical LLMs is acceptable and widely adopted by the community.