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

Authors: Khashayar Gatmiry, Nikunj Saunshi, Sashank J. Reddi, Stefanie Jegelka, Sanjiv Kumar
Primary area: learning theory
Keywords: transformers, attention, looped transformers, task diversity, in-context learning, in-context linear regression, out-of-distribution generalization

Abstract:
The intriguing in-context learning (ICL) abilities of \emph{deep Transformer models} have lately garnered significant attention. By studying in-context linear regression on unimodal Gaussian data, recent empirical and theoretical works have argued that ICL emerges from Transformers' abilities to simulate learning algorithms like gradient descent. However, these works fail to capture the remarkable ability of Transformers to learn \emph{multiple tasks} in context.
To this end, we study in-context learning for linear regression with diverse tasks, characterized by data covariance matrices with condition numbers ranging from $[1, \kappa]$, and highlight the importance of depth in this setting. More specifically, (1) (1) theoretical lower bounds of $\log(\kappa)$ (or $\sqrt{\kappa}$) linear attention layers in the unrestricted (or restricted) attention and (2) we show that the class of {\em multilayer Transformers} can indeed solve such tasks with a number of layers that matches the lower bounds. Furthermore, we show that this expressivity of multilayer Transformer comes at the price of robustness; in particular, multilayer Transformers are not robust to even distributional shifts as small as $O(e^{-L})$ in Wasserstein distance, where $L$ is the depth of the network. We then demonstrate that Looped Transformers ---a special class of multilayer Transformers with weight-sharing--- not only exhibit similar expressive power but are also provably robust under mild assumptions. Besides out-of-distribution generalization, we also show that Looped transformers are the only models that exhibit a monotonic behavior of loss with respect to depth (or number of loops).

Decision: reject
Reviewer scores: 6, 5, 3, 8, 5

Reviews and discussion:
<thread>
creation date:
1729522147640

thread_type:
official_review

<thread_object>
creator: reviewer_ps8e

<content>
summary:
The paper studies the in-context learning (ICL) ability of transformers consisting of multiple attention layers to learn linear regression tasks with diverse covariances where the eigenvalues are assumed to lie in $[\alpha,\beta]$. A lower bound of $\Omega(\sqrt{\beta/\alpha})$ for the constrained case and $\Omega(\log\beta/\alpha)$ for the unconstrained case on the number of layers required to solve all such tasks are established, and transformer constructions achieving matching upper bounds are provided. It is also shown that multilayer Transformers are not robust to task distribution shifts due to overfitting, while looped Transformers can generalize to any distribution, showing that looped models achieve both expressivity and robustness.

strengths:
The task-diverse regression setting constitutes a nice extension to the existing body of literature on ICL of linear regression where learnability can be characterized in terms of the condition number $\beta/\alpha$. The provided upper and lower bounds match, providing a tight characterization of the expressivity of multilayer Transformers in this setting.

weaknesses:
* The architecture under study is quite stylized: the attention mechanism is restricted to be linear or ReLU (robustness results only to apply to the linear case) and feedforward layers or other operations are not considered. While such restrictions are common in the ICL literature, I think this is a significant limitation for this particular paper since the main novelty seems to be the lower bounds. To be specific, the upper bounds in Theorems 4 and 5 (essentially established in previous works) tell us that practical Transformers are capable of executing a variety algorithms in context, even with certain architectural restrictions, and many derivative works have already explored this correspondence. However, the newly provided lower bounds do not shed much light on the limitations of practical Transformers, especially since the proof relies on a polynomial argument which is very 'brittle' and does not give any insight if a non-polynomial aspect such as softmax attention or feedforward layers is added.
* The above weakness could be mitigated via numerical experiments which demonstrate similar deficiencies in softmax Transformers. However, the experiments are limited to restricted linear attention.
* The strong out-of-distribution capabilities of looped Transformers have already been theoretically established in [1]. Their Theorem 4.4 shows that (quote) "a looped Transformer learned on one in-context distribution can generalize to other problem instances with different covariance, owing to the fact that it has learned a good iterative algorithm."
* Unlike [1], the paper does not provide any optimization guarantees for either looped or ordinary Transformers.
* The writing of the paper could be improved. In particular, the Appendix needs to be better structured: the ordering of the proofs does not match the results in the main text, some proofs refer to lemmas that are placed later in the Appendix, and there is often no indication in the main text where the corresponding proofs are, making the logic hard to follow. There are also frequent typos throughout the paper, see below.

