Transformers learn to implement preconditioned gradient descent for in-context learning

By Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, Suvrit Sra. Neural Information Processing Systems, 2023

311 citations recorded by the source corpus.

Several recent works demonstrate that transformers can implement algorithms like gradient descent. By a careful construction of weights, these works show that multiple layers of transformers are expressive enough to simulate iterations of gradient descent. Going beyond the question of expressivity, we ask: Can transformers learn to implement such algorithms by training over random problem instances? To our knowledge, we make the first theoretical progress on this question via an analysis of the loss landscape for linear transformers trained over random instances of linear regression. For a single attention layer, we prove the global minimum of the training objective implements a single iteration of preconditioned gradient descent. Notably, the preconditioning matrix not only adapts to the input distribution but also to the variance induced by data inadequacy. For a transformer with $L$ attention layers, we prove certain critical points of the training objective implement $L$ iterations of preconditioned gradient descent. Our results call for future theoretical studies on learning algorithms by training transformers.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Ocean Engineering

View Paper: https://www.semanticscholar.org/paper/f5e9337477d7a9eb6267d0310549fdefafbb7fe2

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

Citations in the corpus

A Survey on In-context Learning (2022, background): “(2) As ICL is a new paradigm that is dif-1”