On the Role of Depth and Looping for In-Context Learning with Task Diversity

Response to reviewer PS8E

We thank the reviewer for their feedback and pointing out to our typos. Below we have address all their concerns. We have also applied all their comments to the new version of the paper uploaded. In addition, we believe there are some misunderstandings on our contribution and how in compares to the literature that we will clarify below. **Relu in identity for activation is restricted and does not cover all cases that are interesting in practice, and what about non-polynomial activations?** **Reply:** First we note that Relu is a non-polynomial activation. In fact our lower bounding method works after applying feedforward layers or other non-linearities than Relu by approximating the activations with polynomials**. In other words, if the approximation degree is not superlarge, the results are generalizable, but for sake of keeping the presentation of the work simple and avoid generalizations that lack deep novel ideas and follows from more conventional derivations, we decided not to add this sort of generalizations to our current work (Since we have many results including depth lower and upper bounds and robustness aspects etc in the current work). We point out that our work is the first attempt to show any such expressivity lower bound (i.e. limitations) for in-context learning, since most previous works focussed on showing upper bounds for Transformers. We hope this work encourages more study on the limitations of Transformers and the importance of depth in more general settings. **The above weakness could be mitigated via numerical experiments** **Reply:** Our theoretical work was, in fact, motivated by previous empirical works by Garg et al. and Yang et al. which show that Transformers with softmax activations can solve ICL problems with sufficient depth. Since our theoretical analysis was in a slightly different setting, we ran experiments in the restricted linear setting to verify that the findings also hold there. **Out of distribution generalization was studied in Gatmiry et al [1]** **Reply:** Note that the out-of-distribution result in Gatmiry only holds for the global minimizer of the Gaussian case, i.e. when data is sampled iid from a single Gaussian, in which case the data covariance distribution is concentrated on the covariance of the Gaussian distribution. We significantly generalize that result to a much broader class of distributions over the covariance, namely the class of right-spread-out distributions, which can be highly multi-modal. We believe that this is in some sense the minimal assumption on the training distribution that enables one to obtain out-of-distribution generalization. **The paper does not provide an optimization guarantee** **Reply:** The focus of this work is on expressivity of multi-layer and looped transformers and analyzing the global minimizer of distributions that are much more general than Gaussian, and their connections with robustness and out-of-distribution generalization. We believe that we have proved many interesting results in this work around these topics that have to be self-contained in this paper. Proving optimization results is certainly an interesting future direction. **The writing of the paper can be improved** **Reply:** We have applied all the comments of reviewers and improved the presentation of the work in the new version that we have uploaded. Given that all these concerns are addressed, we kindly ask the reviewer to reconsider their score. **Lemma 1 in the main body seems unnecessary** **Reply:** In fact we see Lemma 1 as the core technical result and Theorem 1 as its corollary. Hence, we decided to put Lemma 1 in the main body to just highlight our approach for proving the lower bound in Theorem 1. That being said, the proof in the appendix is really the proof of Lemma 1. We made this point clearer in the new version. Thanks for your feedback. **We also applied all the feedback on typos and minor errors that the reviewer kindly provided to the new uploaded version.** **Questions** **Is there a way to generalize the techniques in this paper to softmax attention or non-linear feedforward networks?** **Reply:** Please see our answer to your first point above **The lower bound for restricted attention is not matched by upper bound Looped transformer, is there a gap in expressivity here?** **Reply:** That is a great point, although we have not shown any formal result here, we suspect that in fact the loop model is also capable of achieving the $\sqrt{\beta/\alpha}$ complexity by a more complicated construction. We leave this interesting direction for future work. **What is the difference between the out-of-distribution upper bound in your work and that of Gatmiry et al. [1]** **Reply:** Please see our answer to your second point above.

Response to Reviewer wVSr

We thank the reviewer for acknowledging our contributions and their valuable feedback. **The result does not apply to softmax, multi-layer perceptron networks, etc** **Reply:** our lower bounding method works after applying feedforward layers or other non-linearities than Relu by approximating the activations with polynomials**. In other words, if the approximation degree is not superlarge, the results are generalizable, but for sake of keeping the presentation of the work simple and avoid generalizations that lack deep novel ideas and follows from more conventional derivations, we decided not to add this sort of generalizations to our current work (Since we have many results including depth lower and upper bounds and robustness aspects etc in the current work). We point out that our work is the first attempt to show any such expressivity lower bound (i.e. limitations) for in-context learning, since most previous works focussed on showing upper bounds for Transformers. We hope this work encourages more study on the limitations of Transformers and the importance of depth in more general settings. **The assumption on the distribution of $w^*$ and $x_q$ although used in previous work is rather contrived* **Reply:** The choice of distribution on $x_q$ is so that it aligns with the distribution of $x_i$’s, since the data covariance matrix defined based on $x_i$’s is going to concentrate around the population covariance. The choice of the distribution of $x^*$ is rather a simplifying assumption for our analysis, and also natural makes the distribution of the output ${w^*}^\top x$ becomes a subexponential whose variance does not depend on $\Sigma$. **The experiments are apparently only run once and there might be some outliers** **Reply:** Each experiment has been run 3 times and the results presented are an average over the 3 runs.

