Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Renpu Liu, Ruida Zhou, Cong Shen, Jing Yang

TL;DR

The paper studies how Transformer-based in-context learning can realize Learning-to-Optimize (L2O) algorithms for sparse recovery. It constructs a LISTA-VM architecture in a $K$-layer Transformer, proving a linear convergence rate in $K$ and the ability to adapt to varying measurement matrices $X$ during inference, while exploiting task structure learned during pre-training. The authors establish a theoretical equivalence between Transformer updates and LISTA-type LISTA-VM steps, with context-dependent updates $D_n^{(k)}$ that depend on the current $X$, enabling generalization beyond fixed training matrices. Empirically, Transformers outperform classical gradient-based solvers and remain competitive with fixed-$X$ LISTA-type methods, while showing strong generalization across demonstration lengths and improved performance when prior support information is available. Overall, the work advances understanding of Transformers as adaptive optimizers in ICL, with practical implications for fast, flexible sparse recovery in varied measurement regimes.

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.

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

TL;DR

The paper studies how Transformer-based in-context learning can realize Learning-to-Optimize (L2O) algorithms for sparse recovery. It constructs a LISTA-VM architecture in a -layer Transformer, proving a linear convergence rate in and the ability to adapt to varying measurement matrices during inference, while exploiting task structure learned during pre-training. The authors establish a theoretical equivalence between Transformer updates and LISTA-type LISTA-VM steps, with context-dependent updates that depend on the current , enabling generalization beyond fixed training matrices. Empirically, Transformers outperform classical gradient-based solvers and remain competitive with fixed- LISTA-type methods, while showing strong generalization across demonstration lengths and improved performance when prior support information is available. Overall, the work advances understanding of Transformers as adaptive optimizers in ICL, with practical implications for fast, flexible sparse recovery in varied measurement regimes.

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.

Paper Structure

This paper contains 32 sections, 14 theorems, 130 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related Works
  3. ICL Mechanism for Transformers.
  4. L2O Algorithms for Sparse Recovery.
  5. Preliminaries
  6. Notations.
  7. Transformer Architecture
  8. In-context Learning by Transformers
  9. Transformer as a LISTA-type Algorithm for In-context Sparse Recovery
  10. In-context Sparse Recovery
  11. Classical Algorithms
  12. Transformer Can Provably Perform LISTA-type Algorithms
  13. Performance of Transformers for In-context Sparse Recovery
  14. Sparse Vector Estimation
  15. Label Prediction
  16. ...and 17 more sections

Key Result

Theorem 4.1

With the Transformer structure described above, under Assumption assump:basic, there exists a set of parameters in the Transformer so that for any $k\in[1:K ]$, $n\in[N]$, we have where $\mathbf{M}^{(k)}\in\mathbb{R}^{d\times d}$ is embedded in the $k$-th Transformer layer.

Figures (3)

  • Figure 1: 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.
  • Figure 2: Experimental results for meta-trained classic LISTA-type algorithms.
  • Figure 3: Experimental results for meta-trained LISTA-VM

Theorems & Definitions (24)

  • Definition 3.1: Masked attention layer
  • Definition 3.2: MLP layer
  • Definition 3.3: Transformer layer
  • Definition 4.1: Linear read-out
  • Definition 4.2: Query read-out
  • Theorem 4.1: Equivalence between ICL and LISTA-VM
  • Remark 1
  • Theorem 5.1: Convergence of ICL
  • Remark 2: Linear convergence rate
  • Remark 3: Generalization across measurement matrix $\mathbf{X}$
  • ...and 14 more