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

## Abstract

Abstract The intriguing in-context learning (ICL) abilities of 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 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 , κ ] 1 𝜅 [1,\kappa] [ 1 , italic_κ ] , and highlight the importance of depth in this setting. More specifically, (a) we show theoretical lower bounds of log ⁡ ( κ ) 𝜅 \log(\kappa) roman_log ( italic_κ ) (or κ 𝜅 \sqrt{\kappa} square-root start_ARG italic_κ end_ARG ) linear attention layers in the unrestricted (or restricted) attention setting and, (b) we show that multilayer Transformers can indeed solve such tasks with a number of layers that matches the lower bounds. However, 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 ) 𝑂 superscript 𝑒 𝐿 O(e^{-L}) italic_O ( italic_e start_POSTSUPERSCRIPT - italic_L end_POSTSUPERSCRIPT ) in Wasserstein distance, where L 𝐿 L italic_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.

## 1 Introduction

Transformer-based language models (Vaswani, ) have demonstrated remarkable in-context learning abilities (Brown et al., ), thereby bypassing the need to fine-tune them for specific tasks. This is a desirable feature for large models because fine-tuning them is often expensive. In particular, given just a few samples of a new learning task presented in the context, the model is able to meta-learn the task and generate accurate predictions without having to update its parameters. In an attempt to study this phenomenon, recent work has shown that Transformers can demonstrate such in-context learning abilities on a variety of learning tasks such as linear or logistic regression, decision trees (; Garg et al., ; Xie et al., ) and, perhaps surprisingly, even training Transformers themselves (Panigrahi et al., ). This has sparked significant interest to theoretically understand the in-context learning phenomenon.
From a more theoretical perspective, there are two crucial aspects to this in-context learning ability: (1) the Transformer architecture is powerful enough to implement iterative algorithms, including first order optimization methods such as gradient or preconditioned gradient descent (Li et al., ; Von Oswald et al., ; Akyürek et al., ), and (2) Transformers with appropriate optimization algorithm can indeed learn to simulate such algorithms (Ahn et al., ; Gatmiry et al., ).
These results provide further insights into how in-context learning could arise within Transformer models. Recent papers (e.g (Ahn et al., ; Gatmiry et al., )) have studied both these aspects for Transformers on a *single task* — linear regression setting where the data is sampled i.i.d Gaussian.

One of the mysterious features of recent foundation models, which separates them from traditional learning systems, is their ability to handle a spectrum of tasks, from question answering within a broad range of topics, to reasoning or in-context learning for various learning setups. The emergence of this ability to handle task diversity naturally necessitates the training data to incorporate a diverse datasets, which inherently makes the training distribution highly multimodal.
Consequently, foundation models need more capacity to be able to solve such diverse tasks.
However, recent work on in-context learning theory focus on unimodal Gaussian data distribution. In such settings, even constant number of layers suffice to achieve any arbitrarily small error.
Thus, while interesting, these works fail to capture the remarkable ability of Transformers to learn *multiple tasks* in context and explicate the role of depth in these settings.

To this end, in this paper we introduce a new in-context learning setting, *task-diverse linear regression*, which specifically examines how well Transformers can in-context learn linear regression problems arising from fairly diverse data. More specifically, in task-diverse linear regression, the data can be sampled from multiple Gaussians where the eigenvalues of its covariance matrix can range from $[1,\kappa]$. For this setting, we ask the following important question:

*For multilayer Transformer, what is the role of depth in task-diverse linear regression? Is it possible to establish upper and lower bounds that depend on task diversity?*

In this paper, we resolve this problem by showing lower bound of $\log\kappa$ and matching upper bounds — indicating a fundamental limit on the depth required to solve the task-diverse linear regression setting. Furthermore, in the special case of attention layers considered in (Von Oswald et al., ), we obtain a lower bound of $\Omega(\sqrt{\kappa})$ instead. Correspondingly, we also show a matching upper bound in this special setting, thereby; providing a comprehensive picture of the representation power of multilayer Transformers in various important settings with task diversity.

While the above results paint a powerful picture of multilayer Transformers, it remains unclear if this comes at any cost. In this paper, we show that this is indeed the case and the price we pay is *robustness*. First, diverging significantly from prior works, we study the problem of *out-of-distribution generalization* on a very general class of distributions (see 2).
Under this general class of distributions, we show that the out-of-distribution generalization of multilayer Transformer can get exponentially worse with depth. Thus, the strong representation power of multilayer Transformers indeed comes at the price of robustness. This begs an important question: is it possible to achieve both robustness and expressivity at the same time?

Recently, there has been interest in the special class of Transformers, called Looped Transformers, which have empirically shown to have much more robust out-of-distribution generalization (Yang et al., ). Looped Transformers are essentially multilayer Transformers that share weights across layers. It is natural to ask if Looped Transformers exhibit similar representation-robustness tradeoff as multilayer Transformers. Perhaps surprisingly, we provide a negative answer to this question. In particular, we show that Looped Transformers can match the lower bound of $\log(\kappa)$ on the representative power, similar to multilayer Transformers. However, this does not come at the cost of robustness. In fact, Looped Transformers can achieve very good out-of-distribution generalization. In light of this discussion, we state the following main contributions of our paper.

- •
We consider the task-diverse in-context linear regression where the eigenvalues of the input covariance matrix has eigenvalues in the range $[1,\kappa]$; to handle this class of instances, we show a $\Omega\left(\sqrt{\kappa}\right)$ lower bound on the number of required attention heads in the restricted attention case (Section 2.3), and a $\Omega\left(\log(\kappa)\right)$ lower bound for the unconstrained case (Theorem 1). The $\Omega\left(\log(\kappa)\right)$ lower bound matches the upper bound presented in Fu et al. (). In particular, our lower bounds hold for linear attention as well as for ReLU non-linearity.
- •
We show a matching upper bound on the constrained case in Equation 6 by proving that the Transformer can implement multiple steps of the Chebyshev iteration algorithm (Theorem 2).
- •
We design an adaptive termination condition for multilayer Transformers based on the norm of the residuals; thereby, providing an flexible and efficient mechanism for early stopping with fewer layers (Theorem 3).
- •
We show that for certain training distributions on the covariance matrix, multilayer Transformers can overfit, in which case even if we are at a global minimizer with zero population loss and the test distribution has exponentially small deviation of $O(Ld^{-1}9^{-L})$ in Wasserstein distance from the training distribution, the test loss can be up to a constant error (Theorem 6).
- •
In contrast to the multilayer case, we show that under a mild condition that the training distribution puts some non-trivial mass on covariance matrices with large eigenvalues, the global minimizer of the training loss for loop Transformers extrapolate to most other distributions that have the same support (Theorem 7).
- •
Our theoretical insights regarding how depth make Transformers more task-diverse, and the robust out-of-distribution generalization of Looped Transformers compared to standard multilayer Transformers are validated through experiments in simple linear regression setting.

### 1.1 Related work

Transformers in Task Diversity.
Raventós et al. () observe a connection when the level of task diversity in the pretraining distribution is above a threshold and the emergence of strong in-context learning abilities for linear regression bu switching from a Bayesian estimator to ridge regression.
Garg et al. (); Akyürek et al. (); Von Oswald et al. () observe that transformers can extrapolate to new test distributions. Min et al. () discover surprising latent task inference capabilities in language models, even when using random labels for in-context demonstrations. Kirsch et al. () empirically show the benefits of task diversity for in-context learning in classification tasks. Chan et al. () investigate which aspects of the pretraining distribution enhance in-context learning and find that a more bursty distributions proves to be more effective than distributions that are closer to uniform. The impact of task diversity has also been studied in the context of meta learning Yin et al. (); Kumar et al. (); Tripuraneni et al. () or transfer learning and fine-tuning Raffel et al. (); Gao et al. ().

Transformers and Looping as Computational Models.
Pérez et al. () Transformers are known to be effective computational models and Turing-complete Pérez et al. (); Giannou et al. (). Garg et al. (); Akyürek et al. (); Von Oswald et al. (); Dai et al. () argue that transformers can implement optimization algorithms like gradient descent by internally forming an appropriate loss function. Allen-Zhu and Li () observe that transformers can generate sentences in Context-free Grammers (CFGs) by utilizing dynamic programming. Other work focuses on how training algorithms for Transformers can recover weights that implement iterative algorithms, such as gradient descent, for in-context learning, assuming simplified distributions for in-context instances Ahn et al. (); Zhang et al. (); Mahankali et al. (). Notably, while most of these studies focus on the single-layer case, recent work by the authors in Gatmiry et al. () provides a global convergence result for training deep Looped Transformers.

## 2 Preliminaries

### 2.1 In-context Linear Regression

We study the role of depth and looping for in-context learning in linear regression. The setting of linear regression for in-context learning has also been studied by Gatmiry et al. (); Ahn et al. (); Von Oswald et al. (); Akyürek et al. (). We briefly describe the setup here.

Linear regression. We consider the linear regression problem $(X,w^{*},y,x_{q})$ with data matrix $X\in\mathbb{R}^{d\times n}$, labels $y$, regressor $w^{*}$, and query vector $x_{q}$. In particular, $X$ consists of the observed data inputs $\{x_{i}\}_{i=1}^{n}$ as its columns:
$X=[x_{1},\dots,x_{n}]$.
In this work, we assume that the instances are *realizable* i.e., $Xw^{*}=y$.
Thus, $w^{*}$ can be obtained by the following pseudo-inverse relation:

$$ $\displaystyle w^{*}=(XX^{\top})^{-1}Xy.$ $$

We define $\Sigma=XX^{\top}$ to be the input covariance since this will play an important role going forward.

### 2.2 Task Diversity & Robustness

To handle task diversity, we need a model that not only makes accurate predictions for a fixed distribution, but is also robust to a class of distributions. Here, we focus on the notion of robustness with respect to various distributions on the covariance matrix
$\Sigma$ of the linear regression instance.
It is unreasonable to expect the model to behave well for all possible distributions, so we focus on the set of instances where the eigenvalues of the covariance matrix are in the range $(\alpha,\beta)$ for $0<\alpha<\beta$.

###### Definition 1 .

We define the set $S_{\alpha,\beta}$ of $(\alpha,\beta)$ “normal” linear regression instances as the set instances with a well-conditioned covariance matrix $X^{\top}X$:

$$ $\displaystyle S_{\alpha,\beta}=\left\{(X,w^{*},y,x_{q}):\alpha I\preccurlyeq XX ^{\top}\preccurlyeq\beta I\right\}.$ $$

In order for the model to handle task diversity on the set of instances in $S_{\alpha,\beta}$, the loss needs to be small for any distribution $P_{\alpha,\beta}$ that is supported on covariances $\alpha I\preccurlyeq\Sigma\preccurlyeq\beta I$, which means the model has to be accurate for all the instances in $S_{\alpha,\beta}$. In this work, we are interested in understanding the ability of Transformers to handle task diversity from two angles: (1) expressiveness of the Transformer architecture is in representing task-diverse models (Section 3) and (2) robustness of these models in terms of out-of-distribution generalization. We first formally define the models used in this paper before exploring these questions.

