On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery

## Abstract

Abstract An intriguing property of the Transformer is its ability to perform in-context learning (ICL), where the Transformer can solve different inference tasks without parameter updating based on the contextual information provided by the corresponding input-output demonstration pairs. It has been theoretically proved that ICL is enabled by the capability of Transformers to perform gradient-descent algorithms ( Von Oswald et al., 2023a, ; Bai et al.,, 2024 ) . This work takes a step further and shows that Transformers can perform learning-to-optimize (L2O) algorithms. Specifically, for the ICL sparse recovery (formulated as LASSO) tasks, we show that a K 𝐾 K italic_K -layer Transformer can perform an L2O algorithm with a provable convergence rate linear in K 𝐾 K italic_K . This provides a new perspective explaining the superior ICL capability of Transformers, even with only a few layers, which cannot be achieved by the standard gradient-descent algorithms.
Moreover, unlike the conventional L2O algorithms that require the measurement matrix involved in training to match that in testing, the trained Transformer is able to solve sparse recovery problems generated with different measurement matrices. Besides, Transformers as an L2O algorithm can leverage structural information embedded in the training tasks to accelerate its convergence during ICL, and generalize across different lengths of demonstration pairs, where conventional L2O algorithms typically struggle or fail. Such theoretical findings are supported by our experimental results.

## 1 Introduction

Since its introduction in Vaswani et al., (), Transformers have become the backbone in various fields such as natural language processing (Radford,, ; Devlin,, ), computer vision (Dosovitskiy,, ) and reinforcement learning (Chen et al.,, ), significantly influencing subsequent research and applications.
A notable capability of Transformers is their good performance for in-context learning (ICL) (Brown et al.,, ), i.e., without further parameter updating, Transformers can perform new inference tasks based on the contextual information embedded in example input-output pairs contained in the prompt. Such ICL capability facilitates state-of-the-art few-shot performances across a multitude of tasks, such as reasoning and language understanding tasks in natural language processing (Chowdhery et al.,, ), in-context dialog generation (Thoppilan et al.,, ) and in-context linear regression (Garg et al.,, ; Fu et al.,, ).

Given the significance of transformers’ ICL capabilities, extensive research efforts have been directed toward understanding the mechanisms behind their ICL performance. In this context, the ICL capability of a pre-trained Transformer is understood as the Transformer’s implicit implementation of learning algorithms during the forward pass. , Dai et al., (), and Bai et al., () suggest that these learning algorithms closely approximate gradient-descent-based optimizers, thus making the Transformer a universal solver for various ICL tasks. Specifically, these works demonstrate that Transformers can approximate gradient descent steps implicitly through some specific constructions of their parameters, enabling them to adapt to new data points during inference without explicit re-training.

However, it is well known that gradient-descent-based algorithms are not efficient in solving complicated optimization problems and thus may not be sufficient to explain the superior performance of Transformers on a plethora of ICL tasks. A recent work (Ahn et al.,, ) suggests that instead of gradient descent, Transformers actually perform pre-conditioned gradient descent during ICL. In other words, it learns a pre-conditioner during pre-training and then utilizes it during ICL to expedite the optimization process. recently demonstrates that the forward pass of a trained transformer can implement meta-optimization algorithms, i.e., it can implicitly define internal objective functions and then optimize these objectives to generate predictions. Similarly, Zhang et al., () show that a mesa-optimizer embedding the covariance matrix of input data can efficiently solve linear regression tasks. Such interpretation of the ICL mechanism shares the same essence as
Learning-to-Optimize (L2O) algorithms, and motivates the following hypothesis:

*Transformer does not simply implement a universal optimization algorithm during ICL. Rather, it extracts useful information from the given dataset during pre-training and then utilizes such information to generate an optimization algorithm that best suits the given ICL task.*

In this paper, we examine this hypothesis through the lens of *in-context sparse recovery*. Sparse recovery is a classical signal processing problem that is of significant practical interest across various domains, such as compressive sensing in medical imaging (Shen et al.,, ) and spectrum sensing (Elad,, ).
Compared with in-context linear regression that is often studied in the literature (Akyürek et al.,, ; ; Bai et al.,, ), the additional sparsity constraint makes the problem much more complicated and requires more sophisticated treatment.
Our main contributions are as follows.

- •
First, we demonstrate that Transformers can implement an L2O algorithm for in-context sparse recovery, and theoretically prove that a $K$-layer Transformer as an L2O algorithm can recover the underlying sparse vector in-context at a convergence rate linear in $K$. Such linear convergence results significantly improve the state-of-the-art convergence results for in-context sparse recovery (Bai et al.,, ; ), and enrich the literature.
- •
Second, we show that the Transformer as an L2O algorithm can actually outperform traditional L2O algorithms for sparse recovery in several aspects: 1) It does not require the measurement matrices involved in training to be the same as those in traditional L2O algorithms (Gregor and LeCun,, ; Chen et al.,, ; Liu and Chen,, ) for sparse recovery, which allows more flexibility to solve various in-context sparse recovery tasks. 2) It allows different numbers of measurements (i.e., prompt length) used for in-context sparse recovery, with guaranteed recovery performance as long as the number of measurements is sufficiently large. 3) It can extract structural properties of the underlying sparse vectors from the training data and utilize them to expedite its ICL convergence.
- •
We compare the ICL performances of Transformers with traditional iterative algorithms and L2O algorithms for sparse recovery empirically. Our experimental results indicate that Transformers substantially outperform traditional gradient-descent-based iterative algorithms, and achieve comparable performances compared with L2O algorithms that are trained and tested using data generated with the same measurement matrix. This supports our claim that Transformers can implement L2O algorithms during ICL. Besides, Transformers also demonstrate remarkable generalization capability when the measurement matrix varies, and achieve accelerated convergence when additional structure is imposed on the underlying sparse vectors, supporting our theoretical findings.

## 2 Related Works

#### ICL Mechanism for Transformers.

Brown et al., () first show that GPT-3, a Transformer-based LLM, can perform new tasks from input-output pairs without parameter updates, suggesting its ICL ability. This intriguing phenomenon of Transformers has attracted many attentions, leading to various interpretations and hypotheses about its underlying mechanism. For example, Han et al., () empirically hypothesize that Transformers perform kernel regression with internal representations when facing in-context examples, and Fu et al., () empirically show that Transformers learn to implement an algorithm similar to iterative Newton’s method for ICL tasks.

To better understand the ICL mechanism in large Transformers, existing works aim to demonstrate the Transformer’s *capability* for ICL by construction, e.g., showing that Transformers can perform gradient-based algorithms to solve ICL tasks by iteratively performing gradient descent layer by layer. In this category, Akyürek et al., () show that by construction, Transformers can implement gradient descent-based algorithms for linear regression problems. construct explicit weights for a Transformer, claiming it can perform gradient descent on linear and non-linear regression tasks. Bai et al., () provide constructions such that Transformers can make selections between different gradient-based algorithms. Ahn et al., () reason the ICL capability of Transformers to their ability to implement a pre-conditioned gradient descent for linear regression tasks, where the pre-condition matrix is learned from pre-training. Recently, demonstrates that the forward pass of a trained transformer can implement meta-optimization algorithms, i.e., it can implicitly define internal objective functions and then optimize these objectives to generate predictions. Similarly, Zhang et al., () show that a mesa-optimizer embedding the covariance matrix of input data can efficiently solve linear regression tasks.
Our work belongs to this category, where we construct a Transformer structure that can implement an L2O algorithm to effectively solve the in-context sparse recovery problem.

#### Training Dynamics of Transformers.

There exist some works aiming to theoretically understand the ICL mechanism in Transformer through their training dynamics.
Ahn et al., (); Mahankali et al., (); Zhang et al., (); Huang et al., () investigate the dynamics of Transformers with a single attention layer and a single head for in-context linear regression tasks.
Cui et al., () prove that Transformers with multi-head attention layers outperform those with single-head attention. Cheng et al., () show that local optimal solutions in Transformers can perform gradient descent in-context for non-linear functions. Kim and Suzuki, () study the non-convex mean-field dynamics of Transformers, and Nichani et al., () characterize the convergence rate for the training loss in learning a causal graph. Additionally, investigate the gradient flow in training multi-head single-layer Transformers for multi-task linear regression. Chen and Li, () propose a supervised training algorithm for multi-head Transformers.
The training dynamics of Transformers for binary classification (; ; Vasudeva et al.,, ; Li et al.,, ; Deora et al.,, ; ) and next-token prediction (NTP) (; ; ; Huang et al.,, ) have also been studied recently.

#### L2O Algorithms for Sparse Recovery.

There is a rich literature on L2O. We discuss the L2O algorithms relevant to sparse recovery here and leave more discussions on general L2O in Appendix A. Sparse recovery, typically formulated as a least absolute shrinkage and selection operator (LASSO), has many important applications like magnetic resonance imaging (Meng et al.,, ) and stock market forecasting (Roy et al.,, ), thus motivates the design of efficient algorithms. E.g., the iterative soft-thresholding algorithm (ISTA) (Daubechies et al.,, ) is proposed to solve LASSO and improves over the standard gradient descent algorithm. Motivated by the ISTA structure, Gregor and LeCun, () introduce the Learned ISTA (LISTA), a feedforward neural network that incorporates trainable matrices into ISTA updates. Chen et al., () and Liu and Chen, () further propose LISTA-Partial Weight Coupling (LISTA-CP) and Analytic LISTA (ALISTA) with fewer trainable parameters, making them easier to train. They also provide theoretical analyses demonstrating a linear convergence rate. However, these existing LISTA-type algorithms only apply to a fixed measurement matrix and fail when the measurement matrices (the inputs of the input-output pairs in ICL) change. In this work, we provide evidence that Transformer can perform a LISTA-type algorithm that tackles this issue and succeeds in in-context sparse recovery tasks, both theoretically in Section 5 and experimentally in Section 6.

## 3 Preliminaries

#### Notations.

For matrix $\mathbf{X}$, we use $[\mathbf{X}]_{p:q,r:s}$ to denote the submatrix that contains rows $p$ to $q$ and columns $r$ to $s$, and we use $[\mathbf{X}]_{:,i}$ and $[\mathbf{X}]_{j,:}$ to denote the $i$-th column and $j$-th row of $\mathbf{X}$ respectively. In some places, we also use $[\mathbf{X}]_{i}$ to denote its $i$-th column for convenience. We use $\|\mathbf{X}\|_{F}$ to denote its Frobenius norm.
For vector $\mathbf{x}$, we use $\|\mathbf{x}\|_{1}$, $\|\mathbf{x}\|$ and $\|\mathbf{x}\|_{\infty}$ to denote its $\ell_{1}$, $\ell_{2}$ and $\ell_{\infty}$ norms, respectively. We denote by $\mathbbm{1}_{d}$ and $\mathbf{0}_{d}$ the $d$-dimensional all-$1$ and all-0 0 column vectors, respectively.
$\mathbbm{1}_{a\times b}$ and $\mathbf{0}_{a\times b}$ denote the all-$1$ and all-0 0 matrices of size $a\times b$, respectively.

### 3.1 Transformer Architecture

In this work, we consider the decoder-based Transformer architecture (Vaswani et al.,, ), where each attention layer is masked by a decoder-based attention mask and followed by a multi-layer perception (MLP) layer.

###### Definition 3.1 (Masked attention layer) .

Denote an $M$-head masked attention layer parameterized by $\{(\mathbf{V}_{m},\mathbf{Q}_{m},\mathbf{K}_{m})_{m\in[M]}\}$ as $\mathrm{Attn}_{\{(\mathbf{V}_{m},\mathbf{Q}_{m},\mathbf{K}_{m})\}}(\cdot)$, where $\mathbf{V}_{m},\mathbf{Q}_{m},\mathbf{K}_{m}\in\mathbb{R}^{D\times D}$, $\forall m\in[M]$. Then, given an input sequence $\mathbf{H}\in\mathbb{R}^{D\times(N+1)}$, the output sequence of the attention layer is

$$ $\displaystyle\mathrm{Attn}_{\{(\mathbf{V}_{m},\mathbf{Q}_{m},\mathbf{K}_{m})\} }(\mathbf{H})=\mathbf{H}+\sum_{m=1}^{M}(\mathbf{V}_{m}\mathbf{H})\times\mathrm {mask}\Big{(}\sigma\big{(}(\mathbf{K}_{m}\mathbf{H})^{\top}(\mathbf{Q}_{m} \mathbf{H})\big{)}\Big{)},$ $$

where $\mathrm{mask}(\cdot)$ is defined as

$$ $\displaystyle\mathrm{mask}(\mathbf{M})=\begin{bmatrix}1&1/2&\cdots&\frac{1}{N+ 1}\\ 0&1/2&\cdots&\frac{1}{N+1}\\ \cdots&\cdots&\cdots&\cdots\\ 0&0&\cdots&\frac{1}{N+1}\\ \end{bmatrix}\odot\mathbf{M}\quad\mbox{for}\quad\mathbf{M}\in\mathbb{R}^{(N+1) \times(N+1)},$ $$

the operator $\odot$ denotes the Hadamard product (element-wise product), and $\sigma(\cdot)$ denotes the activation function. Let $\mathbf{h}_{i}$ denotes the $i$-th column of $\mathbf{H}$. Then, we have

$$ $\displaystyle[\mathrm{Attn}_{\{(\mathbf{V}_{m},\mathbf{Q}_{m},\mathbf{K}_{m}) \}}(\mathbf{H})]_{i}=\mathbf{h}_{i}+\sum_{m=1}^{M}\sum_{j=1}^{i}\frac{1}{i} \sigma((\mathbf{Q}_{m}\mathbf{h}_{i})^{\top}(\mathbf{K}_{m}\mathbf{h}_{j})) \cdot\mathbf{V}_{m}\mathbf{h}_{j}.$ $$

###### Definition 3.2 (MLP layer) .

Given $\mathbf{W}_{1}\in\mathbb{R}^{D^{\prime}\times D}$, $\mathbf{W}_{2}\in\mathbb{R}^{D\times D^{\prime}}$ and a bias vector $\mathbf{b}\in\mathbb{R}^{D^{\prime}}$, an MLP layer following the decoder attention layer, denoted as $\mathrm{MLP}_{\{\mathbf{W}_{1},\mathbf{W}_{2},\mathbf{b}\}}$, maps each token in the input sequence (i.e, each column $\mathbf{h}_{i}$ in $\mathbf{H}\in\mathbb{R}^{D\times N}$) to another token as

$$ $\displaystyle\mathrm{MLP}_{\{\mathbf{W}_{1},\mathbf{W}_{2},\mathbf{b}\}}( \mathbf{h}_{i})=\mathbf{h}_{i}+\mathbf{W}_{2}\sigma(\mathbf{W}_{1}\mathbf{h}_{ i}+\mathbf{b}),$ $$

where $\sigma$ is the non-linear activation function.

In this work, we set the activation function in Definition 3.1 and Definition 3.2 as the element-wise ReLU function. Next, we define the one-layer decoder-based Transformer structure.

###### Definition 3.3 (Transformer layer) .

A one-layer decoder-based Transformer is parameterized by $\mathbf{\Theta}:=\{\mathbf{W}_{1},\mathbf{W}_{2},{\mathbf{b}},(\mathbf{V}_{m},
\mathbf{Q}_{m},\mathbf{K}_{m})_{m\in[M]}\}$, denoted as $\mathrm{TF}_{\mathbf{\Theta}}$. Therefore, give input sequence $\mathbf{H}\in\mathbb{R}^{d\times N}$, the output sequence is:

$$ $\displaystyle\mathrm{TF}_{\mathbf{\Theta}}(\mathbf{H})=\mathrm{MLP}_{\{\mathbf {W}_{1},\mathbf{W}_{2},{\mathbf{b}}\}}\Big{(}\mathrm{Attn}_{\{(\mathbf{V}_{m}, \mathbf{Q}_{m},\mathbf{K}_{m})\}}(\mathbf{H})\Big{)}.$ $$

###### Definition 3.1 (Masked attention layer) .

###### Definition 3.2 (MLP layer) .

###### Definition 3.3 (Transformer layer) .

### 3.2 In-context Learning by Transformers

For an in-context learning (ICL) task, a trained Transformer is given an ICL instance $\mathcal{I}=(\mathcal{D},\mathbf{x}_{N+1})$, where $\mathcal{D}=\{(\mathbf{x}_{i},y_{i})\}_{i\in[N]}$ and $\mathbf{x}_{N+1}$ is a query. Here, $\mathbf{x}_{i}\in\mathbb{R}^{d}$ is an in-context example, and $y_{i}$ is the corresponding label for $\mathbf{x}_{i}$. We assume $y_{i}=f_{\bm{\beta}}(\mathbf{x}_{i})+\epsilon_{i}$, where $\epsilon_{i}$ is an added random noise, and $f_{\bm{\beta}}$ is a deterministic function parameterized by $\bm{\beta}$. Unlike conventional supervised learning, for each ICL instance, $\bm{\beta}\sim P_{\beta}$, i.e., it is randomly sampled from a distribution $P_{\beta}$.

To perform ICL in a Transformer, we first embed the ICL instance into an input sequence $\mathbf{H}\in\mathbb{R}^{D\times N^{\prime}}$. The Transformer then generates an output sequence $\text{TF}(\mathbf{H})$ with the same size as $\mathbf{H}$, based on which a prediction $\widehat{y}_{N+1}$ is generated through a read-out function $F$, i.e., $\widehat{y}_{N+1}=F(\text{TF}(\mathbf{H}))$. The objective of ICL is then to ensure that $\widehat{y}_{N+1}$ closely approximates the target value $y_{N+1}=f_{\bm{\beta}}(\mathbf{x}_{N+1})+\epsilon_{N+1}$ for any ICL instance.

When a Transformer is pre-trained for ICL, it first samples a large set of ICL instances. For each instance, the Transformer generates a prediction $\widehat{y}_{N+1}$ and calculates the prediction loss by comparing it with $y_{N+1}$ using a proper loss function. The training loss is the aggregation of all prediction losses for every ICL instance used in pre-training, and the Transformer is trained to minimize this training loss.

Previous studies about the mechanism of how Transformers performs in-context learning have attracted a lot of attention recently. To start with, it is believed that the ICL capability is due to the Transformer’s implicit implementation of learning algorithms in the forward pass. , Dai et al., (), and Bai et al., () suggest that these learning algorithms closely approximate gradient-descent-based optimizers, thus making the Transformer a universal solver for various ICL tasks. A recent work (Ahn et al.,, ) suggests that instead of gradient descent, Transformers actually perform pre-conditioned gradient descent for in-context least square linear regression. In general, these results corroborate the claim that the mechanism of the Transformer in-context learning is a L2O algorithm. We study this perspective and further provide the evidence supporting this claim by considering a more complicated in-context problem: in-context sparse recovery.

## 4 Transformer as a LISTA-type Algorithm for In-context Sparse Recovery

In this section, we demonstrate that a decoder-based Transformer can implement a novel LISTA-type L2O algorithm, specifically LISTA-VM, as detailed in Theorem 4.1, for in-context sparse recovery. We begin by formally defining the in-context sparse recovery problem.

### 4.1 In-context Sparse Recovery

Sparse recovery is a fundamental problem in fields such as compressed sensing, signal denoising, and statistical model selection. The core concept of sparse recovery is that a high-dimensional sparse signal can be inferred from very few linear observations if certain conditions are satisfied. Specifically, it aims to identify an $S$-sparse vector $\bm{\beta}^{*}\in\mathbb{R}^{d}$ from its noisy linear observations $\mathbf{y}=\mathbf{X}\bm{\beta}^{*}+\bm{\epsilon}$, where $\mathbf{X}\in\mathbb{R}^{N\times d}$ is a measurement matrix, $\bm{\epsilon}\in\mathbb{R}^{N}$ is an isometric Gaussian noise vector with mean vector $\mathbf{0}_{N}$ and covariance matrix $\mathbf{I}_{N\times N}$.
Typically, we assume $d\gg N$, which is the so-called under-determined case. One common assumption for the measurement matrix $\mathbf{X}$ (Pitaval et al.,, ; Ge et al.,, ; Zhu et al.,, ) is that
each row of the matrix is independently sampled from an isometric sub-Gaussian distribution with zero mean and covariance matrix $\mathrm{diag}(\sigma_{1}^{2}\cdots,\sigma_{d}^{2})$, denoted as $P_{\mathbf{x}}$.
which guarantees the critical *restricted isometry property* under mild conditions on the sparsity level (Candes and Tao,, ). In this work, we also assume $\bm{\beta}^{*}$ is randomly sampled from a distribution $P_{\bm{\beta}}$, which admits an $S$-sparse vector with random support.

A popular approach to tackling sparse recovery is the least absolute shrinkage and selection operator (LASSO), which aims to find the optimal sparse vector $\bm{\beta}\in\mathbb{R}^{d}$ that minimizes the following loss:

$$ $\displaystyle\mathcal{L}(\bm{\beta})=\frac{1}{2}\|\mathbf{y}-\mathbf{X}\bm{ \beta}\|^{2}_{2}+\alpha\|\bm{\beta}\|_{1}.$ $$

Here $\alpha$ is a coefficient controlling the sparsity penalty. We denote the transpose of the $i$-th row in $\mathbf{X}$ by $\mathbf{x}_{i}$, i.e., $\mathbf{x}_{i}=[\mathbf{X}^{\top}]_{:,i}$.

