Transformer learns the cross-task prior and regularization for in-context learning

TL;DR

This paper investigates how transformers perform in in-context learning for inverse linear regression, revealing they learn priors and regularization strategies, outperforming traditional methods, especially in ill-posed problems.

Contribution

It introduces a linear transformer model that learns inverse mappings, demonstrating implicit prior and regularization learning, and provides insights into transformers' knowledge extraction mechanisms.

Findings

Transformers outperform ridge regression in inverse linear regression tasks.

The error scales linearly with noise, task dimension ratio, and data condition number.

Low task dimensionality relative to context length is crucial for successful learning.

Abstract

Transformers have shown a remarkable ability for in-context learning (ICL), making predictions based on contextual examples. However, while theoretical analyses have explored this prediction capability, the nature of the inferred context and its utility for downstream predictions remain open questions. This paper aims to address these questions by examining ICL for inverse linear regression (ILR), where context inference can be characterized by unsupervised learning of underlying weight vectors. Focusing on the challenging scenario of rank-deficient inverse problems, where context length is smaller than the number of unknowns in the weight vectors and regularization is necessary, we introduce a linear transformer to learn the inverse mapping from contextual examples to the underlying weight vector. Our findings reveal that the transformer implicitly learns both a prior distribution and an effective regularization strategy, outperforming traditional ridge regression and regularization methods. A key insight is the necessity of low task dimensionality relative to the context length for successful learning. Furthermore, we numerically verify that the error of the transformer estimator scales linearly with the noise level, the ratio of task dimension to context length, and the condition number of the input data. These results not only demonstrate the potential of transformers for solving ill-posed inverse problems, but also provide a new perspective towards understanding the knowledge extraction mechanism within transformers.