Transformers Learn Closed-Form Linear Regression In-Context

In-context learning is a highly studied property of transformers, where models can learn from examples provided within the input sequence without explicit weight updates. While many studies have shown transformers can act as solvers for regression problems, often implicitly implementing gradient descent, it has been less clear if they can acquire analytical, closed-form solutions. This research specifically investigates this by constructing a transformer equipped with linear self-attention. The study demonstrates that this specialized transformer can in-context learn the least squares estimate for a simple linear regression task. Crucially, the paper highlights that the approximate acquisition of this closed-form solution is facilitated by the use of layer normalization, rather than relying on an iterative gradient descent approximation. Experimental examples further illustrate this, particularly when the transformer is trained with L1 regularization and the target output is the least squares estimate, providing a deeper understanding of the mechanisms behind in-context learning for fundamental statistical tasks.