In this work, we study how Transformers solve the sparse recovery problem in context. Specifically, we assume a Transformer has been pre-trained with randomly sampled sparse recovery instances. Then, during ICL, for a given sparse recovery instance with $\bm{\beta}\sim P_{\beta}$, $\{\mathbf{x}_{i}\}_{i=1}^{N}\sim P_{\mathbf{x}}$, $\{\epsilon_{i}\}_{i=1}^{N}\sim P_{\epsilon}$, the Transformer aims to predict $y_{N+1}$ with input $(\mathbf{X},\mathbf{y},\mathbf{x}_{N+1})$, where $y_{N+1}=\mathbf{x}_{N+1}^{\top}\bm{\beta}+\epsilon_{N+1}$.

### 4.2 Classical Algorithms

Gradient descent is known to struggle in solving the LASSO problem due to its inefficiency in effectively handling the sparsity constraint (Chen et al.,, ). This inefficiency has led to the development of more specialized algorithms that can better address the unique challenges posed by the LASSO formulation. A popular approach to solving the LASSO problem is the Iterative Shrinkage Thresholding Algorithm (ISTA). Starting with a fixed initial point $\bm{\beta}^{(1)}$, the update rule in the $k$-th iteration is given by

$$ $\bm{\beta}^{(k+1)}=\mathcal{S}_{\alpha/L}\left(\bm{\beta}^{(k)}-\frac{1}{L} \mathbf{X}^{\top}(\mathbf{X}\bm{\beta}^{(k)}-\mathbf{y})\right).$ $$

Here, $\mathcal{S}_{\alpha/L}$ is the soft-thresholding function defined as $[\mathcal{S}_{\alpha/L}(\mathbf{x})]_{i}=\text{sign}([\mathbf{x}]_{i})\max\{0,
|[\mathbf{x}]_{i}|-\alpha/L\}$, and $L$ is typically chosen as the largest eigenvalue of $\mathbf{X}^{\top}\mathbf{X}$ (Chen et al.,, ; Liu and Chen,, ).

Generally, for any ground-truth sparse vector $\bm{\beta}^{*}$ and any given $\mathbf{X}$, ISTA converges at a sublinear rate (Beck and Teboulle,, ).
The sublinear convergence rate of ISTA is considered inefficient, which has led to the development of various LISTA-type L2O algorithms, such as LISTA (Gregor and LeCun,, ), LISTA-CP (Chen et al.,, ), and ALISTA (Liu and Chen,, ). These algorithms learn the weights in the matrices in ISTA rather than fixing them.

Among them, LISTA-CP is one state-of-the-art (SOTA) method that has been well-studied. The update rule in the $k$-th iteration of LISTA-CP can be expressed as

$$ $\displaystyle\bm{\beta}^{(k+1)}=\mathcal{S}_{\theta^{(k)}}\Big{(}\bm{\beta}^{( k)}-(\mathbf{D}^{(k)})^{\top}(\mathbf{X}\bm{\beta}^{(k)}-\mathbf{y})\Big{)},$ (4.1) $$

where $\{\theta^{(k)},\mathbf{D}^{(k)}\}$ are learnable parameters.
Compared with ISTA with fixed parameters, LISTA-CP obtains $\{\theta^{(k)},\mathbf{D}^{(k)}\}$ through pre-training. Specifically, with a fixed measurement matrix $\mathbf{X}$, it randomly samples $n$ $S$-sparse vectors $\{\bm{\beta}_{j}\}_{j=1}^{n}\sim P_{\beta}$ and generates $\{\mathbf{y}_{j}\}_{j=1}^{n}$, which is then utilized to optimize $\{\theta^{(k)},\mathbf{D}^{(k)}\}$ by minimizing the total predicting loss for $\{\bm{\beta}_{j}\}_{j}$.
Chen et al., () show that, for the same measurement matrix $\mathbf{X}$, given any random instance $(\bm{\beta}^{*},\mathbf{y})$, a pre-trained LISTA-CP will converge to the ground-truth $\bm{\beta}^{*}$ linearly in $K$ under certain necessary conditions on ${\mathbf{X}}$.

### 4.3 Transformer Can Provably Perform LISTA-type Algorithms

Noting that LISTA-type algorithms can efficiently solve sparse recovery problems, in this section, we argue that a trained Transformer can implement a LISTA-type algorithm and efficiently solve a sparse recovery problem in context. To distinguish the algorithm implemented by the Transformer with the classical LISTA-type algorithms, we term it as LISTA with Varying Measurements (LISTA-VM). Towards this end, we provide an explicit construction of a $K$-layer decoder-based Transformer as follows.
A $K$-layer Transformer is the concatenation of $K$ blocks, where each block comprises a self-attention layer followed by an MLP layer. The input to the first self-attention layer, denoted as $\mathbf{H}^{(1)}$, is an embedding of the given in-context sparse recovery instance $\mathcal{I}=(\mathbf{X},\mathbf{y},\mathbf{x}_{N+1})$.

Embedding. Given an in-context sparse recovery instance $\mathcal{I}=(\mathbf{X},\mathbf{y},\mathbf{x}_{N+1})$ we embed the instance into an input sequence $\mathbf{H}^{(1)}\in\mathbb{R}^{(2d+2)\times(2N+1)}$
as follows:

$$ $\displaystyle\mathbf{H}^{(1)}(\mathcal{I})=\begin{bmatrix}\mathbf{x}_{1}& \mathbf{x}_{1}&\cdots&\mathbf{x}_{N}&\mathbf{x}_{N}&\mathbf{x}_{N+1}\\ 0&y_{1}&\cdots&0&y_{N}&0\\ \bm{\beta}^{(1)}_{1}&\bm{\beta}^{(1)}_{2}&\cdots&\bm{\beta}^{(1)}_{2N-1}&\bm{ \beta}^{(1)}_{2N}&\bm{\beta}^{(1)}_{2N+1}\\ 1&0&\cdots&1&0&1\end{bmatrix},$ (4.2) $$

where $\{\bm{\beta}_{i}^{(1)}\}_{i\in[2N+1]}\in\mathbb{R}^{d}$ are implicit parameter vectors initialized as $\mathbf{0}_{d}$, and $\mathbf{x}_{i}$ is the $i$-th column of the transposed measurement matrix, i.e, $[\mathbf{X}^{\top}]_{:,i}$.

Self-attention layer.
The self-attention layer takes as input a sequence of embeddings and outputs a sequence of embeddings of the same length. Let the $K$-layer decoder-based Transformer feature four attention heads uniquely indexed as $+1$, $-1$, $+2$, and $-2$. We construct these heads according to the structure specified in Section C.1. This construction ensures that the self-attention layer performs the $\bm{\beta}$ updating inside the soft-thresholding function in Equation 4.1. Furthermore, with our construction, the learnable matrix $\mathbf{D}^{(k)}$ in LISTA-CP becomes *context-dependent* instead of fixed. In the update rule implemented by the self-attention layer, this matrix becomes $\mathbf{D}^{(k)}=\frac{1}{2N+1}\mathbf{X}(\mathbf{M}^{(k)})^{\top}$, where $\mathbf{M}^{(k)}\in\mathbb{R}^{d\times d}$ is fixed.

MLP layer.
For the MLP layer following the $k$-th self-attention layer, it functions as a feedforward neural network that takes the output of the self-attention layer as its input, and outputs a transformed sequence of the embeddings. Recall that $\mathbf{h}_{i}$ is the $i$-th column in an embedding sequence $\mathbf{H}$. We parameterize $(\mathbf{W}_{1},\mathbf{W}_{2},\mathbf{b})$ in the $k$-th MLP layer to let it function as a partial soft-threshold function:

$$ $\displaystyle\mathrm{MLP}(\mathbf{h}_{i})=\begin{bmatrix}[\mathbf{h}_{i}]_{1:d +1}\\ \mathcal{S}_{\theta^{(k)}}([\mathbf{h}_{i}]_{d+2:2d+1})\\ [\mathbf{h}_{i}]_{2d+2}\end{bmatrix},$ (4.3) $$

where $\mathcal{S}_{\theta^{(k)}}$ is the soft-threshold function.
Essentially, the soft-threshold function is effectively implemented by the MLP layer utilizing the $\mathrm{ReLU}$ activation. This can be realized by combining $-\mathrm{ReLU}(x)$, $\mathrm{ReLU}(-x)$, $\mathrm{ReLU}(x-\theta)$, $-\mathrm{ReLU}(-x+\theta)$, and $x$. The implementation details of the MLP layer can be found in Section C.1.

Read-out function.
Given the output sequence of the Transformer $\mathrm{TF}_{\mathbf{\Theta}}(\mathbf{H}^{(1)})$, to obtain the estimation $\widehat{y}_{N+1}$, it is necessary to *read out* from the output sequence. In this work, we consider two types of read-out functions:

###### Definition 4.1 (Linear read-out) .

$\mathcal{F}_{\text{linear}}$ is defined as the class of linear readout functions such that

$$ $\displaystyle\mathcal{F}_{\text{linear}}=\{F(\cdot)\mid F(\mathbf{h})=\mathbf{ v}^{\top}\mathbf{h},\mathbf{v}\in\mathbb{R}^{D}\}.$ $$

###### Definition 4.2 (Query read-out) .

$\mathcal{F}_{\text{query}}$ is defined as the class of explicit quadratic readout functions such that

$$ $\displaystyle\mathcal{F}_{\text{query}}=\{F(\cdot~{};\widetilde{\mathbf{X}}) \mid F(\mathbf{h},i;\widetilde{\mathbf{X}})=\mathbf{h}^{\top}\mathbf{V} \widetilde{\mathbf{X}}_{\lfloor\frac{i+1}{2}\rfloor,:},\mathbf{V}\in\mathbb{R} ^{D\times d}\},$ $$

where $i\in[N+1]$ is an index number and $\widetilde{\mathbf{X}}=[\mathbf{X}^{\top}~{}\mathbf{x}_{N+1}]^{\top}$.

Given a readout function $F_{\mathbf{v}}\in\mathcal{F}_{\text{linear}}$ parameterized by $\mathbf{v}$ or $F_{\mathbf{V}}\in\mathcal{F}_{\text{query}}$ parameterized by $\mathbf{V}$, the estimation $\widehat{y_{i}}$ obtained by the $K$-layer Transformer is $\widehat{y_{i}}=F_{\mathbf{v}}(\mathbf{h}_{2i+1}^{K+1})$ or $\widehat{y_{i}}=F_{\mathbf{V}}(\mathbf{h}_{2i+1}^{K+1},2i-1)$ respectively.

Before we formally present our main results, we introduce the following assumptions.

###### Assumption 1 .

For $\mathbf{x}\sim P_{\mathbf{x}}$ and $\bm{\beta}^{*}\sim P_{\bm{\beta}}$, we assume $\|\mathbf{x}\|\leq b_{\mathbf{x}}$ and $\|\bm{\beta}^{*}\|_{1}\leq b_{\bm{\beta}}$ almost surely. Besides, we consider the noiseless scenario where $\bm{\epsilon}=\mathbf{0}$.

We note that the boundedness assumption over $\mathbf{x}$ and $\bm{\beta}^{*}$ ensures the robustness of the Transformer and prevents it from blowing up under ill conditions. A similar assumption is adopted in Bai et al., (). The noiseless assumption is for ease of analysis and is common in the analysis of Transformers  (Ahn et al.,, ; Fu et al.,, ; Bai et al.,, ). We note that the following Theorem 4.1 can be straightforwardly extended to the noisy case when the noise is bounded.

We denote the input sequence to the $k$-th self-attention layer as $\mathbf{H}^{(k)}$, and use $\bm{\beta}_{2n+1}^{(k+1)}$ to represent the vector $[\mathbf{H}^{(k+1)}]_{d+1:2d+1,2n+1}$. Then, we state the following theorem.

###### Theorem 4.1 (Equivalence between ICL and LISTA-VM) .

With the Transformer structure described above, under Assumption 1,
there exists a set of parameters in the Transformer so that for any $k\in[1:K]$, $n\in[N]$, we have

$$ $\displaystyle\bm{\beta}_{2n+1}^{(k+1)}=\mathcal{S}_{\theta^{(k)}}\Big{(}\bm{ \beta}^{(k)}_{2n+1}-\frac{1}{2n+1}\mathbf{M}^{(k)}[\mathbf{X}]_{1:n,:}^{\top}( [\mathbf{X}]_{1:n,:}\bm{\beta}^{(k)}_{2n+1}-\mathbf{y}_{1:n})\Big{)},$ (4.4) $$

where $\mathbf{M}^{(k)}\in\mathbb{R}^{d\times d}$ is embedded in the $k$-th Transformer layer.

The proof of Theorem 4.1 is detailed in Section C.2.

###### Remark 1 .

As mentioned above, matrix $\mathbf{D}^{(k)}$ in the update rule of LISTA-CP in Equation 4.1 is learned during pre-training and remains fixed across different in-context sparse recovery instances. As a result, LISTA-CP requires the measurement matrix $\mathbf{X}$ to stay the same during pre-training and inference. In contrast, if we denote $\mathbf{D}^{(k)}_{n}=\frac{1}{2n+1}[\mathbf{X}]_{1:n,:}(\mathbf{M}^{(k)})^{\top}$, then within the update rule of the LISTA-VM algorithm implemented by the Transformer, as detailed in Equation 4.4, the matrix $\mathbf{D}^{(k)}_{n}$ depends on a fixed matrix $\mathbf{M}^{(k)}$ after pre-training as well as on the measurement matrix $\mathbf{X}$ during inference. As a result, the Transformer can adaptively update $\mathbf{D}^{(k)}_{n}$ for instances with different $\mathbf{X}$’s, enabling more flexibility and improved performance for the ICL tasks.

###### Definition 4.1 (Linear read-out) .

###### Definition 4.2 (Query read-out) .

###### Assumption 1 .

###### Theorem 4.1 (Equivalence between ICL and LISTA-VM) .

###### Remark 1 .

## 5 Performance of Transformers for In-context Sparse Recovery

In this section, we demonstrate the effectiveness of the constructed Transformer in implementing the LISTA-VM algorithm and solving in-context sparse recovery problems. We first show that the LISTA-VM algorithm implemented by the Transformer recovers the underlying sparse vector in context at a convergence rate linear in $K$. We then demonstrate that the Transformer can accurately predict ${y}_{N+1}$ at the same time.

### 5.1 Sparse Vector Estimation

###### Theorem 5.1 (Convergence of ICL) .

Let $\delta\in(0,1)$, $N_{0}=8(4S-2)^{2}{\frac{\log d+\log S-\log\delta}{c}}$, $\alpha_{n}=-\log\big{(}1-\frac{2}{3}\gamma+\gamma(2S-1)\sqrt{\frac{\log d-\log
\delta}{nc}}+\sqrt{\frac{\log S-\log\delta}{nc}}\big{)}$, where $c$ is a positive constant and $\gamma$ is a positive constant satisfies $\gamma\leq\frac{3}{2}$.
For a $K$-layer Transformer model with the structure described in Section 4.3, under Assumption 1,
there exists a set of parameters such that for any $n\in[N_{0}:N]$, with probability at least $1-\delta$, we have

$$ $\displaystyle\big{\|}\bm{\beta}^{(K+1)}_{2n+1}-\bm{\beta}^{*}\big{\|}\leq b_{ \bm{\beta}}e^{-\alpha_{n}K}.$ $$

###### Main challenge and key ideas of the proof..

Similar to the proofs in Chen et al., () and Liu and Chen, (), the core step in proving convergence is to ensure that $\mathbf{D}^{(k)}$ exhibits small coherence with $\mathbf{X}$, i.e., $(\mathbf{D}^{(k)})^{\top}\mathbf{X}\approx\mathbf{I}_{d\times d}$.
In Chen et al., () and Liu and Chen, (), such a $\mathbf{D}^{(k)}$ is obtained by minimizing the generalized mutual coherence to a *fixed* $\mathbf{X}$. However, this results in poor generalization across different $\mathbf{X}$’s. In our proof, we consider $\mathbf{X}$ as a random matrix and leverage its sub-Gaussian properties to prove that if $\mathbf{D}^{(k)}=\mathbf{X}(\mathbf{M}^{(k)})^{\top}$, where $\mathbf{M}^{(k)}$ is associated with the covariance of $\mathbf{X}$, then $\mathbf{D}^{(k)}$ will have small coherence with $\mathbf{X}$ with high probability. We defer the detailed proof of Theorem 5.1 to Section D.1.
∎

###### Remark 2 (Generalization across measurement matrix $\mathbf{X}$ ) .

Theorem 5.1 shows that for any $\mathbf{X}$ satisfying Assumption 1, the Transformer can estimate the ground-truth sparse vector $\bm{\beta}^{*}$ in-context at a convergence rate linear in $K$. This is in stark contrast to traditional LISTA-CP type of algorithms, which only work for fixed $\mathbf{X}$. Such generalization is enabled by the input-dependent matrices $\{\mathbf{D}^{(k)}\}_{k}$.
Besides, we also note that $\bm{\beta}^{(K+1)}_{2n+1}$ only depends on $\mathbf{x}_{1},\ldots,\mathbf{x}_{n}$. This implies that even if the measurement matrix $\mathbf{X}$ is of dimension $n\times d$ instead of $N\times d$, when $n\in[N_{0},N]$, the Transformer can still recover $\bm{\beta}^{*}$ accurately. Such results demonstrate the robustness of Transformers to variations in in-context sparse recovery tasks.

###### Remark 3 (Effective utilization of the hidden patterns in ICL tasks) .

We note that the Transformer can be slightly modified to exploit certain hidden structures in the in-context sparse recovery tasks. Specifically, if the support of $\bm{\beta}$ lies in a subset $\mathbb{S}\subset[1:d]$ with $S<|\mathbb{S}|\leq d$, then by slightly modifying the parameters of the Transformer to ensure $[\bm{\beta}^{(k)}]_{i}=0$ for all $i\notin\mathbb{S}$, the ICL performance can be improved by replacing all $d$ involved in Theorem 5.1 by $|\mathbb{S}|$. We defer the corresponding result and analysis to Section D.2.

###### Theorem 5.1 (Convergence of ICL) .

###### Main challenge and key ideas of the proof..

###### Remark 2 (Generalization across measurement matrix 𝐗 𝐗 \mathbf{X} bold_X ) .

###### Remark 3 (Effective utilization of the hidden patterns in ICL tasks) .

### 5.2 Label Prediction

In Section 5.1, we have demonstrated that Transformers can successfully recover the ground-truth sparse vector $\bm{\beta}^{*}$ with linear convergence by implementing a LISTA-type algorithm. In this section, we bridge the gap between this theoretical claim and the explicit objective of in-context sparse recovery, which is to predict $\widehat{y}_{N+1}$ given an in-context instance $(\mathbf{X},\mathbf{y},\mathbf{x}_{N+1})$. This gap might seem trivial at first glance, as given $\mathbf{x}_{N+1}$ and an accurate estimate of $\bm{\beta}$, the label $\widehat{y}_{N+1}$ can be obtained through a simple linear operation. However, we will show that, for a decoder-based Transformer, generating $\widehat{y}_{N+1}$ using the predicted sparse vector $\bm{\beta}$, which is implicitly embedded within the forward-pass sequences, critically depends on the structure of the read-out function.

###### Theorem 5.2 .

Under the same setting as in Section 4.3, with probability at least $1-n\delta-\delta^{\prime}$, we have

$$ $\displaystyle\|y_{n}-\widehat{y}_{n}\|\leq b_{\mathbf{x}}\left(1-\frac{2}{3} \gamma\right)^{K}+\frac{c_{4}K}{\sqrt{n}}\left(1-\frac{2}{3}\gamma\right)^{K-1}$ $$

for a linear read-out function, and with probability at least $1-\delta$, we have

$$ $\displaystyle\|y_{n}-\widehat{y}_{n}\|\leq c_{5}e^{-\alpha_{n}K}$ $$

for a query read-out function, where $c_{4}$, $c_{5}$ are constants.

We defer the proof of Theorem 5.2 to Appendix E.

###### Remark 4 .

Theorem 5.2 indicates that adopting a linear read-out function results in a prediction error of order $\mathcal{O}\big{(}e^{-K}+\frac{K}{\sqrt{n}}e^{-K}\big{)}$, which still exhibits linear convergence with respect to $K$. When a query-based read-out function is employed, the convergence rate improves to $\mathcal{O}(e^{-K})$. However, there exists a gap of order $\mathcal{O}\big{(}\frac{K}{\sqrt{n}}e^{-K}\big{)}$, which diminishes as $n$ becomes large. Empirically, we observe the superiority of using a query-based read-out function in our experimental results, as detailed in Section 6, and we also observe that the gap decreases as $n$ grows.

###### Theorem 5.2 .

###### Remark 4 .

## 6 Experimental Results

Problem setup.
In all experiments, we adhere to the following steps to generate in-context sparse recovery instances.
First, we sample a ground truth sparse vector $\bm{\beta}^{*}$ from a $d=20$ dimensional standard normal distribution, and we fix the sparsity of $\bm{\beta}^{*}$ to be $3$ by randomly setting $17$ entries in $\bm{\beta}^{*}$ to zero. Next, we independently sample
$N=10$ vectors form a $d$ dimensional standard normal distribution and then contract the measurement matrix $\mathbf{X}\in\mathbb{R}^{10\times 20}$ (each sampled $d$ dimensional random vector is a row in $\mathbf{X}$). We also sample an additional $\mathbf{x}_{N+1}$ from the $d$-dimensional standard Gaussian distribution. We follow the noiseless setting in Bai et al., () for sparse recovery, i.e., $\mathbf{y}=\mathbf{X}\bm{\beta}^{*}$.

