Research map / Optimization / Implicit bias and generalization
On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that a K-layer decoder-only Transformer can implement a LISTA-type learning-to-optimize algorithm (LISTA-VM) for in-context sparse recovery, with provable linear-in-K convergence and generalization across measurement matrices.
Problem
Understanding whether Transformers' in-context learning can be explained as learning-to-optimize (L2O) algorithms, specifically for in-context sparse recovery (LASSO), where prior work gives only sublinear convergence and L2O methods require fixed measurement matrices between training and testing.
Approach
Construct a K-layer decoder-based Transformer whose attention layers implement context-dependent LISTA updates with D_n = (1/(2n+1)) X_{1:n}^T M^T, so the learned matrix adapts to each in-context measurement matrix; this gives linear convergence in K and generalization to varying measurement matrices and prompt lengths.
Key results
- Theorem gives ||beta^(K+1) - beta*|| <= b_beta exp(-alpha_n K), i.e., linear convergence in the number of layers K for the constructed Transformer.
- Experiments show Transformers (Small TF, GPT-2) and LISTA-VM outperform ISTA/FISTA, while LISTA, LISTA-CP, ALISTA perform worst when the measurement matrix varies at test time.
- With support restricted to the first 10 entries, LISTA-VM-SS reaches about 0.07 prediction error, roughly 10x better than standard LISTA-VM and comparable to Small TF.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use LISTA-VM instead of classical LISTA-type methods on varying measurement matrices | In-context sparse recovery · Varying X | label prediction loss | Helped |
| Use GPT-2 instead of Small TF | In-context sparse recovery · Varying X | label prediction loss | Helped |
| Test classical LISTA-type algorithms with measurement matrices varying from training | LISTA, LISTA-CP, and ALISTA · Varying X | label prediction loss | Hurt |
| Use the same measurement matrix at test time as during training | Classical LISTA-type algorithms · Fixed X | label prediction loss | Helped |
| Meta-train LISTA and LISTA-CP across measurement matrices rather than train with a fixed matrix | LISTA and LISTA-CP · Varying X | label prediction loss | Hurt |
| Restrict sparse-vector support to the first 10 entries (ablation) | Small TF and GPT-2 · In-context sparse recovery, support restricted to first 10 entries | label prediction loss | Helped |
| Restrict sparse-vector support to the first 10 entries (ablation) | Baseline algorithms · In-context sparse recovery, support restricted to first 10 entries | label prediction loss | No clear effect |
| Use LISTA-VM-SS by zeroing measurement-matrix columns outside the known support (ablation) | LISTA-VM · In-context sparse recovery, support restricted to first 10 entries | prediction error: approximately 0.07; almost 10 × better than the standard version of LISTA-VM | Helped |
| Use query read-out instead of linear read-out (ablation) | Small TF · In-context sparse recovery, label prediction | label prediction error | Helped |
| Increase the number of in-context measurements when comparing linear and query read-outs (ablation) | Small TF · In-context sparse recovery, label prediction | label prediction error | No clear effect |
| Use Query Readout† instead of Query Readout (ablation) | Small TF · In-context sparse recovery, label prediction | label prediction error | No clear effect |
Surprising findings
- Conventional LISTA-type methods, even when meta-trained, fail to generalize to varying measurement matrices, while the Transformer does.
Limitations
- Theory is proved for noiseless observations and bounded assumptions on X and beta*.
- Linear read-out yields a weaker error bound with an extra K/sqrt(n) term than query read-out.
- Experiments are on a small synthetic setting (d=20, N=10, S=3).
Experimental setup
Datasets: synthetic in-context sparse recovery instances
Benchmarks: synthetic in-context sparse recovery (d=20, N=10, S=3)
Metrics: label prediction loss sparse vector estimation error
Baselines: ISTA FISTA LISTA LISTA-CP ALISTA GPT-2
Builds on: LISTA LISTA-CP ALISTA ISTA GPT-2
Closest papers by meaning
- Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.91)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.89)
- In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization (NeurIPS 2024 · similarity 0.88)
- In-Context Deep Learning via Transformer Models (ICML 2025 · similarity 0.88)
- Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context (NeurIPS 2024 · similarity 0.87)
- Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.