Response to Reviewer g1xd

We thank the reviewer for their feedback. Below we will clarify the ambiguities. Line 153: $X w^*$ has to be $X^{\top} w^*$... The index $t$ of $Z$ in line 212 has to be $t+1$... $x_q$ is missing in the definition of $Z^{(0)}$ in line 267 **Reply:** Thank you for pointing out this typo, we fixed it in the new version. **Line 223-228: this is unusual and I could not locate this formula in Ahn et al 2024a** **Reply:** Please see Equation (8) in Ahn et al, right before Lemma 1. So our definition is in fact more general than theirs in that we allow non-zero $u^{(t)}$ while in Ahn et al. assume $u^{(t)} = 0$. **The matrix $Z^{(0)}$ is rescaled in Equation 6 by $\gamma$, which means $x_q$ is also rescaled, which is not the case in the line before.** **Reply:** Sorry for the confusion, in the current version of the proof we only rescale $X$ and $y$ by $\gamma$ as we have denoted in line 262 (so $x_q$ is not scaled). That being said, the $\gamma$ behind $Z^{(0)}$ in Equation (6) is a typo. Thank you very much for pointing this out, we fixed this typo in the new version uploaded. **line 281: w^* should be w:** **Reply:** thank you for spotting this typo. We fixed this typo in the new version. **line 284: “under restricted attention” is typo?** **Reply:**Yes, this is a typo, we fixed it in the new version. **Theorem 2. $\sqrt k$ dependency of the depth is a type of acceleration that is worthy of more discussion** **Reply:** The reviewer is correct, this is indeed a type of acceleration that is achieved using Chebyshev polynomials. We have added a note about its connection with accelerated rates in optimization within the final version (last sentence before section 3.3). **Line 350, k is L ?** **Reply:** Yes thank you for pointing out to this typo, we fixed it in the new uploaded version. **What is meant by "in-distribution" here? Without first clarifying this, it is hard to grasp what "out-of-distribution" refers to. I also cannot figure out what would be the “(supervised) learning tasks” and “population distribution” in this context, given that $X$ is supposed to be fixed while $x_q$ is random.** **Reply:** We understand why the reviewer is confused here. To be superclear, in this paper we study the in-context learning problem on the supervised learning task of linear regression instance; so each instance here is a “learning problem” following the formalization from Garg et al. We consider a “population distributions” on these linear regression instances by assuming a distribution on the data covariance matrix $\Sigma$, the query vector $x_q$ and the regressor $w^*$. This is very similar to the assumptions made in prior work for ICL, except that the inputs do not need to have a Gaussian distribution. Following Gatmiry et al [1], we also assume $w\sim \mathcal N(0,\Sigma^{-1})$. We made changes in the new version in lines 150, 154, and lines 369 and 375 of section 4.2 to make these points clearer. Hopefully this has clarified the confusion for the reviewer. We also resolved the typos in line 884 and 1017 in the new version. **Moreover, it appears that the transformer only sees one context example during training. If that is the case, how can we expect it to generalize?** No, the transformer can see a finite number of samples. Just like any supervised learning problem, even though the model might only see a finite number of samples from a distribution, it can learn to generalize to the population distribution. The situation for ICL should not be any different.

Response to reviewer 6ZLH

We thank the reviewer for valuable comments and feedback. **W1: The presentation has to become more clear…. The role of $\gamma$ in Theorem 1 is not clear** **Reply:** Thank you for your feedback, we try to add more high level discussion to our rigorous proofs to add more intuition. Indeed $\gamma$ has to be removed from Theorem 1. **W2: The lower bound is proved for one instance from the class of covariances with bounded condition number. Is it possible to show lower bounds against a distribution, in expectation or high probability?** **Reply:** No, the result is proved against a range of instances whose covariance has bounded condition number. To show a lower bound, one considers a “normal” class of instances and prove that given any transformer, there is a bad instance in that class on which the transformer behvaes poor, which is exactly what are result states. The reason is, for any fixed instance, and likely for many fixed distributions over instances, one can find a transformer that achieves low loss on that particular instance or distribution, but the transformer does not necessarily behave well for other instances or distributions. **W3: It is not to call matching upper and lower bound as the upper bound has an additional $\og(1/\epsilon)$ factor** **Reply:** Thank you for your feedback on phrasing this part; when mentioning matching upper and lower bound, we mean the upper bound on depth for achieving a constant error, say 0.001, in which the log term is just a constant. We have clarified this in the new version. **W4: Theorem 6 claims the existence of a global minimzier in which out-of-distribution loss exponentially blows up, what about other global minimizers?** **Reply:** Thank you for your feedback on this statement. While we suspect that this is the case for all global minimizers (if any other one exists), we currently do not know how to rule out the existence of another minimizer in which the loss might not blow up, as the corresponding matrix equation is highly complex. That being said, existence of bad global minima on its own shows that training transformers with first-order gradient method can lead to these bad solutions that cannot generalize out of distribution. We have added our conjecture about the global minimizer of the loss right after Theorem 6. **W5: Some parts of proofs are not clear. For example Lemma 3, does the result hold for any polynomial $P(\gamma)$?** **Reply:** Yes, as we state in the statement of Lemma 3, the result holds for any polynomial $P$ with degree $k$. We have added the adjective “arbitrary” before the word “polynomial” to make this clearer. While we appreciate the reviewer’s feedback on clarity of the proof and will try to add more details in the final version, we would like to point out that our vision for presenting the proofs was not to make them unnecessarily lengthy and complex, but rather keep them as short and sharp as possible, conveying the key mathematical ideas. If the reviewer can explicitly point out parts that they think are unclear, that would help us a lot in adding more clarity.

