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

TL;DR

This paper demonstrates that Transformers can perform learning-to-optimize algorithms for sparse recovery tasks, providing theoretical guarantees and showing advantages over traditional methods in generalization, convergence, and flexibility.

Contribution

It proves that Transformers can implement L2O algorithms with linear convergence and generalize across different measurement matrices and demonstration lengths.

Findings

Transformers can perform L2O algorithms with provable linear convergence.

Trained Transformers generalize across different measurement matrices.

Transformers leverage structural information to accelerate convergence.

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.