On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery
Authors: Renpu Liu
Ruida Zhou
Cong Shen
Jing Yang
Venue: ICLR 2025 Poster (2025)
Keywords: Transformer
In-context learning
Learning-to-optimize
Primary area: learning theory

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-layer Transformer can perform an L2O algorithm with a provable convergence rate linear in 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.

Published as a conference paper at ICLR 2025
ON THE LEARN-TO-OPTIMIE CAPABILITIES OF TRANSFORMERS IN IN-CONTEXT SPARSE RECOVERY
Renpu Liu'; Ruida Zhou²; Cong Shen', Jing Yang't
University of Virginia, Univeessity of California, Los Angeles {pzw7bx, cong, ruidaēg acla.edu
ABSTRACT
An intriguing property of the Transformer is its ability to perform in-context learn- ing (ICL), where the Transformer can solve different inference tasks without pa- rameter updating based on the contextual information provided by the correspond- ing input-output demonstration pairs. It has been theoretically proved that ICL is enabled by the capability of Transformers to perform gradient-
algo- rithms (Von Oswald et al. 2023a Bai al. 2024). This work takes step further and shows that Transformers can perform reanningtoo-opimiiee (I.20) algorithms. Specifically, for the ICL sparse recovery (formulated as LASSO) tasks, we show that a K K-layer Transformer can perform an L20 algorithm with a provable conver- gence rate linear K. This provides new perspective explaining the superior ICL capability of Transformers, even with only few layers, which cannot be achieved by the standard gradient-descee algorithms. Moreover, unlike the con- ventional L20 algorithms that require the measurement matrix involved in training to match that testing, the trained Transformer is able to solve sparse recovery problems generated with different measurement matrices. Besides, Transformers as an L20 algorithm can leverage structural information embedded the train- ing tasks to accelerate its convergence during ICI., and generalize across differ- ent lengths of demonstration pairs, where conventional L20 algorithms typically struggle or fail. Such theoretical findings are supported by our experimental re- sults.
INTRODUCTION
Since its introduction in Vaswani et al. (2017). Transformers have become the backbone in various fields such as natural language processing (Radford 2018 Devlin 2018). computer vision (Doso-) vitskiy 2020) and reinforeement learning (Chen al. 2021). significantly influencing subsequent research and applications. A notable capability of Transformers is their good performance for in- context learning (ICL) (Brown al. 2020), 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 ofe----rtttt few-shot perfor- mances across a multitude of tasks, such as reasoning and language understanding tasks in natural language processing (Chowdhery et al., 2023). in-context dialog generation (Thoppilan et al. 2022) and in-context linear regression (Garg et al. 2022 Fu et al., 2023).
Given the significance of transformers' ICL capabilities, extensive research efforts have been di- rected toward understanding the mechanisms behind their ICI, performance. In this context, the ICI. capability pre-trained Transformer is understood as the Transformer's implicit implementation of learning algorithms during the forward pass. Von Oswald et al. (2023a). Dai et al. (2022), and Bai op-
ctt al. (2024) suggest that these learning algorithms closely approximate gradient-descent-bass timizers, thus making the Transformer a universal solver for various ICL tasks. Specifically, these works demonstrate that Transformers can approximate gradient descent steps implicitly through
'These authors contributed equally to this work.
1 "Corrsponniig author.

---

Published as a conference paper at ICLR 2025
some specific constructions of their parameters, enabling them to adapt to new data points during inference without explicit re- re-training.
However, it is well known that seffient-descent-beee plicated optimization problems and thus may not be sufficient to explain the superior performance of Transformers a plethora of ICL tasks. A recent work (Ahn al. 2024) suggests that instead of gradient descent, Transformers actually perform pre-conditioned gradient descent during ICL. In other words, learns pre-conditioner during pre-training and then utilizes it during ICL to expe- dite the optimization process. Von Oswald et al. (2033b) recently demonstrates that the forward pass of a trained transformer can implement meta-op veeta-opimizationn algorithms, i.e, it can implicitly define internal objective functions and then optimize these objectives to generate predictions. Similarly, Zhang et (202)) 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 Searning-to-Optinee (L.20) algorithms, and motivates the following hypothesis:
algorithms not efficient in solving com-
Trunsformer does not simply implement universal optimization algorithm during ICL. Rather, il extracts useful information from the given dalaset 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 classical signal processing problem that is of significant practical interest across various domains, such as compressive sensing in medical imaging (Shen et al. 2017) and spectrum sensing able to implement gradient descent-basee
(Elad 2010). Recent works show that Transformers
algorithms with sublinear convergence rates for in-context sparse recovery (Bai et al.) 2024 Chen
et al. 2024b). However, empirical findings indicate that Transformers
solve in-context sparse recovery more efficiently than gradient bescent--aseee approaches (Bai et al. 2024). Mcanwhile,
there exists plethora of L20 algorithms that solve the classical sparse recovery problem efficiently
with linear convergence guarantees (Gregor and LeCun 2010 Chen et al., 2018 Liu and Chen 2019). Therefore, examining the L20 capabilities of Transformers in solving the in-context sparse recovery task becomes a promising direction and may serve as perfect example to validate our hypothesis. Our main contributions are as follows.
First, we demonstrate that Transformers can implement an L20 algorithm for in-context sparse recovery, and theoretically prove that a K K-layer Transformer as an L2O algorithm can recover the underlying sparse vector in-context at a convergence rate linear in K. The linear convergence results in this work significantly improve the state-of-the-art convergence results for in-context sparse recovery and validate our previous hypothesis.
Second, we show that the Transformer as an L20 algorithm can actually outperform traditional L20 algorithms for sparse recovery several aspects: 1) It does not require the measurement ma- trices involved in training to be the same, which is in stark contrast to traditional L20 algorithms (Gregor and LeCun 2010; Chen et al. 2018, Liu and Chen 2019) for sparse recovery, and 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 struc- tural 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 L20 algorithms for sparse recovery empirically. Our experimental results indicate that Transform- ers substantially outperform traditional gradient-descent-baee comparable performances compared with L20 algorithms that are trained and tested using data generated with the same measurement matrix. This supports our claim that Transformers can implement L20 algorithms during ICL. Besides, Transformers also demonstrate remarkable gen- eralization capability when the measurement matrix varies, and achieve accelerated convergence when additional structure is imposed on the underlying sparse vectors, supporting our theoretical findings.
iterative algorithms, and achieve
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORKS
ICL Mechanism for Transformers. Brown et al. (2020) first show that GPT-3, 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. (2023) empirically hypothesize that Transformers perform kernel regression with internal represen- tations when facing in-context examples, and Fu et al. (2023) 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. (2022) show that by construction, Transformers can imple- ment gradient descent-basee based algorithmss for linear regression problems. [on Oswald et al. (2023a) construct explicit weights for a Transformer, claiming can perform gradient descent on linear and non-linear regression tasks. [Bai et al. (2024) provide constructions such that Transformers can make selections between different gradient-base based algorithms. [Ahn et al. (2024) reason the ICL ca- pability of Transformers to their ability to implement pre-conditioned gradient descent for linear regression tasks, where the pre-condition matrix is learned from pre-training. Recently, Von Oswald et al. (2023b) demonstrates that the forward pass of a trained transformer implement meta- optimization algorithms, i.e., it can implicitly define internal objective functions and then optimize these objectives to generate predictions. Similarly, /Zhang al. (2023) show that mesa-optimizer embedding the covariance matrix of input data can efficiently solve linear regression tasks. Our work belongs this category, where we construct a Transformer structure that can implement an L20 algorithm to effectively solve the in-context sparse recovery problem.
L20 Algorithms for Sparse Recovery. The L20 framework, summarized in the review pa- per (Chen et al. 2022). is an optimization paradigm that develops an optimization method (i.c., a solver) by training across set of similar problems (tasks) sampled from a task distribution. While the training process is often offline and time-consuming, the objective of L20 is to improve the op- timization efficiency and accuracy when the method is deployed online and any new task sampled from the same distribution is encountered.
As general optimization framework, 1,20 has demonstrated its advantage over classical static op- timization frameworks in various optimization problems and applications. In this work, we focus on the L20 algorithms relevant to sparse recovery here and leave more discussions on general L20 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 al.) 2023) and stock market forecasting (Roy et al. 2015). thus motivates the design of efficient algo- rithms. E.g.. the iterative soft- woft-thresholding algorithm (ISTA) (Daubechies et al. (004) proposed to solve LASSO and improves over the standard gradient descent algorithm. Motivated by the ISTA structure, Gregor and LeCun (2010) introduce the Learned ISTA (LISTA), a feedforward neural net- work that incorporates trainable matrices into ISTA updates. Chen et al. (2018) and Liu and Chen (2019) further propose LISTA- Weight Coupling (LISTA-CP CC) and Analytic LISTA (ALISTA) with fewer trainable parameters, making them easier train. They also provide theoretical analyses demonstrating a linear convergence rate.
We note that for these existing LISTA-type algorithms, both the training and testing tasks (instances) are randomly generated by fixing the measuremeni matrix bui varying the underlying sparse vectors. In general, they cannot handle cases where the training instances are generated with varying mea- surement matrices, or the measurement matrices during training and testing do not match.
PRELIMINARIES
Notations. For matrix X, we use [ p.g.r:s to denote the submatrix that contains rows p to g and columns and we use [ and [X], denote the i-th column and j-th row of X respectively. In some places, we also use [] to denote its i-th column for convenience. We use || pr to denote
its Frobenius norm. For vector x, we use |x||,,11 ||x|| and ||x||cc ee to denote its f1, & and for norms,
3

---

Published as a conference paper at ICLR 2025
respectively. We denote by 1, and O2 the d-dimensional all-1 and all-O column vectors, respectively. laxb and Oaxb denote the all-1 and all-O matrices of size x b, respectively.
3.1 TRANSFORMER ARCHITECTURE
In this work, consider the decoder-based Transformer architecture (Vaswani al. 2017). 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 M-head masked attention layer parameterized by {(V Qm, ) (mn) } as (VV (... (.)) where Vms Qms Km DXD, Vm [M].
Then, given an input sequence H E Dx(+1))) the output sequence of the attention layer is
M
Σ (V (V x mask (o (((mH) (QmH))))
A on Q.. (H) H
mm
where mask(M) satisfies maskk(M)], = - MM, S j and (mask(MM). 0 otherwise. (-) denotes the activation function.
Definition 3.2 (MLP layer). Given W1 E D'D W₂ DxD' and bias vector b E RR', MLP layer following the decoder attention layer, denoted as MLP MLP(Ww b2,), maps each token in the input sequence (i.e, each column h, in H E DxN ) to another token as
MLP (W., (b..)) h, W20 i,; b),
where a 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). one-layer decoder-base based Transformer parameterized by O (W1, W2, b, (Vm: m: Qm, Km me(M)) denoted TFe. Therefore, give impul sequence H T dx the output sequence is:
(H) MLP (W.,W,.b) (Atttn(vv...,.Kk..))
3.2 IN-CONTEXT LEARNING BY TRANSFORMERS
For an in-context learning (ICL) task, a trained Transformer is given an ICL instance (D+xN+1), where D (xxx 1/) [EN] and XN+1 is query. Here, x; E Rd is in-context example, and y; is the corresponding label for x;. We assume Vi fx();) 6g, where 6; is an added random noise, and ʃa is a deterministic function parameterized by B. Unlike conventional supervised learning, for each ICL instance, B Pa. i.e., it is randomly sampled from a distribution Ps.
To perform ICL in a Transformer, we first embed the ICL instance into input sequence H DD'N', The Transformer then generates an output sequence TF(H) with the same size based on which a prediction +1 1 is generated through a read-out function F, i.e., ØN+1 = (TF( The objective of ICL then ensure that UN+1 closely approximates the target value yN+1 (x(x)+1) + CN+1 for any ICL instance.
When a Transformer pre-trained for ICL, first samples large set ICL instances. For each instance, the Transformer generates prediction ØN+1 1 and calculates the prediction loss by compar- ing with y using 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 at- tracted 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. Von Oswald al. (2023)). Dai et al. (2022). and Bai et al. (2024) suggest that these learning algorithms closely approx- imate gradient-descent-haee optimizers, thus making the Transformer universal solver for various
4

---

Published as a conference paper at ICLR 2025
ICL tasks. A recent work (Ahn al. 2024) suggests that instead of gradient descent, Transformers actually perform crn-coniito gradient descent for in-context least square linear regression. In general, these results corroborate the claim that the mechanism of the Transformer in-context learn- ing a 1.20 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 LISTA-TYPE ALGORITHM FOR IN-CONTEXT SPARSE RECOVERY
In this section, demonstrate that decoder- Transformer can implement a novel LISTA-type L,00 algorithm, specifically LISTA-VI as detailed in Theorem 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. aims identify an S-sparse vector 3 Id from its noisy linear observations = X,8 + where X E I. N is measurement matrix, e E R is an isometric Gaussian noise vector with mean vector ON and covariance matrix xNxN. Typically, we assume d > N, which is the so-called under-determined case. One common assumption for the measurement matrix X PPitaall al.] 20155) Ge et al.) 20177) Zhh al.] (02T) that each row of the matrix is indepen- dently sampled from isometric sub-Gaussian distribution with zero mean and covariance matrix diag(o? .... o²)), denoted as Px. which guarantees the critical restricted isometry property under mild conditions the sparsity level (Candes and Tao. 2007). In this work, also assume 3 is randomly sampled from a distribution Pg. 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 3 E TRd that minimizes the following loss:
C(3) = ly x + al|31|11
Here is coefficient controlling the sparsity penalty. We denote the transpose of the i-th row in X by x,, i.c., x; = X J
In this work, study how Transformers solve the sparse recovery problem in context. For the pre-training process, set of in-context sparse recovery instances {(X), ,',x//11 ) ) generated according to the relationship v/. = T + + for j E [Nirain] and n E [N + 1]. where @ Ps. x Px, and e), ~ Pe, respectively. Given the set of training instances, pre-trained Transformer is obtained by minimizing a certain loss function.
Non is
After pre- pre-training, during the inference process for ICL, a sparse recovery instance (X is sampled randomly according to the same distributions as in the pre-training, and the Transformer then aims to predict y1+1 using the input (, without any further parameter updating.
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. 2018). This inefficiency has led to the development 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 fixed initial point ((1), the update rule in the k-th iteration is given by g(++1) Sa/L (L
X (x31))-)) -y))
Here, Sa/L is the soft-thresoolding function defined as [Sa/L(x)) = sign( (x))) max(0, | a/L}, and L is typically chosen as the largest eigenvalue of X X CCenn et al. 2018 Liu and Chen 2019
5