###### Definition 1 .

### 2.3 Transformer Model

Self-Attention layer. Central to our paper is the single attention layer, which we define below. First we define the matrix $Z^{(0)}$ as a $(d+1)\times(n+1)$ matrix by combining the data matrix $X$, their labels $y$, and the query vector $x_{q}$ as follows:
$Z^{(0)}=\begin{bmatrix}X&x_{q}\\
y^{\top}&0\end{bmatrix}.$

Given the key, query, and value matrices
$W_{k},W_{q},W_{v}$,
we follow (Ahn et al., ; Von Oswald et al., ) and define the self-attention model $\text{Attn}^{lin}\left({Z};W_{k,q,v}\right)$ with activation function $\sigma$ as:

$$ $\displaystyle\text{Attn}^{lin}\left({Z};W_{k,q,v}\right)\triangleq W_{v}ZM \sigma(Z^{\top}{W_{k}}^{\top}W_{q}Z),$ $\displaystyle M\triangleq\begin{bmatrix}I_{n\times n}&0\\ 0&0\end{bmatrix}\in\mathbb{R}^{(n+1)\times(n+1)},$ $$

where the index $k\times r$ below a matrix determines its dimensions, and by $\sigma(\mathfrak{M})$ for matrix $\mathfrak{M}$ we mean entry-wise application of $\sigma$ on $\mathfrak{M}$.
To simplify the notation for analysis, we combine the key and query matrices into $Q$, as $Q=W_{k}^{\top}W_{q}$. Hence, we have the following alternative parameterization:

$$ $\displaystyle\text{Attn}^{lin}\left({Z};Q,P\right)\triangleq PZM\sigma(Z^{\top }QZ),$ (1) $$

where now the learnable parameters are the matrices $P,Q$. We consider the cases where
$\sigma(x)=Relu(x)$ or $\sigma(x)=x$ is linear. If $\sigma$ is linear, then the model becomes linear self-attention, which was first considered by Ahn et al. (); Von Oswald et al. (); Schlag et al. () to understand the behavior of in-context learning for linear regression. Note that even with linear attention, the output of the attention layer $\text{Attn}^{lin}\left({Z};Q,P\right)$ is a nonlinear map in $(P,Q)$ or $Z$, and hence challenging to analyze. Since we consider multilayer models, it turns into a low-degree polynomial in either $Z$ or in the set of weights $\{P^{(t)},Q^{(t)}\}_{t=0}^{L-1}$; the limits of its expressivity is investigated in this work. Furthermore Ahn et al. () showed that studying linear attention can already provide signal about non-linear attention.
Here, we study the questions of expressivity and robustness.

Transformers. Using $L$ attention heads, we define a Transformer block $\text{TF}({Z^{(0)}},\{P^{(t)},Q^{(t)}\}_{t=0}^{L-1})$. In particular, following the notation in (Ahn et al., ), we define

$$ $\displaystyle Z^{(t)}\triangleq Z^{(t-1)}-\tfrac{1}{n}\text{Attn}^{lin}\left({ Z^{(t-1)}};Q^{(t)},P^{(t)}\right).$ (2) $$

which results in the recursive relation

$$ $\displaystyle Z^{(t+1)}=Z^{(t)}-\tfrac{1}{n}P^{(t)}{Z^{(t)}}M\sigma({Z^{(t)}}^ {\top}Q^{(t)}{Z^{(t)}}).$ (3) $$

Then, the final output of the Transformer is the $((d+1),(n+1))$ entry of $Z^{(t)}$:

$$ $\displaystyle\text{TF}({Z^{(0)}},\{P^{(t)},Q^{(t)}\}_{t=0}^{L-1})=-{Z^{(L)}}_{ (d+1),(n+1)}.$ (4) $$

Note the minus sign in the final output, which is primarily for the ease of our exposition later on.

Looped Transformer Model.
A Looped Transformer model is simply a multilayer Transformer with weight-sharing i.e., we define it as $\text{TF}({Z^{(0)}},P,Q)=-{Z^{(L)}}_{(d+1),(n+1)}$ where

$$ $\displaystyle Z^{(t+1)}=Z^{(t)}-\tfrac{1}{n}P{Z^{(t)}}M\sigma({Z^{(t)}}^{\top} Q{Z^{(t)}}).$ (5) $$

Restricted Attention.
Following Ahn et al. (), in addition to our lower bound in Theorem 1 for the more general case when $\{P^{(t)},Q^{(t)}\}_{t=0}^{L-1}$ can be arbitrary, we also study a relevant special case when the last row and column of $Q^{(t)}$ are zero and only the last row of $P^{(t)}$ is non-zero i.e.,

$$ $\displaystyle Q^{(t)}=\left[\begin{matrix}{}A^{(t)}&0\\ 0&0\end{matrix}\right],\ P^{(t)}=\left[\begin{matrix}{}0_{d\times d}&0\\ {u^{(t)}}^{\top}&1\end{matrix}\right].$ (6) $$

This special case is important as it can implement preconditioned gradient descent, while it keeps the feature matrix $X$ the same going from $Z^{(t-1)}$ to $Z^{(t)}$.
In this case, the left upper $d\times n$ block of the second term in Equation 3 is always zero. This is because we assume the $d\times d$ zero block in $P^{(t)}$. We denote the output of the Transformer in this case by $\text{TF}({Z^{(0)}},\{A^{(t)},u^{(t)}\}_{t=0}^{L-1})$. We denote the first $n$ entries of the last row of $Z^{(t)}$ by vector ${y^{(t)}}^{\top}$, and the last entry by $y^{(t)}_{q}$, i.e. $y^{(t)}_{q}={Z^{(t)}}_{(d+1)\times(n+1)}$. Unrolling Equation 6, we obtain the following recursions for $y^{(t)}$ and $y^{(t)}_{q}$:

$$ $\displaystyle{y^{(t+1)}}^{\top}={y^{(t)}}^{\top}-\frac{1}{n}({y^{(t)}}^{\top}+ {u^{(t)}}^{\top}X)\sigma\Big{(}X^{\top}A^{(t)}X\Big{)},$ $\displaystyle y^{(t+1)}_{q}=y^{(t)}_{q}-\frac{1}{n}({y^{(t)}}^{\top}+{u^{(t)}} ^{\top}X)\sigma\Big{(}X^{\top}A^{(t)}x_{q}\Big{)}.$ $$

Note that the final output of the Transformer, $\text{TF}({Z^{(0)}},\{A^{(t)},u^{(t)}\}_{t=0}^{L-1})$, is equal to $y^{(L)}_{q}$.

## 3 On the Role of Depth in Multilayer Transformers for Task Diversity

In this section, we examine the power and limits of multilayer Transformers. We first present lower bounds for multilayer Transformers in both the general case and the restricted attention case (Section 2) and, then provide matching upper bounds. While this makes a compelling case for the expressivity power of multilayer Transformers, in Section 4.2.1, we show that multilayer Transformers can suffer from overfitting and perform very poorly under distribution shifts.

### 3.1 Lower bound on the power of Multilayer Transformers

In this section, we provide lower bound for multilayer Transformers in solving in-context learning.

###### Lemma 1 (Effect of scaling on accuracy - restricted attention) .

For an $L$-layer Transformer with architecture defined in Equation 3 and 6, with arbitrary weights and positive homogeneous ReLU activation $\sigma$ or simply using linear attention ($\sigma(x)=x$), consider an arbitrary instance of a realizable linear regression $(X,w^{*},y,x_{q})$;

- 1.
there exists a scaling $1\leq\gamma\leq 36L^{2}$ such that the Transformer’s response for the scaled instance $(\gamma X,w^{*},\gamma y,x_{q})$ will be inaccurate by constant:
$\displaystyle\frac{|\text{TF}({\gamma Z^{(0)}},\{A^{(t)},u^{(t)}\}_{t=0}^{L-1}
)-{w^{*}}^{\top}x_{q}|}{|{w^{*}}^{\top}x_{q}|}\geq\frac{1}{4},$
(7)
where $Z^{(0)}=\left[\begin{tabular}[]{c c}$X$&0\\
$y^{\top}$&0\end{tabular}\right]$ is the Transformer input corresponding to the instance $(X,w,y,x_{q})$. Even more, there exists an interval $[a,b]$ inside $[1,36L^{2}]$ with length at least constant (independent of $L$) such that for all $\gamma\in(a,b)$ we have Equation 7.
- 2.
Furthermore, if we do not restrict the weight matrices $\{P^{(t)},Q^{(t)}\}_{t=0}^{L-1}$ to have the form in Equation 6, then there is a scaling $1\leq\gamma\leq 2^{L}$ such that
$\displaystyle\frac{|\text{TF}({\gamma Z^{(0)}},\{P^{(t)},Q^{(t)}\}_{t=0}^{L-1}
)-{w^{*}}^{\top}x_{q}|}{|{w^{*}}^{\top}x_{q}|}\geq\frac{1}{4}.$

###### Theorem 1 (Lower bound on the representation power of Transformers) .

For the Transformer architecture with ReLU or linear activation defined in Equation 4, under the weight restriction Equation 6,
consider the set $\mathcal{S}_{\alpha,\beta}$ of $(\alpha,\beta)$-normal realizable instances of linear regression $(X,w,y,x_{q})$ where $\alpha I\preccurlyeq X^{\top}X\preccurlyeq\beta I$.
Then, given $L\leq\sqrt{\beta/\alpha}$ and for any choice of $\{A^{(t)},u^{(t)}\}_{t=1}^{L}$ and query vector $x_{q}$, there exists a normal instance $(X,w^{*},y,x_{q})\in\mathcal{S}_{\alpha,\beta}$ with that query vector, such that

$$ $\displaystyle\frac{|\text{TF}({\gamma Z^{(0)}},\{A^{(t)},u^{(t)}\}_{t=0}^{L-1} )-w^{\top}x_{q}|}{|w^{\top}x_{q}|}\geq\frac{1}{2}.$ (8) $$

Furthermore, under the unrestricted attention (Equation 1), given $L=O\left(\log\left(\beta/\alpha\right)\right)$ for small enough constant, we incur at least constant error for some instance $(X,w^{*},y,x_{q})\in\mathcal{S}_{\alpha,\beta}$ due to Equation 8.

Theorem 1 asserts that a minimum depth is required for multilayer Transformer to be able to solve all instances of linear regression whose covariance matrix is in the range $[\alpha,\beta]$. In particular, this lower bound depends on $\tfrac{\beta}{\alpha}$, which increases with the diversity of the instances. For the restricted attention case, the lower bound is $\Omega(\sqrt{\beta/\alpha})$ whereas for the general case the lower bound is $\Omega(\log(\beta/\alpha))$.