Response to Reviewer 2duX

We thank the reviewer for a detailed review of the paper and for appreciating the contributions like lower bounds and robustness results. **W1: It’s unclear how well these controlled settings will transfer to real-world problems. This is a broader issue in the literature** **Reply:**We would like to note that most of the results in the paper do not make the Gaussian assumption. In fact our matching upper and lower bounds consider arbitrary linear regression instances with bounded eigenvalues of the covariance. Moreover, our robustness result for looped models holds for a much broader class of distributions (that satisfy the “right-spread-out” assumption) rather than just a single Gaussian as in previous work. **W2 The idea of “diversity” here might be somewhat questionable** **Reply:**We model diversity as different instances that belong to the same learning task, here the learning task is linear regression. While not perfect, this already provides a reasonable test-bed for modeling in-context learning ability. Proving any theoretical result encompassing various learning tasks, although interesting as the reviewer has mentioned, is currently beyond our reach. We believe many fundamental questions about in-context learning of transformers first need to be addressed before that. **W3: Choosing linear/ReLU activation on the self-attention matrix may not map well to reality.** **Reply:** We point out that previous theoretical work on in-context learning mainly considered either linear attention, or at most Relu attention (what references are good?). We do acknowledge that the goal is to move to a stronger model, however proving lower bounds becomes even more difficult. **W4: The authors mention that there is “signal” from linear attention for understanding non-linear attention (citing Ahn et al., 2024b). this is vague** **Reply:** We acknowledge that indeed the behavior of linear transformers can be different, depending on the specific setting and instances, to that of non-linear attention. The reason behind citing the work of Ahn. et al is that they show linear transformers generalize as well as non-linear ones for many complex regimes, so theoretical understanding of linear transformers is also very valuable. Moreover we note that our lower bounds also apply to attentions with Relu activation. Thus theoretically we have shown that adding Relu non-linearity cannot qualitatively improve the performance of linear transformers, which aligns with the empirical observation of Ahn. et al. **Q1: How do the authors justify that linear regression with different covariances sufficiently captures "task diversity"?** **Reply:** The notion of diversity we consider (arbitrary covariances with bounded eigenvalues) is a significant improvement over the settings considered in prior work (unimodal Gaussians). Crucially, for the unimodal Gaussian distribution even a 1-layer model can almost perfectly solve the task. However for the diversity we consider, depth becomes necessary in terms of expressivity as per our results, which makes the setting more appealing. **Q2: Can the authors clarify the signal that linear attention supposedly provides for non-linear tasks, and the feasibility of extending these findings from non-language to language models?** **Reply:** For the first part of your query about linear and non-linear attention, please refer to our response to Q2. For the second part, we believe that a key challenge in proving in-context learning results for language problems lies in understanding how the structure of a given grammar aligns with the expressive power of self-attention. Specifically, this involves questions such as how compactly a verifier can represent a proof that a sentence belongs to a particular grammar. **Q3:implications for practitioners from your work?** **Reply:** This is certainly an important question that we are actively thinking about. There are few takeaways that may be of interest to practitioners Depth matters: For Transformers to solve a diverse set of ICL problems (which language models do), depth plays a very important role. This could be a valuable insight for architecture design, if the focus is ICL abilities. Non-robustness: It was interesting to learn from the theory that multilayer models, if too deep relative to training task diversity, can learn non-robust solutions. It does so by using different layers to solve different tasks seen during training. Thus if we are going deeper, we better make sure to have a more diverse training dataset. Looped models: One fascinating finding is that looped models, although seemingly restrictive, can solve ICL with close to optimal depth in many cases. Furthermore they afford significantly higher robustness since they are forced to learn a more generalizable algorithmic solution, rather than the “shortcuts” mentioned above. This suggests that looped models should be given more attention if such algorithmic abilities are of interest