In-context learning of closed form solution to simple linear regression task using transformer with linear self-attention
Machine LearningArtificial Intelligence
Summary
The authors study how transformers can learn to solve regression problems just by seeing examples in their input (in-context learning). They build a special type of transformer that solves a simple linear regression using a math shortcut (closed-form solution) with layer normalization, instead of relying on the usual gradient descent method. They also show an example where this approach works well when the transformer is trained with l1 regularization to predict the least squares estimate. This helps clarify how transformers might internally solve such problems without explicitly performing gradient descent.
Authors
Katsuyuki Hagiwara
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 solver 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 construct a transformer with 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 approximately obtained by using layer normalization rather than an approximate solution based on gradient descent algorithm. Then, we show an experimental example, in which our implementation is mainly used in the transformer trained with l1 regularization when the target output is the least squares estimate.