###### Remark 3.1 .

(Fu et al., ) show that, given full freedom in the weights of Transformers, one can implement one step of the iterative Newton method with an attention head. Hence, multilayer Transformers with depth at least $\Omega(\ln\left(\beta/\alpha\right))$ are able to accurately solve all instances in $S_{\alpha,\beta}$. Our result essentially shows a matching lower bound for this case in Theorem 1.
Interestingly, their construction uses the same set of weights for all of the attention layers, therefore, their network is indeed a Looped Transformer.
This shows that looped models can also match the lower bound in terms of numbers of layers required.
We will revisit this in Section 4.1.

###### Lemma 1 (Effect of scaling on accuracy - restricted attention) .

###### Theorem 1 (Lower bound on the representation power of Transformers) .

###### Remark 3.1 .

### 3.2 Matching Upper bound for Multilayer Transformer

We show the existence of weights for a linear Transformer that can solve a linear regression instance with bounded condition number using restricted attention defined in Equation 6.

###### Theorem 2 .

For $\sigma(x)=x$, there exists a set of weights $\{A^{(t)},u^{(t)}\}_{t=0}^{L-1}$ for linear Transformer in the restricted case defined in equation 6 with depth at most $L=O(\ln(1/\epsilon)\sqrt{\beta/\alpha})$ such that for every linear regression instance $(X,w^{*},y,x_{q})\in\mathcal{S}_{\alpha,\beta}$ we have

$$ $\displaystyle|\text{TF}({Z^{(0)}},\{A^{(t)},u^{(t)}\}_{t=0}^{L-1})-{w^{*}}^{ \top}x_{q}|\leq\epsilon\|w^{*}\|\|x_{q}\|.$ $$

Notably, the dependency $\sqrt{\beta/\alpha}$ in Theorem 2 matches the lower bound that we show for the same restricted Transformer model in Theorem 1. The square root dependency comes from the fact that one can implement an iterative algorithm for solving a linear regression instance using the properties of Chebyshev polynomials. The proof of Theorem 2 can be found in Section C.3.

###### Theorem 2 .

### 3.3 Adaptive Depth

In this section, we investigate adaptive strategies that can be used for terminating a multilayer Transformer earlier, before passing the input through all the $L$ layers. Surprisingly, our result below, shows convergence of the output of the Transformer if we adaptively terminate based on $\|y^{(\ell)}\|$.

###### Theorem 3 (Termination guarantee) .

For a linear Transformer architecture as defined in Equation 3 with $\sigma(x)=x$, given $u^{(i)}=0$, suppose we wait until $\|y^{(\ell)}\|\leq\frac{\epsilon}{\|x_{q}\|_{XX^{\top}}}$. Then, the output of the Transformer at layer $\ell$ is close to the true label:

$$ $\displaystyle|y_{q}-\text{TF}({Z^{(0)}},\{A^{(t)},u^{(t)}\}_{t=0}^{\ell-1})| \leq\epsilon\|x_{q}\|_{XX^{\top}}.$ $$

This result further highlights the computational benefits by terminating in an adaptive manner.

###### Theorem 3 (Termination guarantee) .

## 4 On the Power & Robustness of Looped Transformers versus Multi-layer Transformers

In the previous section, we studied expressivity of multilayer Transformers for task-diverse linear regression.
In Section 4.1, we discuss how even though Looped Transformers are more restricted compared to multilayer Transformers, their expressive power is just as good for in-context linear regression.
While expressivity is important for task-diversity of our model, we also want the model to be robust to distribution shifts from the training distribution.
In Section 4.2, we show that multilayer Transformers are not very robust to even very tiny perturbations to the training distribution. In contrast, Looped Transformers are provably more robust.
Finally, motivated by the fact that early stopping is a desirable property for neural nets, in Section 4.3 we study the monotonicity behavior of multilayer Transformers with respect to depth and interestingly find that they cannot behave monotonic with respect to all covariance matrices unless the weights of different layers are equal.

### 4.1 Expressivity of Looped Transformers

Looped Transformers are a subclass of multilayer Transformers that uses weight-sharing, and thus, cannot exceed multilayer Transformers in representation power. Despite this restriction, we observe that their ability in handling task-diversity almost remains unchanged. In particular, according to the construction in (Fu et al., ), there exists a Looped transformer in addition to constant number of attention layers, that can achieve small error for all instances in $S_{\alpha,\beta}$ after $\ln(\beta/\alpha)$ number of loops, matching the lower bound in Theorem 1.

###### Theorem 4 (Restatement of Theorem 5.1 in Fu et al. ( 2023 ) ) .

There exists a Looped Transformer with an additional eight attention layers that implements the Newton algorithm i.e., for instance $(X,w^{*},y,x_{q})$, looping $L$ times results in output $\hat{x}_{q}^{\top}w_{L}^{(\text{Newton})}$ where $w_{L}^{(\text{Newton})}\triangleq M_{L}Xy$ where $M_{j}$ is updated as
$M_{j}=2M_{j-1}\Sigma M_{j-1},j\in[1,k],M_{0}=\alpha\Sigma,$
where $\Sigma=XX^{\top}$.

Furthermore, in the restricted attention case as defined in Equation 6, with the implementation of gradient descent by Von Oswald et al. () which also uses weight sharing, it is easy to check that Looped Transformers can still achieve small error on all instances in $S_{\alpha,\beta}$ with depth $O(\beta/\alpha)$ (see Section B); in this case, we see a gap with the $\Omega(\sqrt{\beta/\alpha})$ lower bound in Theorem 1.

###### Theorem 5 (Follows from Proposition 1 in Von Oswald et al. ( 2023 ) ) .

There exists a Looped Transformer architecture in the restricted attention case which for a linear regression instance $S=(X,w^{*},y,x_{q})\in\mathcal{S}_{\alpha,\beta}$ can achieve accuracy $|\text{TF}({Z^{(0)}},\{A,0\}_{t=0}^{L-1})-y|\leq\epsilon$ after $O(\ln(1/\epsilon)\beta/\alpha)$ number of loops.

###### Theorem 4 (Restatement of Theorem 5.1 in Fu et al. ( 2023 ) ) .

###### Theorem 5 (Follows from Proposition 1 in Von Oswald et al. ( 2023 ) ) .

### 4.2 Out-of-Distribution Generalization

So far we examined the crucial role of depth for task-diversity settings. However, it is unclear if the weights that minimize the population loss are robust to distributions shifts. To this end, we study the task diversity of transformers through the lens of out-of-distribution generalization.
We first study the limitations of multi-layer Transformers with respect to out-of-distribution generalization. To enable this study, we consider the population loss with respect to various distributions on the covariance matrix in the restricted attention case defined in Section 2.3. In particular, we focus on distribution shifts on the covariance matrix of the linear regression instances, so we set $u^{(t)}=0$ and define the loss function with respect to a distribution $P$ on the covariance matrix $\Sigma=XX^{\top}$; for every realization of the covariance matrix sampled from $P$, we calculate the loss by averaging over $w^{*}\sim N(0,\Sigma^{-1})$ and $x_{q}\sim N(0,\Sigma)$. This choice of distribution follows (Gatmiry et al., ) and is natural since we expect the distribution of $x_{q}$ to have the same covariance as the data covariance matrix of the observed data, $\Sigma$, and it also simplifies the analysis. Therefore, we define the loss as

$$ $\displaystyle\mathcal{L}^{P}(\{A^{(t)}\}_{t=0}^{L-1})\triangleq\mathbb{E}_{ \Sigma\sim P,w^{*}\sim N(0,\Sigma^{-1}),x_{q}\sim N(0,\Sigma)}\left(\text{TF}( {Z^{(0)}},\{A^{(t)},0\}_{t=0}^{L-1})-y\right)^{2}.$ (9) $$

In particular, for the special case of the looped model (where $A^{(t)}=A$ for $t\in[L-1]$), the loss is:

$$ $\displaystyle\mathcal{L}^{P}(A)\triangleq\mathbb{E}_{\Sigma\sim P,w^{*}\sim N( 0,\Sigma^{-1}),x_{q}\sim N(0,\Sigma)}\left(\text{TF}({Z^{(0)}},\{A,0\}_{t=0}^{ L-1})-y\right)^{2}.$ $$

When the distribution $P$ is a point mass on covariance $\Sigma$, we denote the loss by $\mathcal{L}^{\Sigma}(\{A^{(t)}\}_{t=0}^{L-1})$ and $\mathcal{L}^{\Sigma}(A)$. We denote a distribution on the covariance matrices that is only supported on $\{\alpha I\preccurlyeq\Sigma\preccurlyeq\beta I\}$ by $P_{\alpha,\beta}$, with sub-indices $\alpha,\beta$ showing the interval of the eigenvalues for the support of the covariances.
The key message that we deliver in this section is that loop Transformers are more robust for out-of-distribution generalization compared to multilayer Transformers.
To elucidate the advantage of looped over multilayer Transformers, we need the following definition (Section 4.2.2).

###### Definition 2 .

We say the distribution $P_{\alpha,\beta}$ supported on the covariance matrices of linear regression instances in $S_{\alpha,\beta}$ (or we abbreviate by saying supported on $S_{\alpha,\beta}$) is $(\epsilon,\delta)$ right-spread-out if for every fixed unit vector $v$, $P_{\alpha,\beta}(\|X^{\top}v\|^{2}\geq(1-\delta)\beta)\geq\epsilon$.

The right-spread-out property restricts the distribution on the data covariance $XX^{\top}$ so that for every direction $v\in\mathbb{R}^{d}$, it puts a minimum amount of mass on matrices whose eigenvectors with large eigenvalues, that are close to the right end point of the interval $(\alpha,\beta)$, are close to $v$.

###### Definition 2 .

#### 4.2.1 Weakness of Multilayer Transformers for Out-of-Distribution Generalization

From Section 3.2, recall that there exists a multilayer Transformer of depth $O(\ln(1/\epsilon)\sqrt{\beta/\alpha})$ that can solve any linear regression with covariance $\Sigma$ such that $\alpha I\preccurlyeq\Sigma\preccurlyeq\beta I$. However, it is unclear if minimizing the training loss recovers such a network. Here, we show that this is not the case.

###### Theorem 6 (Multilayer Transformers blow up out of distribution) .

Given $\frac{\beta}{\alpha}\geq 10$ and for any $\epsilon>0$, there exists a $(\frac{1}{L},0)$ right-spread-out distribution $P_{\alpha,\beta}$ such that there exists a global minimizer $\{{A^{(t)}}^{*}\}_{t=0}^{L-1}$ of the loss $\mathcal{L}^{P_{\alpha,\beta}}(\{A^{(t)}\}_{t=0}^{L-1},0)$ for linear Transformer defined in equation 3, so that for any other distribution $\tilde{P}_{\alpha,\beta}$ on $\Sigma$ which has at least $\epsilon$ mass supported on $8\alpha I\preccurlyeq\Sigma\preccurlyeq(1-\delta^{\prime})\beta I$, the out of distribution loss $\mathcal{L}^{(\tilde{P}_{\alpha,\beta})}$ exponentially blows up at $\{{A^{(t)}}^{*}\}_{t=0}^{L-1}$ with the depth:

$$ $\displaystyle\mathcal{L}^{(\tilde{P}_{\alpha,\beta})}(\{{A^{(t)}}^{*}\}_{t=0}^ {L-1},0)\geq\epsilon\delta^{\prime}d9^{L-1}.$ $$

In particular, Theorem 6 shows that multilayer Transformers are essentially prone to overfiting, so that with a slight deviation of the test distribution from the train distribution (even $O(\tfrac{L}{d9^{L-1}})$) in Wasserstein distance, the test loss will incur constant error for the global minimizer of the train loss.

###### Theorem 6 (Multilayer Transformers blow up out of distribution) .

#### 4.2.2 Out-of-distribution generalization of Looped Transformers

Our result in Theorem 6 states that multi-layer Transformers can behave poorly for out-of-distribution generalization. Here, we show that Looped Transformers can indeed circumvent this issue by restricting the model with weight sharing when the training distribution is right-spread-out.

###### Theorem 7 (Looped Transformer is robust out-of-distribution) .

Given an $(\epsilon,\delta)$ right-spread-out distribution $P(\alpha,\beta)$ on $S_{\alpha,\beta}$, let ${A^{*}}$ be the global minimizer of $\mathcal{L}^{(P_{\alpha,\beta})}(A,0)$ for the linear Looped Transformer defined in equation 5 with $\sigma(x)=x$ in the region $\alpha I\preccurlyeq A^{-1}\preccurlyeq\beta I$. Then, given any $\epsilon^{\prime}\leq\epsilon$ and any other distribution $\tilde{P}_{\alpha,(1-\delta^{\prime})\beta}$ that is supported on $(\alpha,(1-\delta^{\prime})\beta)$, for

$\delta^{\prime}=\delta+\frac{\ln(d/\epsilon^{\prime})}{L}$ such that $\delta^{\prime}\leq 1-\frac{\alpha}{\beta}$, we have
$\mathcal{L}^{(\tilde{P}_{\alpha,(1-\delta^{\prime})\beta})}(A^{*},0)\leq
\epsilon^{\prime}\wedge\left(1-\frac{\alpha}{\beta}\right)^{2L}.$

Theorem 7 demonstrates that if the training distribution is sufficiently spread out, the global minimizer in a looped model remains robust to a wide range of distributional shifts in the test loss. Specifically, the test distribution can be any arbitrary distribution on the covariance matrix, as long as its eigenvalues lie within the slightly narrowed interval of $(\alpha,(1-\delta^{\prime})\beta)$, where $\delta^{\prime}=\delta+\gamma\ln(d)$ approaches $\delta$, the spread-out parameter of the training distribution, as the number of loops increases. This contrasts with the behavior of multilayer Transformers, which, as shown in Theorem 6, can overfit when trained on right-spread-out distributions.

###### Theorem 7 (Looped Transformer is robust out-of-distribution) .

### 4.3 Non-monotonicity of the loss in Multilayer Transformer

In Section 3.3, we discussed how one can adaptively pick a depth based on problem difficulty. A related idea is that of early-exiting (Teerapittayanon et al., ), where the goal is to exit the model at an earlier layer for easier example. This naturally provides inference efficiency. An important consideration in the early exit literature is monotonicity with respect to layers, i.e., the model’s error decreases as the layer index increases. This has been studied in detail (Baldock et al., ; ) and there is also work on enforcing such monotonicity (Schuster et al., ; Jazbec et al., ).
In the linear regression incontext learning setting considered in this paper, we make an intriguing observation: if a multilayer model has monotonically decreasing error with depth for a diverse set of distributions, then it must be a looped model.
This suggests that looped models are naturally suited for early-exiting strategies. We believe this phenomenon deserves further exploration in future work. The following theorem formalizes this idea.

###### Theorem 8 (Monotonic behavior w.r.t depth $\rightarrow$ equal weights) .

For multi-layer Transformer $\text{TF}({\{A^{(t)},0\}_{t=0}^{L-1}},\{A^{(t)},u^{(t)}\}_{t=0}^{L-1})$ if the average value of the loss is monotonic for $0\leq t\leq L-1$ for every covariance matrix $\Sigma$, then we have $A^{(1)}=\dots=A^{(L)}$. Moreover, for the loop model and for every covariance $\Sigma$, there exists an $L_{0}$ which depends on $\Sigma$, such that for all $L\geq L_{0}$, the behavior of the loss becomes monotonic with respect to $L$.

Proof of Theorem 8 can be found in Section F. At a high level, Theorem 8 states that in order to have a multilayer Transformer in the restricted attention case (Section 2) for which the population loss behaves monotonically for an arbitrary choice of distribution on the covariance matrix, then the only possibility is for the Transformer to be looped. Furthermore, Theorem 8 states that this monotonicity property indeed holds for the looped model for large enough depth. The proof of Theorem 8 uses the fact that for any covariance distribution, the loss function can be related to the spectrum of the weight matrices $A^{(i)}$’s, and that having unequal $A^{(i)}$’s, one can construct an adversarial distribution on the covariance for which the loss does not behave monotonically with respect to the depth.

###### Theorem 8 (Monotonic behavior w.r.t depth → → \rightarrow → equal weights) .

## 5 Experiments

Figure: Figure 1: We evaluate the test loss of looped models and multilayer model as the function of loops and depth respectively. In (a), we plot the test loss as a function of number of layers for three different covariance ranges. As predicted by theory, the larger the range ($\kappa$), the more layers are needed to get a small loss. In (b) we observe that the number of loops required to solve the problem is very close to the number of layers required for multilayer model.
Refer to caption: extracted/5960819/Media/difficulty_multilayer_plot.png

In this section, we run experiments on the in-context learning linear regression problem to validate the theoretical results. In particular, we would like to demonstrate: (1) role of depth (in both multilayer and looped Transformers) with increasing task diversity and (2) robustness of Looped Transformer for out-of-domain generalization.
We use the codebase and experimental setup from (Ahn et al., ) for all our linear regression experiments. In particular, we work with $d=10$ dimensional inputs and each input instance consists of $n=20$ pairs of $(x,y)$ in the context.
We train with $L$ attention layer models for multilayer training and 1 layer attention model looped $L$ times, as described in Section 2.3. For simplicity we use the restricted linear attention setting. In all experiments, we follow the same data distribution as described in Section 4.2.
For all experiments, the covariances are sampled from training and test distributions as follows: train distribution is $\Sigma=sI_{d}$, $s\sim\mathcal{D}_{\text{train}}$ and test distribution: $\Sigma=sI_{d}$, $s\sim\mathcal{D}_{\text{test}}$. In each experimental setting we will specify $\mathcal{D}_{\text{train}}$ and $\mathcal{D}_{\text{test}}$.

Role of Depth. We demonstrate the importance of depth in presence of task diversity. Consider the following train and test distributions $\mathcal{D}_{\text{train}}=\mathcal{D}_{\text{test}}=\text{Unif}([1,\kappa])$.
Task diversity is controlled by varying $\kappa\in[2,4,8]$, with higher value corresponding to more diverse tasks. In Figure 1, we see that that more diverse tasks (larger $\kappa$) require more layers (or loops) in order to achieve low loss. This aligns with the theoretical results from Section 3.1. Furthermore, we find that looped model with $L$ loops can be very competitive to multilayer model with $L$ layers even in the most diverse setting.

Figure: Figure 2: We evaluate the robustness of looped models and multilayer model as the level of deviation of the test distribution from the train distribution increases. The three plots differ in the number of distinct covariances used in the train distribution. The test distribution is supported on a window of size $w$ around the train covariances, and $w$ is varied on the x-axis. We see that in all settings the looped model is more robust than the multilayer model, especially on the left plot where there is least diversity in the training distribution, as predicted by the theory.
Refer to caption: extracted/5960819/Media/ood_cov2.png

Out-of-distribution generalization.
For this set of experiments, we train models using a distribution comprising of $k$ covariances. The precise form the train and test distributions are:

- •
Train distribution: $\mathcal{D}_{\text{train}}=\text{Unif}(\{s_{1},\dots,s_{k}\})$ where $S_{k}=\{s_{1},\dots,s_{k}\}$ is selected to be $k$ values uniformly spread out in the range $[1,8]$.
- •
Test distribution: $\mathcal{D}_{\text{train}}=\text{Unif}(\cup_{i=1}^{k}[s_{i}-w,s_{i}+w])$ for some deviation $w\in\mathcal{R}$.

Note that the train distribution is a mixture of $k$ tasks, whereas the test distribution has a slight deviation using a window size of $w$ around the training covariances.
For evaluations, $w$ is varied from 0, corresponding to in-distribution evaluation, to a max value of 0.5.
In Figure 2, we plot the robustness of multilayer and looped model for different training distributions corresponding to $k=2,5$ and $9$. Consistent with our theory in Section 4.2, we observe that Looped Transformers have much better out-of-distribution generalization compared to multilayer Transformers.
This also aligns with the empirical results from Yang et al. () on various incontext learning problems.
Furthermore, as the number of train covariances increase, the looped model becomes perfectly robust.

## 6 Conclusion

In this paper, we study a more realistic in-context setting than prior works — task-diverse linear regression, which we believe is more reflective of the incontext abilities of Transformer based foundation models. For this setting, we provide lower bounds on the depth and show Transformers match these bounds; thereby, providing a comprehensive picture about number of layers required to solve the problem. While multilayer Transformers demonstrate this powerful representation power, we show they have weak out-of-distribution generalization, which gets exponentially worse with respect to depth. Surprisingly, we show the looped Transformers exhibit similar representative power with much better robustness.
We validate our theoretical findings through experiments on linear regression setting.
Finally, we touch upon aspects like adaptive depth and monotonicity of loss in layers, and we believe that these are very interesting future directions to pursue to understand the power of looped models.
While the incontext linear regression setting is quite simplified, it already provides a lot of interesting insights in the role of depth and task diversity. Extending this study to other settings, like reasoning, is a very interesting future direction.

## Appendix A Related work