**Typos and errors**

* Theorem 1: the second part should refer to non-restricted rather than restricted attention. The bound $L\le \sqrt{\beta/\alpha}$ should be $O(\sqrt{\beta/\alpha})$. Equation (7) should not have the $\gamma$ factor.
* Lemma 1 has no proof and having both Lemma 1 and Theorem 1 in the main body of text is redundant. Capitalization of "1. there" should be fixed.
* Theorem 4: The statement of Newton's method is wrong.
* Theorem 7: $P(\alpha,\beta)$ needs to be $P_{\alpha,\beta}$.
* Proof of "Lemma 7": needs to be Theorem 7. In the same paragraph, $A$ needs to be defined first and Lemma 4 should come before this.
* The next sentence reads "we get $I-\beta^{-1}\Sigma\le I$, hence $(I-\beta^{-1}\Sigma)\le I$"
* $1-\delta'$ instead of $1\delta'$ after (12)
* The sentence before Lemma 3 does not make sense and should be deleted.
* Need clarification: what is $\lambda_{min},\lambda_{max}$ in Lemma 3?
* In the proof of Lemma 3: lipschitz, chebyshev (need capitalization); miaximized, conlcude, taht, legnth (typos)
* Capitalization such as looped/Looped, equation/Equation is inconsistent.
* There are many instances of "Equation equation X" which need to be fixed.
* There are many instances of $\mathbb{E}_{XYZ}$ being displayed as $\mathbb{E}XYZ$.

[1] Gatmiry et al. Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? ICML 2024.

questions:
* Is there any way for the results or techniques to be extended to, or suggest new directions regarding, softmax attention or nonlinear feedforward layers?
* The lower bound for restricted attention is not matched by looped Transformers in Theorem 5. Is there an expressivity gap in this setting?
* Regarding out-of-distribution learning capability, what is the main difference with the result in [1]? Is the reason for robustness due to learning a good iterative algorithm as [1] suggests?
</content>
</thread_object>

</thread>
<thread>
creation date:
1729984237406

thread_type:
official_review

<thread_object>
creator: reviewer_wvsr

<content>
summary:
This paper follows a line of work studying the ability of transformers to learn to solve linear regression problems in-context. The key novelty is the consideration of the training tasks having condition number within an interval. The paper theoretically examines transformers with ReLU- and identity-activated attention (as opposed to softmax) and without an MLP or layer normalization, and shows tight upper and lower bounds on the number of layers required to solve tasks with condition number within this interval. The paper also shows that minimizers of the training loss do not generalize to new tasks whose condition number lies outside the interval for the training tasks, whereas a simple and previously-studied modification of the transformer, namely the looped transformer, which repeats attention parameters across layers, indeed generalizes. Toy experiments are conducted to support the theory.

strengths:
- The paper addresses an area of interest to the theoretical community -- understanding the ability of (simplified) transformers to solve linear regression in-context -- and provides novel, decently surprising results on this topic. As the authors point out, considering the case in which the transformers are trained on tasks with varying condition number is an improvement over prior work. The separation between multilayer and looped transformers on OOD data is especially interesting.

- The theory is rigorous, comprehensive, and clearly discussed.

- The paper is very well-written overall.

- The experiments support the theory.

weaknesses:
- The results are limited to simplified transformers that use ReLU or identity activation instead of softmax in the attention block, and do not use MLPs or layer normalization. This is common in the literature, but should be acknowledged in the introduction and abstract.

- Although it has been used previously in the literature, the assumption in the OOD case that $w^*$ and $x_q$ have inverse covariances is quite contrived and limited.

- The experiments are apparently only run once -- they should be repeated over several trials with error bars shown, especially since the curves seem to have large variance, e.g. it is strange that looped transformers generalize worse than multilayer transformers in part of Figure 2b, this could be an outlying trial. Running more trials is thus important for our understanding.

- (Minor) Figure 1a should be log-scaled.

questions:
n/a
</content>
</thread_object>

</thread>
<thread>
creation date:
1730678319509

thread_type:
official_review

<thread_object>
creator: reviewer_g1xd

