In-context learning of solution to simple linear regression task using transformer with masked linear self-attention

Abstract

In-context learning is a remarkable property of transformers and has recently received a lot of interest. In many studies of in-context learning, it has been shown that transformers are capable of implementing solvers for linear and non-linear regression problems, in which the most of them implement gradient descent algorithm. However, it is still unclear whether those implementations have actually been acquired through training. In this paper, we constructed a transformer with masked linear self-attention, which in-context learns the least squares estimate in a simple regression task. The point here is that the closed form (analytical) solution is obtained by using layer normalization rather than an approximate solution based on the gradient descent algorithm. Then, we showed an experimental example, in which, by taking the non-uniqueness of the transformer representation into account, our implementation for a closed form solution is used in the transformer trained with the l1 regularization when the target output is the least squares estimate.