Transformers in Task Diversity.
Raventós et al. [] empirically study the relationship between task diversity in the pretraining distribution and the emergence of in-context learning abilities for linear regression. They observe an intriguing phase transition, where as task diversity increases, the learned model shifts from acting as a Bayesian estimator to implementing ridge regression. Garg et al. [] investigate in-context learning across various tasks, where both the inputs and the link function for the in-context instances are sampled from an underlying distribution. For in-context linear regression, Garg et al. [], Akyürek et al. [], Von Oswald et al. [] further observe that transformers are capable of extrapolating to new test distributions, both in terms of the inputs and the regressor. Min et al. [] discover surprising latent task inference capabilities in language models; even when using random labels for in-context demonstrations (i.e., in the test distribution), the model can still produce accurate predictions, provided the input distribution in the test instances is similar to that of the pretraining data. The impact of task diversity has also been studied in the context of meta learning Yin et al. [], Kumar et al. [], Tripuraneni et al. []. Kirsch et al. [] empirically demonstrate the benefits of task diversity for in-context learning in classification tasks. Chan et al. [] investigate which aspects of the pretraining distribution enhance in-context learning and find that a more bursty distribution, with potentially a large number of rare clusters, proves to be more effective than distributions that are closer to uniform. Raffel et al. [], Gao et al. [] explore the role of task diversity in the pretraining distribution for improvements in the context of transfer learning and fine-tuning on downstream tasks.

Transformers as Computational Models.
Pérez et al. [] Transformers are known to be powerful computational models, Pérez et al. [] demonstrating their Turing-completeness. Giannou et al. [] propose that Looped Transformers can function as programmable computers without scaling the size of its architecture with the runtime logic such as program loops. Akyürek et al. [], Von Oswald et al. [], Dai et al. [] argue that transformers are expressive enough to implement optimization algorithms like gradient descent by internally forming an appropriate loss function. Garg et al. [] empirically show that transformers can perform in-context learning across various learning tasks, while Allen-Zhu and Li [] observe that transformers can learn and generate sentences in Context-free Grammers (CFGs) by utilizing dynamic programming. Other work focuses on how training algorithms for Transformers can recover weights that implement iterative algorithms, such as gradient descent, for in-context learning, assuming simplified distributions for in-context instances Ahn et al. [], Zhang et al. [], Mahankali et al. []. Notably, while most of these studies focus on the single-layer case, recent work by the authors in Gatmiry et al. [] provides a global convergence result for training deep Looped Transformers.

## Appendix B Remaining Proofs

###### Proof of Lemma 7 .

First, note that the global minimizer $A^{*}$ of $\mathcal{L}$ should satisfy $\mathcal{L}^{(P_{\alpha,\beta})}(A^{*},0)\leq\mathcal{L}^{(P_{\alpha,\beta})}(
A,0)\leq 1$. To see this, note that $\mathcal{L}^{(P_{\alpha,\beta})}(A^{*},0)\leq\mathcal{L}^{(P_{\alpha,\beta})}(
A,0)$ for $A=\beta^{-1}I$, as from the formula in Lemma 4:

$$ $\displaystyle\mathcal{L}(A,0)$ $\displaystyle=\mathbb{E}{\Sigma\sim P_{\alpha,\beta}}{\text{trace}{\left((I- \Sigma^{1/2}A\Sigma^{1/2})^{2L}\right)}}$ $\displaystyle=\mathbb{E}{\Sigma\sim P_{\alpha,\beta}}{\text{trace}{\left((I- \beta^{-1}\Sigma)^{2L}\right)}}.$ $$

But since $\Sigma\leq\beta I$, we get $I-\beta^{-1}\Sigma\preccurlyeq I$, hence $(I-\beta^{-1}\Sigma)\preccurlyeq I$, which implies

$$ $\displaystyle\mathcal{L}(A^{*})\leq d.$ (10) $$

Now we show that if given this upper bound on the optimal loss, we should necessarily have

$$ $\displaystyle{A^{*}}^{-1}\succcurlyeq\frac{(1-\delta)\beta}{2}(1-\frac{\ln(1/ \epsilon)}{L})I.$ $$

Suppose this is not the case; Then, there is unit direction $v$ such that

$$ $\displaystyle v^{\top}{A^{*}}^{-1}v<\frac{(1-\delta)\beta}{2}(1-\frac{\ln(1/ \epsilon)}{L}).$ (11) $$

Now from the $(\epsilon,\delta)$-right-spread-out assumption, with probability at least $\epsilon$ we have

$$ $\displaystyle v^{\top}\Sigma v=v^{\top}XX^{\top}v\geq(1-\delta)\beta.$ (12) $$

Now we claim that under this event we should necessarily have

$$ $\displaystyle\left\lVert I-\Sigma^{1/2}A^{*}\Sigma^{1/2}\right\rVert>1+\frac{ \ln(d/\epsilon)}{L}.$ $$

Suppose this is not the case. Then

$$ $\displaystyle\Sigma^{1/2}A^{*}\Sigma^{1/2}\leq(2+\frac{\ln(d/\epsilon)}{L})I.$ $$

But this implies

$$ $\displaystyle(2+\frac{\ln(d/\epsilon)}{L}){A^{*}}^{-1}\geq\Sigma,$ $$

which gives

$$ $\displaystyle(2+\frac{\ln(d/\epsilon)}{L})v^{\top}{A^{*}}^{-1}v\geq v^{\top} \Sigma v.$ $$

But this contradicts Equations equation 11 and equation 12 as

$$ $\displaystyle(2+\frac{\ln(d/\epsilon)}{L})(1-\frac{\ln(d/\epsilon)}{L})(1- \delta)\beta=(1-\frac{\ln(d/\epsilon)}{L}-\left(\frac{\ln(d/\epsilon)}{L} \right)^{2})(1-\delta)\beta<(1-\delta)\beta.$ $$

Hence, we should have

$$ $\displaystyle\left\lVert I-\Sigma^{1/2}A^{*}\Sigma^{1/2}\right\rVert>1+\frac{ \ln(d/\epsilon)}{L}.$ $$

But this means that under this event, which happens with probability at least $\epsilon$ under $P_{\alpha,\beta}$, we should also have

$$ $\displaystyle\left(I-\Sigma^{1/2}A^{*}\Sigma^{1/2}\right)^{2L}\geq\left(1+ \frac{\ln(d/\epsilon)}{L}\right)^{L}>\frac{d}{\epsilon}.$ $$

This implies

$$ $\displaystyle\mathcal{L}(A^{*},0)$ $\displaystyle=\mathbb{E}{\Sigma\sim P_{\alpha,\beta}}{\text{trace}{\left((I- \Sigma^{1/2}A^{*}\Sigma^{1/2})^{2L}\right)}}$ $\displaystyle\geq\Pr(v^{\top}\Sigma v\geq(1-\delta)\beta)\times\frac{d}{ \epsilon}>d,$ $$

which contradicts Equation equation 10. Therefore, we should have

$$ $\displaystyle{A^{*}}^{-1}\succcurlyeq\frac{(1-\delta)\beta}{2}(1-\frac{\ln(d/ \epsilon)}{L})I.$ (13) $$

Now note that for every $\Sigma$ sampled from $P_{\alpha,(1\delta^{\prime})\beta}$, according to definition we have

$$ $\displaystyle\Sigma\preccurlyeq(1-\delta^{\prime})\beta I,$ $$

which combined with Equation equation 13 implies

$$ $\displaystyle\Sigma^{1/2}A^{*}\Sigma^{1/2}$ $\displaystyle\leq\frac{(1-\delta^{\prime})\beta}{\frac{(1-\delta)\beta}{2}(1- \frac{\ln(d/\epsilon^{\prime})}{L})}=$ $\displaystyle\leq 2(1-\frac{\ln(d/\epsilon)}{L}).$ $$

The last inequality is because

$$ $\displaystyle 1-\delta^{\prime}=1-\delta-2\frac{\ln(d/\epsilon^{\prime})}{L} \leq 1-\delta-2(1-\delta)\frac{\ln(d/\epsilon^{\prime})}{L}+(1-\delta)\left( \frac{\ln(d/\epsilon^{\prime})}{L}\right)^{2}=\left(1-\delta\right)\left(1- \frac{\ln(d/\epsilon^{\prime})}{L}\right)^{2}.$ $$

On the other hand, note that for any $\Sigma$ in the support of $P_{\alpha,(1-\delta^{\prime})\beta}$, we have

$$ $\displaystyle\alpha I\preccurlyeq\Sigma.$ $$

Therefore,

$$ $\displaystyle\left(\frac{\alpha}{\beta}-1\right)I\preccurlyeq\Sigma^{1/2}A^{*} \Sigma^{1/2}-I\preccurlyeq\left(1-\frac{2\ln(d/\epsilon^{\prime})}{L}\right)I,$ $$

which implies

$$ $\displaystyle\left\lVert\left(I-\Sigma^{1/2}A^{*}\Sigma^{1/2}\right)^{2L} \right\rVert\leq\left(1-\left(\frac{2\ln(d/\epsilon^{\prime})}{L}\wedge\frac{ \alpha}{\beta}\right)\right)^{2L}\leq\frac{\epsilon^{\prime}}{d}\vee\left(1- \frac{\alpha}{\beta}\right)^{L}.$ $$

Therefore

$$ $\displaystyle\mathcal{L}^{(\tilde{P}_{\alpha,(1-\delta^{\prime})\beta})}(A^{*} ,0)$ $\displaystyle=\mathbb{E}{\Sigma\sim P_{\alpha,\beta}}{\text{trace}{\left((I- \Sigma^{1/2}A^{*}\Sigma^{1/2})^{2L}\right)}}\leq\epsilon^{\prime}\vee d\left(1 -\frac{\alpha}{\beta}\right)^{L}.$ $$

∎

###### Proof of Theorem 5 .

As described in Appendix A.1 in Von Oswald et al. [], we can construct one iteration of gradient descent for the linear regression quadratic loss in the restricted attention case, which takes the following form:

$$ $\displaystyle{y^{(t+1)}}^{\top}={y^{(t)}}^{\top}-\eta{y^{(t)}}^{\top}X^{\top}X,$ $\displaystyle y^{(t+1)}_{q}=y^{(t)}_{q}-\eta{y^{(t)}}^{\top}X^{\top}x_{q}.$ $$

Then, taking the step size $\eta=\frac{1}{2\beta}$, for the linear regression error we get

$$ $\displaystyle|\text{TF}({Z^{(0)}},\{A,0\}_{t=0}^{L-1})-y|={w^{*}}^{\top}\Big{( }I-\eta(XX^{\top})\Big{)}^{L}x_{q}\leq\|w^{*}\|\|x_{q}\|(1-\frac{\alpha}{2 \beta})^{L},$ $$

which completes the proof.
∎

###### Proof of Lemma 7 .

###### Proof of Theorem 5 .

## Appendix C Lower bound

### C.1 A formula for the loss in the multilayer case

###### Lemma 2 .

[Output with non-linearity]
We have the following relations