Baselines. The baselines for our experiments include traditional iterative algorithms such as ISTA and FISTA (Beck and Teboulle,, ). We also evaluate three classical LISTA-type L2O algorithms: LISTA, LISTA-CP, ALISTA (Gregor and LeCun,, ; Chen et al.,, ; Liu and Chen,, ). For all of these algorithms, we set the number of iterations $K=12$. We generate a single fixed measurement matrix $\mathbf{X}$. For each training epoch, we create $N=50,000$ instances from 50,000 randomly generated sparse vectors. We also evaluate the LISTA-VM algorithm introduced in Theorem 4.1, where we set the number of iterations $K=12$ as well. For each training epoch, we randomly sample 100 measurement matrices, each generating 500 instances from 500 randomly generated sparse vectors, which results in a total of 50,000 instances. For comparison, we also meta-train LISTA and LISTA-CP using the same training method as LISTA-VM. We do not perform meta-training for ALISTA, as the training process of ALISTA involves solving a non-convex optimization problem for each different measurement matrix $\mathbf{X}$, which makes meta-training for ALISTA unrealistic. For all baseline algorithms, we minimize the sparse vector prediction loss $\sum_{j=1}^{N}\|\widehat{\bm{\beta}}_{j}-\bm{\beta}_{j}\|^{2}$ using gradient descent for each epoch. We run all baseline experiments for 340 epochs.

Transformer structure. We consider two Transformer models, i.e., a small Transformer model (denoted as Small TF) and GPT-2. Small TF has $12$ layers, each containing a self-attention layer followed by an MLP layer. Each self-attention layer has 4 attention heads. We set the embedding dimension to $D=42$, and the embedding structure according to Equation 4.2. For GPT-2, we employ $12$ layers, and set the embedding dimension to $256$ and the number of attention heads per layer to $8$.
In order to train Small TF and GPT-2, we randomly generate $64$ instances per epoch and train the algorithms for $10^{6}$ epochs.
The training process minimizes the label prediction loss $\sum_{j=1}^{T}(y_{N+1,j}-\widehat{y}_{N+1,j})^{2}$.
We run the experiments for Small TF and GPT-2 on an NVIDIA RTX A5000 GPU with 24G memory. The training time for Small TF is approximately $8$ hours, while the training time for GPT-2 is around $12$ hours.

Figure: (a) Without support constraints
Refer to caption: x1.png

Results. We test the prediction performance of the baseline algorithms and Transformers on a sparse recovery instance $(\mathbf{X},\mathbf{y},\mathbf{x}_{N+1})$, and plot the label prediction loss in Figure 1.

We first do not impose any support constraint on $\bm{\beta}$. We start with the general setting where the testing instance is randomly generated (Varying $\mathbf{X}$). As shown in Figure 1(a), GPT-2 outperforms Small TF, followed by LISTA-VM, which outperforms iterative algorithms FISTA and ISTA, while the classical LISTA-type algorithms LISTA, LISTA-CP, and ALISTA perform the worst. Such results highlight the efficiency of LISTA-VM, Small TF and GPT-2 in solving ICL sparse recovery problems, corroborating our theoretical result in Theorem 5.1. Meanwhile, classical LISTA-type algorithms cannot handle mismatches between the measurement matrices in pre-training and testing, leading to poor prediction performance. When $\mathbf{X}$ during testing is fixed to be the same as that in pre-training, all LISTA-type algorithms achieve performances comparable with Small TF. For meta-trained LISTA and LISTA-CP, the corresponding prediction loss (denoted by red marks in Figure 1(a)) are still much higher than that under LISTA-VM. This indicates that meta-training cannot help those classical LISTA-type algorithms achieve performances comparable with LISTA-VM, further corroborating the strong generalization capability induced by the constructed Transformer structure.

Next, we impose an additional constraint on the support of the sparse vectors. For in-context sparse recovery instances used in both pre-training and testing, we set the support of the sparse vector $\bm{\beta}$ to the first $10$ entries, i.e., $\mathbb{S}=\{1,2,\cdots,10\}$. As observed in Figure 1(b), Small TF and GPT-2 significantly improve their performances in Figure 1(a), while other baseline algorithms do not exhibit significant performance improvements.
In Figure 1(b), we present the experimental results for a support-selected version of LISTA-VM, referred to as LISTA-VM-SS. This algorithm is a simple variation of LISTA-VM, where we incorporate prior knowledge of the support by setting all columns in $\mathbf{X}$ whose indices are not in the prior support set to be zero vectors.
As we claim in Remark 3 and Corollary D.1, a Transformer could perform this LISTA-VM-SS by utilizing prior knowledge of the support. Our results show that the LISTA-VM-SS achieves an in-context prediction error of approximately $0.07$, which is almost 10 times better than the standard version of LISTA-VM and is comparable to the prediction error of Small TF. This empirical finding corroborates our claim in Remark 3.

Finally, we examine the label prediction loss under Small TF with three types of read-out functions, i.e., linear read-out function (*Linear Readout*), query read-out function (*Query Readout*), and another quadratic read-out function with parameters selected according to the proof of Theorem 5.2 (*Query Readout^†*). In Figure 1(c), we observe that the label prediction error is lower with the query read-out functions than with the linear read-out function. When $n$ becomes large, the gap between the linear read-out function and the other two types of query read-out functions becomes insignificant, which is consistent with our theoretical result in Theorem 5.2. Meanwhile, those two query read-out functions behave very similarly.

## 7 Conclusion

In this work, we demonstrated that Transformers’ known ICL capabilities could be understood as performing L2O algorithms. Specifically, we showed that for sparse recovery tasks, Transformers can execute the LISTA-VM algorithm with a provable linear convergence rate. Our results highlight that, unlike existing LISTA-type algorithms, which are limited to solving individual sparse recovery problems with fixed measurement matrices, Transformers can address a general class of sparse recovery problems with varying measurement matrices without requiring parameter updates. Experimentally, we demonstrated that Transformers can leverage prior knowledge from training tasks and generalize effectively across different lengths of demonstration pairs, where traditional L2O methods typically fail. These insights not only deepen our understanding of Transformers’ capabilities in ICL but also suggest new potential applications for utilizing Transformers in other optimization tasks.

## Appendix A Additional Related Works

#### General L2O Techniques.

L2O leverages machine learning to develop optimization algorithms, aiming to improve existing methods and innovate new ones. As highlighted by Sucker et al., () and Chen et al., (), L2O intersects with meta-learning (also known as “learning-to-learn”) and automated machine learning (AutoML).

Unlike meta-learning, which focuses on enabling models to quickly adapt to new tasks with minimal data by leveraging prior knowledge from diverse tasks (Finn et al.,, ; Hospedales et al.,, ), L2O aims to improve the optimization process itself by developing adaptive algorithms tailored to specific tasks, leading to faster convergence and enhanced performance in model training (Andrychowicz et al.,, ; Li and Malik,, ). Thus, while meta-learning enhances task adaptability, L2O refines the efficiency of the optimization process. In contrast, AutoML focuses on model selection, optimization algorithm selection, and hyperparameter tuning (Yao et al.,, ); L2O distinguishes itself by its ability to generate new optimization techniques through learned models.

L2O has demonstrated significant potential across various optimization fields and applications. For instance, Andrychowicz et al., () introduced a method where optimization algorithms are learned using recurrent neural networks trained to optimize specific classes of functions. Li and Malik, () proposed learning optimization algorithms through reinforcement learning, utilizing guided policy search to develop optimization strategies. Furthermore, Hruby et al., () applied L2O to address the “minimal problem”, a common challenge in computer vision characterized by the presence of many spurious solutions. They trained a multilayer perceptronmodel to predict initial problem solutions, significantly reducing computation time.

## Appendix B Table of Notations

| Notation | Definition |
| --- | --- |
| $\mathbf{X}$ | Measurement Matrix |
| $\bm{\beta}^{*}$ | Ground-truth sparse vector |
| $\epsilon$ | Noise Vector |
| $Y$ | $Y=\mathbf{X}\bm{\beta}^{*}+\epsilon$ |
| $d$ | Number of columns of $\mathbf{X}$ |
| $N$ | Number of Measurement Vectors, i.e., number of rows of $\mathbf{X}$ |
| $S$ | Sparsity of $\bm{\beta}^{*}$, i.e., $\|\bm{\beta}^{*}\|_{0}=S$ |
| $\mathbb{S}$ | Support set of $\bm{\beta}^{*}$ |
| $D$ | Dimension in the Self-attention Layer |
| $M$ | Number of Heads in the Self-attention Layer |
| $D^{\prime}$ | Hidden Dimension in the MLP Layer |
| $K$ | Number of Layers in Transformer |
| $\mathbf{Q}_{i}^{(k)}$ | Query Matrix of the $k$-th Layer of Transformer’s $i$-th Head |
| $\mathbf{K}_{i}^{(k)}$ | Key Matrix of the $k$-th Layer of Transformer’s $i$-th Head |
| $\mathbf{V}_{i}^{(k)}$ | Value Matrix of the $k$-th Layer of Transformer’s $i$-th Head |
| $\mathbf{H}^{(k)}$ | Input Sequence of the $k$-th Layer of Transformer |
| $\mathbf{h}_{i}^{(k)}$ | $[\mathbf{H}^{(k)}]_{:,i}$ |
| $\text{sign}(x)$ | Sign Function: $\text{sign}(x)=|x|/x$ if $x\neq 0$, $\text{sign}(x)=0$ if $x=0$ |
| $\mathcal{S}_{\theta}(x)$ | Soft Thresholding Function: $\mathcal{S}_{\theta}(x)=\text{sign}(x)\max\{0,|x|-\theta\}$ |
| $\sigma:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ | ReLU Function:$[\sigma(x)]_{i}=x_{i}$ if $x_{i}\geq 0$, $[\sigma(x)]_{i}=0$ if $x_{i}<0$ |
| $\neg\mathcal{E}$ | Complement of an event $\mathcal{E}$ |

## Appendix C Deferred Proofs in Section 4.3

### C.1 Transformer Structure

#### Attention layer.

Consider a model consisting of $K$ Transformer layers, where each layer is equipped with four attention heads. These heads are uniquely indexed as $+1$, $-1$, $+2$, and $-2$ to distinguish their specific roles within the layer.

$$ $\displaystyle\mathbf{Q}_{\pm 1}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{(d+1)\times(d+1)}&\mathbf{0}_{(d+1) \times d}&\mathbf{0}_{d+1}\\ \mathbf{0}_{d\times(d+1)}&\mathbf{M}^{Q,(k)}_{\pm 1}&\mathbf{0}_{d}\\ \mathbf{0}_{1\times(d+1)}&\mathbf{0}_{1\times d}&-B\end{bmatrix},$ $\displaystyle\mathbf{Q}_{\pm 2}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{d\times(2d+1)}&\mathbf{0}_{d}\\ \mathbf{0}_{1\times(2d+1)}&m_{\pm 2}^{Q,(k)}\\ \mathbf{0}_{(d+1)\times(2d+1)}&\mathbf{0}_{d+1}\end{bmatrix}$ $\displaystyle\mathbf{K}_{\pm 1}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{(d+1)\times d}&\mathbf{0}_{(d+1) \times(d+1)}&\mathbf{0}_{d+1}\\ \mathbf{I}_{d\times d}&\mathbf{0}_{d\times(d+1)}&\mathbf{0}_{d}\\ \mathbf{0}_{1\times d}&\mathbf{0}_{1\times(d+1)}&{1}\end{bmatrix},$ $\displaystyle\mathbf{K}_{\pm 2}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{d\times d}&\mathbf{0}_{d}&\mathbf{0}_ {d\times d+1}\\ \mathbf{0}_{1\times d}&1&\mathbf{0}_{1\times(d+1)}\\ \mathbf{0}_{(d+1)\times d}&\mathbf{0}_{d+1}&\mathbf{0}_{(d+1)\times(d+1)}\end{bmatrix}$ $\displaystyle\mathbf{V}_{\pm 1}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{(d+1)\times d}&\mathbf{0}_{(d+1) \times(d+2)}\\ \mathbf{M}^{V,(k)}_{\pm 1}&\mathbf{0}_{d\times(d+2)}\\ \mathbf{0}_{{1}\times d}&\mathbf{0}_{{1}\times(d+2)}\end{bmatrix},$ $\displaystyle\mathbf{V}_{\pm 2}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{(d+1)\times d}&\mathbf{0}_{(d+1) \times d+2}\\ \mathbf{M}^{V,(k)}_{\pm 2}&\mathbf{0}_{d\times(d+2)}\\ \mathbf{0}_{{1}\times d}&\mathbf{0}_{{1}\times(d+2)}\end{bmatrix},$ (C.1) $$

where $\mathbf{M}^{Q,(k)}_{+1}$, $\mathbf{M}^{Q,(k)}_{-1}$, $\mathbf{M}^{V,(k)}_{+1}$, $\mathbf{M}^{V,(k)}_{-1}$, $\mathbf{M}^{V,(k)}_{+2}$, and $\mathbf{M}^{V,(k)}_{-2}$ are all $d\times d$ matrices, and $m_{+2}^{Q,(k)}$, $m_{-2}^{Q,(k)}$ are scalars.

#### MLP layer.

For the MLP layer following the $k$-th self-attention layer, we set

$$ $\displaystyle\mathbf{W}_{1}=\begin{bmatrix}\mathbf{W}_{1,\text{sub}}\\ -\mathbf{W}_{1,\text{sub}}\\ \mathbf{W}_{1,\text{sub}}\\ -\mathbf{W}_{1,\text{sub}}\end{bmatrix},\mathbf{W}_{2}^{\top}=\begin{bmatrix}- \mathbf{I}_{(2d+2)\times(2d+2)}\\ \mathbf{I}_{(2d+2)\times(2d+2)}\\ \mathbf{I}_{(2d+2)\times(2d+2)}\\ -\mathbf{I}_{(2d+2)\times(2d+2)}\end{bmatrix},\mathbf{b}^{(k)}=\begin{bmatrix} \mathbf{0}_{5d+5}\\ -\theta^{(k)}\cdot\mathbf{1}_{d}\\ \mathbf{0}_{d+2}\\ \theta^{(k)}\cdot\mathbf{1}_{d}\\ 0\end{bmatrix}.$ (C.2) $$

where the submatrix $\mathbf{W}_{1,\text{sub}}$ is defined as

$$ $\displaystyle\mathbf{W}_{1,\text{sub}}=\begin{bmatrix}\mathbf{0}_{(d+1)\times( d+1)}&\mathbf{0}_{(d+1)\times d}&\mathbf{0}_{d+1}\\ \mathbf{0}_{d\times(d+1)}&\mathbf{I}_{d\times d}&\mathbf{0}_{d}\\ \mathbf{0}_{1\times(d+1)}&\mathbf{0}_{1\times d}&0\end{bmatrix}.$ $$

Therefore, we obtain that the output of the MLP layer is

$$ $\displaystyle\mathrm{MLP}(\mathbf{h}_{i})=\begin{bmatrix}[\mathbf{h}_{i}]_{1:d +1}\\ \mathcal{S}_{\theta^{(k)}}([\mathbf{h}_{i}]_{d+2:2d+1})\\ [\mathbf{h}_{i}]_{2d+2}\end{bmatrix}.$ (C.3) $$

where $\mathcal{S}_{\theta^{(k)}}$ is the soft-thresholding function parameterized by $\theta^{(k)}$.

### C.2 Proof of Theorem 4.1

We start by stating an equivalent form of
Theorem 4.1 below, where we specific $\mathbf{M}^{(k)}$ in Theorem 4.1 to be $\gamma^{(k)}\mathbf{M}^{V}$.

###### Theorem C.1 (Equivalent form Theorem 4.1 ) .

Suppose Assumption 1 holds.
For a Transformer with $K$ layers as described in Section 4.3, set the input sequence as:

$$ $\displaystyle\mathbf{H}^{(1)}=\begin{bmatrix}[\mathbf{X}^{\top}]_{:,1}&[ \mathbf{X}^{\top}]_{:,1}&\cdots&[\mathbf{X}^{\top}]_{:,N}&[\mathbf{X}^{\top}]_ {:,N}&\mathbf{x}_{N+1}\\ 0&y_{1}&\cdots&0&y_{N}&0\\ \bm{\beta}^{(1)}_{1}&\bm{\beta}^{(1)}_{2}&\cdots&\bm{\beta}^{(1)}_{2N-1}&\bm{ \beta}^{(1)}_{2N}&\bm{\beta}^{(1)}_{2N+1}\\ 1&0&\cdots&1&0&1\end{bmatrix}.$ (C.4) $$

Denote $\mathbf{H}^{(k+1)}$ as the output of the $K$-th layer of the Transformer and define $\bm{\beta}_{2n+1}^{(K+1)}=\mathbf{H}^{(k+1)}_{d+2:2d+1,2n+1}$.
There exists a set of parameters within the Transformer such that for all $k\in[1,K]$, we have:

$$ $\displaystyle\bm{\beta}_{2n+1}^{(k+1)}=\mathcal{S}_{\theta^{(k)}}\Big{(}\bm{ \beta}^{(k)}_{2n+1}-\gamma^{(k)}(\mathbf{D}_{n})^{\top}([\mathbf{X}]_{1:n,:} \bm{\beta}^{(k)}_{2n+1}-\mathbf{y}_{1:n})\Big{)},$ $$

where $\mathbf{D}_{n}=\frac{1}{2n+1}[\mathbf{X}]_{1:n,:}(\mathbf{M}^{V})^{\top}$ and $\mathbf{M}^{V}\in\mathbb{R}^{d\times d}$ is embedded in the $k$-th Transformer lsyer.

Proving the theorem is equivalent to demonstrating the existence of a Transformer for which the following proposition holds for any $k\geq 2$:

###### Proposition 1 .

Suppose Assumption 1 holds.
For a $K$ layers Transformer with structure described in Section 4.3, the input sequence of the $k$-th layer of the Transformer satisfies:

$$ $\displaystyle\mathbf{H}^{(k)}=\begin{bmatrix}[\mathbf{X}^{\top}]_{:,1}&[ \mathbf{X}^{\top}]_{:,1}&\cdots&[\mathbf{X}^{\top}]_{:,N}&[\mathbf{X}^{\top}]_ {:,N}&\mathbf{x}_{N+1}\\ 0&y_{1}&\cdots&0&y_{N}&0\\ \bm{\beta}^{(k)}_{1}&\bm{\beta}^{(k)}_{2}&\cdots&\bm{\beta}^{(k)}_{2N-1}&\bm{ \beta}^{(k)}_{2N}&\bm{\beta}^{(k)}_{2N+1}\\ 1&0&\cdots&1&0&1\end{bmatrix},$ $$

where it holds that $\bm{\beta}_{2n+1}^{(k)}=\mathcal{S}_{\theta^{(k-1)}}\Big{(}\bm{\beta}^{(k-1)}_
{2n+1}-\gamma^{(k-1)}(\mathbf{D}_{n})^{\top}([\mathbf{X}]_{1:n,:}\bm{\beta}^{(
k)}_{2n+1}-\mathbf{y}_{1:n})\Big{)}$, $\mathbf{D}_{n}=\frac{1}{2n+1}[\mathbf{X}]_{1:n,:}(\mathbf{M}^{V})^{\top}$ and $|\bm{\beta}_{2n+1}^{(k-1)}|\leq C_{\bm{\beta}}^{k}$ for Constant $C_{\bm{\beta}}$.

###### Proof of Proposition 1 .

We prove Proposition 1 is true for all $k\geq 2$ by induction. First, $\mathbf{H}^{(k)}$ for $k=2$ satisfies the condition automatically; therefore, Proposition 1 is true when $k=2$. Then, we demonstrate that if Proposition 1 is true for $k-1$, it remains valid for $k$. For odd values of $i$, the token-wise outputs of the $k$-th attention layer corresponding to the first and second heads satisfy:

$$ $\displaystyle\mathbf{Q}_{\pm 1}^{(k)}\mathbf{h}_{i}^{(k)}=\begin{bmatrix} \mathbf{0}_{d+1}\\ \mathbf{M}^{Q,(k)}_{\pm 1}\bm{\beta}_{i}^{(k)}\\ -B\end{bmatrix};\quad\mathbf{K}_{\pm 1}^{(k)}\mathbf{h}_{i}^{(k)}=\begin{ bmatrix}\mathbf{0}_{d+1}\\ [\mathbf{X}^{\top}]_{:,\lfloor\frac{i+1}{2}\rfloor}\\ 1\end{bmatrix};\quad\mathbf{V}_{\pm 1}^{(k)}\mathbf{h}_{i}^{(k)}=\begin{ bmatrix}\mathbf{0}_{d+1}\\ \mathbf{M}^{V,(k)}_{\pm 1}[\mathbf{X}^{\top}]_{:,\lfloor\frac{i+1}{2}\rfloor} \\ 0\end{bmatrix}.$ $$

The token-wise outputs of the third and fourth heads for odd $i$ satisfy

$$ $\displaystyle\mathbf{Q}_{\pm 2}^{(k)}\mathbf{h}_{i}^{(k)}=\begin{bmatrix} \mathbf{0}_{d}\\ m^{Q,(k)}_{\pm 2}\\ \mathbf{0}_{d+1}\end{bmatrix};\quad\mathbf{K}_{\pm 2}^{(k)}\mathbf{h}_{i}^{(k) }=\mathbf{0}_{2d+2};\quad\mathbf{V}_{\pm 2}^{(k)}\mathbf{h}_{i}^{(k)}=\begin{ bmatrix}\mathbf{0}_{d+1}\\ \mathbf{M}^{V,(k)}_{\pm 2}[\mathbf{X}^{\top}]_{:,\lfloor\frac{i+1}{2}\rfloor} \\ 0\end{bmatrix}.$ $$

Besides, for any $i$ that is an even number, the token-wise outputs of the $k$-th attention layer corresponding to the first and second heads satisfy:

$$ $\displaystyle\mathbf{Q}_{\pm 1}^{(k)}\mathbf{h}_{i}^{(k)}=\begin{bmatrix} \mathbf{0}_{d+1}\\ \mathbf{M}^{Q,(k)}_{\pm 1}\bm{\beta}_{i}^{(k)}\\ 0\end{bmatrix};\quad\mathbf{K}_{\pm 1}^{(k)}\mathbf{h}_{i}^{(k)}=\begin{ bmatrix}\mathbf{0}_{d+1}\\ [\mathbf{X}^{\top}]_{:,\lfloor\frac{i+1}{2}\rfloor}\\ 0\end{bmatrix};\quad\mathbf{V}_{\pm 1}^{(k)}\mathbf{h}_{i}^{(k)}=\begin{ bmatrix}\mathbf{0}_{d+1}\\ \mathbf{M}^{V,(k)}_{\pm 1}[\mathbf{X}^{\top}]_{:,\lfloor\frac{i+1}{2}\rfloor} \\ 0\end{bmatrix}.$ $$

The token-wise outputs of the third and fourth heads for even $i$ satisfy

$$ $\displaystyle\mathbf{Q}_{\pm 2}^{(k)}\mathbf{h}_{i}^{(k)}=\mathbf{0}_{2d+2}; \quad\mathbf{K}_{\pm 2}^{(k)}\mathbf{h}_{i}^{(k)}=\begin{bmatrix}\mathbf{0}_{d }\\ y_{\frac{i}{2}}\\ \mathbf{0}_{d+1}\end{bmatrix};\quad\mathbf{V}_{\pm 2}^{(k)}\mathbf{h}_{i}^{(k) }=\begin{bmatrix}\mathbf{0}_{d+1}\\ \mathbf{M}^{V,(k)}_{\pm 2}[\mathbf{X}^{\top}]_{:,\lfloor\frac{i+1}{2}\rfloor} \\ 0\end{bmatrix}.$ $$

Therefore, for $\mathbf{h}_{i}$ with an odd index and head with index $u\in\{1,-1\}$, we have

$$ $\displaystyle\sigma\left(\Big{\langle}\mathbf{Q}^{(k)}_{u}\mathbf{h}_{i}^{k}, \mathbf{K}^{(k)}_{u}\mathbf{h}_{j}^{k}\Big{\rangle}\right)\cdot\mathbf{V}_{u}^ {(k)}\mathbf{h}^{(k)}_{j}$ $\displaystyle=\sigma\left((\bm{\beta}^{(k)}_{i})^{\top}(\mathbf{M}_{u}^{Q})^{ \top}[\mathbf{X}^{\top}]_{:,\lfloor\frac{j+1}{2}\rfloor}-\mathbbm{1}_{\{j\%2=1 \}}(j)B\right)\cdot\begin{bmatrix}\mathbf{0}_{d+1}\\ \mathbf{M}^{V,(k)}_{u}[\mathbf{X}^{\top}]_{:,\lfloor\frac{j+1}{2}\rfloor}\\ 0\end{bmatrix}.$ (C.5) $$

Also, for $\mathbf{h}_{i}$ with an odd index and head with index $u\in\{2,-2\}$, we have

$$ $\displaystyle\sigma\left(\Big{\langle}\mathbf{Q}^{(k)}_{u}\mathbf{h}_{i}^{k}, \mathbf{K}^{(k)}_{u}\mathbf{h}_{j}^{k}\Big{\rangle}\right)\cdot\mathbf{V}_{u}^ {(k)}\mathbf{h}^{(k)}_{j}=\mathbbm{1}_{{\{j\%2=0\}}}\cdot\sigma\left(m_{u}^{Q, (k)}\cdot y_{\frac{j}{2}}\right)\cdot\begin{bmatrix}\mathbf{0}_{d+1}\\ \mathbf{M}^{V,(k)}_{u}[\mathbf{X}^{\top}]_{:,\lfloor\frac{j+1}{2}\rfloor}\\ 0\end{bmatrix}.$ (C.6) $$

We specify the parameters of the self-attention layer as

$$ $\displaystyle\mathbf{M}_{1}^{Q,(k)}=-\mathbf{I}_{d\times d},\quad\mathbf{M}_{1 }^{V,(k)}=\gamma^{(k)}\mathbf{M}^{V},$ $\displaystyle\mathbf{M}_{-1}^{Q,(k)}=\mathbf{I}_{d\times d},\quad\mathbf{M}_{- 1}^{V,(k)}=-\gamma^{(k)}\mathbf{M}^{V},$ $\displaystyle m_{2}^{q,(k)}=1,\quad\mathbf{M}_{2}^{V,(k)}=\gamma^{(k)}\mathbf{ M}^{V},$ $\displaystyle m_{-2}^{q,(k)}=-1,\quad\mathbf{M}_{-2}^{V,(k)}=-\gamma^{(k)} \mathbf{M}^{V}.$ $$

Then, for head index $u\in\{1,-1\}$ we have

$$ $\displaystyle(\bm{\beta}^{(k)}_{i})^{\top}(\mathbf{M}_{u}^{Q})^{\top}[\mathbf{ X}^{\top}]_{:,\lfloor\frac{j+1}{2}\rfloor}-B$ $\displaystyle\leq\big{|}(\bm{\beta}^{(k)}_{i})^{\top}[\mathbf{X}^{\top}]_{:, \lfloor\frac{j+1}{2}\rfloor}\big{|}-B$ $\displaystyle\leq\big{\|}\bm{\beta}^{(k)}_{i}\big{\|}_{1}\big{\|}[\mathbf{X}^{ \top}]_{:,\lfloor\frac{j+1}{2}\rfloor}\big{\|}_{\infty}-B$ (C.7) $\displaystyle\leq\big{\|}\bm{\beta}^{(k)}_{i}-\bm{\beta}^{*}\big{\|}_{1}\big{ \|}[\mathbf{X}^{\top}]_{:,\lfloor\frac{j+1}{2}\rfloor}\big{\|}_{\infty}+b_{\bm {\beta}}b_{\mathbf{x}}-B,$ $$

where Equation C.7 is given by Hölder’s Inequality. Therefore, combining this with the assumption $\|\mathbf{x}\|_{\infty}\leq b_{\mathbf{x}}$ and $\|\bm{\beta}^{*}\|_{1}\leq b_{\bm{\beta}}$, we obtain

$$ $\displaystyle(\bm{\beta}^{(k)}_{i})^{\top}(\mathbf{M}_{u}^{Q})^{\top}[\mathbf{ X}^{\top}]_{:,\lfloor\frac{j+1}{2}\rfloor}-B$ $\displaystyle\leq b_{\mathbf{x}}\big{\|}\bm{\beta}^{(k)}_{i}-\bm{\beta}^{*} \big{\|}_{1}+b_{\bm{\beta}}b_{\mathbf{x}}-B.$ (C.8) $$

Recall that for odd $i$, we assume that Proposition 1 is true for $k-1$, then $\bm{\beta}_{i}^{(k)}$ satisfies

$$ $\displaystyle\bm{\beta}_{i}^{(k)}=\mathcal{S}_{\theta^{(k-1)}}\Big{(}\bm{\beta }^{(k-1)}_{i}-\gamma^{(k-1)}(\mathbf{D}_{n})^{\top}([\mathbf{X}]_{1:\lfloor \frac{i+1}{2}\rfloor,:}\bm{\beta}^{(k)}_{i}-\mathbf{y}_{1:\lfloor\frac{i+1}{2} \rfloor})\Big{)}$ $$

and $|\bm{\beta}_{2n+1}^{(k-1)}|\leq C_{\bm{\beta}}^{k}$. Therefore, we have

$$ $\displaystyle|\bm{\beta}_{i}^{(k)}|$ $\displaystyle\leq|\bm{\beta}_{i}^{(k-1)}|+|\gamma^{(k-1)}|\|\mathbf{M}^{V}\| \left\|\frac{1}{i}[\mathbf{X}]_{1:\lfloor\frac{i+1}{2}\rfloor,:}^{\top}[ \mathbf{X}]_{1:\lfloor\frac{i+1}{2}\rfloor,:}\right\||\bm{\beta}_{i}^{(k-1)}|+ b_{\bm{\beta}}|\gamma^{(k-1)}|\|\mathbf{M}^{V}\|\left\|\frac{1}{i}[\mathbf{X}] _{1:\lfloor\frac{i+1}{2}\rfloor,:}\right\|$ $\displaystyle\leq C_{\bm{\beta}}^{k}+b_{\mathbf{x}}^{2}d|\gamma^{(k-1)}|\| \mathbf{M}^{V}\|\frac{\lfloor\frac{i+1}{2}\rfloor}{i}C_{\bm{\beta}}^{k}+b_{\bm {\beta}}b_{\mathbf{x}}|\gamma^{(k-1)}|\|\mathbf{M}^{V}\|\frac{\lfloor\frac{i+1 }{2}\rfloor}{i}$ $\displaystyle\leq C_{\bm{\beta}}^{k}+b_{\mathbf{x}}^{2}d|\gamma^{(k-1)}|\| \mathbf{M}^{V}\|C_{\bm{\beta}}^{k}+b_{\bm{\beta}}b_{\mathbf{x}}|\gamma^{(k-1)} |\|\mathbf{M}^{V}\|.$ (C.9) $$

Equation C.9 arises from the fact that $\lfloor\frac{i+1}{2}\rfloor/i\leq 1$. Therefore, denoting $\gamma_{\max}=\max_{k}\|\gamma^{(k)}\|$ and letting $C_{\bm{\beta}}=1+b_{\mathbf{x}}^{2}d\gamma_{\max}\|\mathbf{M}^{V}\|+b_{\bm{
\beta}}b_{\mathbf{x}}\gamma_{\max}\|\mathbf{M}^{V}\|$, we have

$$ $\displaystyle|\bm{\beta}_{i}^{(k)}|\leq C_{\bm{\beta}}^{k}+b_{\mathbf{x}}^{2}d |\gamma^{(k-1)}|\|\mathbf{M}^{V}\|C_{\bm{\beta}}^{k}+b_{\bm{\beta}}b_{\mathbf{ x}}|\gamma^{(k-1)}|\|\mathbf{M}^{V}\|\leq C_{\bm{\beta}}^{k+1}.$ $$

Therefore we have

$$ $\displaystyle b_{\mathbf{x}}\big{\|}\bm{\beta}^{(k)}_{i}-\bm{\beta}^{*}\big{\| }_{1}\leq b_{\mathbf{x}}\sqrt{d}(b_{\bm{\beta}}+C_{\bm{\beta}}^{k}).$ $$

Then, by setting $B\geq b_{\bm{\beta}}b_{\mathbf{x}}+b_{\mathbf{x}}\sqrt{d}(b_{\bm{\beta}}+C_{
\bm{\beta}}^{K})$, we obtain that $B>\big{|}(\bm{\beta}^{(k)}_{i})^{\top}[\mathbf{X}^{\top}]_{:,\lfloor\frac{j+1}
{2}\rfloor}\big{|}$ for all $k\leq K$ and any odd $i$. Then, combining with Equation C.8 we have

$$ $\displaystyle\sigma\left((\bm{\beta}^{(k)}_{i})^{\top}(\mathbf{M}_{u}^{Q})^{ \top}[\mathbf{X}^{\top}]_{:,\lfloor\frac{j+1}{2}\rfloor}-B\right)=0.$ $$

Further, from Equation C.5, for $i$ is odd number and $j\in[1,3,\cdots,i]$ we have

$$ $\displaystyle\sigma\left(\Big{\langle}\mathbf{Q}^{(k)}_{u}\mathbf{h}_{i}^{k}, \mathbf{K}^{(k)}_{u}\mathbf{h}_{j}^{k}\Big{\rangle}\right)\cdot\mathbf{V}_{u}^ {(k)}\mathbf{h}^{(k)}_{j}$ $\displaystyle=\mathbf{0}_{2d+2}.$ $$

Besides, for $i$ being an odd number and $j\in[2,4,\cdots,i-1]$, we have

$$ $\displaystyle\sigma\left(\Big{\langle}\mathbf{Q}^{(k)}_{u}\mathbf{h}_{i}^{k}, \mathbf{K}^{(k)}_{u}\mathbf{h}_{j}^{k}\Big{\rangle}\right)\cdot\mathbf{V}_{u}^ {(k)}\mathbf{h}^{(k)}_{j}$ $\displaystyle=\sigma\left((\bm{\beta}^{(k)}_{i})^{\top}(\mathbf{M}_{u}^{Q})^{ \top}[\mathbf{X}^{\top}]_{:,\lfloor\frac{i+1}{2}\rfloor}\right)\cdot\begin{ bmatrix}\mathbf{0}_{d+1}\\ \mathbf{M}^{V,(k)}_{u}[\mathbf{X}^{\top}]_{:,\lfloor\frac{j+1}{2}\rfloor}\\ 2\end{bmatrix}.$ $$

Thus, when $h_{i}$ is of odd index, summing over all $j\leq i$ and $u\in\{1,-1\}$, we obtain:

$$ $\displaystyle\sum_{u\in\{-1,+1\}}\sum_{j=1}^{i}\sigma\left(\Big{\langle} \mathbf{Q}^{(k)}_{u}\mathbf{h}_{i}^{k},\mathbf{K}^{(k)}_{u}\mathbf{h}_{j}^{k} \Big{\rangle}\right)\cdot\mathbf{V}_{u}^{(k)}\mathbf{h}^{(k)}_{j}$ $\displaystyle=\sum_{u\in\{-1,+1\}}\sum_{j=2}^{i-1}\sigma\left(\Big{\langle} \mathbf{Q}^{(k)}_{u}\mathbf{h}_{i}^{k},\mathbf{K}^{(k)}_{u}\mathbf{h}_{j}^{k} \Big{\rangle}\right)\cdot\mathbf{V}_{u}^{(k)}\mathbf{h}^{(k)}_{j}$ $\displaystyle=\sum_{j=2}^{i-1}-(\bm{\beta}^{(k)}_{i})^{\top}[\mathbf{X}^{\top} ]_{:,\lfloor\frac{j+1}{2}\rfloor}\cdot\begin{bmatrix}\mathbf{0}_{d+1}\\ \gamma^{(k)}\mathbf{M}^{V}[\mathbf{X}^{\top}]_{:,\lfloor\frac{j+1}{2}\rfloor} \\ 0\end{bmatrix}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{d+1}\\ -\gamma^{(k)}\mathbf{M}^{V}[\mathbf{X}^{\top}]_{:,1:\frac{i-1}{2}}[\mathbf{X}] _{1:\frac{i-1}{2},:}\bm{\beta}^{(k)}_{i}\\ 0\end{bmatrix}.$ (C.10) $$

Next, we consider heads with indexes $m\in\{2,-2\}$. From Equation C.6 we obtain

$$ $\displaystyle\sum_{u\in\{-2,+2\}}\sum_{j=1}^{i}\sigma\left(\Big{\langle} \mathbf{Q}^{(k)}_{u}\mathbf{h}_{i}^{k},\mathbf{K}^{(k)}_{u}\mathbf{h}_{j}^{k} \Big{\rangle}\right)\cdot\mathbf{V}_{u}^{(k)}\mathbf{h}^{(k)}_{j}$ $\displaystyle=\sum_{u\in\{-2,+2\}}\sum_{j=2}^{i-1}\sigma\left(\Big{\langle} \mathbf{Q}^{(k)}_{u}\mathbf{h}_{i}^{k},\mathbf{K}^{(k)}_{u}\mathbf{h}_{j}^{k} \Big{\rangle}\right)\cdot\mathbf{V}_{u}^{(k)}\mathbf{h}^{(k)}_{j}$ $\displaystyle=\sum_{j=2}^{i-1}y_{\frac{j}{2}}\cdot\begin{bmatrix}\mathbf{0}_{d +1}\\ \gamma^{(k)}\mathbf{M}^{V}[\mathbf{X}^{\top}]_{:,\frac{j}{2}}\\ \mathbf{0}_{2}\end{bmatrix}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{d+1}\\ \gamma^{(k)}\mathbf{M}^{V}[\mathbf{X}^{\top}]_{:,1:\frac{i-1}{2}}\mathbf{y}_{1 :\frac{i-1}{2}}\\ \mathbf{0}_{2}\end{bmatrix}.$ (C.11) $$

Combining Equation C.10 and Equation C.11, we obtain the following equation for $\mathbf{h}_{i}$ with an odd index:

$$ $\displaystyle\mathbf{h}_{i}^{(k+1)}$ $\displaystyle=\mathrm{MLP}\left(\mathbf{h}_{i}^{(k)}+\frac{1}{i}\sum_{u\in\{ \pm 1,\pm 2\}}\sum_{j=1}^{i}\sigma\left(\Big{\langle}\mathbf{Q}^{(k)}_{u} \mathbf{h}_{i}^{k},\mathbf{K}^{(k)}_{u}\mathbf{h}_{j}^{k}\Big{\rangle}\right) \cdot\mathbf{V}_{u}^{(k)}\mathbf{h}^{(k)}_{j}\right)=\begin{bmatrix}[\mathbf{X }^{\top}]_{:,\lfloor\frac{i+1}{2}\rfloor}\\ 0\\ \bm{\beta}^{k+1}_{i}\\ 1\end{bmatrix},$ $$

where

$$ $\displaystyle\bm{\beta}^{k+1}_{i}$ $\displaystyle=\mathcal{S}_{\theta^{(k)}}\left(\bm{\beta}_{i}^{(k)}-\frac{ \gamma^{(k)}}{i}\mathbf{M}^{V}[\mathbf{X}^{\top}]_{:,1:\frac{i-1}{2}}[\mathbf{ X}]_{1:\frac{i-1}{2},:}\bm{\beta}^{(k)}_{i}+\gamma^{(k)}\mathbf{M}^{V}[\mathbf {X}^{\top}]_{:,1:\frac{i-1}{2}}\mathbf{y}_{1:\frac{i-1}{2}}\right)$ $\displaystyle=\mathcal{S}_{\theta^{(k)}}\left(\bm{\beta}_{i}^{(k)}-\frac{ \gamma^{(k)}}{i}\mathbf{M}^{V}[\mathbf{X}^{\top}]_{:,1:\frac{i-1}{2}}\left([ \mathbf{X}]_{1:\frac{i-1}{2},:}\bm{\beta}^{(k)}_{i}-\mathbf{y}_{1:\frac{i-1}{2 }}\right)\right).$ (C.12) $$

Similarly, we obtain the following equation for $\mathbf{h}_{i}$ with an even index:

$$ $\displaystyle\mathbf{h}_{i}^{(k+1)}$ $\displaystyle=\mathrm{MLP}\left(\mathbf{h}_{i}^{(k+1)}+\frac{1}{i}\sum_{u\in\{ \pm 1,\pm 2\}}\sum_{j=1}^{i}\sigma\left(\Big{\langle}\mathbf{Q}^{(k)}_{u} \mathbf{h}_{i}^{k},\mathbf{K}^{(k)}_{u}\mathbf{h}_{j}^{k}\Big{\rangle}\right) \cdot\mathbf{V}_{u}^{(k)}\mathbf{h}^{(k)}_{j}\right)=\begin{bmatrix}[\mathbf{X }^{\top}]_{:,\lfloor\frac{i+1}{2}\rfloor}\\ 0\\ \bm{\beta}^{k+1}_{i}\\ 0\end{bmatrix}.$ (C.13) $$

Note that an explicit formulation of $\bm{\beta}_{i}^{(k+1)}$ is not required when $i$ is even. Next, combining Equation C.12 and Equation C.13 gives

$$ $\displaystyle\mathbf{H}^{(k+1)}=\begin{bmatrix}[\mathbf{X}^{\top}]_{:,1}&[ \mathbf{X}^{\top}]_{:,1}&\cdots&[\mathbf{X}^{\top}]_{:,N}&[\mathbf{X}^{\top}]_ {:,N}&\mathbf{x}_{N+1}\\ 0&y_{1}&\cdots&0&y_{N}&0\\ \bm{\beta}^{(k+1)}_{1}&\bm{\beta}^{(k+1)}_{2}&\cdots&\bm{\beta}^{(k+1)}_{2N-1} &\bm{\beta}^{(k+1)}_{2N}&\bm{\beta}^{(k+1)}_{2N+1}\\ 1&0&\cdots&1&0&1\end{bmatrix}.$ $$

The proof for Proposition 1 is thus complete.
∎

###### Theorem C.1 (Equivalent form Theorem 4.1 ) .

###### Proposition 1 .

###### Proof of Proposition 1 .

## Appendix D Deferred Proofs in Section 5.1

### D.1 Proof of Theorem 5.1

In Section 4.1, we assume that each row of the measurement matrix $\mathbf{X}$ is independently sampled from an isometric sub-Gaussian distribution with zero mean and covariance $\mathrm{diag}(\sigma_{1}^{2},\cdots,\sigma_{d}^{2})$. Without loss of generality, let $\sigma_{1}\geq\sigma_{2}\geq\cdots\geq\sigma_{d}$. We start the proof of Theorem 5.1 by introducing the auxiliary lemmas: Lemma 1, Lemma 2, and Lemma 3.

###### Lemma 1 .

Assume $\bm{\beta}^{(k)}$ and $\bm{\beta}^{(k+1)}$ satisfy:

$$ $\displaystyle\bm{\beta}^{(k+1)}$ $\displaystyle=\mathcal{S}_{\theta^{(k)}}\Big{(}\bm{\beta}^{(k)}-\gamma\mathbf{ D}^{\top}(\mathbf{X}\bm{\beta}^{(k)}-\mathbf{y})\Big{)}.$ (D.1) $$