<content>
summary:
This paper examines the expressivity of multi-layer ReLU or linear transformers in performing in-context learning (ICL) for linear regression tasks with varying condition numbers (from $1$ to $\kappa$). The authors derive the necessary and sufficient number of layers required for the transformers to solve the ICL problem as a function of $\kappa$. Additionally, the paper compares the out-of-distribution performance of a standard transformer with that of a looped transformer, where the QKV matrices are shared.

strengths:
There might be some merits in the results. But the writing of the paper is so below standard that I am unsure if I understand those results correctly. Therefore I am unwilling to make potentially incorrect assertions about the strength of the paper.

weaknesses:
I find the writing of the paper to be significantly below standard, to the point where it impairs the clarity and understanding of the main results. Below, I outline some potential issues (though I believe there are likely more) and my questions:

- **Line 153:** Based on the context, it should be $X^\top w^* = y$ since $X$ is $d \times n$.
- **Line 212:** According to equation (1), some index $t$ in equation (2) should be $t-1$.
- **Lines 223–228:** The formula in equation (5) is unusual, particularly with the inclusion of $(u^{(t)}, 1)$ in the last row and many zeros in $P^{(t)}$. I could not locate this formula in Ahn et al. 2024a—could you point out where this setting is considered in Ahn et al. 2024a? Additionally, what is the role of the non-zero $u^{(t)}$?
- **Line 267:** In the definition of $Z^{(0)}$, should one of the zeros be $x_q$?
- **Equation (6) in Lemma 1:** I find this equation difficult to interpret. The input to the transformer is rescaled by $\gamma$, which implies that $x_q$ in the input is also rescaled by $\gamma$ (as per the previous question). However, the “label” $(w^*)^\top x_q$ is not rescaled by $\gamma$. If my interpretation is correct, this lemma seems trivial since the transformer output and the label are not on the same scale. I have a similar concern with Theorem 1.
- **Line 281:** $w$ should be $w^*$.
- **Line 284:** “…under restricted attention…”—is this a typo?
- **Theorem 2:** The $\sqrt{\kappa} = \sqrt{\beta/\alpha}$ dependence in the number of layers suggests that the transformers are implementing accelerated gradient descent (AGD) rather than gradient descent (GD). I believe this point is worth a more detailed discussion.
- **Line 350:** Should $k$ be $L$?
- **Section 4.2:** This section is particularly difficult to follow. 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. 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?

The appendix is even more problematic. I only skimmed it but found it nearly unreadable.

- **Line 884:** The sentence is incomplete.
- **Line 1017:** $\lambda_i$ should be $\theta_i$.

In its current state, the submission requires substantial revision and polish. I recommend rejecting the paper in its present form, but I encourage the authors to resubmit after addressing these issues and thoroughly refining the manuscript.

questions:
See above please.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730711895698

thread_type:
official_review

<thread_object>
creator: reviewer_6zlh

<content>
summary:
The study explores the in-context learning (ICL) capabilities of deep Transformer models for linear regression tasks with diverse data covariance conditions. It highlights the significance of model depth, showing that multilayer Transformers can solve such tasks with a number of layers matching theoretical lower bounds. Furthermore, the authors show that this expressivity comes at the cost of robustness,  these models are sensitive to minor distributional shifts. The research also studies the Looped Transformers and shows that such a simplified transformer model can be more robust to the distributional shift.

strengths:
The strengths of this paper are summarized as follows:

* This paper considers task-diverse in-context linear regression where the in-context data may have different covariance matrices. Then, this paper proves the lower bound on the number of required attention heads to handle the task-diverse tasks.
* This paper also develops the upper bound on a class of transformer models, which matches the lower bound results.
* This paper also investigates the robustness of the multilayer transformer regarding the distributional shift and reveals its connection with the model depth, and the structural of the transformer (e.g., looped transformer vs. standard transformer).

weaknesses:
The weaknesses of this paper are summarized as follows:

* The presentation of this paper needs improvements. In particular, in Theorem 1, the role of $gamma$ is not clear to me. Why does one need to involve $\gamma$ in this theorem, is there any condition for $\gamma$ (may be the same condition in Lemma 1)? 

* In Theorem 1, the statement is made on one particular instance $(X, w^*, y, x_q)$, it is indeed possible that the model cannot provide accurate prediction for all instances. However, in practice one may expect the model to have good performance in expectation or with high probability, would it be possible to extend the lower bound results to the expectation or high probability cases?