$$ $\displaystyle{y^{(t+1)}}^{\top}={y}^{\top}\prod_{i=0}^{t}\Big{(}I-\frac{1}{n} \sigma\Big{(}X^{\top}A^{(i)}X\Big{)}\Big{)},$ $\displaystyle y^{(t+1)}_{q}=\sum_{j=1}^{t-1}y^{\top}\Big{(}I-\prod_{i=0}^{j-1} \Big{(}I-\frac{1}{n}\sigma\Big{(}X^{\top}A^{(i)}X\Big{)}\Big{)}\Big{)}\sigma \Big{(}X^{\top}A^{(j)}x_{q}\Big{)}+\prod_{i=0}^{t-1}\Big{(}I-\frac{1}{n}\sigma \Big{(}X^{\top}A^{(i)}X\Big{)}\Big{)}\sigma\Big{(}X^{\top}A^{(t)}x_{q}\Big{)}.$ $$

###### Proof.

Assuming the formula for $y^{(t)}$ and from the update equation for $y^{(t+1)}$,

$$ $\displaystyle{y^{(t+1)}}^{\top}$ $\displaystyle={y^{(t)}}^{\top}-\frac{1}{n}{y^{(t)}}^{\top}\sigma\Big{(}X^{\top }A^{(t)}X\Big{)}$ $\displaystyle={y^{(t)}}^{\top}\Big{(}I-\frac{1}{n}\sigma\Big{(}X^{\top}A^{(t)} X\Big{)}\Big{)}$ $\displaystyle={y}^{\top}\prod_{i=0}^{t}\Big{(}I-\frac{1}{n}\sigma\Big{(}X^{ \top}A^{(i)}X\Big{)}\Big{)}.$ $$

Therefore, for $y^{(t)}_{q}$ we have

$$ $\displaystyle y^{(t+1)}_{q}$ $\displaystyle=y^{(t)}_{q}+\frac{1}{n}{y^{(t)}}^{\top}\sigma\Big{(}X^{\top}A^{( t)}x_{q}\Big{)}$ $\displaystyle=\sum_{j=1}^{t-1}y^{\top}\Big{(}I-\prod_{i=0}^{j-1}\Big{(}I-\frac {1}{n}\sigma\Big{(}X^{\top}A^{(i)}X\Big{)}\Big{)}\Big{)}\sigma\Big{(}X^{\top}A ^{(j)}x_{q}\Big{)}+\prod_{i=0}^{t-1}\Big{(}I-\frac{1}{n}\sigma\Big{(}X^{\top}A ^{(i)}X\Big{)}\Big{)}\sigma\Big{(}X^{\top}A^{(t)}x_{q}\Big{)}.$ $$

∎

where recall $\Sigma=\frac{1}{n}\sum_{i=1}^{n}x_{i}{x_{i}}^{\top}$. Now we substitute $\Sigma=$

###### Lemma 3 (Polynomial blow up) .

Suppose $P(\gamma)$ is a polynomial of degree $k$ with $P(0)=0$. Then, given $k\leq\frac{\sqrt{\lambda_{max}}}{6\sqrt{\lambda_{min}}}$, for an arbitrary choice of $\alpha\in\mathbb{R}$, there exists a choice of $\tilde{\gamma}\in(\lambda_{min},\lambda_{max})$ such that

$$ $\displaystyle\frac{|P(\tilde{\gamma})-\alpha|}{|\alpha|}\geq\frac{1}{4}.$ $$

Furthermore, there exists an interval $(a,b)$ of length at least $\frac{\lambda_{\max}-\lambda_{\min}}{64k^{2}}$ such that for all $\gamma\in(a,b)$:

$$ $\displaystyle\frac{|P(\gamma)-\alpha|}{|\alpha|}\geq\frac{1}{8}.$ $$

###### Proof.

Let $\kappa=\frac{\lambda_{min}}{\lambda_{max}}$. Consider the following scaling of the $k$th Chebyshev polynomial $T_{k}(\gamma)$:

$$ $\displaystyle Q_{k}(\gamma)=T_{k}(\frac{2\gamma-(\lambda_{min}+\lambda_{max})} {\lambda_{max}-\lambda_{min}}).$ $$

Now from the fact that $T_{k}(x)=(x+\sqrt{x^{2}-1})^{k}+(x-\sqrt{x^{2}-1})^{k}$ outside the interval $[-1,1]$, we get for $x=-1-\alpha$:

$$ $\displaystyle(1+\sqrt{\alpha})^{k}\leq|T_{k}(x)|\leq 1+(1+3\sqrt{\alpha})^{k}.$ $$

Therefore, for $\kappa=\lambda_{min}/\lambda_{max}$ we have

$$ $\displaystyle Q_{k}(0)=T_{k}(1-\frac{2\lambda_{min}}{\lambda_{max}-\lambda_{ min}})\leq 1+(1+3\sqrt{4\kappa})^{k},$ $$

and using the fact that $k\leq 1/(6\sqrt{\kappa})$, we have $Q_{k}(0)\leq 4$. On the other hand, from the property of Chebyshev polynomials, we get that $Q_{k}(0)$ alternates between $-1$ and $1$ $k$ times. This means that if we define $R(\gamma)=Q_{k}(\gamma)/Q_{k}(0)$, we have that $R(0)=1$, and that $R$ alternates above and below the $y=\frac{1}{4}$ and $y=-\frac{1}{4}$ lines $k$ times. Now suppose the absolute value of $P(\gamma)$ is always within the interval $(\alpha-\alpha/4,\alpha+\alpha/4)$. Then for the polynomial $\tilde{P}(\gamma)=1-P(\gamma)/\alpha$ we have $\tilde{P}(0)=1$ and $\tilde{P}(\gamma)<\frac{1}{4}$ for all points $\gamma$ in the interval $[\lambda_{min},\lambda_{max}]$. Therefore, if we consider the polynomial $S(\gamma)=\tilde{P}(\gamma)-R(\gamma)$, then $S$ has degree at most $k$, $S(0)=0$, and $S(\gamma)$ alternates $t$ times between positive and negative numbers in the interval $[\lambda_{min},\lambda_{max}]$. But this means that $S$ has at least $t+1$ real roots. The contradiction finishes the proof for the first part.
Proof for the second part.
To show the second part, note that the $k$th Chebyshev polynomial $T_{k}(x)$ is lipschitz around its roots; In particular, consider the trigonomic property $T_{k}(\cos(\theta))=\cos(k\theta)$, now defining $x(\theta)=\cos(\theta)$, we see that

$$ $\displaystyle\frac{d}{dx}T_{k}(x(\theta))$ $\displaystyle=\frac{d}{d\theta}T_{k}(x(\theta))\frac{d\theta}{dx}$ $\displaystyle=\frac{d}{d\theta}\cos(k\theta)\Big{(}\frac{dx}{d\theta}\Big{)}^{ -1}$ $\displaystyle=k\sin(k\theta)/\sin(\theta).$ $$

Namely, the extremal points of $T_{K}(\cos(\theta))$ are at $\theta=\frac{i\pi}{k}$ for $i=0,\dots,k-1$. Now for $|\Delta\theta|\leq\frac{\pi}{4k}$,

$$ $\displaystyle T_{k}(\cos(i\frac{\pi}{k}+\Delta\theta))=\cos(\pi k+k\Delta)\geq \cos(\frac{\pi}{4})=\frac{1}{\sqrt{2}}\geq\frac{1}{2}.$ $$

Therefore, we conclude that the $i$th root $\cos(\frac{\pi i}{k})$ of $T_{k}$, for all $x$ in the interval $(\cos(\frac{\pi i}{k}-\frac{\pi}{4k}),\cos(\frac{\pi i}{k}+\frac{\pi}{4k}))$ we have

$$ $\displaystyle T_{k}(x)\geq\frac{1}{2}.$ $$

Moreover, it is easy to see that the length of this interval is miaximized for $i=0$ with length $2(1-\cos(\frac{\pi}{k}))$. Therefore,
for all $0\leq i<k$, if we define the intervals

$$ $\displaystyle(a_{i},b_{i})$ $\displaystyle=(\cos(\frac{\pi i}{k}-\frac{\pi}{4k})(\frac{\lambda_{\max}- \lambda_{\min}}{2})+\frac{\lambda_{\min}+\lambda_{\max}}{2},\cos(\frac{\pi i}{ k}+\frac{\pi}{4k})(\frac{\lambda_{\max}-\lambda_{\min}}{2})+\frac{\lambda_{ \min}+\lambda_{\max}}{2}),$ $$

then the scaled chebyshev polynomials $Q_{k}$ that we defined for the previous part have the property that for all $0\leq i<k$ and $\gamma\in(a_{i},b_{i})$, $Q_{k}(\gamma)\geq\frac{1}{2}$, and therefore $R(\gamma)\geq\frac{1}{8}$. Now suppose for every $0\leq i<k$, there exists $\gamma_{i}\in(a_{i},b_{i})$ such that $P(\gamma_{i})<\frac{1}{8}$. Then, for the polynomial $S(\gamma)=\tilde{P}(\gamma)-R(\gamma)$ we see that again it alternates $t$ times between positive and negative at $\gamma_{i}$’s, hence has at least $k+1$ roots, which is again a contradiction. Therefore, we conlcude that there exists $0\leq i<k$ such taht for all $\gamma\in(a_{i},b_{i})$:

$$ $\displaystyle\frac{|P(\tilde{\gamma})-\alpha|}{|\alpha|}\geq\frac{1}{8}.$ $$

Finally, note that the length of the intervals $(a_{i},b_{i})$ is minimized for $i=0$, in which case the legnth is

$$ $\displaystyle(1-\cos(\frac{\pi}{4k}))(\lambda_{\max}-\lambda_{\min})\geq\big{| }1-\sqrt{1-\frac{\pi^{2}}{16k^{2}}}(\lambda_{\max}-\lambda_{\min})\big{|}\geq \frac{\pi^{2}}{64k^{2}}(\lambda_{\max}-\lambda_{\min}).$ $$

∎

###### Lemma 2 .

###### Proof.

###### Lemma 3 (Polynomial blow up) .

###### Proof.

### C.2 Proof of Theorem 1

###### Proof of Theorem 1 .

First we consider the case when the weight matrices are restricted to $A^{(t)},u^{(t)}$ as in equation 6. We substitute $X=\gamma U$ and $x_{q}=\theta x$ with parameters $\gamma,\theta$ and arbitrary fixed matrix $U$ and vector $x$. The key idea is to look at $y^{(t)}_{q}$ as a polynomial of $\gamma$ and $\theta$, $P_{t}(\gamma)$, where the degree of $P_{t}(\gamma)$ is at most $2t$ and $P_{t}(0)=0$.

We show this by induction. In particular, we strengthen the argument by also proving that each entry of $y^{(t)}$ is a polynomial of degree at most $2t+1$ of $\gamma$. The base of induction is clear since $y_{q}^{(0)}=0,y^{(0)}=y$ are both degree zero polynomials of $\gamma$. For the step of induction, suppose the argument holds for $t$ and we want to prove it for $t+1$. We write the update rule for $y^{(t)}_{q}$:

$$ $\displaystyle y^{(t+1)}_{q}$ $\displaystyle=y^{(t)}_{q}-\tfrac{1}{n}({y^{(t)}}^{\top}+{u^{(t)}}^{\top}X) \sigma\Big{(}X^{\top}A^{(t)}x_{q}\Big{)}.$ (14) $$

Now from the hypothesis of induction, we know that $y^{(t)}_{q}$ is represented by a polynomial of degrees at most $2t$ of $\gamma$.
On the other hand, note that from the positive scaling property of ReLU

$$ $\displaystyle\sigma\Big{(}X^{\top}A^{(t)}x_{q}\Big{)}=\sigma\Big{(}\gamma U^{ \top}A^{(t)}x_{q}\Big{)}=\gamma\sigma\Big{(}U^{\top}A^{(t)}x_{q}\Big{)}.$ $$

Therefore,
the second part of the second term in equation 14, i.e. ${u^{(t)}}^{\top}X\sigma\Big{(}X^{\top}A^{(t)}x_{q}\Big{)}$ has degree exactly $2$ in $\gamma$. On the other hand, from the hypothesis of induction we know that each entry of $y^{(t)}$ is a polynomial of degree at most $2t+1$. Therefore,
the first part of the second term, $\frac{1}{n}{y^{(t)}}^{\top}\sigma\Big{(}X^{\top}A^{(t)}x_{q}\Big{)}$, has degree at most $2t+2$. Hence, overall $y^{(t+1)}_{q}$ is a polynomial of degree at most $2t+2$ of $\gamma$.
Next, we show the step of induction for $y^{(t+1)}$. We have the update rule

$$ $\displaystyle{y^{(t+1)}}^{\top}={y^{(t)}}^{\top}-\tfrac{1}{n}({y^{(t)}}^{\top} +{u^{(t)}}^{\top}X)\sigma\Big{(}X^{\top}A^{(t)}X\Big{)}.$ $$

Now again from the scale property of the activation

$$ $\displaystyle\sigma\Big{(}X^{\top}A^{(t)}X\Big{)}$ $\displaystyle=\sigma\Big{(}(\gamma U)^{\top}A^{(t)}(\gamma U)\Big{)}=\gamma^{2 }\sigma\Big{(}U^{\top}A^{(t)}U\Big{)}.$ $$

Hence, from the hypothesis of induction, $y^{(t)}$ is a polynomial of degree at most $2t+1$, ${y^{(t)}}^{\top}\sigma\Big{(}X^{\top}A^{(t)}X\Big{)}$ is of degree at most $2t+3$, and ${u^{(t)}}^{\top}X\sigma\Big{(}X^{\top}A^{(t)}X\Big{)}$ is of degree at most three. Overall we conclude that $y^{(t+1)}$ is of degree at most $2t+3$ which finishes the step of induction.
The result then follows from Lemma 3.

Next, we show the second argument without the weight restriction equation 6 with a similar strategy, namely we substitute $X=\gamma U$ and bound the degree of polynomial that entries of $Z^{(t)}$ are of variable $\gamma$. In particular we claim that each entry of $Z^{(t)}$ is a polynomial of degree at most $3^{t}$ of $\gamma$. The base of induction is trivial as $Z^{(1)}$ is a polynomial of degree at most $1$ of $\gamma$. Now for step of induction, given that $Z^{(t)}$ is a polynomial of degree at most $3^{t}$, in light of the recursive relation equation 3, $Z^{(t)}$ and the one-homogeneity of ReLU, we get that $Z^{(t+1)}$ is of degree at most $3^{t+1}$, proving the step of induction. The result follows again from Lemma 3.
∎

###### Proof of Theorem 1 .

### C.3 Proof of Theorem 2

###### Proof of Theorem 2 .

Using Lemma 4, setting $u^{(i)}=0,\forall i\leq L$ we have

$$ $\displaystyle|\text{TF}({Z^{(0)}},\{A^{(t)},u^{(t)}\}_{t=0}^{L-1})-{w^{*}}^{ \top}x_{q}|$ $\displaystyle=|y_{q}-{w^{*}}^{\top}x_{q}|$ $\displaystyle=|{w^{*}}^{\top}\prod_{i=0}^{L-1}(I-\Sigma A^{(i)})x_{q}|.$ $$

Now pick $A^{(i)}=\theta_{i}^{-1}I$, where $\theta_{i}=\left(\frac{\beta-\alpha}{2}\lambda_{i}+\frac{\beta+\alpha}{2}
\right)^{-1}$ is the $i$th root of the Chebyshev polynomial scaled into the interval $(\alpha,\beta)$. This way

$$ $\displaystyle\prod_{i=0}^{L-1}(I-\Sigma A^{(i)})$ $\displaystyle=\prod_{i=0}^{t-1}(I-\Sigma/\lambda_{i})$ $\displaystyle=C\prod_{i=0}^{t-1}(\lambda_{i}I-\Sigma)$ $\displaystyle=CP^{L,(\alpha,\beta)}(\Sigma),$ $$

where $P^{L,(\alpha,\beta)}(x)=T_{L}(\frac{2x-(\alpha+\beta)}{\beta-\alpha})$ is the degree $L$ Chebyshev polynomial whose roots are scaled into interval $[\alpha,\beta]$ and $C$ is a constant. Note that the constant $C$ is picked in a way that $P^{L,(\alpha,\beta)}(0)=1$. From , this condition implies that $C\leq\frac{1}{2^{L\sqrt{\alpha/\beta}-1}}$, hence picking degree $L=O(\ln(1/\epsilon)\sqrt{\beta/\alpha})$ implies that the value of $P^{L,(\alpha,\beta)}$ on the interval $[\alpha,\beta]$ is at most $\epsilon$. Now looking at the eigenvalues, this implies

$$ $\displaystyle-\epsilon I\leq\prod_{i=0}^{t-1}(I-\Sigma/\lambda_{i})\leq \epsilon I.$ $$

Finally from Cauchy-Schwarz