---

Published as a conference paper at ICLR 2025
Generally, for any ground-truth sparse vector 3 and any given X, ISTA converges sublinear rate (Beck and Teboulle 2009). The sublinear convergence rate ISTA considered inefficient, which has led to the development of various LISTA-type L2O algorithms, such as LISTA (Gre- gor and LeCun, 2010), LISTA-CP (Chen et al. 2018), and ALISTA (Liu and Chen, 2019). 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
g(++1) = Spmn (B (D( (x,3() - y))
(4.1)
where {B(k), ,D()) } are learnable parameters. Compared with ISTA with fixed parameters, LISTA- CP obtains {Ø(k) D(k) } through pre-training. Specifically, with fixed measurement matrix X. it randomly samples n S-sparse vectors {B, ~ P3 and generates (y, } which then utilized Jheh to optimize { { } by minimizing the total predicting loss for {,3, 13- Chen al. (2018) show that, for the same measurement matrix X, given any random instance (B",y), .y), a pre-trained LISTA- CP will converge the ground-truth 3 linearly in K under certain necessary conditions 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 trained Transformer can implement LISTA-type algorithm and efficiently solve sparse recovery problem in context. To distinguish the algorithm implemented by the Transformer with the classical LISTA-type algorithms, we term as LISTA with Varying Measurements (LISTA- VM). Towards this end, we provide an explicit construction of a K K-layer secoder-based Transformer as follows. K-layer Transformer the concatenation of blocks, where each block comprises a self-attention layer followed by an MLP layer. The input to the first self-attention layer, denoted as (1), is an embedding of the given in-context sparse recovery instance I (X, y, XN-
Embedding. Given an in-context sparse recovery instance I (X, y, xN+1) we embed the in- stance into an input sequence ( IR (x22+2)) follows:
XN 0
XN-1
(Z)
3
3
V1
3
1)
UN
B (1)
(4.2)
2N-1
2N
2N+1
he|aN++ Rdd are implicit parameter vectors initialized as O22 and x; is the i-th col- umn of the transposed measurement matrix, i.e, [X]] I We note that a similar embedding structure is adopted in Bai et al. 2024)..
where (B(0)
Self-attention layer. The self-attention layer takes as input a sequence of embeddings and outputs sequence of embeddings of the same length. Let the K Kllayer 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 Appendix [.. This construction ensures that the self-attention layer performs the 33 updating inside the soft-thresholding function in Equation (4. Furthermore, with our construction, the learnable matrix ( in LISTA-CP becomes context- deeeed-de instead of fixed. In the update rule implemented by the self-attention layer, this matrix becomes D
2N+1 ((M where (()) E pdxd 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 transformed sequence of the embeddings. Recall that h; is the i-th column in an embedding sequence H. We pa- rameterize (W1, W2, b) in the k-th MLP layer to let it function as a partial soft-threshold function:
[h [hs]1.441
MLP(h,) = Sgun [hy (d42:2d-1)
(4.3)
where Sac is the soft-threshold function. Essentially, the soft-threshold function is effectively implemented by the MLP layer utilizing the ReLU activation. This can be realized by combining

---

Published as a conference paper at ICLR 2025
- (x), ReLU(-=2), ReLU(( 0), -ReLU(- (-LU +). and The implementation details of the MLP layer can found Appendix C.1
Read-out function. Given the output sequence of the Transformer TF (H()), to obtain the esti- mation I, 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). Flinear is defined as the class of linear readout functions such that
Flinear {F(-) = h
Definition 4.2 (Query read-out). Fquery is defined as the class of explicit quadratic readout functions such that
Fquery = (F(-; F(h, i; h VX
1 VERDxd)
where E [N +1] is an index number and X [ XN+1I'.
Given a readout function Fy E Finear parameterized by v or Fv E Fquery parameterized by V, the
K+1 estimation Yi obtained by the K-layer Transformer is Øi Fr 2++1 or Di Fv 21-11 2i 1)
K+1
respectively.
Before we formally present our main results, we introduce the following assumptions.
Assumption 1. For x Px and 3 Pg. we assume ||x|| & bx and || < b9 almost surely. Besides, consider the noiseless scenario where 0.
We note that the boundedness assumption over x and 3 ensures the robustness of the Transformer and prevents from blowing up under ill conditions.
similar assumption is adopted in Bai et al. (2024). The noiseless assumption is for ease of analysis and is common in the analysis of Trans- formers (Ahn al. 2024. Fu al. [223] Bai al. al.] (022)). 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 the k-th self-attention layer H()), and use Blkt1 represent the vector [ (k+ 1|441:241.2241 Then, we state the following theorem.
Theorem 4.1 (Equivalence between ICL and LISTA-VM), With the Transformer structure de- scribed above, under Assumption [ there exists set of parameters in the Transformer so that for any l [1: K]. E [N], we have
B Spcce (B/N+=
M (k)
[
([X)
1.n, B v1m)).
(4.4)
2n +1
1:n,
where (R)
(k) E Bdd is embedded in the k-th Transformer layer.
The proof of Theorem 4.1 detailed in Appendix C.2)
Remark 1. As mentioned above, matrix ((k) in the update rule of LISTA-CP in Equation 4411) is learned during pre-training and remains fixed across different in-context sparse recovery instances. As a result, LISTA-CP requires the measurement matrix X. stay the same during pre-training and inference. In contrast, if we denote ( 2n+1 [Xim: (M))) then within the update rule of the
LISTA-VM algorithm implemented by the Transformer, as detailed in Equation 4.4), the matrix ( depends on a fixed matrix ((k) after pre-training well the measurement matrix X during inference. As a result, the Transformer can adaptively update D th for instances with differrnt X's, enabling more flexibility and improved performance for the ICL tasks.
PERFORMANCE OF TRANSFORMERS FOR IN-CONTEXT SPARSE RECOVERY
In this section, we demonstrate the effectiveness of the constructed Transformer in implementing the LISTA-VNS 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 convergence rate linear in K. We then demonstrate that the Transformer can accurately predict + the same time.
7

---

Published as a conference paper at ICLR 2025
5.1 SPARSE VECTOR ESTIMATION
Theorem 5.1 (Convergence of ICL), Let (0, No = 8(4S - 2) log 8 8-lng On = - log (1
3 125 v log logd-log. +y
/ Bolog-) where e positive constant and y is positive constant
log
satisfies 7 3. For a K K-layer Transformer model with the structure described in Section 4.3 under
Assumption [ there exists set of parameters such that for any randomly generated sparse recovery instance and any n E [No: N,, with probability at least - 6, we have
[ - -3 S hize K..
Main challenge and key ideas of the proof. Similar to the proofs in Chen et al. (0018) and Liu and Chen (2019). the core step in proving convergence is to ensure that ( exhibits small coherence with X. i.e. (D(t)) T X R Laxd- In Chen et al. (2018) and Liu and Chen (2019), such a D(k) is obtained by minimizing the generalized mutual coherence to a fixed X. However, this results in poor generalization across different X's. In our proof, we consider X as a random matrix and leverage its aab-Gaussian properties to prove that if ( X(M(+)) where M(k) is associated with the covariance of X, then (8) will have small coherence with X with high probability. We defer the detailed proof of Theorem 5.1 to Appendix D.I
Remark (Linear convergence rate). The linear convergence rate demonstrated in Theorem 5.1 is due to the incorporation of curvature information into the update rule. Specifically, the learned matrices (()) serve as approximations of the inverse Hessian. By leveraging the statistical proper- ties of the problem, these matrices effectively accelerate convergence, enabling the Transformer to mimic second-order optimization methods. This allows the Transformer to overcome the traditional limitations of first-order methods, which are typically restricted to a sublinear O(1/ (//K) convergence rate unless additional assumptions are introduced (Nesterov) 2005 Nemirovskij and Yudin) (983).
Remark (Generalization across measurement matrix X). Theorem 5.1 shows that for any X sat- isfying Assumption 7 the Transformer can estimate the ground-truth sparse vector 3 in-context at a convergence rate linear in K. This is in stark contrast to traditional LISTA-CP type of algo- rithums, which only work for fixed X. Such generalization is enabled by the input-depentee
matrices
(D()), J.- Besides, we also note that (3 only depends on X1,..., Xn. This implies that even if the measurement matrix X is of dimension x d instead of N x d, when n E [No, N). the Trans- former can still recover 3 accurately. Such results demonstrate the robustness of Transformers to variations in in-context sparse recovery tasks.
Remark 4 (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 B lies in subset S C [1: d) with S < |S] < d, then by slightly modify- ing the parameters of the Transformer to ensure [g(k) 0 for all i & S, the ICL performance can be improved by replacing all d involved in Theorem 5.7 by [S]. We defer the corresponding result and analysis to Appendix D.2
5.2 LABEL PREDICTION
have demonstrated that Transformers can successfully recover the ground-truth sparse vector 3 with lincar convergence by implementing LISTA-type algorithm. In this section, we bridge the gap between this theoretical claim and the explicit objective of in-context sparse recov- ery, which is to predict UN+1 given an in-context instance (X, y, 1)) This gap might seem trivial at first glance, as given XN+11 and accurate estimate of B3, the label ØN+1 can be obtained through a simple linear operation. However, we will show that, for decoder-based Transformer, generating ØN+1 using the predicted sparse vector B3, which is implicitly embedded within the forward-pass pass sequences, critically depends on the structure of the read-out function.
In Section 5.1
Theorem 5.2. Under the same setting as in Section probability at least nō 8¹, we have
4.3 for any n. E [No + 1: N + +1] with
lll. Tnll < bx (1 1
2 3
K
2
K-1
+
1
nn
3
8

---

Published as a conference paper at ICLR 2025
for linear read-out function, and with probability at least - 6, we have 11 Tnll I caek for a query read-out function, where Ca, Cs are constants.
We defer the proof of Theorem 5.2 to Appendix E
Remark 5. Theerem 5.2 indicates that adopting linear read-out function results in prediction error of order Oie K + K e=="), which still exhibits linear convergence with respect to K. When query-based read-out function is employed, the convergence rate improves to O(e However, there exists a gap of order O ( e==," 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 O
and we also observe that the gap decreases grows.
6 EXPERIMENTAL RESULTS
Problem setup. In all experiments, we adhere to the following steps to generate in-context sparse recovery instances. First, we sample ground truth sparse vector 3 from a d 20 dimensional standard normal distribution, and we fix the sparsity of 3 to be 3 by randomly setting 17 entries in 3 to zero. Next, we independently sample N 10 vectors form a d dimensional standard normal distribution and then contract the measurement matrix X E R 100020 (each sampled d dimensional random vector is row in X). We also sample an additional X from the d-dimensional standard Gaussian distribution. We follow the noiseless setting in Bai et al. (2024) for sparse recovery, i.e.,
y -x,"
Baselines. The baselines for our experiments include traditional iterative algorithms such as ISTA and FISTA (Beck and Teboulle 2009). We also evaluate three classical LISTA-type L20 algorithms: LISTA, LISTA-CP, ALISTA (Gregor and LeCun 2010; Chen et al. 2018: Liu and Chen 2019). For all of these algorithms, we set the number of iterations = 12. We generate a single fixed measure- ment matrix X. For each training epoch, we create / 50, 000 instances from 50, 000 randomly generated sparse vectors. During inference, we evaluate the pre-trained LISTA-type models under two settings: (1) when the measurement matrix remains identical to that used during pre-training, reported as "Fixed X", and (2) when the measurement matrix is varying through random sampling, reported "Varying X"
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 ma- trices, each generating 500 instances from 500 randomly generated sparse vectors, which results in 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 non-conved optimization problem for each different measure- ment matrix X, which makes meta-training for ALISTA unrealistic. For all baseline algorithms, minimize the sparse vector prediction loss E/= [| 3 using gradient descent for each epoch. We run all baseline experiments for 340 epochs.
Transformer structure. We consider (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 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 attention heads per layer to 8. We note that the Small TF model shares the same 4-head configuration and ReL.U activation function as described in Theorem 4.1 while used in practical applications. In order to train Small TF and GPT-2, we randomly generate 64 instances per epoch and train the algorithms for 10 epochs. The training process minimizes the label prediction loss E (リ₁ 3. 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 hours, while the training time for GPT-2 is around 12 hours.
Transformer models, i.e., small Transformer model
configuration of GPT-2 commonly
Results. We test the prediction performance of the baseline algorithms and Transformers on a sparse recovery instance (X. y, XN+1), and plot the label prediction loss in Figure
We first do not impose any support constraint on B. We start with the general setting where the testing instance is randomly generated (Varying X). As shown in Figure Ta GPT-2 outperforms

---

Published as a conference paper at ICLR 2025
1971
1514
C
AUSTA
LSTA-W
LSTIUN JST
Smal77
Vary F
Venying
Faod
03
0.3
0.5
(a) Without support constraints
(b) With support constraints
(c) Without support constraints
Figure I: Experimental results for sparse recovery. (a) S = 3. (b) S = 3, and the support is restricted to be within the first 10 entries. (c) Prediction with different read-outs functions.
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 re- sults 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 test- ing, leading to poor prediction performance. When 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 mecta-trained LISTA and LISTA-CP, the corresponding prediction loss (denoted by red marks in Figure lab are still much higher than that under LISTA-VM. VM. This indicates that meta-training can- not help those classical LIST, LISTA-type Tt-type algorithms achieve performances comparable with LISTA- -MM, further corroborating the strong generalization capability induced by the constructed Transformer.
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, set the support of the sparse vector & to the first 10 entries, i.e., = (1,2,... 10}. As observed in Figure 1b Small TFF and GPT- 2 significantly improve their performances in Figure Ial while other bascline algorithms do not exhibit significant performance improvements. In Figure lb we present the experimental results for support-selected version of LISTA- LISTA-VM, referred as LISTA- VM-SS. This algorithm is simple variation of LISTA-VM, where we incorporate prior knowledge of the support by setting all columns in X whose indices are not the prior support set to zero vectors. As we claim in Remark 4 and Corollary D.I 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 10x better than the standard version of LISTA-VAS and is comparable to the prediction error of Small TF. This empirical finding corroborates Remark [
Finally, we examine the label prediction loss under Small TF with three types of read-out functions, ie., linear read-out function (Linear Readout), query read-out function (Query Readout), and an- other quadratic read-out function with parameters selected according to the proof of Theorem [.2 (Query Readou("). In Figure Ic 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 insignif- icant, which 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 L20 algorithms. Specifically, we showed that for in-context 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 individ- ual sparse recovery problems with fixed measurement matrices, Transformers can address general class of sparse recovery problems with varying measurement matrices during inference without re- quiring 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 L20 methods typically fail.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The work of R. Liu and J. Yang was supported in part by the U.S. National Science Foundation un- der the grants ECCS-2133170 and ECCS-2318759, The work of R. Zhou was supported in part by NSF grants 2139304, 2146838 and the Army Research Laboratory grant under Cooperative Agree- ment W91 INF- 17-2-0196. The work of C. Shen was supported in part by the U.S. National Sci- ence Foundation under the grants CNS-2002902, ECCS-2029978, CCCS-2143359, ECCS-2033671, CPS-2313110, and ECCS-2332060.
REFERENCES
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. (2024). Transformers learn to implement pre- conditioned gradient descent for in-context learning. Advances Neural Information Processing Systems, 36.
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. (2022). What learning algorithm is in-context learning? investigations with linear models. arXiv preprint urXiv:2211.1500 151661.
Andrychowicz, M., Denil, M., Gomez, S., Hoffman, M. W., Pfau, D., Schaul, T., Shillingford, B., and de Freitas, N. (2016). Learning to learn by gradient descent by gradient descent.
Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. (2024). Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances neural information process- ing systems, 36.
Beck, A. and Teboulle, M. (2009). A fast iterative shrinkage-thresholding algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1):183-202
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901.
Candes, E. and Tao, T. (2007). The Dantzig selector: Statistical estimation when p is much larger than n. The Annals of Statistics, 35(6):2313-22 2351.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mor- datch, I. (2021). Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 15084-15550977
Chen, S. and Li, Y. (2024). Provably learning multi-head attention layer. arXiv preprint arXiv:2402.0408% 04004.
Chen, S., Sheen, H., Wang, T., and Yang, Z. (2024a). Training dynamics multi-head soft- attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.19442
Chen, T., Chen, X., Chen, W., Heaton, H., Liu, J., Wang. Z., and Yin, W. (2022). Learning to optimize: A primer and benchmark. Journal of Machine Learning Research, 23(189):1-59.
Chen, X., Liu, J., Wang, Z., and Yin, W. (2018). Theoretical linear convergence of unfolded ista and its practical weights and thresholds. Advances in Neural Information Processing Systems, 31.
Chen, X., Zhao, L., and Zou, D. (2024b). How transformers utilize multi-heas attention in-context learning? study sparse linear regression. arXiv preprint arXiv:2408.0455 04532..
Cheng, X., Chen, Y., and Sra, S. (2023). Transformers implement functional gradient descent to learn non-linear functions in context. arXiv preprint arXiv:23/2.0652.
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung,
H. W., Sutton, C., Gehrmann, S., et al. (2023). Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113.
Cui, Y., Ren, J., He, P., Tang. J., and Xing, Y. (2024). Superiority of multi-hear attention in- context linear regression. arXiv preprint arXiv:2401.14420 17426.

---

Published as a conference paper at ICLR 2025
Dai, D., Sun, Y., Dong, L., Hao, Y., Ma. S., Sui, Z., and Wei, F. (2022). Why gpt learn in- context? language models implicitly perform gradient descent as meta-optimizers. arXiv preprint arXiv:2212.10559,
Daubechies, I,, Defrise, M., and De Mol, C. (2004). An iterative thresholding algorithm for linear inverse problems with sparsity constraint. Communications on Pure and Applied Mathematicy: Journal Issued by the Courant Institute of Mathematical Sciences, 57(1 57:113:34331445
Deora, P., Ghaderi, R., Taheri, H., and Thrampoulidis, C. (2023). On the optimization and general- ization of multi-head attention. arXiv preprint arXiv:2310.12680.
Devlin, J. (2018), Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0480 048005.
Dosovitskiy, A. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929
Elad, M. (2010). Sparse and redundant representations: from theory to applications in signal and image processing. Springer Science & Business Media.
Finn, C., Abbeel, P., and Levine, (2017). Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning. pages 1126-1135. PMLR.
Fu, D., Chen, T.-Q., Jia, R., and Sharan, V. (2023). Transformers learn higher-order optimization methods for in-context learning: study with linear models. arXiv preprint arXiv:2310. 17086.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. (2022). What can transformers learn in-context? case study of simple function classes. Advances in Neural Information Processing Systems, 15:30583-30593
Ge, R., Jin, C., and Zheng. Y. (2017). No spurious local minima in nonconvex low rank problems: A unified geometrie analysis. In International Conference Machine Learning, pages 133--124.. PMLR.
Gregor, K. and LeCun, Y. (2010). Learning fast approximations of sparse coding. In Proceedings of the 27th international conference international conference machine learning, pages 399- 406.
Han, C., Wang, Z., Zhao, H., and Ji, H. (2023). Explaining emergent in-context learning as kernel regression. arXiv preprint arXiv:2305.127666
Hospedales, T., Antoniou, A., Micaelli, P., and Storkey, A. (2021). Meta-learingg in neural networks: A survey. IEEE transactions on pattern analysis and machine intelligence, 44(9):5149-516,
Hruby, P., Duff, T., Leykin, A., and Pajdla, T. (2022). Learning to solve hard minimal problems. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5532-5542.
Huang, R., Liang, Y., and Yang. J. (2024). Non-asymptotic convergence of training transformers for next-token prediction. arXiv preprint arXiv:2409.17333.
Huang, Y., Cheng, Y., and Liang, Y. (2023). In-context convergence of transformers. arXiv preprint arXiv:2310.05249.
Kim, J. and Suzuki, T. (2024). Transformers learn nonlinear features in context: Nonconvex mean- field dynamics on the attention landscape. arXiv preprint arXiv:2402.01258.
Li, H., Wang, M., Liu, S., and Chen, P.-Y. (2023). A thcoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arxiv:2302.06015.
Li, H., Wang, M.. Lu, S., Cui, X., and Chen, P.-Y. (2024a). Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint 2XXv:24002.16607.
Li, K. and Malik, J. (2016). Learning optimize, arXiv preprint 16Xiv:1606.01888
12

---

Published as a conference paper at ICLR 2025
Li, Y., Huang, Y., Ildiz, M. E., Rawat, A. S., and Oymak. S. (2024b). Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics, pages 685-693, PMLR.
Liu, J. and Chen, X. (2019). Alista: Analytic weights are as good as learned weights in lista. In International Conference Learning Representations (ICIR).
Mahankali, A., Hashimoto, T. B., and Ma, T. (2023). One step of gradient descent is provably the op- timal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576.
Meng, Y., Yang, Y., Hu, M.. Zhang, Z., and Zhou, X. (2023). Artificial intelligence- basel radiomics in bone tumors: Technical advances and clinical application. In Seminars in Cancer Biology. Elsevier.
Nemirovskij. A. S. and Yudin, D. B. (1983). Problem complexity and method efficiency in opti- mization.
Nesterov, Y. (2005). Smooth minimization non-smooth functions. Mathematical programming, 103:127-152.
Nichani, E., Damian, A., and Lee, J. D. (2024). How transformers learn causal structure with gradient descent. arXiv preprint 2rxiv:2402.14111 /4735.
Pitaval, R.-A., Dai, W., and Tirkkonen, O. (2015). Convergence gradient descent for low-rank matrix approximation. IEEE Transactions Information Theory, (1()444551-455
Radford, A. (2018). Improving language understanding by generative pre-training.
Roy, S. S., Mittal, D., Basu, A., and Abraham, A. (2015). Stock market forecasting using lasso
linear regression model. In of the First International Afro-European Conference for Industrial Advancement AECIA 2014, pages 37 1-381. Springer
Conference for Industrial Advancement: Proceedings
Shen, D., Wu, G., and Suk, H.-I. (2017). Deep learning in medical image analysis. Annual review of biomedical engineering, 19(1):221-245
Sucker, M., Fadili, J., and Ochs, P. (2024). Learning-to-optimii with pac-bayesian guarantees:
Theoretical considerations and practical implementation. arXiv preprint arXiv:2404.03::
Tarzanagh, D. A., Li, Y., Thrampoulidis, C., and Oymak, S. (2023a). Transformers as support vector machines. arXiv preprint 2rxiv:2308.1688
Tarzanagh, D. A., Li, Y., Zhang, X., and Oymak, S. (2023b). Max-margin token selection in attention mechanism. In Thirry-sevnnth Conference Neural Information Processing Systems.
Thoppilan, R., De Freitas, D., Hall, J., Shazeer, N., Kulshreshtha, A., Cheng, H.-T., Jin, A., Bos, T., Baker, L., Du, Y., et al. (2022). Lamda: Language models for dialog applications. arXiv preprint arX1v:2201.08239
Tian, Y., Wang, Y., Chen, B., and Du, S. (2023a). Scan and snap: Understanding training dynam- ics and token composition in 1-layer transformer. Advances in Neural Information Processing Systems, 16779711-719777
Tian, Y., Wang, Y., Zhang, Z., Chen, B., and Du, S. (2023b). Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention. arXiv preprint arXiv:23 10. 00535.
Vasudeva, B., Deora, P., and Thrampoulidis, C. (2024). Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.05738.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L.,, and Polo-
sukhin, I. (2017). Attention all you need. Advances in neural information processing systems, 30.
13

---

Published as a conference paper at ICLR 2025
Von Oswald, J., Niklasson, E., Randazzz, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. (2023a). Transformers learn in-context by gradient descent. In International Conference Machine Learning, pages 35151-35111 PMLR.
Von Oswald, Niklasson, E., Schlegel, M., Kobayashi, S., Zucchet, N., Scherrer, N., Miller, N., Sandler, M., Vladymyrov, M., Pascanu, R., et al. (2023b). Uncovering mesa-optimization algo- rithms in transformers. arXiv preprint arXiv:2309.05888.
Yao, Q.. Wang, M., Chen, Dai, W., Li, Y.-F., Tu, W.-W., Yang, Q., and Yu, Y. (2018). Taking human out of learning applications: A survey on automated machine learning, arXiv preprint arXiv: 1810.180JJ.0 13306.
Zhang, R., Frei, S., and Bartlett, P. L. (2023). Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.099 09927.
Zhu, Z., Li, Q. Tang. G., and Wakin, M. B. (2021). The global optimization geometry of low-rank matrix optimization. IEEE Transactions on Information Theory, 6782):1308-131.
14

---

Published as a conference paper at ICLR 2025
Supplementary Materials
CONTENTS
Introduction
Related Works
Preliminaries
Transformer Architecture
3.2
In-context Learning by Transformers
4
Transformer as a LISTA-type Algorithm for In-context Sparse Recovery
In-contexi Sparse Recovery
4.2
Classical Algorithms
4.3
Transformer Can Provably Perform LISTA LISTA-type Algorithms
5
Performance of Transformers for In-context Sparse Recovery
5.1
Sparse Vector Estimation
5.2
Label Prediction
8
Experimental Results
Conclusion
10
Additional Related Works
16
Table of Notations
17
Deferred Proofs in Section 4.3
17
Transformer Structure
17
Proof of Theorem 4.
18
D Deferred Proofs in S.c
21
Proof of Theorem
D.2 Proof of Remark 4
27
Deferred Proofs in Section 5.2
29
Proof of Theorem 5.2
29
Additional Experiment Results
35
15

---

Published as a conference paper at ICLR 2025
A ADDITIONAL RELATED WORKS
General L20 Techniques. I.20 leverages machine learning to develop optimization algorithms, aiming improve existing methods and innovate new ones. As highlighted by SSucker et al. (2024) and [heen et al. 2022), L20 intersects with meta-learning (also known as "learning-to-leari ") and automated machine learning (AutoML).
Unlike meta-learning, which focuses on enabling models to quickly adapt to new tasks with min- imal data by leveraging prior knowledge from diverse tasks (Finn al., 2017 Hospedales et al. 2021), I.20 aims to improve the optimization process itself by developing adaptive algorithms tai- lored specific tasks, leading faster convergence and enhanced performance in model train- ing (Andrychowicz et al.) [016); [Li and Malik 2016). Thus, while meta-learning enhances task adaptability, L20 refines the efficiency of the optimization process. In contrast, AutoML focuses on model selection, optimization algorithm selection, and hyperparameter tuning (aao et al. 2018); L,20 distinguishes itself by its ability to generate new optimization techniques through learned mod- els.
L20 has demonstrated significant potential across various optimization fields and applications. For (estance,AAdrychoovvee et al. (2016) introduced a method where optimization algorithms learned using recurrent neural networks trained to optimize specific classes of functions. Li and Malik (2016) proposed learning optimization algorithms through reinforcement learning, utilizing guided policy search to develop optimization strategies. Furthermore, Hruby et al. (2022) applied L20 address the "minimal problem", common challenge in computer viston characterized by the presence of many spurious solutions. They trained a multilayer perceptron model to predict initial problem solutions, significantly reducing computation time.
Training Dynamics of Transformers. There exist some works aiming theoretically understand the ICL mechanism in Transformer through their training dynamics. [Ahn et al.) (2024); Mahankali et al. (202);; Zhang al. (2023;; Huang al. (2023) investigate the dynamics of Transformers with a single attention layer and a single head for in-context linear regression tasks. [Cui et al. (2024) prove that Transformers with multi-head attention layers outperform those with single-head attention. Cheng et al.] (2033) show that local optimal solutions in Transformers can perform gradient descent in-context for non-linear functions. Kim and Suzuki (2024) study the non-convex mean- field dynamics of Transformers, and Nichani et al. (2024) characterize the convergence rate for the training loss in learning causal graph. Additionally, ly. Chen et al. (2024a) investigate the gradient flow in training multi-head single-layer Transformers for multi-task linear regression. Chen and Li (2024) propose supervised training algorithm for multi-heat Transformers.
The training dynamics of Transformers for binary classification (Tarzanagh et al.
2023b a Vasudeva
al. 2024 et al., 2023 Deora et al. 2023 [i al., 2024a) and next-token prediction (NTP)
(Tian et al. 2023a Li al. 2024b Huang et al. 2024) have also been studied recently.
16

---

Published as a conference paper at ICLR 2025
B
TABLE OF NOTATIONS
Notation
Definition
Measurement Matrix
3
Ground-truth sparse vector
Noise Vector
=33 + +
Number of columns of X
N
Number of Mcasurement Vectors, i.c., number of rows of X
Sparsity of 3 f.e., | 10
Support set of 3
D
Dimension in the Self-attention Layer
Number of Heads in the Self-attention Layer
D'
Hidden Dimension the MLP Layer
K
Number of Layers in Transformer
QU4)
Query Matrix of the k-th Layer of Transformer's 's i-th Head
K ((A)
Key Matrix of the k- th Layer of Transformer's i-th Head
V ((k
Value Matrix of the k-th Layer of Transformer's i-th Head
H())
Input Sequence of the k-th Layer of Transformer
[H(k)
(agn(ar)
Sign Function: sign( /
0, sign(z) = 0
(x)
Soft Thresholding Function: So = sign(2) max (0, [ 0]
0
Rd
12d
ReLU Function: (a (x)]; = Iz xr 2 0. [o(x)) = 0 <0
Complement of event
C DEFERRED PROOFS IN SECTION 4.3
C.I TRANSFORMER STRUCTURE
Attention layer. Consider model consisting of K Transformer layers, where each layer equipped with four attention heads. These heads are uniquely indexed as +1, -1, +2, and -2 distinguish their specific roles within the layer.
(d- (x-1)x(++1)
1i+ (1)
0
(2x+2)11
O₂ Q.(k)
Q (R)
(1x(4+1)
0a
0 (1x+2)+1)
m442
01x(d+1)
Oxxd
((d+1)x(24+))
0
0 (d-1)xx
(d+1)x
0
Odxd
0a
Ouxd+1
(k)
i1
Laxd
0dx(d+1)
0₂
01xd
01x(++1)
0 O1xd
O 01x(d+1)
1
0 ((d+1)xd
0
0 ((a+1)x(++1)
(d-1)xd
(d+1)x((+2)
(d+1))dd
0(441)xa42
(k)
il
M (.)
O Oxx(d+2)
(k) 12
M Ki(k))
Oax(d+2)
(C.I)
0 Oxxd
O1x(d42)
Oxxd
O 0xx(d+2
where 2..K)
9.)))
Vi(k) V.(k), K.)) +1
and M (k()) are all matrices, and m
Q.(k)
+1
12
m (k)k)
are scalars.
MLP layer. For the MLP layer following the k-th self-attention layer, we set
05d45 g(h) 1d
Wi 1.suub
((2042)x(2442)
W,
W 1,sub
W =
(22442)x
(()==
Oa-2
(C.2)
W L.ub
[2d42) (24+2)
g(h) 1,
W W.suh
(224+2)((2422)
17

---

Published as a conference paper at ICLR 2025
where the submatrix W 1,ub is defined as Wi,sub (41)((()+))(((+)))) Iax: Laxdi 0). Therefore, the output of the MLP layer is
[i
ML,(h;)
Sgms
[h (d+22dd-1)
(C.3)
[h:l2d42
where Socks the soft-thresholding function parameterized by (())
C.2 PROOF OF THEOREM
We start by stating an equivalent form of Theorem 4.1 below, where we specific M in Theo-
rem 4.1 to be (K)
Theorem C.1 (Equivalent form Theorem
Suppose Assumption 7
holds. For Transformer
with layers as described Section
4.3
set the imput sequence as:
X
X
1:1
X
J.N
X
I:N
XN+1
H
B
1
3
y
B
11)
(
NN
(1)
(C.4)
2N-1
2N
2N+1
Denote H(++)) as the oulput of the K-th layer of the Transformer and define (K+1)
((k+1)
042.2411,224 There exists a set of parameters within the Transformer such that for all k. [1, K], we have:
(k+1) 2n+ Spen (B - (D,,)) ([X) Iom, 32 V1..)..
where D 2n+1 (....MM) and MY E gdxd is embedded in the k-th Transformer lsyer.
Proving the theorem is equivalent to demonstrating the existence of Transformer for which the following proposition holds for any h: > 2:
Proposition Suppose Assumption holds. For a K layers Transformer with structure described the inpul sequence of the k-th layer of the Transformer satisfies:
in Section 4.3
X
1
X
X
X
J.NN
XN+1
(R)
UN
3
B
B
(A)
3
(k)
2N-1
22++1
where it holds that B(h) Spce-1)
(B
182+t1
((--1)(,,)) ([X]
(M)
2n+1 Misn)) D,,
X Xlm (MV)T and 1// < C2ty for Constant Cj..
Proof of Proposition We prove Proposition true for all 2 by induction. First, (k) for k 2 satisfies the condition automatically; therefore, Proposition [ is true when & = 2. Then, we demonstrate that if Proposition [ is true for & 1, remains valid for k. For odd values of the token-wise outputs of the k-thh attention layer corresponding to the first and second heads satisfy:
0
04+1
Od+1
(k)
M
((k)
K
11
(K)
x
1
;
(b) (k)
M
Vi(k)
B
The token-wise outputs of the third and fourth heads for odd satisfy
0g
0441
m 12
(k)
K O2442i
M
Vi(0)) +2
X
J.
0
18

---

Published as a conference paper at ICLR 2025
Besides, for any i that is an number, the token-wist outputs of the k-th attention layer corre- sponding to the first and second heads satisfy:
O M +1
0++1
Oa+1
g( ;
K
(k)
x
;
(k)
M
Vi(k)
X
0
0
The token-wiss wise outputs of the third and fourth heads for even satisfy
0₂
Oa+1
D4 O24423
(R)
M (R)
;
=2
1.11 1
0a+1
Therefore, for h, with an odd index and head with index a E {1,-1) we have
((Q! ( )) n(
Oc+1 M
T (MM) [x -
1..1
(C.5)
Also, for h, with an odd index and head with index u E (2, -2}, we have
0a+1
a (( n k())))) w( (+) 102==00 a ( (m (...000.4.)
MY
(Kk)
X
1.
(C.6)
We specify the parameters of the self-attention layer as
(. = -Iaxd: M. = -00 MY
M Laxdd M Vi(k) --000 MY,
m2 (k) =1, My (.()) (K)
7.()) = -,, MV(() (V
m=
Then, for head index u E -1, we have
( X 1:: B s (1((h)) ) X 1 B
< |, || X
x B
(C.7)
< || B01
1,ر + bah B,
where Equation
C.7
is given by Hölder's 's Inequality. Therefore, combining this with the assumption
|x| bx and
3 ll bg. we obtain
(g ( T (M?) T [X
B < bx bx - 33111, + bigbx - B.
(C.8)
Recall that for odd i, we assume that Proposition I is true for - 1. then 3/(h) satisfies
B kk) S (B(k-1)) (K- (D..) (
1: 16 )
and [3( S Ch. Then, we obtain
(
<18 (1a (1(K-D))11
| -بب J.. [ ببهة |1044 1 1)
X *
C3 + 1p211 1hhp
C3 + baاي| (k- 1>
M|
1 + Ech
V Ich + ( (k-1) 1||||| MM
(C.9)
19

---

Published as a conference paper at ICLR 2025
Equation C.9 arises from the fact that
1/t / Σ 1. Denoting Ymax max; || and letting
Cg
bx mmax
|MV|
M || we have
1B/401
S C + dd
10 ||| | ( S Catt!
Therefore, have
bx|ls!!) - 3011, S bx bxvd(bs + Cp).
Then, by setting B 2 babx + bx Vdb3 + C3 we obtain that B > ((3 ) 1. for all
k S K and any odd i. Then, combining with Equation (C.8) we have
(M))
I., 1 T
B 0.
Further, from Equation C.5 for i is odd number and j E [1,3, ... .d] we have
(( h K nf))
v: h 02a+2. (k)
Besides, for being an odd number and j [2, 42, , 1). we have
0
a (( 1 h) V( = ((M)) 1.141)
MY X
I., 1
Thus, when h2 is of odd index, summing over all j $ and a E {1, -1}, we obtain:
Σ
h h bt)) n
wE{
Σ
σ h b+))
n
-E(-1+1))=2
i-1
0
Σ
T
1
2=2
04+1
M
X
1:
X 1:
11
11
(C.10)
0
Next, consider heads with indexes
{2, 2].
From Equation (C.6) obtain
Σ Σ
(2)
K (()))))
v h
(k)
wEI 2+(2)++)
Σ
Σ 0
(k) h
v( (8) K
wEt 26([- j=2
O
Σ U44
(A)
M
X
02
04+1
7 (1) MV
X
1. -
11:
(C.11)
02
20

---

Published as a conference paper at ICLR 2025
Combining Equation (C.0 and Equation [C.11. odd index:
obtain the following equation for h, with an
h
(++1)
=
MLP
h
(k)
Σ
Σ
o (( h
k())
(k)
( #2) 1=1
[
3
where
7(k)
0/+t1 Sg(x)
(8!hh)
M
[X
1:11:
[X] 1: B((b) 7 [X
11 111:
(B(())
7(h)
= =Saw
(C.12)
Similarly, we obtain the following equation for h, with an even index:
h h
(k+1)
MLP
(k+1)
Σ 2) f=1
Σ
a
((09 (8)
"C
111 s+1
0
(C.13)
gk+1
0
Note that explicit formulation of G tion (C.12) and Equation (C.13)
(N++1
is not required when is even. Next, combining Equa-
gives
X
J..1
X
X
JAN
X
Jn,
XN+1
(k-1)
V1
UN
0
B
(++1)
B
B
B 1
((++1)
2
2N
The proof for Proposition iis thus complete.
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 X is independently sampled from isometric sub-Gaussiaa distribution with zero and covariance diag(of.... aQ)
Without loss generality, let σ 2 02 2 ... 2 σa- We start the proof of Theorem 5.1 by in-
troducing the auxiliary lemmas: Lemma 0 Lemma 2 and Lemma 3
Lemma 1. Assume (ik) and g(++1) satisfy:
glk+1) Spacs (B(K) -D D (x))()) y)) -
(D.1)
Define mmin = mines D D.T. Fmax max D. x. . and let g(k) = m (ea(k-2) 2), Then, if (1( Ba1 S caa-)
where = || l and log 7(25-1)omms +I1 m7
and D, 7, g(k) satisfy the following conditions:
722 - 1)0max + +11 7amin 1
(D.2)
0 D x 51 Vi [d],
then ((++1) satisfies
[|(0+11)
Bi & c (2(k-1)
21

---

Published as a conference paper at ICLR 2025
Proof of Lemmal 7 Note that are considering the noiseless model: y XB3 We begin by rewriting Equation (D.I:
B(++1)= = Sg(t) (3(6) (+) T (x3(() y))
= Sgres (3(k) ( T x(3(6) -3))
Consider entries of 3 ((K+1) in the support of 3
B 1) Spms (p!") (b) (D.s) X (3 (1) 3b)
E E (k) - 710 (D.s)) x (B (k) (b) ((() de, (BBthtt)),,
where (g) is the sub-gradient of |3):
if [3], + 0, [0L] 1-1 if [3): = O.
{sign([3])
Then, for any E S, we have
gl(+1) E -/0)
(D.,)) TT x { - B3)
glk) ale) (1t
Note that 3 (k) - (D),)) (D.,i) (3 (k) - (8 can be rewritten as
3 (k) 7 (Dr.i) Tx (0B (k) (3)
s
= = (k) - -7(h (k) Σ (D X (0) - 33) - (b) T (3()) - -):)
jEsij
= B - (k)
(D...) x x (8 - -3)) + (1-,(6((D.)))
(D.,)) T (8! - -31)
JEB.j
Therefore,
3 (((+11)
3 E (A (8)
(D...)),,,, -))+ (1-+0111000 JES,J
(1-1(0)) (k) (D...) (B, - -3 - ge,
Based on the definition of Oli, we derive the following inequality:
( Bils
Σ (D...) x (10 - 371 + 1
(2) (D...) Xx. (1( - -371 + 0(3)
jE3,j
Recall that 0 < (k) (D...) X < 1, we obtain that
[ -3 S g (k) max
Σ
18/0)
331 + ---11 (k)
0 min |10(4)= -B + B(K) +(0)
JES
From Lemma 2 we have support (g(++)) E S, thus, 1|3(h++1) - -B"112 = 11,a -33111 Then,
[|((++1] - 1 S
(K)
18((0) -B,1 + -|1-1 mmin |13 - -3 + (1k)
T
max
iss
JE
S (1) - m)x
(/ - + |1 - Tmia |1(hh) 3: + (g())
(
EES
(s - 1)0mux + 1 - (k) mmin |1((K) B: ||| + +Sg(A))
< (x()) 1)0max + || (100 (1) Tmin
)e1e-oa(--1)
Combining with 02 log 711) - 1)+ 7(0) 000mm amin
we obtain
[| - -B <cee -22k
Therefore the proof is complete.
22

---

Published as a conference paper at ICLR 2025
Lemma 2. Suppose all conditions mentioned in Theorem 5.1 hold. For lo E N, if
support(3()) E and || B S (2(k))))-1)
then holds that
support(3(kt))) ES.
Proof of Lemma 2 have
For a fixed k, support(3(h)) e and || Bbll S C₁C cie o2(k-1), then we
(k+11) Sgiu (B -(1h) (D ,00 X x (b)
Sgie) (-g(k)) (D...e X ( BB))
where g0 = [s Then, for all E SC, we obtain
glkt2) = Soics (-)(0)
X., 87))
JEB
Note that
(() Σ (D..) X (3) -37)
S 711 (k) max [1(((k) 30
jES
S 5 ma - c2(k-11
Given ((k) defined in Theorem 1, we obtain:
-000 (K)
T (3) B3)) s((
JES
Consequently, it follows that (,x0) 0 for all i E SC.
Lemma 3. Suppose Assumption 1 holds and let D such that the following inequalities hold with probubility at least 1 - 35:
2N41 I X((V) Then, there exists an MV
N
log S log
min D ES
N 2
Nc
N
log log
max
(D..)
1+
seS
N 2
Ne
log d log
max *$
(D...)
T X S
Nc
Proof of Lemma 3 Recall that 01 2 02 2 2 od. Let M
2
Lixd. To prove the lemma, we
first introduce a lower bound for min;es
x Note that
N
2
( X
M
X
X
2N +1
2N +1
2N
Note that X is sub-exponential random variable. Then, from the tail bound for sub-exponential random variables, there exists a constant e such that
- N
} 21 21-ca --exp Nemin {3,5)
23

---

Published as a conference paper at ICLR 2025
Consider 0 1, we have
- exp (-N(-- Ncs²),
follows that
P
min
1
Σ
x2
ES 01 j=1
P
min
2 --
N
-
PP
X
EE
JE
N
1 -
1
IE3
j=1
1 exp
Nes
Let 8
log S d Ne
Then, with probability at least 1 - 6, we have
N
log S log 8
x2
EE3
Ne
j=1
Therefore, with probability least 8, holds that
2
N
log S log 5
min (D.,)
X.
min
1-
2N +
N
+
Nc
Similarly, we have the following inequality holding with probability at least - 6:
N
log S log 5
max
(D..)
X
S
1+
N +
Nc
Next, we will provide high probability upper bound for maxi4j i + j. we have
(D..i) X ا
Observe that when
2
(D.)) 2N (* X
X Xk.j
2N
where XR. is a centered sub-exponential random variable. Then, from the tail bound for sub- exponential random variables, there exist constants and s < 1 such that
xx.. N t=l
p
> 1 exp (- --Ncs²). S 5
24

---

Published as a conference paper at ICLR 2025
Then, have
2
P
X Xk.)
2
2N
Od
k=
N
2 PP
max
X
Xk.))
5
$F
0d
k=1
- P
(k.s)
=3}
k=
N
21
Σ
Xk.j)
28
h=1
21 Σ
Xx
Xx.i) >8
28
od
21 d exp Nes
Let
dogd-log Ne
thus with probability at least - 6, we obtain
log d log 5
max [
(D..) X
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 7 hhols. Let ā E (0,1). For K K-layer Transformer model with the structure described in Section 4.3 there exists set of parameters in the Transformer such that for any n E [No, NY. with probability at least 1 - 36, have
3 S a..k a.
where (KK+)) = [H+++))
1,1:2441,2041 and
On -lg ( 1
(1- 2
log d log 5 ne
log S log 5 nc
3 3 + 7 1-2 1)
(D.3)
11 and N, = 8(45 2) log dlgg 5 Here, c is a positive constant and can be any positive constant less than 1.6.
To prove Theorem D.1 first, for D,, 2 an event:
X XI1: (M defined in Equation 4.4 we define
log S - log
A =
min (DDn)..) X X1mm les
2 n
1
+ 2
max iES
(Dn)..) X Ximma
n1
1+
log log 8
n + 2
log dl log 8
max (Dal..) T Xims/
T
S
(D.4)
nc
Next, we introduce proposition as following.
Proposition 2. Suppose Assumption I 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 set of parameters such that if [No, N), we have
|| +1)
B S -oe-
25

---

Published as a conference paper at ICLR 2025
where 3 is given in Equation (D3)) than 1.6.
[H(K+1)
d+ 1,244 ei = =111%,,, and No 8(4S 2) log log de S-1og Oln Here, e is a positive constant and 7 can be any positive constant less
Proof of Proposition 2 We prove it by induction. First, for k 1, we initialize 3 02, thus || 10 3 || Note that ex 11, so Proposition 2 holds for k = 1.
Next, assume Proposition 2 holds for k. From the proof of Theorem 5.1 the (k: + 1)-th Transformer layer of the following structure:
the input sequence of
X 1.1
X
X
JIN
X 1
XN+1 0 (k+1)
H
UN
(((t2)
B
(k+1)
(++1)
β 2N
(k+1)
2N-1
22N+1 1
1
where
( Sgce (3 700 (D,) ([X) 1:n, B -V1sn)
From Lemma 1 let 7(())))))))) then for 61 1133 C2 =
log 1²1 S cie cx(k-1 if the following conditions hold:
we have
7(25 7225-1)0ma -11m0-l
11-00min
||
725 max +11 70 S
(D.5)
((Dn).)) X 11m, S 1, Vi E [d].
(D.6)
First,
2)2 log d log S log 8
n 2 No where No 8(4S- 2)²
(D.7)
By rearranging Inequality D.7 we have
log + log S - log d
Vn 2 2)
log d log log
V2(48 2)
c
log d log
log S log
2 (48 - 2)
log d log 8
log
log
2 (48 2)
(D.8)
c
Inequality D.8 is equivalent to
1
log d log 8
1 2
log S log 5
1
(2S -1)
+
vn
2
log d - log 6
log S log 8
(28
1)
1
mc
log d - log 6
log S log 8
(28- 1)1
1
(D.9)
NC
Combining Inequality D.9 with Equation (D4) we have
log d log 6
mmxx (2S - 1) < (28 1)V
log S log 8
1
n +2 2
nc
min-
(D.10)
26

---

Published as a conference paper at ICLR 2025
Next, recall that MV axd and (Dnl..)) X 2n-11 (Xims,) MM X Therefore, for any i [d], we have do
Xim,, 0if 2 0. Note that if satisfies the following condition:
1
7 min
(D.11)
Omin max ((Dn)..) X X1m0,
Then, Inequality D.6 will hold and
-Ymmin 2 0. Besides, we have
728 max
|1 min S mi |1 70
(D.12)
min + 70
(D.13)
1,
where Incquality D.12 derived from Inequality D.10 and Equation (.13) is obtained using Equa- tion (D.1T). Consequently, Inequalities DD5 and D.6 hold if event occurs. Furthermore, under the condition that event A occurs, by setting
(D.14)
log log ne
1
n
have satisfies Equation D.I1I Note that assume
No. Thus, have
2
2 1.6.
log S S-log
1+
log S S-----lla log Noc
+
8(4S-2)
Therefore, any choice of y such that y S 1.6 will ensure that Inequality (D.14) satisfied. Next, noting that event A holds and y satisfies Equation (DIIT. we obtain
|13
-3"
k-1
$ || lla (7(25- - 1)00 max |1 (YTmmil
7(25 -11 log d log 8
log S - log
< || l
+
1
k-1
log dl log 8 ne
log S
log
< ||
722 1
+
+
n2 n 2
n2 + 2
nc
h-1
log d log 5
2
log S - log 5
7225 - 1 $ ||
+1
+ +1
nc
nc
= = (t-1).
Therefore, the proof is complete.
Note that Lemma demonstrates that event A occurs with probability of at least 1 1-20 26. Therefore, the proof of Theorem D.I completed by combining Lemmal with Proposition 2
D.2 PROOF OF REMARK
Note that in Appendix D.2 assume there is constraint the support of ( its support lies within a subset c [1: dj. with S < [S] < d, where S denotes the sparsity of 3 We establish the corollary by demonstrating the following corollary:
Corollary D.1 (Equivalent statement of Remark Suppose Assumption 1 holds and S < [S] < 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 of parameters such that if 2 NS. with probability at least 8, we have
[,
|1 -K
27

---

Published as a conference paper at ICLR 2025
where Bgnt1 [H(K+1)]
-d 1.2224+1,,24+ and
a loo (1-2 (1-37 + 725 11)
2
log 3
log
log S log 6
(D.15)
nc
nc
< = -1 11, and No = 8(45 2)2
log +log S-log
Here, e is a positive constant, and y can be
any positive constant less than 1.6.
Proof of Corollary D.1 based on Equation (C.2
W.L.og,, assume S = f1,2.... [)). First, we modify the MLP layers
We maintain structure of Wa as described in Equation and
reconstruct W1 and b
We define the following submatrices:
(d++))((++1
x[+1)x
(d-1)x (8))
0a+1
(d++1
|||
(d-1118)
W (,,(((())
0 (d-18))x(d+1)
id-18)) 8
(dd11)))((-19))
(d-1
0 (xx(+11
0 x
0 (d-119))
(d+1)x(d+1)
xd41)x
(d-1)x(d-18)
0a+1
(d+x)
(d-1(1)
0
Wi siusb((2)
0 (d- (dd119))((++11
(d x 8]
I (d-
(d- [S])
0 01x(d+1)
0 1
0 (d((-18)
0 ((+1))(d+1))
(2+1))x S]
0 (d+1)x(d-18))
0a+1
(1+1)
|8|
0 (i-x(d-111))
0
W (.su(33)
0 (1-18))xd++1)
(dd11)))(d-181)
0 (d-119))
(xx++))
xx
0 xxx(d-181)
(d+1)x((+1)
(d-
0a+1
(1sxx(++1)
0 (1
0
W (.)bb(1)
(8)+))))(++1))
(it [S]) 8]
(d-) (91)
0
(a-118)) [8])
0 (x(++11
01x [S]
0 (d-(818
0
To simplify notations, let
O 05d45
W (,s((())
((R) 1s8
W
W
1.sub(2)
.b(0)
02d (kk) 1 |s]
W 1.sub(3)
s.b((44)_
0 --d [3]
The output of the modified MLP layer
h M1/1:
Sgas ([h)
(42.0++11+11
MLP(h,)
Og-13
[h:2a44
Next, also modify the parameters in the self-attention layers. For convenience, we introduce projection matrix I (iag([[g S), O (14))))))))1))) The reconstructed self-attention layers are described below:
O 0(+1)x((+1)
0 O(4+1)xd
04+1
(2x+1)))
Oa
(A) +1
Q.(k)
Odx(d+1)
0,
01x(2d41)
01x(441)
01xd
B
((d+1)x
0
O(a+1)xd
(d+1)x(d+1))
0a+1
Oaxa
Oa
Oaxd41
K(C
1
(dx(d+1)
Oa
K(A
Oxxd
1
01x(d+1)
Oxxd
01x(d+1)
1
O O1411)xd
Oa41
(d++1x((d-1)
O(4+))xd
0 ((41)xx(442)
0 (d4+))xd
(a+1)xxd22
(k)
I V.(
(k)
I
0dx(2+2)
+2
Oax(d+2)
Ozxd
02x(d+2)
O2xd
0 02x(d+2)
28

---

Published as a conference paper at ICLR 2025
Therefore, when the imput sequence is
X
1,,
X
X
[
J.N
[XT
JINN
XN+1
H
0
311
UN
3
k)
3
B (R)
3
((k)
2N-1
2N+1
0
and the corresponding output of the self-attention layer satisfies
(k+1)
diag ,,,,,,,,,, ,11,,,1)HH (k+())))
Then, it is equivalent consider the input sequence
1
X
1
X
11,
[xN+1]118]
UN
kk
1:|1
1
[
11:11
1
(R)
[e
2N 1:
(k)
Consequently, by applying Theorem 5.1 to this sub-problem, we derive the corresponding corollary.
E DEFERRED PROOFS IN SECTION 5.2
E. 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 || for every and k: 2 2.
Lemma ((k) satisfies Equation C.12 and [| S Cx. then we have [|g(k+2)|| S Ox.
(N++) where Cx amin(X (X X)
Proof of Lemma
From Equation
C. 12
we derive that
(k)
g. = =Sgo0)
M
:1:: (x)
3( -y
(1- (R)
(()
-ggr
M X
1:
1:
3/
M 1 X
1111::11 (1,47
Since 11 (x:] < | it follows that
(k)
7(0)
1,
MV
1. X| - (3/10)
M X
1. Y
1:44
(b)
(K)
i o
Umin
[
1,11 X
(k ||
i. 0.02
|| [
1: 1
(()
(R)
VIE
(N +1+ Omin (xx)
113
t dbxbs.
The satisfaction of Inequality (i) is achieved by setting M
2 Laxds while Inequality (ii) is oz
verified through the application of the Cauchy interlacing theorem. Consequently, by integrating the expression Cx
db. Fg. (N+1) with the aforementioned inequality, obtain
Jmin (X XX)
||
(N o Tmin (x xxxx
(A)
dba (N (N +1) amin (XX) o2
(()) dbxb3 < Cxi
Thus, the proof is complete.
29

---

Published as a conference paper at ICLR 2025
Next, prove the following theorem to demonstrate the prediction loss for the linear read-out function.
Theorem E.1. Suppose Assumption holds. Consider 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 2 No. there exists set of parameters in the Transformer such that the following inequality holds with probability at least 1 - -8' nå:
2 3
K
CsK
2
K-1
11 Tnll S bx
1
1
?
Vn
3
where
No 8(45 2) log + log S
log
c
2v 2bxn log d + log S log
2B' 2Cxbx
C1
+
VP
vD
vp
where e is positive constant, Y) is defined in Equation Appendix E.7
E.7
8' B', are constants defined in
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 and >> No, and denote [X]]] "1,, X,. For the first attention head, its structure is as follows:
0₂² 1
O.
0₂
Jast (K-1)=
K Klaat (k+1)=
B'
Vlast (K+1)
3 ((K+1)
(6%2-1
Xn 0
0a+1
Therefore, we obtain
(Questy Just h (K (K+1))TT (K/ (k+1)) = = B'.
Let p P( (f B'), which is the probability that the inner product x, 3 (K+1) exceeds the threshold B'. This probability is the same for any due to the symmetry of the distribution. We introduce a random variable oj E {0, 1} to represent the following: we set 0, = 0 if x B (k+1)
B' and = 1 1 xf (K+ 2 B'. We define the following events:
& = Σ 0 = 0).
Eg(E) = {mm- p - -
Σ 0j
p 2
S pn
+
2
iSNo
52n
Assume E, (E( holds. Then, we obtain
1
Jast h
(K+1) k- h (K+2))
V (K+1)
O2
1
Σ
[TTg(((1)) - B¹)
(1922=1)
j=1
Oa+1
Og
xTg[(++)) B²)
x
(1922=1)
2n
0
5:0z=1
Let the last MLP layer satisfies
0₂
0a
ML.P
27 + B'
p
02+1-
0,+11
30

---

Published as a conference paper at ICLR 2025
it follows that
MLP
Q Olaat h
(K+1)
KK/ath (x+1))
Vlasth (1+1))
0g T 04+1
=100== 2np-p p E
[((k+1)
20 B' B' 2np-p
Denote 8n 2mp-p Ej:0,=1 X (
2 0z) B' + B', therefore 2mp-p p
1/ 7nl = T
(g(k+
2(E
B' B'
x
2np P jong m1
2np
< |
3 2np P
Σ
x (3(
+
22 03) B'
2np p
j:0==1
S bx
gr
Σ
g(++1)
+ 2e B'I.
(E.1)
2mp P j:0==1
2np p
We randomly select set f0/ j 2 2No 1} such that |o/H satisfying the following conditions:
If Ho 0f = =
then
2'
Hoj 0g = 13 1o/31= = pm z.
[o, 0j =1] pn
2" then
0 pm
We define the set O as follows:
foz: 00 = 1} 1o5)) if [o; :0; 1}} 2
2
(o;:0; = (d),
if [1o0:00 :0; = 1} < pm
2
Therefore, we have
3 2np p j=a==1
Σ
x xTTg((+
3 mp
Σ
3 3(K+1)
+
Σ
xTT((+1))
n1
J:0
0 up
Σ EO
B
(k+1)
Cxbxe
(E.2)
mp
where Equation
E.2 follows from Lemma
Observe that
0 np
Σ
(((+1)
Σ
3 ((
P
2 E:
np
2 f:0, E
31

---

Published as a conference paper at ICLR 2025
Then, from Theorem 5.1
the following inequality holds with probability at least 1 - 3nnd:
Σ
||3 - 2i+1
1 1
Σ 11 - - 11a S
Σ
| l aaek o
(E.3)
100
E02441E0
Eo2+1E0
where
( 11-2 2
log d log
log S log 6
04 = -loog
1 1-37+ 37 + 725 -
-
+
ic
ic
From Seahhy-Schwazz Inequality we have:
log d logó
log S log
log + log 5
log
7(2 V)
+
< 2 -1)y
(E.4)
ic
ic
ic
Combining Inequalities
H.3 and
E.4
gives
18-
2i41
|
Σ R00++1E0
1
log d + log S - log 8 3 2y(25 V)
2
ic
Note that, by induction, it straightforward to prove that the following inequality holds if 3 S 1:
2 1 3 27(28 111/
1log + logg S - llog 8
x 1 1 37
2
+ 7
1
37
2 K-1
log d + log S
log 6
(E.S)
ic
when Y7 satisfies
(1 -13
IN 2
(E.6)
[1 - }5) (K 1)
log dog log
which can be satisfied by simply setting
(1 - 7))
1 =
(E.7)
(1 7) (K - vV
log d' Sllog Nac
Therefore, Inequality E..6) holds for any 2 Ng. Based on Inequality
E.5
we have
1 -- 11- 2
log d + log S log 5
K
Σ
3 27 -01 1)
Not
K
2
log d log S log 8
1- 1- 37
+7222 -111
ic
Noo+1
Notnp-R f Σ
2 K
2
K-1 1 No &
log + log S - log
1 37
+ nk
-
1
-7
3
ic
No+1
N=
< (np-
P
1-
2 K
+ mk
--
1
2
K-1
log d + log S - log
No Not
Σ
1
2
3
V8
im No+1
< (np
P
2
(
3
k
+ nk
1
3
2
K-1
log d + logg S logg
c
2(V
Nooo
No
np
p
VNo
No
(F.8)
32

---

Published as a conference paper at ICLR 2025
By combining Inequalities (E.3)
and E.
we obtain
Σ
11
-
6o2c+160 EO
P.
(1
2
K
2
R-1
log d + log S - logd
$ (np-
2
3
1
37
P.
2
K
2 3
K-1
log d + log S - log 6
S (mp-
2
3
+
nt
-
1
Therefore,
K
2/2n
2
2
K-1
log d + log S log
113
2i+1
?
+
1
1
3
npp
Vnp
No+1
(E.9)
Combining Equation
Equation
E.2 and Equation
E.9
we obtain
1/3/0 1n
2
2 2b22
K-1
log d log S
log 8
2e
2xx
< 0x
1 37
1 (1-37
B'
3
Vpp
c
np
np
Set c = K /np(11 - 27/3)* K we obtain
)
2
C.K
2
K-1
11 Fnll S 1x
1
1
1
3
vn
3
where
2v 2bxl
log d + log S - log 5
B 2Cxbx
C4
+
vp
vp
vp
Then, the theorem follows by denoting 8' as 8' = - + P( /np(1 22/22)
Finally, establish the result concerning the label prediction loss presented in Theorem 5.2 through the following corollary:
Corollary E.1. Suppose Assumption 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 such that, if No, the following inequality holds with probability at least 1 ō:
11 Dn+ill S CgE a.
where Gg = Vdbyci, and ci and an are defined in Theorem 5.1
Proof of Corollary E.
We set the matrix in the read out function Fy as
0 (at1)dd
V
Losa
O1sd
From Proposition
when is an odd index, h
(K+1)
is of the following structure:
X
h (K+1)
g g(K+ (K+1)
33

---

Published as a conference paper at ICLR 2025
Therefore, the output after the read-out function is
T = (r, (
(v( = [X] 141
3. ((++1)
2
Note that | + S
vdbx, we have
3 s|k+il||
3 - -3(
X
Combining with Theorem 5.1
the proof is thus complete.
Therefore, by combining Corollary
E.1 and Theorem
we arrive at the conclusion of Theo-
rem 5.2
34

---

Published as a conference paper at ICLR 2025
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 fixed measurement matrix X. In this section, we further show that unlike LISTA-VM, these classical LISTA-type algorithms fail 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 iter- ations K 12. During each epoch, we randomly sample 100 measurement matrices, each gen- erating 500 instances from 500 randomly generated sparse vectors, resulting in total of 50,000 in- stances per epoch. For each epoch, minimize the sparse vector prediction loss Σ | - /|? 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 LISTA and LISTA-CP remains around 3. As illustrated the figure, there no observable trend indicating improvement in the error throughout the training process.
observe that the prediction error on varying X for meta-trained
Prediction Error for Meta-Tra ined LISTA
Prediction Error for Meta-Trained LISTA-CP
Orgen
233
22
Tredtt
3.0
2.9
100
252
50
15c
200
300
350
Epoch
Enoch
(a) LISTA
(b) LISAT-CT CP
Figure 2: Experimental results for meta-trained classic LISTA-type type algorithms.
For comparison, in Figure 3 we also provide the results of the prediction error on varying 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.
Prediction Error for LISTA-VM
Orginal
1.1
2.0
8.
100
250
250
300
350
Epoch
Figure 3: Experimental results for meta-trained LISTA-VM
35


Reviews:
[{"review_id": "3kV3FuMkoU", "reviewer": "Reviewer_iC5a", "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.", "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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "**Strengths:**\n\n1. **Strong Theoretical Contributions:**  \n   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.\n\n2. **Clarity and Rigor of Proofs:**  \n   The proofs are well-written and easy to follow, both in the main text and in the appendix.\n\n3. **Robust Empirical Validation:**  \n   The empirical experiments are thorough and effectively corroborate the theoretical findings", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FxX9LcYo1N", "reviewer": "Reviewer_bS8s", "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.", "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). \n\n- 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. \n\n- 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. \n\n- 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?\n\n- 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? \n\n- 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.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "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. \n\n2. The convergence rate results (for sparse estimation and prediction) place the proposed approach on a concrete theoretical footing. \n\n3. 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.", "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). \n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wSnqzjDesV", "reviewer": "Reviewer_7MPy", "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.", "questions": "See weakness", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The manuscript is articulated with clarity and is accessible for comprehension. \nThe proposed hypothesis presents a valuable avenue for in-depth analysis. \nPreliminary concepts are provided with comprehensive equations and definitions. \nElaborate proofs are accessible in the appendix. \nThe quality of English is commendable and articulated naturally.", "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.\n1. 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. \n2. 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.\n3. 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. \n4. 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.\n5. 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. \n6. 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\n[1] Zero-shot image restoration using denoising diffusion null-space model\n[2] Metalearning-based alternating minimization algorithm for nonconvex optimization\n[3] Invertible Diffusion Models for Compressed Sensing\n[4] OmniSSR: Zero-shot Omnidirectional Image Super-Resolution using Stable Diffusion Model\n[5] Blind Super-Resolution Via Meta-Learning and Markov Chain Monte Carlo Simulation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YxnF07TsL9", "reviewer": "Reviewer_FsZo", "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.", "questions": "See in the comments in Weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "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.", "weaknesses": "I have two concerns which should be better explained in the paper:\n- 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?\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)