Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?

TL;DR

This paper demonstrates that multi-layer, looped Transformers can learn to implement multi-step gradient descent algorithms for in-context linear regression, showing convergence and data-adaptive preconditioning through theoretical analysis and experiments.

Contribution

It provides the first theoretical analysis showing that looped Transformers can learn and converge to multi-step gradient descent algorithms in linear regression tasks.

Findings

Transformers can implement multi-step preconditioned gradient descent.

Gradient flow converges despite non-convex landscape.

Theoretical validation through synthetic experiments.

Abstract

The remarkable capability of Transformers to do reasoning and few-shot learning, without any fine-tuning, is widely conjectured to stem from their ability to implicitly simulate a multi-step algorithms -- such as gradient descent -- with their weights in a single forward pass. Recently, there has been progress in understanding this complex phenomenon from an expressivity point of view, by demonstrating that Transformers can express such multi-step algorithms. However, our knowledge about the more fundamental aspect of its learnability, beyond single layer models, is very limited. In particular, can training Transformers enable convergence to algorithmic solutions? In this work we resolve this for in-context linear regression with linear looped Transformers -- a multi-layer model with weight sharing that is conjectured to have an inductive bias to learn fix-point iterative algorithms. More specifically, for this setting we show that the global minimizer of the population training loss implements multi-step preconditioned gradient descent, with a preconditioner that adapts to the data distribution. Furthermore, we show a fast convergence for gradient flow on the regression loss, despite the non-convexity of the landscape, by proving a novel gradient dominance condition. To our knowledge, this is the first theoretical analysis for multi-layer Transformer in this setting. We further validate our theoretical findings through synthetic experiments.