$$ $\displaystyle|{w^{*}}^{\top}\prod_{i=0}^{L-1}(I-\Sigma A^{(i)})x_{q}|$ $\displaystyle\leq\|w^{*}\|\|x_{q}\|\|\prod_{i=0}^{L-1}(I-\Sigma A^{(i)}\|\leq \epsilon\|w^{*}\|\|x_{q}\|,$ $$

which completes the proof.
∎

###### Proof of Theorem 2 .

## Appendix D Robustness of Looped Transformers

###### Lemma 4 (Transformer and loss formulas) .

Transformer output can be calculated as

$$ $\displaystyle\text{TF}({Z^{(0)}},\{A^{(t)},u^{(t)}\}_{t=0}^{L-1})={y^{(L)}}^{ \top}={y_{q}}-{w^{*}}^{\top}\prod_{i=0}^{t-1}(I-\Sigma A^{(i)})x_{q}.$ $$

Furthermore, the loss can be written as

$$ $\displaystyle\mathcal{L}(\{A^{(t)},0\}_{t=0}^{L-1})=\mathbb{E}_{X}{\text{trace }{\left(\prod_{t=0}^{L-1}(I-\Sigma^{1/2}A^{(t)}\Sigma^{1/2})\prod_{t=L-1}^{0}( I-\Sigma^{1/2}A^{(t)}\Sigma^{1/2})\right)}}.$ $$

###### Proof.

First, we show the following recursive relations:

$$ $\displaystyle{y^{(t)}}^{\top}={w^{*}}^{\top}\prod_{i=0}^{t-1}(I-\Sigma A^{(i)} )X~{}$ (15) $\displaystyle{y^{(t)}_{q}}^{\top}={y_{q}}-{w^{*}}^{\top}\prod_{i=0}^{t-1}(I- \Sigma A^{(i)})x_{q},$ (16) $$

We show this by induction on $t$. For the base case, we have ${w^{*}}^{\top}X=y^{(0)}$ and $y^{(0)}_{q}=0=y_{q}-{w^{*}}^{\top}x_{q}$. For the step of induction, we have Equations equation 15 and equation 16 for $t-1$. We then open the following update rule:

$$ $\displaystyle Z^{(t)}=Z^{(t-1)}-\frac{1}{n}P^{(t)}MZ^{(t)}A^{(t)}{Z^{(t)}}^{ \top},$ $$

as

$$ $\displaystyle{y^{(t+1)}}^{\top}$ $\displaystyle={y^{(t)}}^{\top}-\frac{1}{n}{y^{(t)}}^{\top}X^{\top}A^{(t)}X$ $\displaystyle={w^{*}}^{\top}\prod_{i=0}^{t-1}(I-\Sigma A^{(i)})X-{w^{*}}^{\top }\prod_{i=0}^{t-1}(I-\Sigma A^{(i)})(XX^{\top})A^{(i)}X$ $\displaystyle={w^{*}}^{\top}\prod_{i=0}^{t}(I-\Sigma A^{(t)})X,$ $$

where we used the fact that $XX^{\top}=\Sigma$. Moreover

$$ $\displaystyle y^{(t+1)}_{q}$ $\displaystyle=y^{(t)}_{q}-\frac{1}{n}{y^{(t)}}^{\top}X^{\top}Qx_{q}$ $\displaystyle=y_{q}-{w^{*}}^{\top}\prod_{i=0}^{t-1}(I-\Sigma A^{(i)})x_{q}$ $\displaystyle-{w^{*}}^{\top}\prod_{i=0}^{t-1}(I-\Sigma A^{(i)})(XX^{\top})A^{( i)}x_{q}$ $\displaystyle=y_{q}-{w^{*}}^{\top}\prod_{i=0}^{t}(I-\Sigma A^{(i)})x_{q}.$ $$

Therefore

$$ $\displaystyle\mathbb{E}{X}{(y_{q}^{(L)}-y_{q})^{2}}$ $\displaystyle=\mathbb{E}{X}{\Big{(}{w^{*}}^{\top}\prod_{i=0}^{L-1}(I-\Sigma A^ {(i)})x_{q}\Big{)}^{2}}$ $\displaystyle=\mathbb{E}{X}{\Big{(}{w^{*}}^{\top}\Sigma^{1/2}\prod_{i=0}^{L-1} (I-\Sigma^{1/2}A^{(i)}\Sigma^{1/2}){\Sigma^{-1/2}}x_{q}\Big{)}^{2}}.$ $$

Now taking expectation with respect to $w^{*}\sim N(0,\Sigma^{-1}),x_{q}\sim N(0,{\Sigma})$:

$$ $\displaystyle\mathcal{L}(\{A^{(t)},0\}_{t=0}^{L-1})=\mathbb{E}{X}{\text{trace} {\left(\prod_{i=0}^{L-1}(I-\Sigma^{1/2}A^{(i)}\Sigma^{1/2})\prod_{i=L-1}^{0}(I -\Sigma^{1/2}A^{(i)}\Sigma^{1/2})\right)}}.$ $$

∎

###### Lemma 4 (Transformer and loss formulas) .

###### Proof.

### D.1 Proof of Theorem 6

###### Proof of Theorem 6 .

Consider a set of diagonal matrices $\{{B^{(t)}}\}_{t=0}^{L-1}$ with a set of $Ld$ distinct diagonal entries that satisfy $\alpha I\preccurlyeq{{B^{(t)}}}^{-1}\preccurlyeq 2\alpha I$ for all $0\leq t\leq L-2$ and ${B^{(L-1)}}^{-1}=\beta I$, and define the distribution

$$ $\displaystyle P_{\alpha,\beta}\triangleq\frac{1}{L}\sum_{t=0}^{L-1}\delta_{{B^ {(t)}}},$ $$

where $\delta_{{B^{(t)}}}$ is a point mass on ${B^{(t)}}$.
Moreover, for parameter $\delta^{\prime}$, consider an alternative distribution $\tilde{P}_{\alpha,\beta}$ on $\Sigma$ where it has at least $\epsilon$ amount of mass on matrices $4\alpha I\preccurlyeq\Sigma\preccurlyeq(1-\delta^{\prime})\beta I$.
First, note that the global minimum of $\mathcal{L}^{P_{\alpha,\beta}}$ is zero and $\mathcal{L}^{P_{\alpha,\beta}}(\{{A^{(t)}}^{*}\}_{t=0}^{L-1},0)$ for ${A^{(t)}}^{*}\triangleq{{B^{(t)}}}^{-1}\forall 0\leq t\leq L-1$.

Now for an arbitrary matrix $\Sigma$ such that $8\alpha I\preccurlyeq\Sigma\preccurlyeq(1-\delta^{\prime})\beta I$, we have for all $0\leq t\leq L-2$:

$$ $\displaystyle\Sigma^{1/2}{A^{(t)}}^{*}\Sigma^{1/2}$ $\displaystyle=\Sigma^{1/2}{{B^{(t)}}}^{-1}\Sigma^{1/2}$ $\displaystyle\succcurlyeq\Sigma^{1/2}{2\alpha I}^{-1}\Sigma^{1/2}$ $\displaystyle\succcurlyeq 4I,$ $$

which implies

$$ $\displaystyle I-\Sigma^{1/2}{A^{(t)}}^{*}\Sigma^{1/2}\preccurlyeq 3I.$ $$

Therefore, for all $0\leq t\leq L-2$,

$$ $\displaystyle\left(I-\Sigma^{1/2}{A^{(t)}}^{*}\Sigma^{1/2}\right)^{2}\geq 9I.$ $$

On the other hand, for $t=L-1$:

$$ $\displaystyle\Sigma^{1/2}{A^{(L-1)}}^{*}\Sigma^{1/2}=\frac{1}{\beta}\Sigma\leq (1-\delta^{\prime})I,$ $$

which implies

$$ $\displaystyle I-\Sigma^{1/2}{A^{(L-1)}}^{*}\Sigma^{1/2}\succcurlyeq\delta^{ \prime}I.$ $$

Hence,

$$ $\displaystyle\left(I-\Sigma^{1/2}{A^{(L-1)}}^{*}\Sigma^{1/2}\right)^{2}\geq \delta^{\prime 2}I.$ $$

Therefore, by repeating the inequality $CBC^{\top}\geq CDC^{\top}$ for arbitrary matrices $B,C,D$ where $B\succcurlyeq D$, we get

$$ $\displaystyle\prod_{t=0}^{L-1}(I-\Sigma^{1/2}{A^{(t)}}^{*}\Sigma^{1/2})\prod_{ t=L-1}^{0}(I-\Sigma^{1/2}{A^{(t)}}^{*}\Sigma^{1/2})\succcurlyeq\delta^{\prime} 9^{L-1}I.$ $$

Therefore

$$ $\displaystyle\text{trace}{\left(\prod_{t=0}^{L-1}(I-\Sigma^{1/2}{A^{(t)}}^{*} \Sigma^{1/2})\prod_{t=L-1}^{0}(I-\Sigma^{1/2}{A^{(t)}}^{*}\Sigma^{1/2})\right) }\geq d\delta^{\prime}9^{L-1}.$ $$

On the other hand, note that from our assumption the distribution $\tilde{P}_{\alpha,\beta}$ has at least $\epsilon$ mass on matrices $\Sigma$ with $4\alpha I\preccurlyeq\Sigma\preccurlyeq(1-\delta^{\prime})\beta I$. Therefore

$$ $\displaystyle\mathcal{L}^{\tilde{P}_{\alpha,\beta}}(\{{A^{(t)}}^{*}\}_{t=0}^{L -1},0)\geq\epsilon\delta^{\prime}d9^{L-1}.$ $$

∎

###### Proof of Theorem 6 .

## Appendix E Proof of Theorem 3

###### Proof.

From Lemma 4 (for $u^{(t)}=0),\ \forall t\leq L$, we have

$$ $\displaystyle{y^{(t)}}^{\top}={w^{*}}^{\top}\prod_{i=0}^{t-1}(I-\Sigma A^{(i)} )X,$ $\displaystyle y^{(t)}_{q}=y_{q}-{w^{*}}^{\top}\prod_{i=0}^{t-1}(I-\Sigma A^{(i )})x_{q}.$ $$

Therefore

$$ $\displaystyle\|y^{(L)}\|^{2}$ $\displaystyle=\|X^{\top}\prod_{i=0}^{L-1}(I-\Sigma A^{(i)})w^{*}\|^{2}$ $\displaystyle=\|\prod_{i=0}^{L-1}(I-\Sigma A^{(i)})w^{*}\|^{2}_{XX^{\top}}.$ $$

Therefore,

$$ $\displaystyle|y_{q}-y^{(L)}_{q}|$ $\displaystyle=|{w^{*}}^{\top}\prod_{i=0}^{L-1}(I-\Sigma A^{(i)})x_{q}|$ $\displaystyle\leq\|x_{q}\|_{(XX^{\top})^{-1}}\|\prod_{i=0}^{L-1}(I-\Sigma A^{( i)})w^{*}\|_{XX^{\top}}$ $\displaystyle\leq\epsilon.$ $$

∎

###### Proof.

## Appendix F Proof of Theorem 8

###### Proof of Theorem 8 .

Suppose $\Sigma^{1/2}=vv^{\top}$ is rank one and define

$$ $\displaystyle v^{\top}A^{(i)}v\left\lVert v\right\rVert^{2}=\gamma_{i}.$ $$

Then

$$ $\displaystyle\Sigma^{1/2}A^{(i)}\Sigma^{1/2}=vv^{\top}A^{(i)}vv^{\top}=\gamma_ {i}\tilde{v}{\tilde{v}}^{\top},$ $$

for $\gamma_{i}>2$, where $\tilde{v}$ is the normalized version of $v$. This implies

$$ $\displaystyle\left(I-\Sigma^{1/2}A^{(i)}\Sigma^{1/2}\right)=(1-\gamma_{i}) \tilde{v}{\tilde{v}}^{\top}+\left(I-\tilde{v}{\tilde{v}}^{\top}\right),$ $$

which then gives

$$ $\displaystyle\mathcal{L}^{\Sigma}(\{A^{(i)}\}_{i=0}^{L},0)$ $\displaystyle=\text{trace}{\left(\prod_{i=0}^{L-1}(I-\Sigma^{1/2}A^{(i)}\Sigma ^{1/2})\prod_{i=L-1}^{0}(I-\Sigma^{1/2}A^{(i)}\Sigma^{1/2})\right)}$ $\displaystyle=\prod_{i=0}^{L-1}(\gamma_{i}-1)^{2}+(d-1).$ $$

Therefore, the loss is increasing with depth if $\gamma_{t}>2$ and decreasing if $\gamma_{t}<2$. Therefore, for the behavior of the loss wrt to depth to be either increasing or decreasing for all $t$, for every $v$, we need to have that either all $\gamma_{i}$’s are great than or equal to $2$, or less than or equal to $2$. But this implies that for all $v$ and all $i,j$:

$$ $\displaystyle v^{\top}A^{(i)}v=v^{\top}A^{(j)}v.$ $$

Suppose this is not the case for some $v,i,j$, namely

$$ $\displaystyle v^{\top}A^{(i)}v<v^{\top}A^{(j)}v.$ $$

Then, we can scale $v$ so that $v^{\top}A^{(i)}v<2$ and $v^{\top}A^{(i)}v>2$ which contradicts the monotonicity of the loss wrt to depth. Therefore, for all $v$, and all $i,j$, we showed

$$ $\displaystyle v^{\top}A^{(i)}v=v^{\top}A^{(j)}v,$ $$

which implies for all $i,j\in[L]$ we should have

$$ $\displaystyle A^{(i)}=A^{(j)}.$ $$

Next, we show the argument for loop models. Namely, suppose $A^{(1)}=\dots=A^{(L)}=A$. The loss in this case becomes

$$ $\displaystyle\mathcal{L}^{\Sigma}(\{A\}_{i=0}^{L},0)$ $\displaystyle=\text{trace}{\left((I-\Sigma^{1/2}A\Sigma^{1/2})^{2L}\right)}.$ (17) $$

Now let $\beta_{1}<\dots<\beta_{d-1}$ are the eigenvalues of $\Sigma^{1/2}A\Sigma^{1/2}$. Then

$$ $\displaystyle\mathcal{L}^{\Sigma}(\{A\}_{i=0}^{L},0)=\sum_{i=0}^{d_{1}}(1- \beta_{i})^{2L}.$ (18) $$

Now if $|1-\beta_{i}|\leq 1$ for all $0\leq i\leq d-1$, then each term $(1-\beta_{i})^{2L}$ is decreasing in $L$. Otherwise, suppose there is $j$ such that $|1-\beta_{j}|>1$. Then, defining $L_{0}=\ln(d/(|1-\beta_{j}|-1))/\ln(|1-\beta_{j}|)$, then for $L\geq L_{0}$, we have

$$ $\displaystyle(1-\beta_{j})^{2L}\geq\frac{d}{|1-\beta_{j}|-1}.$ $$

Hence,

$$ $\displaystyle(1-\beta_{j})^{2(L+1)}$ $\displaystyle\geq(1-\beta_{j})^{2L}+(1-\beta_{j})^{2L}\times(|1-\beta_{j}|-1)$ (19) $\displaystyle\geq(1-\beta_{j})^{2L}+d.$ (20) $$

This means going from $L$ to $L+1$, the loss increase by at least $d$, while the potential decrease in the loss from the other $d-1$ terms in equation 18 is at most $d-1$. Hence, we showed that for all $L\geq L_{0}$, the loss is increasing in this case, which compeletes the proof.

∎

###### Proof of Theorem 8 .

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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?
soundness: 3
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)
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Please see the weakness section.
soundness: 3
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).
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.
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)$?

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: See above please.
soundness: 2
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.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: n/a
soundness: 4
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
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?
soundness: 3
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.
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.
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.

Meta-review: 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.

Decision: Reject