Define $\sigma_{\min}=\min_{i\in\mathbb{S}}\left|\mathbf{D}_{:,i}^{\top}\mathbf{X}_{:,
i}\right|$, $\sigma_{\max}=\max_{i\neq j}\left|\mathbf{D}_{:,i}^{\top}\mathbf{X}_{:,j}\right|$ and let $\theta^{(k)}=\gamma\sigma_{\max}c_{1}e^{-c_{2}(k-2)}$, where $c_{1}=\|\bm{\beta}^{*}\|_{1}$ and $c_{2}=\log\left(\frac{1}{\gamma(2S-1)\sigma_{\max}+\left|1-\gamma\sigma_{\min}
\right|}\right)$. Then, if $\|\bm{\beta}^{(k)}-\bm{\beta}^{*}\|_{1}\leq c_{1}e^{-c_{2}(k-2)}$ and $\mathbf{D}$, $\gamma$, $\theta^{(k)}$ satisfy the following conditions:

$$ $\begin{gathered}\gamma(2S-1)\sigma_{\max}+\left|1-\gamma\sigma_{\min}\right| \leq 1,\\ 0\leq\gamma\mathbf{D}_{:,i}^{\top}\mathbf{X}_{:,i}\leq 1\quad\forall i\in[d], \end{gathered}$ (D.2) $$

then $\bm{\beta}^{(k+1)}$ satisfies

$$ $\displaystyle\|\bm{\beta}^{(k+1)}-\bm{\beta}^{*}\|_{1}\leq c_{1}e^{-c_{2}(k-1)}.$ $$

###### Proof of Lemma 1 .

Note that we are considering the noiseless model: $\mathbf{y}=\mathbf{X}\bm{\beta}^{*}$. We begin by rewriting
Equation D.1:

$$ $\displaystyle\bm{\beta}^{(k+1)}$ $\displaystyle=\mathcal{S}_{\theta^{(k)}}\Big{(}\bm{\beta}^{(k)}-\gamma^{(k)} \mathbf{D}^{\top}(\mathbf{X}\bm{\beta}^{(k)}-\mathbf{y})\Big{)}$ $\displaystyle=\mathcal{S}_{\theta^{(k)}}\Big{(}\bm{\beta}^{(k)}-\gamma^{(k)} \mathbf{D}^{\top}\mathbf{X}(\bm{\beta}^{(k)}-\bm{\beta}^{*})\Big{)}.$ $$

Consider entries of $\bm{\beta}^{(K+1)}$ in the support of $\bm{\beta}^{*}$:

$$ $\displaystyle\bm{\beta}^{(k+1)}_{\mathbb{S}}$ $\displaystyle=\mathcal{S}_{\theta^{(k)}}\Big{(}\bm{\beta}^{(k)}_{\mathbb{S}}- \gamma^{(k)}(\mathbf{D}_{:,\mathbb{S}})^{\top}\mathbf{X}_{:,\mathbb{S}}(\bm{ \beta}^{(k)}_{\mathbb{S}}-\bm{\beta}^{*}_{\mathbb{S}})\Big{)}$ $\displaystyle\in\bm{\beta}^{(k)}_{\mathbb{S}}-\gamma^{(k)}(\mathbf{D}_{:, \mathbb{S}})^{\top}\mathbf{X}_{:,\mathbb{S}}(\bm{\beta}^{(k)}_{\mathbb{S}}-\bm {\beta}^{*}_{\mathbb{S}})-\theta^{(k)}\partial\ell_{1}(\bm{\beta}_{\mathbb{S}} ^{(k+1)}),$ $$

where $\partial\ell_{1}(\bm{\beta})$ is the sub-gradient of $|\bm{\beta}|$:

$$ $[\partial\ell_{1}(\bm{\beta})]_{i}=\begin{cases}\{\text{sign}([\bm{\beta}]_{i} )\}&\text{if }[\bm{\beta}]_{i}\neq 0,\\ {[-1,1]}&\text{if }[\bm{\beta}]_{i}=0.\end{cases}$ $$

Then, for any $i\in\mathbb{S}$, we have

$$ $\displaystyle\bm{\beta}^{(k+1)}_{i}\in\bm{\beta}^{(k)}_{i}-\gamma^{(k)}( \mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,\mathbb{S}}(\bm{\beta}^{(k)}_{\mathbb{S} }-\bm{\beta}^{*}_{\mathbb{S}})-\theta^{(k)}\partial\ell_{1}(\bm{\beta}_{ \mathbb{S}}^{(k+1)}).$ $$

Note that $\bm{\beta}^{(k)}_{i}-\gamma^{(k)}(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,
\mathbb{S}}(\bm{\beta}^{(k)}_{\mathbb{S}}-\bm{\beta}^{*}_{\mathbb{S}})$ can be rewritten as

$$ $\displaystyle\bm{\beta}^{(k)}_{i}-\gamma^{(k)}(\mathbf{D}_{:,i})^{\top}\mathbf {X}_{:,\mathbb{S}}(\bm{\beta}^{(k)}_{\mathbb{S}}-\bm{\beta}^{*}_{\mathbb{S}})$ $\displaystyle=\bm{\beta}^{(k)}_{i}-\gamma^{(k)}\sum_{j\in\mathbb{S},j\neq i}( \mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,j}(\bm{\beta}^{(k)}_{j}-\bm{\beta}^{*}_{ j})-\gamma^{(k)}(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,i}(\bm{\beta}^{(k)}_{i} -\bm{\beta}^{*}_{i})$ $\displaystyle=\bm{\beta}^{*}_{i}-\gamma^{(k)}\sum_{j\in\mathbb{S},j\neq i}( \mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,j}(\bm{\beta}^{(k)}_{j}-\bm{\beta}^{*}_{ j})+\big{(}1-\gamma^{(k)}(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,i}\big{)}(\bm{ \beta}^{(k)}_{i}-\bm{\beta}^{*}_{i}).$ $$

Therefore,

$$ $\displaystyle\bm{\beta}^{(k+1)}_{i}-\bm{\beta}^{*}_{i}\in-\gamma^{(k)}\sum_{j \in\mathbb{S},j\neq i}(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,j}(\bm{\beta}^{(k )}_{j}-\bm{\beta}^{*}_{j})+\big{(}1-\gamma^{(k)}(\mathbf{D}_{:,i})^{\top} \mathbf{X}_{:,i}\big{)}(\bm{\beta}^{(k)}_{i}-\bm{\beta}^{*}_{i})-\theta^{(k)} \partial\ell_{1}(\bm{\beta}_{i}^{(k+1)}).$ $$

Based on the definition of $\partial\ell_{1}$, we derive the following inequality:

$$ $\displaystyle|\bm{\beta}^{(k+1)}_{i}-\bm{\beta}^{*}_{i}|\leq\gamma^{(k)}\sum_{ j\in\mathbb{S},j\neq i}\left|(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,j}\right|| \bm{\beta}^{(k)}_{j}-\bm{\beta}^{*}_{j}|+\left|1-\gamma^{(k)}(\mathbf{D}_{:,i} )^{\top}\mathbf{X}_{:,i}\right||\bm{\beta}^{(k)}_{i}-\bm{\beta}^{*}_{i}|+ \theta^{(k)}.$ $$

Recall that $0<\gamma^{(k)}(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,i}<1$, we obtain that

$$ $\displaystyle|\bm{\beta}^{(k+1)}_{i}-\bm{\beta}^{*}_{i}|\leq\gamma^{(k)}\sigma _{\max}\sum_{j\in\mathbb{S},j\neq i}|\bm{\beta}^{(k)}_{j}-\bm{\beta}^{*}_{j}|+ \left|1-\gamma^{(k)}\sigma_{\min}\right||\bm{\beta}^{(k)}_{i}-\bm{\beta}^{*}_{ i}|+\theta^{(k)}.$ $$

From Lemma 2, we have $\mathrm{support}(\bm{\beta}^{(k+1)})\in\mathbb{S}$, thus, $\|\bm{\beta}^{(k+1)}-\bm{\beta}^{*}\|_{1}=\|\bm{\beta}^{(k+1)}_{\mathbb{S}}-
\bm{\beta}^{*}_{\mathbb{S}}\|_{1}$.
Then,

$$ $\displaystyle\|\bm{\beta}^{(k+1)}-\bm{\beta}^{*}\|_{1}$ $\displaystyle\leq\sum_{i\in\mathbb{S}}\left(\gamma^{(k)}\sigma_{\max}\sum_{j \in\mathbb{S},j\neq i}|\bm{\beta}^{(k)}_{j}-\bm{\beta}^{*}_{j}|+\left|1-\gamma ^{(k)}\sigma_{\min}\right||\bm{\beta}^{(k)}_{i}-\bm{\beta}^{*}_{i}|+\theta^{(k )}\right)$ $\displaystyle\leq\gamma^{(k)}(S-1)\sigma_{\max}\sum_{i\in\mathbb{S}}|\bm{\beta }^{(k)}_{i}-\bm{\beta}^{*}_{i}|+\left|1-\gamma^{(k)}\sigma_{\min}\right|\left \|\bm{\beta}^{(k)}_{i}-\bm{\beta}^{*}_{i}\right\|_{1}+S\theta^{(k)}$ $\displaystyle=\left(\gamma^{(k)}(S-1)\sigma_{\max}+\left|1-\gamma^{(k)}\sigma_ {\min}\right|\right)\left\|\bm{\beta}^{(k)}_{i}-\bm{\beta}^{*}_{i}\right\|_{1} +S\theta^{(k)}$ $\displaystyle\leq\left(\gamma^{(k)}(2S-1)\sigma_{\max}+\left|1-\gamma^{(k)} \sigma_{\min}\right|\right)c_{1}e^{-c_{2}(k-1)}.$ $$

Recall that $c_{2}=\log\left(\frac{1}{\gamma^{(k)}\sigma_{\max}(2S-1)+\left|1-\gamma^{(k)}
\sigma_{\min}\right|}\right)$, we obtain

$$ $\displaystyle\|\bm{\beta}^{(k+1)}-\bm{\beta}^{*}\|_{1}\leq c_{1}e^{-c_{2}k}.$ $$

Therefore the proof is complete.
∎

###### Lemma 2 .

Suppose all conditions mentioned in Theorem 5.1 hold. For $k\in\mathbb{N}$, if

$$ $\displaystyle\mathrm{support}(\bm{\beta}^{(k)})\in\mathbb{S}\quad\text{and} \quad\|\bm{\beta}_{\mathbb{S}}^{(k)}-\bm{\beta}_{\mathbb{S}}^{*}\|\leq c_{1}e^ {-c_{2}(k-1)},$ $$

then it holds that

$$ $\displaystyle\mathrm{support}(\bm{\beta}^{(k+1)})\in\mathbb{S}.$ $$

###### Proof of Lemma 2 .

For a fixed $k$, if $\mathrm{support}(\bm{\beta}^{(k)})\in\mathbb{S}$ and $\|\bm{\beta}_{\mathbb{S}}^{(k)}-\bm{\beta}_{\mathbb{S}}^{*}\|\leq c_{1}e^{-c_{
2}(k-1)}$, then we have

$$ $\displaystyle\bm{\beta}_{\mathbb{S}^{C}}^{(k+1)}$ $\displaystyle=\mathcal{S}_{\theta^{(k)}}\Big{(}\bm{\beta}_{\mathbb{S}^{C}}^{(k )}-\gamma^{(k)}(\mathbf{D}_{:,\mathbb{S}^{C}})^{\top}\mathbf{X}_{:,\mathbb{S}} (\bm{\beta}^{k}_{\mathbb{S}}-\bm{\beta}^{*}_{\mathbb{S}})\Big{)}$ $\displaystyle=\mathcal{S}_{\theta^{(k)}}\Big{(}-\gamma^{(k)}(\mathbf{D}_{:, \mathbb{S}^{C}})^{\top}\mathbf{X}_{:,\mathbb{S}}(\bm{\beta}^{k}_{\mathbb{S}}- \bm{\beta}^{*}_{\mathbb{S}})\Big{)},$ $$

where $\mathbb{S}^{C}=[d^{\prime}]\backslash\mathbb{S}$. Then, for all $i\in\mathbb{S}^{C}$, we obtain

$$ $\displaystyle\bm{\beta}_{i}^{(k+1)}=\mathcal{S}_{\theta^{(k)}}\Big{(}-\gamma^{ (k)}\sum_{j\in\mathbb{S}}(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,j}(\bm{\beta}^ {k}_{j}-\bm{\beta}^{*}_{j})\Big{)}.$ $$

Note that

$$ $\displaystyle\left|-\gamma^{(k)}\sum_{j\in\mathbb{S}}(\mathbf{D}_{:,i})^{\top} \mathbf{X}_{:,j}(\bm{\beta}^{k}_{j}-\bm{\beta}^{*}_{j})\right|$ $\displaystyle\leq\gamma^{(k)}\sigma_{\max}\|\bm{\beta}^{(k)}-\bm{\beta}^{*}\|_ {1}$ $\displaystyle\leq\gamma^{(k)}\sigma_{\max}c_{1}e^{-c_{2}(k-1)},$ $$

Given $\theta^{(k)}$ as defined in Theorem 1, we obtain:

$$ $\displaystyle\left|-\gamma^{(k)}\sum_{j\in\mathbb{S}}(\mathbf{D}_{:,i})^{\top} \mathbf{X}_{:,j}(\bm{\beta}^{(k)}_{j}-\bm{\beta}^{*}_{j})\right|\leq\theta^{(k )}.$ $$

Consequently, it follows that $\bm{\beta}_{i}^{(k)}=0$ for all $i\in\mathbb{S}^{C}$.
∎

###### Lemma 3 .

Suppose Assumption 1 holds and let $\mathbf{D}=\frac{1}{2N+1}\mathbf{X}(\mathbf{M}^{V})^{\top}$. Then, there exists an $\mathbf{M}^{V}$ such that the following inequalities hold with probability at least $1-3\delta$:

$$ $\displaystyle\min_{i\in\mathbb{S}}\left|(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{: ,i}\right|\geq\frac{N}{N+2}\left(1-\sqrt{\frac{\log S-\log\delta}{Nc}}\right),$ $\displaystyle\max_{i\in\mathbb{S}}\left|(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{: ,i}\right|\leq\frac{N}{N+2}\left(1+\sqrt{\frac{\log S-\log\delta}{Nc}}\right),$ $\displaystyle\max_{i\neq j}\left|(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,i} \right|\leq\sqrt{\frac{\log d-\log\delta}{Nc}}.$ $$

###### Proof of Lemma 3 .

Recall that $\sigma_{1}\geq\sigma_{2}\geq\cdots\geq\sigma_{d}$. Let $\mathbf{M}^{V}=\frac{2}{\sigma_{d}^{2}}\mathbf{I}_{d\times d}$. To prove the lemma, we first introduce a lower bound for $\min_{i\in\mathbb{S}}\left|\mathbf{D}_{:,i}^{\top}\mathbf{X}_{:,i}\right|$. Note that

$$ $\displaystyle(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,i}=\frac{1}{2N+1}\mathbf{M }^{V}(\mathbf{X}_{:,i})^{\top}\mathbf{X}_{:,i}=\frac{2}{2N+1}\cdot\frac{1}{ \sigma^{2}_{d}}\sum_{j=1}^{N}\mathbf{X}_{j,i}^{2}\geq\frac{2}{2N+1}\sum_{j=1}^ {N}\frac{1}{\sigma^{2}_{i}}\mathbf{X}_{j,i}^{2}.$ $$

Note that $\mathbf{X}_{j,i}^{2}$ is a sub-exponential random variable. Then, from the tail bound for sub-exponential random variables, there exists a constant $c$ such that

$$ $\displaystyle\mathbb{P}\Big{\{}\frac{1}{N}\sum_{j=1}^{N}(\frac{1}{\sigma_{i}} \mathbf{X}^{2}_{j,i})\geq 1-s\Big{\}}\geq 1-\exp\big{(}-Nc\min\big{\{}s^{2},s \big{\}}\big{)}.$ $$

Consider $0<s\leq 1$, we have

$$ $\displaystyle\mathbb{P}\Big{\{}\frac{1}{N}\sum_{j=1}^{N}(\frac{1}{\sigma_{i}} \mathbf{X}^{2}_{j,i})\geq 1-s\Big{\}}\geq 1-\exp\big{(}-Ncs^{2}\big{)},$ $$

it follows that

$$ $\displaystyle\mathbb{P}\Big{\{}\frac{1}{N}\min_{i\in\mathbb{S}}\frac{1}{\sigma ^{2}_{d}}\sum_{j=1}^{N}\mathbf{X}_{j,i}^{2}\geq 1-s\Big{\}}$ $\displaystyle\geq\mathbb{P}\Big{\{}\frac{1}{N}\min_{i\in\mathbb{S}}\sum_{j=1}^ {N}\frac{1}{\sigma^{2}_{i}}\mathbf{X}_{j,i}^{2}\geq 1-s\Big{\}}$ $\displaystyle=1-\mathbb{P}\Bigg{\{}\bigcup_{i\in\mathbb{S}}\Big{\{}\frac{1}{N} \sum_{j\in[1]}^{N}\frac{1}{\sigma^{2}_{i}}\mathbf{X}_{j,i}^{2}<1-s\Big{\}} \Bigg{\}}$ $\displaystyle\geq 1-\sum_{i\in\mathbb{S}}\mathbb{P}\Big{\{}\frac{1}{N}\sum_{j= 1}^{N}\frac{1}{\sigma^{2}_{i}}\mathbf{X}_{j,i}^{2}<1-s\Big{\}}$ $\displaystyle\geq 1-S\exp\big{(}-Ncs^{2}\big{)}.$ $$

Let $s=\sqrt{\frac{\log S-\log\delta}{Nc}}$. Then, with probability at least $1-\delta$, we have

$$ $\displaystyle\frac{1}{N}\min_{i\in\mathbb{S}}\frac{1}{\sigma^{2}_{d}}\sum_{j=1 }^{N}\mathbf{X}_{j,i}^{2}\geq 1-\sqrt{\frac{\log S-\log\delta}{Nc}}.$ $$

Therefore, with probability at least $1-\delta$, it holds that

$$ $\displaystyle\min_{i\in\mathbb{S}}\left|(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{: ,i}\right|=\frac{2}{2N+1}\min_{i\in\mathbb{S}}\frac{1}{\sigma^{2}_{d}}\sum_{j= 1}^{N}\mathbf{X}_{j,i}^{2}\geq\frac{N}{N+2}\left(1-\sqrt{\frac{\log S-\log \delta}{Nc}}\right).$ $$

Similarly, we have the following inequality holding with probability at least $1-\delta$:

$$ $\displaystyle\max_{i\in\mathbb{S}}\left|(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{: ,i}\right|\leq\frac{N}{N+2}\left(1+\sqrt{\frac{\log S-\log\delta}{Nc}}\right).$ $$

Next, we will provide a high probability upper bound for $\max_{i\neq j}\left|(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,j}\right|$. Observe that when $i\neq j$, we have

$$ $\displaystyle(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,j}=\frac{1}{2N+1}\mathbf{M }^{Q}(\mathbf{X}_{:,i})^{\top}\mathbf{X}_{:,j}=\frac{2}{2N+1}\sum_{k=1}^{N} \frac{1}{\sigma^{2}_{d}}\mathbf{X}_{k,i}\mathbf{X}_{k,j},$ $$

where $\mathbf{X}_{k,i}\mathbf{X}_{k,j}$ is a centered sub-exponential random variable. Then, from the tail bound for sub-exponential random variables, there exist constants $c$ and $0<s\leq 1$ such that

$$ $\displaystyle\mathbb{P}\Bigg{\{}\Big{|}\frac{1}{N}\sum_{k=1}^{N}(\frac{1}{ \sigma_{d}^{2}}\mathbf{X}_{k,i}\mathbf{X}_{k,j})\Big{|}\leq s\Bigg{\}}\geq 1- \exp\big{(}-Ncs^{2}\big{)}.$ $$

Then we have

$$ $\displaystyle\mathbb{P}\Bigg{\{}\max_{i\neq j}\Big{|}\frac{2}{2+2N}\sum_{k=1}^ {N}(\frac{1}{\sigma_{d}}\mathbf{X}_{k,i}\mathbf{X}_{k,j})\Big{|}\leq s\Bigg{\}}$ $\displaystyle\geq\mathbb{P}\Bigg{\{}\max_{i\neq j}\Big{|}\frac{1}{N}\sum_{k=1} ^{N}(\frac{1}{\sigma_{d}}\mathbf{X}_{k,i}\mathbf{X}_{k,j})\Big{|}\leq s\Bigg{\}}$ $\displaystyle=1-\mathbb{P}\Bigg{\{}\bigcup_{i,j:i\neq j}\Big{\{}\Big{|}\frac{1 }{N}\sum_{k=1}^{N}(\frac{1}{\sigma_{d}}\mathbf{X}_{k,i}\mathbf{X}_{k,j})\Big{| }>s\Big{\}}\Bigg{\}}$ $\displaystyle\geq 1-\sum_{i,j:i\neq j}\mathbb{P}\Bigg{\{}\Big{|}\frac{1}{N} \sum_{k=1}^{N}(\frac{1}{\sigma_{d}}\mathbf{X}_{k,i}\mathbf{X}_{k,j})\Big{|}>s \Bigg{\}}$ $\displaystyle\geq 1-\sum_{i,j}\mathbb{P}\Bigg{\{}\Big{|}\frac{1}{N}\sum_{k=1}^ {N}(\frac{1}{\sigma_{d}}\mathbf{X}_{k,i}\mathbf{X}_{k,j})\Big{|}>s\Bigg{\}}$ $\displaystyle\geq 1-d^{2}\exp\big{(}-Ncs^{2}\big{)}.$ $$

Let $s=\sqrt{\frac{\log d-\log\delta}{Nc}}$, thus with probability at least $1-\delta$, we obtain

$$ $\displaystyle\max_{i\neq j}\left|(\mathbf{D}_{:,i})^{\top}\mathbf{X}_{:,i} \right|\leq\sqrt{\frac{\log d-\log\delta}{Nc}}.$ $$

Then, the proof is complete.
∎

To prove Theorem 5.1, we first state an equivalent theorem here:

###### Theorem D.1 (Equivalent form of Theorem 5.1 ) .

Suppose Assumption 1 holds. Let $\delta\in(0,1)$.
For a $K$-layer Transformer model with the structure described in Section 4.3, there exists a set of parameters in the Transformer such that for any $n\in[N_{0},N]$, with probability at least $1-3\delta$, we have

$$ $\displaystyle\|\bm{\beta}^{(K+1)}_{2n+1}-\bm{\beta}^{*}\|\leq c_{1}e^{-\alpha_ {n}K},$ $$

where $\bm{\beta}^{(K+1)}_{2n+1}=[\mathbf{H}^{(K+1)}]_{d+1:2d+1,2n+1}$ and

$$ $\displaystyle\alpha_{n}=-\log\left(1-\frac{2}{3}\gamma+\gamma(2S-1)\sqrt{\frac {\log d-\log\delta}{nc}}+\gamma\sqrt{\frac{\log S-\log\delta}{nc}}\right),$ (D.3) $$

$c_{1}=\|\bm{\beta}^{*}\|_{1}$ and $N_{0}=8(4S-2)^{2}\frac{\log d+\log S-\log\delta}{c}$. Here, $c$ is a positive constant and $\gamma$ can be any positive constant less than $1.6$.

To prove Theorem D.1, first, for $\mathbf{D}_{n}=\frac{1}{2n+1}[\mathbf{X}]_{1:n,:}(\mathbf{M}^{V})^{\top}$ defined in Equation 4.4, we define an event:

$$ $\displaystyle A=\Bigg{\{}$ $\displaystyle\min_{i\in\mathbb{S}}\left|([\mathbf{D}_{n}]_{:,i})^{\top}\mathbf {X}_{1:n,i}\right|\geq\frac{n}{n+2}\left(1-\sqrt{\frac{\log S-\log\delta}{nc}} \right),\quad$ $\displaystyle\max_{i\in\mathbb{S}}\left|([\mathbf{D}_{n}]_{:,i})^{\top}\mathbf {X}_{1:n,i}\right|\leq\frac{n}{n+2}\left(1+\sqrt{\frac{\log S-\log\delta}{nc}} \right),$ $\displaystyle\max_{i\neq j}\left|([\mathbf{D}_{n}]_{:,i})^{\top}\mathbf{X}_{1: n,i}\right|\leq\sqrt{\frac{\log d-\log\delta}{nc}}\Bigg{\}}.$ (D.4) $$

Next, we introduce Proposition 2.

###### Proposition 2 .

Suppose Assumption 1 and event $A$ defined in Equation D.4 hold. For a $K$-layer Transformer model with the structure described in Section 4.3, and let the input sequence satisfy Equation C.4. Then, there exists a set of parameters such that if $n\in[N_{0},N]$, we have

$$ $\displaystyle\|\bm{\beta}^{(K+1)}_{2n+1}-\bm{\beta}^{*}\|\leq c_{1}e^{-\alpha_ {n}K},$ $$

where $\bm{\beta}^{(K+1)}_{2n+1}=[\mathbf{H}^{(K+1)}]_{d+1:2d+1,2n+1}$,
$c_{1}=\|\bm{\beta}^{*}\|_{1}$, and $N_{0}=8(4S-2)^{2}\frac{\log d+\log S-\log\delta}{c}$. $\alpha_{n}$ is given in Equation D.3. Here, $c$ is a positive constant and $\gamma$ can be any positive constant less than $1.6$.

###### Proof of Proposition 2 .

We prove it by induction. First, for $k=1$, we initialize $\bm{\beta}_{2n+1}^{(1)}=\mathbf{0}_{d}$, thus $\|\bm{\beta}^{(1)}_{2n+1}-\bm{\beta}^{*}\|_{1}=\|\bm{\beta}^{*}\|_{1}$. Note that $c_{1}=e^{\alpha}\|\bm{\beta}^{*}\|_{1}$, so Proposition 2 holds for $k=1$.

Next, we assume Proposition 2 holds for $k$. From the proof of Theorem 5.1, the input sequence of the $(k+1)$-th Transformer layer is of the following structure:

$$ $\displaystyle\mathbf{H}^{(k+1)}=\begin{bmatrix}[\mathbf{X}^{\top}]_{:,1}&[ \mathbf{X}^{\top}]_{:,1}&\cdots&[\mathbf{X}^{\top}]_{:,N}&[\mathbf{X}^{\top}]_ {:,N}&\mathbf{x}_{N+1}\\ 0&y_{1}&\cdots&0&y_{N}&0\\ \bm{\beta}^{(k+1)}_{1}&\bm{\beta}^{(k+1)}_{2}&\cdots&\bm{\beta}^{(k+1)}_{2N-1} &\bm{\beta}^{(k+1)}_{2N}&\bm{\beta}^{(k+1)}_{2N+1}\\ 1&0&\cdots&1&0&1\end{bmatrix},$ $$

where

$$ $\displaystyle\bm{\beta}_{2n+1}^{(k+1)}=\mathcal{S}_{\theta^{(k)}}\Big{(}\bm{ \beta}^{(k)}_{2n+1}-\gamma^{(k)}(\mathbf{D}_{n})^{\top}([\mathbf{X}]_{1:n,:} \bm{\beta}^{(k)}_{2n+1}-\mathbf{y}_{1:n})\Big{)},$ $$

From Lemma 1, let $\gamma^{(k)}=\gamma$, then for $c_{1}=\|\bm{\beta}^{*}\|_{1}$, $c_{2}=\log\left(\frac{1}{\gamma(2S-1)\sigma_{\max}+\left|1-\gamma\sigma_{\min}
\right|}\right)$, we have $\|\bm{\beta}^{(k)}_{2n+1}-\bm{\beta}^{*}\|\leq c_{1}e^{-c_{2}(k-1)}$ if the following conditions hold:

$$ $\displaystyle\gamma(2S-1)\sigma_{\max}+\left|1-\gamma\sigma_{\min}\right|\leq 1,$ (D.5) $\displaystyle 0\leq\gamma([\mathbf{D}_{n}]_{:,i})^{\top}\mathbf{X}_{1:n,i}\leq 1 ,\quad\forall i\in[d].$ (D.6) $$

First, let

$$ $\displaystyle n\geq N_{0}\quad\text{where}\quad N_{0}=8(4S-2)^{2}{\frac{\log d +\log S-\log\delta}{c}}.$ (D.7) $$

By rearranging Inequality D.7, we have

$$ $\displaystyle\sqrt{n}$ $\displaystyle\geq 2\sqrt{2}(4S-2)\sqrt{\frac{\log d+\log S-\log\delta}{c}}$ $\displaystyle\geq\sqrt{2}(4S-2)\sqrt{\frac{\log d+\log S-2\log\delta}{c}}$ $\displaystyle\geq(4S-2)\left(\sqrt{\frac{\log d-\log\delta}{c}}+\sqrt{\frac{ \log S-\log\delta}{c}}\right)$ $\displaystyle\geq(4S-2)\sqrt{\frac{\log d-\log\delta}{c}}+\sqrt{\frac{\log S- \log\delta}{c}}.$ (D.8) $$

Inequality D.8 is equivalent to

$$ $\displaystyle\frac{1}{\sqrt{n}}\left((2S-1)\sqrt{\frac{\log d-\log\delta}{c}}+ \frac{1}{2}\sqrt{\frac{\log S-\log\delta}{c}}\right)\leq\frac{1}{2}$ $\displaystyle\Longrightarrow$ $\displaystyle\quad(2S-1)\sqrt{\frac{\log d-\log\delta}{nc}}\leq\frac{1}{2} \left(1-\sqrt{\frac{\log S-\log\delta}{nc}}\right)$ $\displaystyle\Longrightarrow$ $\displaystyle\quad(2S-1)\sqrt{\frac{\log d-\log\delta}{nc}}\leq\frac{n}{n+2} \left(1-\sqrt{\frac{\log S-\log\delta}{nc}}\right).$ (D.9) $$

Combining Inequality D.9 with Equation D.4 we have

$$ $\displaystyle\sigma_{\max}(2S-1)$ $\displaystyle\leq(2S-1)\sqrt{\frac{\log d-\log\delta}{nc}}$ $\displaystyle\leq\frac{n}{n+2}\left(1-\sqrt{\frac{\log S-\log\delta}{nc}}\right)$ $\displaystyle\leq\sigma_{\min}.$ (D.10) $$

Next, recall that $\mathbf{M}^{V}=\frac{2}{\sigma_{d}^{2}}\mathbf{I}_{d\times d}$ and $([\mathbf{D}_{n}]_{:,i})^{\top}\mathbf{X}_{1:n,i}=\frac{1}{2n+1}(\mathbf{X}_{1
:n,i})^{\top}\mathbf{M}^{V}\mathbf{X}_{1:n,i}$. Therefore, for any $i\in[d]$, we have $\gamma([\mathbf{D}_{n}]_{:,i})^{\top}\mathbf{X}_{1:n,i}\geq 0$ if $\gamma\geq 0$. Note that if $\gamma$ satisfies the following condition:

$$ $\displaystyle\gamma\leq\min\left\{\frac{1}{\sigma_{\min}},\frac{1}{\max\{([ \mathbf{D}_{n}]_{:,i})^{\top}\mathbf{X}_{1:n,i}\}}\right\},$ (D.11) $$

then Inequality D.6 will hold and $1-\gamma\sigma_{\min}\geq 0$.
Besides, we have

$$ $\displaystyle\gamma(2S-1)\sigma_{\max}+\left|1-\gamma\sigma_{\min}\right|$ $\displaystyle\leq\gamma\sigma_{\min}+|1-\gamma\sigma_{\min}|$ (D.12) $\displaystyle=\gamma\sigma_{\min}+1-\gamma\sigma_{\min}$ (D.13) $\displaystyle=1,$ $$

where Inequality D.12 is derived from Inequality D.10, and Equation D.13 is obtained using Equation D.11. Consequently, Inequalities D.5 and D.6 hold if event $A$ occurs. Furthermore, under the condition that event $A$ occurs, by setting

$$ $\displaystyle\gamma\leq\frac{1}{\frac{n}{n+2}\left(1+\sqrt{\frac{\log S-\log \delta}{nc}}\right)},$ (D.14) $$

we have $\gamma$ satisfies Equation D.11. Note that we assume $n\geq N_{0}$. Therefore,

$$ $\displaystyle\frac{1}{\frac{n}{n+2}\left(1+\sqrt{\frac{\log S-\log\delta}{nc}} \right)}$ $\displaystyle\geq\frac{1}{1+\sqrt{\frac{\log S-\log\delta}{N_{0}c}}}\geq\frac{ 2}{1+\sqrt{\frac{1}{8(4S-2)}}}\geq 1.6.$ $$

Therefore, any choice of $\gamma$ such that $\gamma\leq 1.6$ will ensure that Inequality (D.14) is satisfied. Next, noting that event $A$ holds and $\gamma$ satisfies Equation D.11, we obtain

$$ $\displaystyle\|\bm{\beta}^{(k)}_{2n+1}-\bm{\beta}^{*}\|$ $\displaystyle\leq\|\bm{\beta}^{*}\|_{1}\Big{(}\gamma(2S-1)\sigma_{\max}+\left| 1-\gamma\sigma_{\min}\right|\Big{)}^{k-1}$ $\displaystyle\leq\|\bm{\beta}^{*}\|_{1}\left(\gamma(2S-1)\sqrt{\frac{\log d- \log\delta}{nc}}+1-\gamma\frac{n}{n+2}\left(1-\sqrt{\frac{\log S-\log\delta}{ nc}}\right)\right)^{k-1}$ $\displaystyle\leq\|\bm{\beta}^{*}\|_{1}\left(\gamma(2S-1)\sqrt{\frac{\log d- \log\delta}{nc}}+1-\gamma\frac{n}{n+2}+\gamma\frac{n}{n+2}\sqrt{\frac{\log S- \log\delta}{nc}}\right)^{k-1}$ $\displaystyle\leq\|\bm{\beta}^{*}\|_{1}\left(\gamma(2S-1)\sqrt{\frac{\log d- \log\delta}{nc}}+1-\frac{2}{3}\gamma+\gamma\sqrt{\frac{\log S-\log\delta}{nc}} \right)^{k-1}$ $\displaystyle=c_{1}e^{-\alpha_{n}(k-1)}.$ $$

The proof is complete.
∎

Note that Lemma 3 demonstrates that event $A$ occurs with a probability of at least $1-2\delta$. Therefore, the proof of Theorem D.1 is completed by combining Lemma 3 with Proposition 2.

###### Lemma 1 .

###### Proof of Lemma 1 .

###### Lemma 2 .

###### Proof of Lemma 2 .

###### Lemma 3 .

###### Proof of Lemma 3 .

###### Theorem D.1 (Equivalent form of Theorem 5.1 ) .

###### Proposition 2 .

###### Proof of Proposition 2 .

### D.2 Proof of Remark 3

Note that in Section D.2, we assume there is a constraint on the support of $\bm{\beta}^{*}$: its support set lies within a subset $\mathbb{S}\subset[1:d]$, with $S<|\mathbb{S}|\leq d$, where $S$ denotes the sparsity of $\bm{\beta}^{*}$.We establish the corollary by demonstrating the following corollary:

###### Corollary D.1 (Equivalent statement of Remark 3 .) .

Suppose Assumption 1 holds and $S\leq|\mathbb{S}|\leq d$.
For a $K$-layer Transformer model with the structure described in Section 4.3, and let the input sequence satisfy Equation C.4. Then, there exists a set of parameters such that if $n\geq N_{0}^{S}$, with probability at least $1-\delta$, we have

$$ $\displaystyle\|\bm{\beta}^{(K)}_{2n+1}-\bm{\beta}^{*}\|\leq c_{1}e^{-\alpha_{n }^{S}K},$ $$

where $\bm{\beta}^{(K)}_{2n+1}=[\mathbf{H}^{(K+1)}]_{d+1:2d+1,2N+1}$ and

$$ $\displaystyle\alpha_{n}^{S}=-\log\Bigg{(}1-\frac{2}{3}\gamma+\gamma(2S-1)\sqrt {\frac{\log|\mathbb{S}|-\log\delta}{nc}}+\gamma\sqrt{\frac{\log S-\log\delta}{ nc}}\Bigg{)},$ (D.15) $$

$c_{1}=\|\bm{\beta}^{*}\|_{1}$, and $N_{0}^{p}=8(4S-2)^{2}\frac{\log|\mathbb{S}|+\log S-\log\delta}{c}$. Here, $c$ is a positive constant, and $\gamma$ can be any positive constant less than $1.6$.

###### Proof of Corollary D.1 .

W.l.o.g., assume $\mathbb{S}=\{1,2,\ldots,|\mathbb{S}|\}$. First, we modify the MLP layers based on Equation C.2. We maintain the structure of $\mathbf{W}_{2}$ as described in Equation C.2, and reconstruct $\mathbf{W}_{1}$ and $\mathbf{b}^{(k)}$. We define the following submatrices:

$$ $\displaystyle\mathbf{W}_{1,\text{sub}(1)}=\begin{bmatrix}\mathbf{0}_{(d+1) \times(d+1)}&\mathbf{0}_{(d+1)\times|\mathbb{S}|}&\mathbf{0}_{(d+1)\times(d-| \mathbb{S}|)}&\mathbf{0}_{d+1}\\ \mathbf{0}_{|\mathbb{S}|\times(d+1)}&\mathbf{I}_{|\mathbb{S}|\times|\mathbb{S} |}&\mathbf{0}_{|\mathbb{S}|\times(d-|\mathbb{S}|)}&\mathbf{0}_{|\mathbb{S}|}\\ \mathbf{0}_{(d-|\mathbb{S}|)\times(d+1)}&\mathbf{0}_{(d-|\mathbb{S}|)\times| \mathbb{S}|}&\mathbf{I}_{(d-|\mathbb{S}|)\times(d-|\mathbb{S}|)}&\mathbf{0}_{( d-|\mathbb{S}|)}\\ \mathbf{0}_{1\times(d+1)}&\mathbf{0}_{1\times|\mathbb{S}|}&\mathbf{0}_{1\times (d-|\mathbb{S}|)}&0\\ \end{bmatrix},$ $\displaystyle\mathbf{W}_{1,\text{sub}(2)}=\begin{bmatrix}\mathbf{0}_{(d+1) \times(d+1)}&\mathbf{0}_{(d+1)\times|\mathbb{S}|}&\mathbf{0}_{(d+1)\times(d-| \mathbb{S}|)}&\mathbf{0}_{d+1}\\ \mathbf{0}_{|\mathbb{S}|\times(d+1)}&-\mathbf{I}_{|\mathbb{S}|\times|\mathbb{S }|}&\mathbf{0}_{|\mathbb{S}|\times(d-|\mathbb{S}|)}&\mathbf{0}_{|\mathbb{S}|} \\ \mathbf{0}_{(d-|\mathbb{S}|)\times(d+1)}&\mathbf{0}_{(d-|\mathbb{S}|)\times| \mathbb{S}|}&\mathbf{I}_{(d-|\mathbb{S}|)\times(d-|\mathbb{S}|)}&\mathbf{0}_{( d-|\mathbb{S}|)}\\ \mathbf{0}_{1\times(d+1)}&\mathbf{0}_{1\times|\mathbb{S}|}&\mathbf{0}_{1\times (d-|\mathbb{S}|)}&0\\ \end{bmatrix},$ $\displaystyle\mathbf{W}_{1,\text{sub}(3)}=\begin{bmatrix}\mathbf{0}_{(d+1) \times(d+1)}&\mathbf{0}_{(d+1)\times|\mathbb{S}|}&\mathbf{0}_{(d+1)\times(d-| \mathbb{S}|)}&\mathbf{0}_{d+1}\\ \mathbf{0}_{|\mathbb{S}|\times(d+1)}&\mathbf{I}_{|\mathbb{S}|\times|\mathbb{S} |}&\mathbf{0}_{|\mathbb{S}|\times(d-|\mathbb{S}|)}&\mathbf{0}_{|\mathbb{S}|}\\ \mathbf{0}_{(d-|\mathbb{S}|)\times(d+1)}&\mathbf{0}_{(d-|\mathbb{S}|)\times| \mathbb{S}|}&-\mathbf{I}_{(d-|\mathbb{S}|)\times(d-|\mathbb{S}|)}&\mathbf{0}_{ (d-|\mathbb{S}|)}\\ \mathbf{0}_{1\times(d+1)}&\mathbf{0}_{1\times|\mathbb{S}|}&\mathbf{0}_{1\times (d-|\mathbb{S}|)}&0\\ \end{bmatrix},$ $\displaystyle\mathbf{W}_{1,\text{sub}(4)}=\begin{bmatrix}\mathbf{0}_{(d+1) \times(d+1)}&\mathbf{0}_{(d+1)\times|\mathbb{S}|}&\mathbf{0}_{(d+1)\times(d-| \mathbb{S}|)}&\mathbf{0}_{d+1}\\ \mathbf{0}_{|\mathbb{S}|\times(d+1)}&-\mathbf{I}_{|\mathbb{S}|\times|\mathbb{S }|}&\mathbf{0}_{|\mathbb{S}|\times(d-|\mathbb{S}|)}&\mathbf{0}_{|\mathbb{S}|} \\ \mathbf{0}_{(d-|\mathbb{S}|)\times(d+1)}&\mathbf{0}_{(d-|\mathbb{S}|)\times| \mathbb{S}|}&\mathbf{I}_{(d-|\mathbb{S}|)\times(d-|\mathbb{S}|)}&\mathbf{0}_{( d-|\mathbb{S}|)}\\ \mathbf{0}_{1\times(d+1)}&\mathbf{0}_{1\times|\mathbb{S}|}&\mathbf{0}_{1\times (d-|\mathbb{S}|)}&0\\ \end{bmatrix}.$ $$

Let

$$ $\displaystyle\mathbf{W}_{1}=\begin{bmatrix}\mathbf{W}_{1,\text{sub}(1)}\\ \mathbf{W}_{1,\text{sub}(2)}\\ \mathbf{W}_{1,\text{sub}(3)}\\ \mathbf{W}_{1,\text{sub}(4)}\end{bmatrix},\mathbf{b}^{(k)}=\begin{bmatrix} \mathbf{0}_{5d+5}\\ -\theta^{(k)}\cdot\mathbf{1}_{|\mathbb{S}|}\\ \mathbf{0}_{2d+2-|\mathbb{S}|}\\ \theta^{(k)}\cdot\mathbf{1}_{|\mathbb{S}|}\\ \mathbf{0}_{1+d-|\mathbb{S}|}\end{bmatrix}.$ $$

Therefore, the output of the modified MLP layer is

$$ $\displaystyle\mathrm{MLP}(\mathbf{h}_{i})=\begin{bmatrix}[\mathbf{h}_{i}]_{1:d +1}\\ \mathcal{S}_{\theta^{(k)}}([\mathbf{h}_{i}]_{d+2:d+|\mathbb{S}|+1})\\ \mathbf{0}_{d-|\mathbb{S}|}\\ [\mathbf{h}_{i}]_{2d+2}\end{bmatrix}.$ $$

Next, we also modify the parameters in the self-attention layers. For convenience, we introduce a projection matrix $\mathbf{I}^{S}=\mathrm{diag}(\mathbf{I}_{|\mathbb{S}|\times|\mathbb{S}|},
\mathbf{0}_{(d-|\mathbb{S}|)\times(d-|\mathbb{S}|)})$. The reconstructed self-attention layers are described below:

$$ $\displaystyle\mathbf{Q}_{\pm 1}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{(d+1)\times(d+1)}&\mathbf{0}_{(d+1) \times d}&\mathbf{0}_{d+1}\\ \mathbf{0}_{d\times(d+1)}&\mathbf{I}^{S}\mathbf{M}^{Q,(k)}_{\pm 1}&\mathbf{0}_ {d}\\ \mathbf{0}_{1\times(d+1)}&\mathbf{0}_{1\times d}&-B\end{bmatrix},$ $\displaystyle\mathbf{Q}_{\pm 2}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{d\times(2d+1)}&\mathbf{0}_{d}\\ \mathbf{0}_{1\times(2d+1)}&m_{\pm 2}^{Q,(k)}\\ \mathbf{0}_{(d+1)\times(2d+1)}&\mathbf{0}_{d+1}\end{bmatrix}$ $\displaystyle\mathbf{K}_{\pm 1}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{(d+1)\times d}&\mathbf{0}_{(d+1) \times(d+1)}&\mathbf{0}_{d+1}\\ \mathbf{I}^{S}&\mathbf{0}_{d\times(d+1)}&\mathbf{0}_{d}\\ \mathbf{0}_{1\times d}&\mathbf{0}_{1\times(d+1)}&{1}\end{bmatrix},$ $\displaystyle\mathbf{K}_{\pm 2}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{d\times d}&\mathbf{0}_{d}&\mathbf{0}_ {d\times d+1}\\ \mathbf{0}_{1\times d}&1&\mathbf{0}_{1\times(d+1)}\\ \mathbf{0}_{(d+1)\times d}&\mathbf{0}_{d+1}&\mathbf{0}_{(d+1)\times(d+1)}\end{bmatrix}$ $\displaystyle\mathbf{V}_{\pm 1}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{(d+1)\times d}&\mathbf{0}_{(d+1) \times(d+2)}\\ \mathbf{I}^{S}\mathbf{M}^{V,(k)}_{\pm 1}&\mathbf{0}_{d\times(d+2)}\\ \mathbf{0}_{2\times d}&\mathbf{0}_{2\times(d+2)}\end{bmatrix},$ $\displaystyle\mathbf{V}_{\pm 2}^{(k)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{(d+1)\times d}&\mathbf{0}_{(d+1) \times d+2}\\ \mathbf{I}^{S}\mathbf{M}^{V,(k)}_{\pm 2}&\mathbf{0}_{d\times(d+2)}\\ \mathbf{0}_{2\times d}&\mathbf{0}_{2\times(d+2)}\end{bmatrix}.$ $$

Therefore, when the input sequence is

$$ $\displaystyle\mathbf{H}^{(k)}=\begin{bmatrix}[\mathbf{X}^{\top}]_{:,1}&[ \mathbf{X}^{\top}]_{:,1}&\cdots&[\mathbf{X}^{\top}]_{:,N}&[\mathbf{X}^{\top}]_ {:,N}&\mathbf{x}_{N+1}\\ 0&y_{1}&\cdots&0&y_{N}&0\\ \bm{\beta}^{(k)}_{1}&\bm{\beta}^{(k)}_{2}&\cdots&\bm{\beta}^{(k)}_{2N-1}&\bm{ \beta}^{(k)}_{2N}&\bm{\beta}^{(k)}_{2N+1}\\ 1&0&\cdots&1&0&1\end{bmatrix},$ $$

the corresponding output of the self-attention layer satisfies

$$ $\displaystyle\mathbf{H}^{(k+1)}=\mathrm{diag}(\mathbf{I}^{S},1,\mathbf{I}^{S}, 1)\mathbf{H}^{(k+1)}.$ $$

Then, it is equivalent to consider the input sequence as

$$ $\displaystyle\begin{bmatrix}[\mathbf{X}^{\top}]_{1:|\mathbb{S}|,1}&[\mathbf{X} ^{\top}]_{1:|\mathbb{S}|,2}&\cdots&[\mathbf{X}^{\top}]_{1:|\mathbb{S}|,N}&[ \mathbf{X}^{\top}]_{1:|\mathbb{S}|,N}&[\mathbf{x}_{N+1}]_{1:|\mathbb{S}|}\\ 0&y_{1}&\cdots&0&y_{N}&0\\ [\bm{\beta}^{(k)}_{1}]_{1:|\mathbb{S}|}&[\bm{\beta}^{(k)}_{2}]_{1:|\mathbb{S}| }&\cdots&[\bm{\beta}^{(k)}_{2N-1}]_{1:|\mathbb{S}|}&[\bm{\beta}^{(k)}_{2N}]_{1 :|\mathbb{S}|}&[\bm{\beta}^{(k)}_{2N+1}]_{1:|\mathbb{S}|}\\ 1&0&\cdots&1&0&1\end{bmatrix}.$ $$

Consequently, by applying Theorem 5.1 to this sub-problem, we derive the corresponding corollary.
∎

###### Corollary D.1 (Equivalent statement of Remark 3 .) .

###### Proof of Corollary D.1 .

## Appendix E Deferred Proofs in Section 5.2

### E.1 Proof of Theorem 5.2

Note that in Theorem 5.2, we have two results: first, the label prediction loss for the linear read-out function, and second, the label prediction loss for the quadratic read-out function.

We start by introducing the following lemma, which provides an upper bound on $\|\bm{\beta}^{(k)}_{j}\|$ for every $j$ and $k\geq 2$.

###### Lemma 4 .

If $\bm{\beta}^{(k)}$ satisfies Equation C.12 and $\|\bm{\beta}^{(k)}\|\leq C_{\mathbf{x}}$, then we have $\|\bm{\beta}^{(k+1)}\|\leq C_{\mathbf{x}}$, where $C_{\mathbf{x}}=\frac{db_{\mathbf{x}}b_{\bm{\beta}}(N+1)}{\sigma_{\min}(\mathbf
{X}^{\top}\mathbf{X})}$.

###### Proof of Lemma 4 .

From Equation C.12, we derive that

$$ $\displaystyle\bm{\beta}^{k+1}_{i}$ $\displaystyle=\mathcal{S}_{\theta^{(k)}}\left(\bm{\beta}_{i}^{(k)}-\frac{ \gamma^{(k)}}{i}\mathbf{M}^{V}[\mathbf{X}^{\top}]_{:,1:\frac{i-1}{2}}\left([ \mathbf{X}]_{1:\frac{i-1}{2},:}\bm{\beta}^{(k)}_{i}-Y_{1:\frac{i-1}{2}}\right)\right)$ $\displaystyle=\mathcal{S}_{\theta^{(k)}}\left(\left(\mathbf{I}-\frac{\gamma^{( k)}}{i}\mathbf{M}^{V}[\mathbf{X}^{\top}]_{:,1:\frac{i-1}{2}}[\mathbf{X}]_{1: \frac{i-1}{2},:}\right)\bm{\beta}_{i}^{(k)}-\frac{\gamma^{(k)}}{i}\mathbf{M}^{ V}[\mathbf{X}^{\top}]_{:,1:\frac{i-1}{2}}Y_{1:\frac{i-1}{2}}\right).$ $$

Since $\|\mathcal{S}_{\theta^{(k)}}(x)\|\leq\|x\|$, it follows that

$$ $\displaystyle\|\bm{\beta}^{k+1}_{i}\|$ $\displaystyle\leq\left\|\left(\mathbf{I}-\frac{\gamma^{(k)}}{i}\mathbf{M}^{V}[ \mathbf{X}^{\top}]_{:,1:\frac{i-1}{2}}[\mathbf{X}]_{1:\frac{i-1}{2},:}\right) \bm{\beta}_{i}^{(k)}-\frac{\gamma^{(k)}}{i}\mathbf{M}^{V}[\mathbf{X}^{\top}]_{ :,1:\frac{i-1}{2}}Y_{1:\frac{i-1}{2}}\right\|$ $\displaystyle\mathrel{\underset{(\text{i})}{\leq}}\left(1-\frac{\gamma^{(k)}}{ i\cdot\sigma_{d}^{2}}\sigma_{\min}\left([\mathbf{X}^{\top}]_{:,1:\frac{i-1}{2} }[\mathbf{X}]_{1:\frac{i-1}{2},:}\right)\right)\|\bm{\beta}_{i}^{(k)}\|+\frac{ \gamma^{(k)}}{i\cdot\sigma_{d}^{2}}\|[\mathbf{X}^{\top}]_{:,1:\frac{i-1}{2}}Y_ {1:\frac{i-1}{2}}\|$ $\displaystyle\mathrel{\underset{(\text{ii})}{\leq}}\left(1-\frac{\gamma^{(k)}} {(N+1)\cdot\sigma_{d}^{2}}\sigma_{\min}\left(\mathbf{X}^{\top}\mathbf{X}\right )\right)\|\bm{\beta}_{i}^{(k)}\|+\frac{\gamma^{(k)}}{\sigma_{d}^{2}}db_{ \mathbf{x}}b_{\bm{\beta}}.$ $$

The satisfaction of Inequality (i) is achieved by setting $\mathbf{M}^{V}=\frac{2}{\sigma_{d}^{2}}\mathbf{I}_{d\times d}$, while Inequality (ii) is verified through the application of the Cauchy interlacing theorem. Consequently, by integrating the expression $C_{\mathbf{x}}=\frac{db_{\mathbf{x}}b_{\bm{\beta}}(N+1)}{\sigma_{\min}(\mathbf
{X}^{\top}\mathbf{X})}$ with the aforementioned inequality, we obtain

$$ $\displaystyle\|\bm{\beta}^{k+1}_{i}\|\leq$ $\displaystyle\left(1-\frac{\gamma^{(k)}}{(N+1)\cdot\sigma_{d}^{2}}\sigma_{\min }\left(\mathbf{X}^{\top}\mathbf{X}\right)\right)\frac{db_{\mathbf{x}}b_{\bm{ \beta}}(N+1)}{\sigma_{\min}\left(\mathbf{X}^{\top}\mathbf{X}\right)}+\frac{ \gamma^{(k)}}{\sigma_{d}^{2}}db_{\mathbf{x}}b_{\bm{\beta}}\leq C_{\mathbf{x}},$ $$

Therefore, the proof is complete.
∎

Next, we prove the following theorem to demonstrate the prediction loss for the linear read-out function.

###### Theorem E.1 .

Suppose Assumption 1 holds.
Consider a Transformer with $K+1$ layers, where the first $K$ layers have the structure described in Section 4.3, and the input sequence is as defined in Equation C.4. Then, if $n\geq N_{0}$, there exists a set of parameters in the Transformer such that the following inequality holds with probability at least $1-\delta^{\prime}-n\delta$:

$$ $\displaystyle\|y_{n}-\widehat{y}_{n}\|\leq b_{\mathbf{x}}\left(1-\frac{2}{3} \gamma\right)^{K}+\frac{c_{4}K}{\sqrt{n}}\left(1-\frac{2}{3}\gamma\right)^{K-1},$ $$

where

$$ $\displaystyle N_{0}=8(4S-2)^{2}\frac{\log d+\log S-\log\delta}{c},$ $\displaystyle c_{4}=\frac{2\sqrt{2}b_{\mathbf{x}}\eta}{\sqrt{p}}\sqrt{\frac{ \log d+\log S-\log\delta}{c}}+\frac{2B^{\prime}}{\sqrt{p}}+\frac{2C_{\mathbf{x }}b_{\mathbf{x}}}{\sqrt{p}},$ $$

where $c$ is a positive constant, $\eta$ is defined in Equation E.7, $\delta^{\prime}$ $B^{\prime}$, $p$ are constants defined in Section E.1.

We are now prepared to demonstrate Theorem E.1.

###### Proof.

For this layer, we assign one attention head to follow the structure described below. Consider the case where $i=2n-1$ and $n\gg N_{0}$, and denote $[\mathbf{X}^{\top}]_{:,n}$ as $\mathbf{x}_{n}$. For the first attention head, its structure is as follows:

$$ $\displaystyle\mathbf{Q}_{1}^{\mathrm{last}}\mathbf{h}_{i}^{(K+1)}$ $\displaystyle=\begin{bmatrix}\mathbf{0}_{d}\\ 1\\ \mathbf{x}_{n}\\ 0\end{bmatrix},\quad\mathbf{K}_{1}^{\mathrm{last}}\mathbf{h}_{i}^{(K+1)}= \begin{bmatrix}\mathbf{0}_{d}\\ -B^{\prime}\\ \bm{\beta}^{(K+1)}_{i}\\ 0\end{bmatrix},\quad\mathbf{V}_{1}^{\mathrm{last}}\mathbf{h}_{i}^{(K+1)}= \begin{bmatrix}\mathbf{0}_{d}\\ \mathbbm{1}_{\{i\%2=1\}}\\ \mathbf{0}_{d+1}\end{bmatrix}.$ $$

Therefore, we obtain

$$ $\displaystyle(\mathbf{Q}_{1}^{\mathrm{last}}\mathbf{h}_{i}^{(K+1)})^{\top}( \mathbf{K}_{1}^{\mathrm{last}}\mathbf{h}_{j}^{(K+1)})=\mathbf{x}_{n}^{\top}\bm {\beta}^{(K+1)}_{j}-B^{\prime}.$ $$

Let $p=\mathbb{P}\{\mathbf{x}_{n}^{\top}\bm{\beta}^{(K+1)}_{j}\geq B^{\prime}\}$, which is the probability that the inner product $\mathbf{x}_{n}^{\top}\bm{\beta}^{(K+1)}_{j}$ exceeds the threshold $B^{\prime}$. This probability is the same for any $n$ due to the symmetry of the distribution. We introduce a random variable $o_{j}\in\{0,1\}$ to represent the following: we set $o_{j}=0$ if $\mathbf{x}_{n}^{\top}\bm{\beta}^{(K+1)}_{j}<B^{\prime}$ and $o_{j}=1$ if $\mathbf{x}_{n}^{\top}\bm{\beta}^{(K+1)}_{j}\geq B^{\prime}$. We define the following events:

$$ $\displaystyle\mathcal{E}_{1}=\Big{\{}\sum_{j\leq N_{0}}o_{j}=0\Big{\}},\qquad \mathcal{E}_{2}(\epsilon)=\Big{\{}pn-\frac{p}{2}-\epsilon\leq\sum_{j\leq n}o_{ j}\leq pn-\frac{p}{2}+\epsilon\Big{\}}.$ $$

Assume $\mathcal{E}_{1}\cap\mathcal{E}_{2}(\epsilon)$ holds. Then, we obtain

$$ $\displaystyle\frac{1}{i}\sum_{j=1}^{i}\left\langle\mathbf{Q}_{1}^{\mathrm{last }}\mathbf{h}_{i}^{(K+1)},\mathbf{K}_{1}^{\mathrm{last}}\mathbf{h}_{j}^{(K+1)} \right\rangle\cdot\mathbf{V}_{1}^{\mathrm{last}}\mathbf{h}_{j}^{(K+1)}$ $\displaystyle=\frac{1}{i}\sum_{j=1}^{i}\sigma\left(\mathbf{x}_{n}^{\top}\bm{ \beta}^{(K+1)}_{j}-B^{\prime}\right)\cdot\begin{bmatrix}\mathbf{0}_{d}\\ \mathbbm{1}_{\{i\%2=1\}}\\ \mathbf{0}_{d+1}\end{bmatrix}$ $\displaystyle=\frac{1}{2n-1}\sum_{j:o_{j}=1}\left(\mathbf{x}_{n}^{\top}\bm{ \beta}^{(K+1)}_{j}-B^{\prime}\right)\cdot\begin{bmatrix}\mathbf{0}_{d}\\ \mathbbm{1}_{\{i\%2=1\}}\\ \mathbf{0}_{d+1}.\end{bmatrix}.$ $$

Let the last MLP layer satisfies

$$ $\displaystyle\mathrm{MLP}\begin{bmatrix}\mathbf{0}_{d}\\ \widehat{y}\\ \mathbf{0}_{d+1}.\end{bmatrix}=\begin{bmatrix}\mathbf{0}_{d}\\ \frac{2}{p}\widehat{y}+B^{\prime}\\ \mathbf{0}_{d+1},\end{bmatrix},$ $$

it follows that

$$ $\displaystyle\mathrm{MLP}\left(\frac{1}{i}\sum_{j=1}^{i}\left\langle\mathbf{Q} _{1}^{\mathrm{last}}\mathbf{h}_{i}^{(K+1)},\mathbf{K}_{1}^{\mathrm{last}} \mathbf{h}_{j}^{(K+1)}\right\rangle\cdot\mathbf{V}_{1}^{\mathrm{last}}\mathbf{ h}_{j}^{(K+1)}\right)=\begin{bmatrix}\mathbf{0}_{d}\\ \frac{2}{2np-p}\sum_{j:o_{j}=1}\mathbf{x}_{n}^{\top}\bm{\beta}^{(K+1)}_{j}- \frac{2(\sum o_{j})}{2np-p}B^{\prime}+B^{\prime}\\ \mathbf{0}_{d+1}.\end{bmatrix}.$ $$

Denote $\widehat{y}_{n}=\frac{2}{2np-p}\sum_{j:o_{j}=1}\mathbf{x}_{n}^{\top}\bm{\beta}
^{(K+1)}_{j}-\frac{2(\sum o_{j})}{2np-p}B^{\prime}+B^{\prime}$, therefore

$$ $\displaystyle|y_{n}-\widehat{y}_{n}|$ $\displaystyle=\left|\mathbf{x}_{n}^{\top}\bm{\beta}^{*}-\frac{2}{2np-p}\sum_{j :o_{j}=1}\mathbf{x}_{n}^{\top}\bm{\beta}^{(K+1)}_{j}+\frac{2(\sum o_{j})}{2np- p}B^{\prime}-B^{\prime}\right|$ $\displaystyle\leq\|\mathbf{x}_{n}\|\left\|\bm{\beta}^{*}-\frac{2}{2np-p}\sum_{ j:o_{j}=1}\mathbf{x}_{n}^{\top}\bm{\beta}^{(K+1)}_{j}\right\|+\left|\frac{2( \sum o_{j})}{2np-p}B^{\prime}-B^{\prime}\right|$ $\displaystyle\leq b_{\mathbf{x}}\left\|\bm{\beta}^{*}-\frac{2}{2np-p}\sum_{j:o _{j}=1}\bm{\beta}^{(K+1)}_{j}\right\|+\frac{2\epsilon}{2np-p}|B^{\prime}|.$ (E.1) $$

We randomly select a set $\{o_{j}^{\prime}:j\geq 2N_{0}-1\}$ such that $|\{o_{j}^{\prime}\}|\leq\epsilon$, satisfying the following conditions:

- •
If $|\{o_{j}:o_{j}=1\}|\geq pn-\dfrac{r}{2}$, then
$\left|\{o_{j}:o_{j}=1\}\setminus\{o_{j}^{\prime}\}\right|=pn-\dfrac{r}{2};$
- •
If $|\{o_{j}:o_{j}=1\}|<pn-\dfrac{r}{2}$, then
$\left|\{o_{j}:o_{j}=1\}\cup\{o_{j}^{\prime}\}\right|=pn-\dfrac{r}{2}.$

We define the set $\mathcal{O}$ as follows:

$$ $\mathcal{O}=\begin{cases}\{o_{j}:o_{j}=1\}\setminus\{o_{j}^{\prime}\},&\text{ if }\left|\{o_{j}:o_{j}=1\}\right|\geq pn-\dfrac{r}{2},\[8.61108pt] \{o_{j}:o_{j}=1\}\cup\{o_{j}^{\prime}\},&\text{if }\left|\{o_{j}:o_{j}=1\} \right|<pn-\dfrac{r}{2}.\end{cases}$ $$

Therefore, we have

$$ $\displaystyle\left\|\bm{\beta}^{*}-\frac{2}{2np-p}\sum_{j:o_{j}=1}\mathbf{x}_{ n}^{\top}\bm{\beta}^{(K+1)}_{j}\right\|\leq$ $\displaystyle\left\|\bm{\beta}^{*}-\frac{1}{np-\frac{p}{2}}\sum_{j:o_{j}\in \mathcal{O}}\bm{\beta}^{(K+1)}_{j}\right\|+\left\|\frac{1}{np-\frac{p}{2}}\sum _{j:o_{j}\in\{o_{j}^{\prime}\}}\mathbf{x}_{n}^{\top}\bm{\beta}^{(K+1)}_{j}\right\|$ $\displaystyle\leq\left\|\bm{\beta}^{*}-\frac{1}{np-\frac{p}{2}}\sum_{j:o_{j} \in\mathcal{O}}\bm{\beta}^{(K+1)}_{j}\right\|+\frac{C_{\mathbf{x}}b_{\mathbf{x }}\epsilon}{np-\frac{p}{2}},$ (E.2) $$

where Equation E.2 follows from Lemma 4. Observe that

$$ $\displaystyle\left\|\bm{\beta}^{*}-\frac{1}{np-\frac{p}{2}}\sum_{j:o_{j}\in \mathcal{O}}\bm{\beta}^{(K+1)}_{j}\right\|\leq\frac{1}{np-\frac{p}{2}}\sum_{j: o_{j}\in\mathcal{O}}\left\|\bm{\beta}^{*}-\bm{\beta}^{(K+1)}_{j}\right\|.$ $$

Then, from Theorem 5.1, the following inequality holds with probability at least $1-3n\delta$:

$$ $\displaystyle\sum_{i:o_{2i+1}\in\mathcal{O}}\|\bm{\beta}^{*}-\bm{\beta}^{(K+1) }_{2i+1}\|\leq\sum_{i:o_{2i+1}\in\mathcal{O}}\|\bm{\beta}^{*}-\bm{\beta}^{(K+1 )}_{2i+1}\|_{1}\leq\sum_{i:o_{2i+1}\in\mathcal{O}}\|\bm{\beta}^{*}\|_{1}e^{- \alpha_{i}K},$ (E.3) $$

where

$$ $\displaystyle\alpha_{i}=-\log\Bigg{(}1-\frac{2}{3}\gamma+\gamma(2S-1)\sqrt{ \frac{\log d-\log\delta}{ic}}+\gamma\sqrt{\frac{\log S-\log\delta}{ic}}\Bigg{)}.$ $$

From Cauchy–Schwarz Inequality we have:

$$ $\displaystyle\gamma(2S-1)\sqrt{\frac{\log d-\log\delta}{ic}}+\gamma\sqrt{\frac {\log S-\log\delta}{ic}}$ $\displaystyle\leq 2\gamma(2S-1)\sqrt{\frac{\log d+\log S-\log\delta}{ic}}.$ (E.4) $$

Combining Inequalities E.3 and E.4 gives

$$ $\displaystyle\sum_{i:o_{2i+1}\in\mathcal{O}}\|\bm{\beta}^{*}-\bm{\beta}^{(K+1) }_{2i+1}\|\leq\|\bm{\beta}^{*}\|_{1}\sum_{i:o_{2i+1}\in\mathcal{O}}\left(1- \frac{2}{3}\gamma+2\gamma(2S-1)\sqrt{\frac{\log d+\log S-\log\delta}{ic}} \right)^{K}.$ $$

Note that, by induction, it is straightforward to prove that the following inequality holds if $1-\frac{2}{3}\gamma\leq 1$:

$$ $\displaystyle\left(1-\frac{2}{3}\gamma+2\gamma(2S-1)\sqrt{\frac{\log d+\log S- \log\delta}{ic}}\right)^{K}$ $\displaystyle\leq\left(1-\frac{2}{3}\gamma\right)^{K}+\eta_{i}K\left(1-\frac{2 }{3}\gamma\right)^{K-1}\sqrt{\frac{\log d+\log S-\log\delta}{ic}},$ (E.5) $$

when $\eta_{i}$ satisfies

$$ $\displaystyle\eta_{i}\geq\frac{\left(1-\frac{2}{3}\gamma\right)}{\left(1-\frac {2}{3}\gamma\right)-(K-1)\sqrt{\frac{\log d+\log S-\log\delta}{ic}}},$ (E.6) $$

which can be satisfied by simply setting

$$ $\displaystyle\eta_{i}=\eta=\frac{\left(1-\frac{2}{3}\gamma\right)}{\left(1- \frac{2}{3}\gamma\right)-(K-1)\sqrt{\frac{\log d+\log S-\log\delta}{N_{0}c}}}.$ (E.7) $$

Therefore, Inequality (E.6) holds for any $i\geq N_{0}$. Based on Inequality (E.5), we have

$$ $\displaystyle\sum_{i:o_{2i+1}\in\mathcal{O}}\left(1-\frac{2}{3}\gamma+2\gamma( 2S-1)\sqrt{\frac{\log d+\log S-\log\delta}{ic}}\right)^{K}$ $\displaystyle\quad\leq\sum_{i=N_{0}+1}^{N_{0}+np-\frac{p}{2}}\left(1-\frac{2}{ 3}\gamma+2\gamma(2S-1)\sqrt{\frac{\log d+\log S-\log\delta}{ic}}\right)^{K}$ $\displaystyle\quad\leq\sum_{i=N_{0}+1}^{N_{0}+np-\frac{p}{2}}\left(1-\frac{2}{ 3}\gamma\right)^{K}+\eta K\left(1-\frac{2}{3}\gamma\right)^{K-1}\sum_{i=N_{0}+ 1}^{N_{0}+np-\frac{p}{2}}\sqrt{\frac{\log d+\log S-\log\delta}{ic}}$ $\displaystyle\quad\leq(np-\frac{p}{2})\left(1-\frac{2}{3}\gamma\right)^{K}+ \eta K\left(1-\frac{2}{3}\gamma\right)^{K-1}\sqrt{\frac{\log d+\log S-\log \delta}{c}}\sum_{i=N_{0}+1}^{N_{0}+np-\frac{p}{2}}\frac{1}{\sqrt{i}}$ $\displaystyle\quad\leq(np-\frac{p}{2})\left(1-\frac{2}{3}\gamma\right)^{K}+ \eta K\left(1-\frac{2}{3}\gamma\right)^{K-1}\sqrt{\frac{\log d+\log S-\log \delta}{c}}2\left(\sqrt{N_{0}+np-\frac{p}{2}}-\sqrt{N_{0}}\right).$ (E.8) $$

By combining Inequalities (E.3) and (E.7), we obtain

$$ $\displaystyle\sum_{i:o_{2i+1}\in\mathcal{O}}\|\bm{\beta}^{*}-\bm{\beta}^{(K+1) }_{2i+1}\|$ $\displaystyle\leq(np-\frac{p}{2})\left(1-\frac{2}{3}\gamma\right)^{K}+\eta K \left(1-\frac{2}{3}\gamma\right)^{K-1}\sqrt{\frac{\log d+\log S-\log\delta}{c}}$ $\displaystyle\leq(np-\frac{p}{2})\left(1-\frac{2}{3}\gamma\right)^{K}+2\eta K \sqrt{np-\frac{p}{2}}\left(1-\frac{2}{3}\gamma\right)^{K-1}\sqrt{\frac{\log d+ \log S-\log\delta}{c}}.$ $$

Therefore,

$$ $\displaystyle\frac{1}{np-\frac{p}{2}}\sum_{i=N_{0}+1}^{n}\|\bm{\beta}^{*}-\bm{ \beta}^{(K+1)}_{2i+1}\|\leq\left(1-\frac{2}{3}\gamma\right)^{K}+\frac{2\sqrt{2 }\eta K}{\sqrt{np}}\left(1-\frac{2}{3}\gamma\right)^{K-1}\sqrt{\frac{\log d+ \log S-\log\delta}{c}}.$ (E.9) $$

Combining Equation E.1, Equation E.2 and Equation E.9 we obtain

$$ $\displaystyle\|y_{n}-\widehat{y}_{n}\|$ $\displaystyle\leq b_{\mathbf{x}}\left(1-\frac{2}{3}\gamma\right)^{K}+\frac{2 \sqrt{2}b_{\mathbf{x}}\eta K}{\sqrt{np}}\left(1-\frac{2}{3}\gamma\right)^{K-1} \sqrt{\frac{\log d+\log S-\log\delta}{c}}+\frac{2\epsilon}{np}|B^{\prime}|+ \frac{2C_{\mathbf{x}}b_{\mathbf{x}}\epsilon}{np}.$ $$

Set $\epsilon=K\sqrt{np}(1-2\gamma/3)^{K}$, we obtain

$$ $\displaystyle\|y_{n}-\widehat{y}_{n}\|\leq b_{\mathbf{x}}\left(1-\frac{2}{3} \gamma\right)^{K}+\frac{c_{4}K}{\sqrt{n}}\left(1-\frac{2}{3}\gamma\right)^{K-1},$ $$

where

$$ $\displaystyle c_{4}=\frac{2\sqrt{2}b_{\mathbf{x}}\eta}{\sqrt{p}}\sqrt{\frac{ \log d+\log S-\log\delta}{c}}+\frac{2|B^{\prime}|}{\sqrt{p}}+\frac{2C_{\mathbf {x}}b_{\mathbf{x}}}{\sqrt{p}}.$ $$

Denote $\delta^{\prime}$ as $\delta^{\prime}=\mathbb{P}\{\neg\mathcal{E}_{1}\}+\mathbb{P}\{\neg\mathcal{E}_
{2}(K\sqrt{np}(1-2\gamma/2)^{K})$, then the theorem follows.

∎

Next, we establish the result concerning the label prediction loss presented in Theorem 5.2 through the following corollary:

###### Corollary E.1 .

Suppose Assumption 1 holds.
For a Transformer with $K$ layers, where all layer structures are described in Section 4.3, and the input sequence is set as in Equation C.4, there exists a set of parameters in the Transformer and an explicitly quadratic readout function $F$ such that, if $n\geq N_{0}$, the following inequality holds with probability at least $1-\delta$:

$$ $\displaystyle\|y_{n+1}-\widehat{y}_{n+1}\|\leq c_{6}e^{-\alpha_{n}K},$ $$

where $c_{6}=\sqrt{d}b_{\mathbf{x}}c_{1}$, and $c_{1}$ and $\alpha_{n}$ are defined in Theorem 5.1.

###### Proof of Corollary E.1 .

We set the $\mathbf{V}$ matrix in the read out function $F_{\mathbf{V}}$ as

$$ $\displaystyle\mathbf{V}=\begin{bmatrix}\mathbf{0}_{(d\!+\!1)\!\times\!d}\\ \mathbf{I}_{d\!\times\!d}\\ \mathbf{0}_{1\!\times\!d}\end{bmatrix}.$ $$

From Proposition 1, when $i$ is an odd index, $\mathbf{h}_{i}^{(K+1)}$ is of the following structure:

$$ $\displaystyle\mathbf{h}_{i}^{(K+1)}=\begin{bmatrix}[\mathbf{X}^{\top}]_{:, \frac{i+1}{2}}\\ 0\\ \bm{\beta}^{(K+1)}_{i}\\ 1\end{bmatrix}.$ $$

Therefore, the output after the read-out function is

$$ $\displaystyle\widehat{y}_{\frac{i+1}{2}}=F_{\mathbf{V}}(\mathbf{h}_{i}^{(K+1)} )=[\mathbf{X}]_{\frac{i+1}{2},:}\bm{\beta}^{(K+1)}_{i}.$ $$

Note that $\left\|[\mathbf{X}]_{\frac{i+1}{2},:}\right\|\leq\sqrt{d}b_{\mathbf{x}}$, we have

$$ $\displaystyle\left\|y_{\frac{i+1}{2}}-\widehat{y}_{\frac{i+1}{2}}\right\|\leq \left\|[\mathbf{X}]_{\frac{i+1}{2},:}\right\|\left\|\bm{\beta}^{*}-\bm{\beta}^ {(K+1)}_{i}\right\|\leq\sqrt{d}b_{\mathbf{x}}\left\|\bm{\beta}^{*}-\bm{\beta}^ {(K+1)}_{i}\right\|.$ $$

Combining with Theorem 5.1, the proof is thus complete.
∎

Therefore, by combining Corollary E.1 and Theorem E.1, we arrive at the conclusion of Theorem 5.2.

###### Lemma 4 .

###### Proof of Lemma 4 .

###### Theorem E.1 .

###### Proof.

###### Corollary E.1 .

###### Proof of Corollary E.1 .

## Appendix F Additional Experiment Results

In Section 6, we demonstrate that classical LISTA-type algorithms, such as LISTA, perform poorly when applied to varying measurement matrices after being trained on a fixed measurement matrix $\mathbf{X}$. In this section, we further show that unlike LISTA-VM, these classical LISTA-type algorithms fail to handle in-context sparse recovery problems, even when trained on varying measurements.

The training setup is identical to how we train LISTA-VM. Specifically, we set the number of iterations to $K=12$. During each epoch, we randomly sample 100 measurement matrices, each generating 500 instances from 500 randomly generated sparse vectors, resulting in a total of 50,000 instances per epoch. For each epoch, we minimize the sparse vector prediction loss $\sum_{j=1}^{N}\|\widehat{\bm{\beta}}_{j}-\bm{\beta}_{j}\|^{2}$ using gradient descent. We train the model for a total of 340 epochs and conduct in-context testing at the end of every epoch.

In the results shown in Figure 2, we observe that the prediction error on varying $\mathbf{X}$ for meta-trained LISTA and LISTA-CP remains around $3$. As illustrated in the figure, there is no observable trend indicating improvement in the error throughout the training process.

Figure: (a) LISTA
Refer to caption: x4.png

For comparison, in Figure 3, we also provide the results of the prediction error on varying $\mathbf{X}$ for meta-trained LISTA-VM. The final prediction error is around $0.68$, which is significantly more promising compared to classic LISTA-type algorithms.

Figure: Figure 3: Experimental results for meta-trained LISTA-VM
Refer to caption: x6.png

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: 1). Could the authors clarify their rationale for choosing to study this problem in the context of sparse vector recovery and explain how it relates to the hypothesis presented in this work?
soundness: 4
strengths: **Strengths:**

1. **Strong Theoretical Contributions:**  
   The paper presents a compelling theoretical claim addressing a crucial aspect of machine learning with significant potential for diverse applications. Moreover, the proof techniques introduced are versatile and can be extended to other domains beyond sparse vector recovery, such as compressed sensing and various inverse problems.

2. **Clarity and Rigor of Proofs:**  
   The proofs are well-written and easy to follow, both in the main text and in the appendix.

3. **Robust Empirical Validation:**  
   The empirical experiments are thorough and effectively corroborate the theoretical findings
summary: This paper presents a theoretical study of transformers to address the hypothesis regarding in-context learning for sparse recovery problems, as discussed in lines 62-65. Transformers are an integral part of modern deep learning architectures, and their contributions cannot be overstated. In particular, this work focuses on making theoretical claims that transformers can perform LISTA-type algorithms. The authors support this with a rigorous theoretical claim in Theorem 5.1, which demonstrates with high probability that transformers can recover the optimal sparse vector. Lastly, the authors empirically validate their claim in Figure 1.
weaknesses: The hypothesis outlined in lines 62-65 appears somewhat disconnected from the application to sparse recovery. The connection between in-context learning (ICL) and sparse recovery is not clearly established, making it challenging to understand why this particular relationship is being explored.

---

confidence: 2
contribution: 2
initial_rating: 6
presentation: 2
questions: - In my opinion, the phrase L2O is not appropriate in the context of this paper (or any approach for constructing a learned network to solve an inverse problem for that matter). These methods essentially try to approximate the conditional mean of the parameter of interest given the measured data and do not find an approximation to a variational reconstruction problem such as LASSO (not without any constraints on the learnable units in the architecture, at least). 

- How does a K-layer pre-trained transformer perform if more layers are added during the inference time (by extending the learned layers using some reasonable strategy)? If this leads to divergence, it would not be appropriate to claim that this approach learns to approximate the LASSO solution. This just gives us a more powerful and expressive architecture that does a better job of approximating the conditional mean estimator of the underlying sparse vector from its linear measurement. 

- Could you explain whether Figure 1 considers the same sensing matrix during training and inference? Maybe I missed it, but I was not able to figure it out from Section 6. 

- How is the support constraint incorporated into your framework? In particular, why does the performance of the proposed method improve with these constraints, while the other baseline methods perform roughly the same?

- While first-order optimization methods for LASSO can only achieve a sublinear rate (and not anything faster, provably), the result in Theorem 5.1 appears paradoxical. I would suggest highlighting the key differences in the setting and assumptions made in the optimization literature and this work to make this bit clear. I understand that the rate in Theorem 5.1 holds with high probability, but is it possible to construct a sparse signal for which a pre-trained transformer fails to achieve a linear rate? 

- Assumption 1 seems somewhat non-standard to what is generally assumed about the sensing matrix in the compressed sensing literature. For instance, a Gaussian sensing matric would not satisfy this assumption.
soundness: 3
strengths: 1. The theoretical insights on the versatility of the transformer architecture are interesting and useful for a broader class of problems. For example, the fact that LISTA-type algorithms can be implemented using a transformer opens up the possibility of using transformers for building more powerful learned reconstruction operators for more general inverse problems. 

2. The convergence rate results (for sparse estimation and prediction) place the proposed approach on a concrete theoretical footing. 

3. Generalizability to a different sensing matrix during the inference time is a very powerful property of a learned approach, which (in my knowledge) is not generally possessed by standard learned reconstruction operators trained on a specific measurement matrix.
summary: The authors leverage the ability of Transformers to perform in-context learning (ICL) for solving the sparse recovery problem. It is shown that the Transformers can implement first-order optimization algorithms (such as ISTA) and their learned variants (e.g., LISTA) for sparse approximation. Moreover, it is shown theoretically that a K-layer transformer has a linear convergence rate in K (with high probability). The ability of transformers to solve the sparse recovery problem corresponding to a different sensing matrix than what it was trained on is also demonstrated.
weaknesses: 1. The notion of ICL in the context of sparse recovery is not very well explained. I would suggest rewriting Sec. 4.1 with a better explanation of the setup (i.e., exactly what the model is trained on and what exactly is given as input to the model during inference). 

2. The convergence result (Theorem 5.1) merely ascertains the existence of a set of parameters such that the recovery is accurate with high probability, and does not state anything about the convergence of a pre-trained transformer on a new problem instance.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: See weakness
soundness: 3
strengths: The manuscript is articulated with clarity and is accessible for comprehension. 
The proposed hypothesis presents a valuable avenue for in-depth analysis. 
Preliminary concepts are provided with comprehensive equations and definitions. 
Elaborate proofs are accessible in the appendix. 
The quality of English is commendable and articulated naturally.
summary: This paper threw an interesting hypothesis: the trained transformers essentially play the role of the L2O algorithm when solving the inverse problem. Enriched theory analysis and derivation are proposed, which are almost correct. This hypothesis is validated underlying a sparse recovery problem with respect to ISTA-based alternating algorithms. The fundamental idea can be summarized as proving the convergence and accuracy follow a similar criterion in ISTA-variant algorithms. Simulations on the trained transformer and ISTA-based algorithms for solving sparse recovery tasks are presented, and the blind inverse problem is particularly picked out to show the stronger performance of the transformer, and thus leads to the L2O capability.
weaknesses: I made a rejection at this stage as the following shortbacks, I am willing to change my score after the rebuttal phase with respect to the author's responses.
1. The entire work is divided. The primary concept, prove that transformer has the L2O approximation capability in forward process, is isolated to the main task, sparse recovery problem. In this version ,this paper tends to mash-up the L2O approximation on transformer and sparse recovery problem. Essentially, each of them deserves to publish an independent paper. 
2. Followed by 1, this causes the following issues: i) no analysis and illustrations on the necessity of L2O for sparse recovery tasks; ii) the bridge between the fundamental proofs and evaluation tasks is fragile; iii) the increments of this paper compared to . Von Oswald et al. (2023a) and  Bai et al. (2024) are invisible, though the authors claim the difference is the L2O consideration.
3. The writing of this paper is also a mixture. The logical flow of the introduction is saltatory. The fundamental motivation is not clear enough to convince me that the L2O algorithmic principle is the explanation of the success of the transformer because of there are L2O-based algorithms that achieve similar performance. 
4. The whole section 3 is overwhelming. I don't think it deserves so much context to present details network layer symbols and definition. Besides, Definition 3.1-3.3 is kind of a definition abuse. I highly suggest a concise and shorter version of this section.
5. The key parts, Sections 4.3, 5.1, and 5.2, show a trivial way of proving the author's hypothesis. Performing LISTA-type algorithms do not provide a direct derivation of the L2O of the transformer forward process, at least to me. 
6. Last but the most fatal issue lies in simulation. The authors tend to show the superior performance of a trained transformer compared to those alternating minimization-based ISTA-type algorithms, concluding the L2O capability of transformer in the blind inverse problem of sparse recovery where matrix X is not given. This is unfair and not reasonable. The trained transformer obviously surpasses those model-based algorithms which are designed for matrix X is given. The transformer solves the blind inverse problem due to the data-driven prior from the training process, instead of L2O capability that focuses on improving the optimization strategy during the iterations. A fair comparison should be set among algorithms with learning ability and capable of solving blind inverse problems, such as
[1] Zero-shot image restoration using denoising diffusion null-space model
[2] Metalearning-based alternating minimization algorithm for nonconvex optimization
[3] Invertible Diffusion Models for Compressed Sensing
[4] OmniSSR: Zero-shot Omnidirectional Image Super-Resolution using Stable Diffusion Model
[5] Blind Super-Resolution Via Meta-Learning and Markov Chain Monte Carlo Simulation

---

confidence: 4
contribution: 4
initial_rating: 6
presentation: 4
questions: See in the comments in Weakness.
soundness: 3
strengths: The paper is well-written, with clearly presented theoretical findings and compelling experimental results. The insights are both intriguing and hold valuable potential for contributing to the scientific community.
summary: The authors of this paper study sparse recovery using transformers in an in-context scenario. They provided theoretical results and conducted experiments regarding the ability of transformers to execute an L2O algorithm.
weaknesses: I have two concerns which should be better explained in the paper:
- The embedding form in Equation 4.2 appears somewhat ad hoc; do you have any references supporting this choice? I understand that this choice is closely tied to the design of the four attention heads in Appendix C1, so I wonder: are alternative forms feasible, and if so, how might they impact the theoretical results?

- In the experimental section, the authors test with GPT-2 using 8 heads, which differs from the 4-head configuration in the theoretical analysis. What motivates this discrepancy? Additionally, I encourage the authors to discuss how varying the number of heads might affect the theoretical findings.

Meta-review: The authors show that sparse recovery can be solved using in-context learning (ICL) with transformer models. Specifically, they show that transformers with ICL can implement LISTA-style sparse recovery algorithms, and that the convergence rate of the recovery algorithm is linear in the depth of the transformer model. There also appear to be numerical benefits compared to older algorithms such as LISTA and ALISTA.  

The paper is very nicely written and the contributions are clear. The results are not terribly surprising and follow a fairly logical next step pointed by the existing literature: there have been a host of recent papers showing that transformers can do gradient descent with ICL; LISTA is just gradient descent with learnable parameters interleaved with soft thresholding; and soft thresholding can be implemented using the ReLU layers in the MLPs of transformers. The results on convergence rates seem to also be mostly a consequence from previous papers. Nonetheless, it takes effort to make everything work, and the authors do a good job putting all the building blocks together in a clear manner.

Decision: Accept (Poster)