* It is not to call "matching upper and lower bound" as the upper bound has an additional $\log(1/\epsilon)$ factor. A more rigorous claim should be "the upper bound matches the lower one up to logarithmic factors."

* Theorem 6 is a bit weird. It claims that there exists a global minimizer that will have very bad robustness as $L$ increases. I just do not quite understand why this argument is made on one existing global minimizer, is it possible that for other global minimizers, the robustness can be better? This should be made very clear.

* The proof is extremely not well organized. Many proofs do not have clean logic and are very hard to follow, thus making it hard to rigorously check the correctness of the proof. For instance, in Lemma 3, does the result hold for any polynomial function $P(\gamma)$?

questions:
Please see the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730995800018

thread_type:
official_review

<thread_object>
creator: reviewer_2dux

<content>
summary:
This paper studies in-context learning (ICL) for multiple regression tasks, focusing on "task diversity" by varying the covariances of linear regression tasks. As far as I understood this corresponds to the standard ICL setting where all n tuples (x_i,y_i) are sampled from the same "task", but the model has to be robust to changes of the covariance of the gaussian used to sample x_i. The authors present a set of theoretical results to show that 1) Transformers can do this with depth 2) they are not robust in out of distribution tasks (eg there's a price to pay for depth) 3) looped transformers (eg TFs with weight sharing) are robust to this issue (nice!). The work introduces theoretical lower bounds on the depth of Transformers required for such tasks, and also constructive proofs on how these are matched. It's interesting to see lower bounds in this context, as they are as far as I am aware, pretty rare. Some of their results are supported by experimental evidence.

strengths:
+ The paper establishes architectural lower bounds for ICL problems, which is interesting and relatively rare in this literature.
+ It highlights the lack of robustness in out-of-distribution generalization for standard*  multilayer Transformers, which is valuable. 
+ Looped Transformers show promise as they are robust to these issues, aligning well with previous work.
+ Basic experiments support the theoretical claims, adding some evidence to back up the theory.

*(caveating this that the activations are linear/relu)

weaknesses:
- It’s unclear how well these controlled settings will transfer to real-world problems. This is a broader issue in the literature on simple in-context learning (ICL) tasks where x is gaussian and y a regressor label, not just this paper.

- The idea of “diversity” here might be somewhat questionable; using the same task (linear regression) but with different covariances could still mean we’re looking at effectively the same task?

- Choosing linear/ReLU activation on the self-attention matrix may not map well to reality. 

- The authors mention that there is “signal” from linear attention for understanding non-linear attention (citing Ahn et al., 2024b). This is vague, and it’s unclear how the signal carries over from linear to standard attention, especially for language tasks. It seems like an added leap of faith without much support. (eg non-language ICL on linear transformer -> non-language ICL on standard transformer -> language ICL on linear transformer). I get that the theory becomes really complicated with other activations, but the amount of relevance of these results to real settings is still unclear.

questions:
- How do the authors justify that linear regression with different covariances sufficiently captures "task diversity"?

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

- This is a "harder" question, and you do not have to answer it for me to increase my score, but a good one to ponder on: In what ways do these results change/further inform the way that a practitioner thinks about transformers and ICL?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732627259922

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response

comment:
****We thank the reviewers for their valuable feedback to improve our work. We deliberately applied all of the comments to the paper, and uploaded the new version. We have left the changes compared to the previous version in blue. We kindly ask the reviewers to reconsider their score based on the improved version that we uploaded.****
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper considers learning ability of Transformers to perform in-context learning where the covariance of the input can differ in each task. It is shown that Transformers can implement multi-step gradient descent. In addition to that, the authors discuss robustness of Transformers, that is, the multi-layer Transformers are much less robust than Looped Transformers with parameter sharing.   

This paper somehow lacks its novelty. The presented upper-bound is already indicated by existing work. The novelty mainly lies in the lower bound analysis. However, the lower bound is not really about Transformers but its proof relies on a polynomial argument.
In addition to that, the proof is not well organized. The logic in the proofs is not clearly described. It requires substantial revision.

As a summary, the results in this paper are not so strong and the writing should go through substantial revision. Thus, I do not recommend acceptance of this paper.